From 5b927b7ceb92d9b1998e22d50cdc4ee400db08d0 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:40:01 +0800 Subject: [PATCH 001/209] [CI/Build] Diff-aware source_file_dependencies for CUDA/NPU pipelines (#6597) Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: Cursor --- .../common/ci_source_file_dependencies.yml | 438 ++++++++++++++++++ .buildkite/common/scripts/upload_pipeline.py | 263 +++++++++-- .buildkite/cuda/bootstrap-upload-steps.yml | 28 +- .buildkite/cuda/test-merge.yml | 129 +----- .buildkite/cuda/test-nightly.yml | 104 +++-- .buildkite/cuda/test-ready.yml | 146 +----- .buildkite/cuda/test-weekly.yml | 288 +++++++++++- .buildkite/npu/bootstrap-upload-steps.yml | 9 +- .buildkite/npu/pipeline-npu-a3.yml | 8 +- .buildkite/npu/test-npu-nightly.yml | 42 +- .claude/skills/vllm-omni-test/SKILL.md | 8 +- .../vllm-omni-test/references/test-routing.md | 4 +- docs/contributing/ci/ci_settings.md | 54 ++- tests/buildkite/test_upload_pipeline.py | 334 ++++++++++++- tests/e2e/accuracy/test_hidream_o1_image.py | 4 +- .../test_lingbot_world_v2.py | 4 +- .../test_minimax_h3_dlo_dp2_fl2va.py | 2 +- .../test_minimax_h3_dlo_dp2_ref2va.py | 2 +- .../minimax_h3/test_minimax_h3_fasth3_lora.py | 2 +- .../minimax_h3/test_minimax_h3_turbo_lora.py | 2 +- .../e2e/online_serving/test_lingbot_video.py | 6 +- .../online_serving/test_lingbot_video_moe.py | 6 +- tests/e2e/online_serving/test_magi2.py | 5 +- .../test_minimax_h3_dlo_dp2_t2va.py | 1 - tools/pre_commit/check_buildkite.py | 1 + 25 files changed, 1479 insertions(+), 411 deletions(-) create mode 100644 .buildkite/common/ci_source_file_dependencies.yml diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml new file mode 100644 index 000000000..0455922c8 --- /dev/null +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -0,0 +1,438 @@ +# Uploader-only path-prefix presets for CUDA L2–L5 and NPU L4 test YAML. +# Referenced as ``source_file_dependencies: omni_qwen3_omni_function`` (or a list of +# keys) on a step; ``upload_pipeline.py`` expands the key(s) only. +# +# YAML anchors above the registry hold model business-code paths. Job keys +# ({model}_function / _perf / _accuracy / _reliability / _doc / _cov) alias +# those anchors and add the scripts the job actually runs. Feature keys below +# are unchanged. ``source_filter_fallback`` is not a job key: when no job-key +# prefix matches, a matching diff keeps every job for that upload. +# +# YAML anchors are nested lists; the loader flattens them. + +x-qwen3-tts: &qwen3_tts + - vllm_omni/worker/gpu_model_runner.py + - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/common/alias_free_activation.py + - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py + - vllm_omni/deploy/qwen3_tts.yaml + - vllm_omni/entrypoints/openai/serving_speech.py + - vllm_omni/entrypoints/openai/tts_adapters/ + - vllm_omni/engine/stage_pool.py + +x-minicpmo-4-5: &minicpmo_4_5 + - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py + - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py + - vllm_omni/deploy/minicpmo_4_5.yaml + - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja + +x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex + - vllm_omni/deploy/minicpmo_4_5_duplex.yaml + - vllm_omni/engine/duplex/ + - vllm_omni/entrypoints/duplex/ + - vllm_omni/clients/ + - vllm_omni/entrypoints/duplex_request_client.py + - vllm_omni/outputs/duplex.py + - vllm_omni/model_executor/duplex_sampling.py + - vllm_omni/experimental/fullduplex/ + - vllm_omni/core/sched/ + - vllm_omni/outputs/multimodal_accumulation.py + - vllm_omni/engine/stage_init_utils.py + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/engine/orchestrator.py + - vllm_omni/entrypoints/cli/serve.py + - vllm_omni/entrypoints/openai/api_server.py + - vllm_omni/entrypoints/openai/duplex_capability.py + - vllm_omni/config/stage_config.py + +x-qwen2-5-omni: &qwen2_5_omni + - vllm_omni/model_executor/models/qwen2_5_omni/ + - vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py + - vllm_omni/deploy/qwen2_5_omni.yaml + +x-qwen3-omni: &qwen3_omni + - vllm_omni/entrypoints/openai/serving_chat.py + - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py + - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py + - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py + - vllm_omni/deploy/qwen3_omni_moe.yaml + - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml + +x-wan22: &wan22 + - vllm_omni/diffusion/models/wan2_2/ + - vllm_omni/diffusion/models/dmd2/ + +x-cosmos3: &cosmos3 + - vllm_omni/diffusion/models/cosmos3/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py + - vllm_omni/experimental/world_models/ + - vllm_omni/model_extras/cosmos3.py + +x-qwen-image: &qwen_image + - vllm_omni/diffusion/models/qwen_image/ + - vllm_omni/diffusion/models/dmd2/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py + +x-hunyuan-image3: &hunyuan_image3 + - vllm_omni/diffusion/models/hunyuan_image3/ + - vllm_omni/model_executor/models/hunyuan_image3/ + - vllm_omni/deploy/hunyuan_image3_dit.yaml + +x-minimax-h3: &minimax_h3 + - vllm_omni/diffusion/models/minimax_h3/ + +x-hunyuan-video: &hunyuan_video + - vllm_omni/diffusion/models/hunyuan_video/ + - vllm_omni/model_executor/models/hunyuan_video/ + +x-ltx2: <x2 + - vllm_omni/diffusion/models/ltx2/ + - vllm_omni/model_extras/ltx2.py + +x-sana-video: &sana_video + - vllm_omni/diffusion/models/sana_video/ + - vllm_omni/model_extras/sana_video.py + +x-z-image: &z_image + - vllm_omni/diffusion/models/z_image/ + +x-stable-audio: &stable_audio + - vllm_omni/diffusion/models/stable_audio/ + +x-lingbot-video: &lingbot_video + - vllm_omni/diffusion/models/lingbot_video/ + - vllm_omni/model_extras/lingbot_video.py + +x-lingbot-world: &lingbot_world + - vllm_omni/diffusion/models/lingbot_world/ + + +source_file_dependencies: + # Shared uploader/registry paths. Not a job key — do not attach this to a step. + # Used only when no job-key prefix matches: then a hit keeps every job. + source_filter_fallback: + - .buildkite/common/scripts/upload_pipeline.py + - .buildkite/common/ci_source_file_dependencies.yml + - .buildkite/common/ci_mirror_hardwares.yml + - tests/helpers/ + + # --- features --- + diffusion_distributed_attention: + - tests/diffusion/distributed/ + - tests/diffusion/attention/test_attention_sp.py + - vllm_omni/diffusion/distributed/ + - vllm_omni/diffusion/attention/ + - vllm_omni/diffusion/data.py + - vllm_omni/diffusion/models/ + + custom_pipeline: + - 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/ + + rlhf_verl_omni: + - 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/ + - vllm_omni/inputs/ + - vllm_omni/outputs/ + + # Tiny-model diffusion jobs (shared across many model trees under models/). + diffusion_tiny_model: + - tests/model_tests/diffusion/ + - vllm_omni/diffusion/models/ + + diffusion_offloader: + - tests/diffusion/offloader/ + - vllm_omni/diffusion/offloader/ + + diffusion_quantization: + - tests/diffusion/quantization/ + - vllm_omni/diffusion/quantization/ + + # --- models --- + # --- omni / tts jobs --- + omni_qwen3_omni_function: + - *qwen3_omni + - tests/e2e/offline_inference/test_qwen3_omni.py + - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py + - tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py + - tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py + - tests/e2e/online_serving/test_qwen3_omni.py + - tests/e2e/online_serving/test_qwen3_omni_expansion.py + - tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py + + omni_qwen3_omni_accuracy: + - *qwen3_omni + - tests/e2e/accuracy/qwen3_omni/ + + omni_qwen3_omni_perf: + - *qwen3_omni + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json + - tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json + - tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json + - tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json + + omni_qwen3_omni_reliability: + - *qwen3_omni + - tests/dfx/reliability/test_reliability_qwen3_omni.py + + omni_qwen3_omni_doc: + - *qwen3_omni + - examples/online_serving/qwen3_omni/ + - examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py + - tests/examples/online_serving/test_qwen3_omni.py + + omni_qwen2_5_omni_doc: + - *qwen2_5_omni + - examples/online_serving/qwen2_5_omni/ + - examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py + - tests/examples/online_serving/test_qwen2_5_omni.py + + omni_minicpmo_4_5_function: + - *minicpmo_4_5 + - tests/e2e/offline_inference/test_minicpmo_4_5.py + - tests/e2e/online_serving/test_minicpmo_4_5.py + - tests/e2e/online_serving/test_minicpmo_4_5_expansion.py + + omni_minicpmo_4_5_accuracy: + - *minicpmo_4_5 + - tests/e2e/accuracy/minicpmo_4_5/ + + omni_minicpmo_4_5_perf: + - *minicpmo_4_5 + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_minicpmo_4_5_ready.json + - tests/dfx/perf/tests/test_minicpmo_4_5.json + + omni_minicpmo_4_5_doc: + - *minicpmo_4_5 + - examples/online_serving/minicpmo/realtime_web/ + - examples/online_serving/minicpmo/realtime_duplex_demo.py + - tests/examples/test_minicpmo_realtime_web_server.py + - tests/examples/test_minicpmo_realtime_web_static.py + - tests/examples/test_minicpmo_realtime_duplex_simple_demo.py + + omni_minicpmo_4_5_duplex_function: + - *minicpmo_4_5 + - *minicpmo_4_5_duplex + - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py + - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + - tests/e2e/online_serving/test_duplex_client_live.py + - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py + - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py + - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py + - tests/assets/minicpmo_4_5/ + - tests/clients/ + - tests/engine/duplex/ + - tests/entrypoints/duplex/ + - tests/model_executor/models/minicpmo_4_5/duplex/ + + omni_minicpmo_4_5_duplex_perf: + - *minicpmo_4_5 + - *minicpmo_4_5_duplex + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json + + tts_qwen3_tts_function: + - *qwen3_tts + - tests/e2e/features/prefill_phase/ + - tests/e2e/offline_inference/test_qwen3_tts_base.py + - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py + - tests/e2e/online_serving/test_qwen3_tts_base.py + - tests/e2e/online_serving/test_qwen3_tts_base_expansion.py + - tests/e2e/online_serving/test_qwen3_tts_customvoice.py + - tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py + + tts_qwen3_tts_perf: + - *qwen3_tts + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_tts.json + + tts_qwen3_tts_cov: + - *qwen3_tts + - .buildkite/common/scripts/run_cov_split.sh + - pyproject.toml + - tests/e2e/offline_inference/test_qwen3_tts_base.py + - tests/e2e/online_serving/test_qwen3_tts_base.py + + # --- diffusion jobs --- + diffusion_wan22_function: + - *wan22 + - tests/e2e/offline_inference/test_wan22_t2v.py + - tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py + - tests/e2e/online_serving/test_wan22_t2v.py + - tests/e2e/online_serving/test_wan22_expansion.py + + diffusion_wan22_accuracy: + - *wan22 + - tests/e2e/accuracy/wan22_i2v/ + - tests/e2e/accuracy/test_diffusers_backend_similarity.py + + diffusion_wan22_perf: + - *wan22 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json + + diffusion_wan22_reliability: + - *wan22 + - tests/dfx/reliability/test_reliability_wan22.py + + diffusion_image_to_video_doc: + - *wan22 + - *hunyuan_video + - *ltx2 + - *sana_video + - *lingbot_video + - examples/offline_inference/image_to_video/ + - tests/examples/offline_inference/test_image_to_video.py + + diffusion_cosmos3_function: + - *cosmos3 + - tests/e2e/online_serving/test_cosmos3.py + + diffusion_cosmos3_accuracy: + - *cosmos3 + - tests/e2e/accuracy/test_cosmos3_similarity.py + + diffusion_cosmos3_perf: + - *cosmos3 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_cosmos3_vllm_omni.json + + diffusion_qwen_image_function: + - *qwen_image + - tests/e2e/online_serving/test_qwen_image.py + - tests/e2e/online_serving/test_qwen_image_expansion.py + + diffusion_qwen_image_accuracy: + - *qwen_image + - tests/e2e/accuracy/test_qwen_image.py + - tests/e2e/accuracy/test_diffusers_backend_similarity.py + + diffusion_qwen_image_perf: + - *qwen_image + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_qwen_image_vllm_omni.json + + diffusion_text_to_image_doc: + - *qwen_image + - *z_image + - examples/offline_inference/text_to_image/ + - examples/online_serving/text_to_image/ + - tests/examples/offline_inference/test_text_to_image.py + - tests/examples/online_serving/test_text_to_image.py + - vllm_omni/diffusion/models/nextstep_1_1/ + - vllm_omni/diffusion/models/longcat_image/ + - vllm_omni/diffusion/models/ovis_image/ + - vllm_omni/diffusion/models/omnigen2/ + - vllm_omni/diffusion/models/sd3/ + - vllm_omni/diffusion/models/flux/ + - vllm_omni/diffusion/models/flux2/ + - vllm_omni/diffusion/models/flux2_klein/ + - vllm_omni/diffusion/models/hidream_image/ + - vllm_omni/diffusion/models/krea2/ + - vllm_omni/diffusion/models/lingbot_video/ + + diffusion_hunyuan_image3_function: + - *hunyuan_image3 + - tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py + + diffusion_hunyuan_image3_accuracy: + - *hunyuan_image3 + - tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py + + diffusion_hunyuan_image3_perf: + - *hunyuan_image3 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json + - tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json + - tests/dfx/perf/tests/test_hunyuan_image_tp4.json + + diffusion_hunyuan_image3_reliability: + - *hunyuan_image3 + - tests/dfx/reliability/test_reliability_hunyuan_image.py + + diffusion_minimax_h3_function: + - *minimax_h3 + - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py + - tests/e2e/online_serving/minimax_h3/ + - tests/diffusion/models/minimax_h3/ + - vllm_omni/diffusion/lora/ + - vllm_omni/diffusion/offloader/ + - vllm_omni/diffusion/distributed/ + - vllm_omni/entrypoints/openai/ + + diffusion_minimax_h3_accuracy: + - *minimax_h3 + - tests/e2e/accuracy/minimax_h3/ + + diffusion_minimax_h3_perf: + - *minimax_h3 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json + + diffusion_hunyuan_video_function: + - *hunyuan_video + - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py + + diffusion_hunyuan_video_accuracy: + - *hunyuan_video + - tests/e2e/accuracy/hunyuanvideo15_t2v/ + - tests/e2e/accuracy/hunyuanvideo15_i2v/ + + diffusion_hunyuan_video_perf: + - *hunyuan_video + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json + - tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json + + diffusion_z_image_function: + - *z_image + - tests/e2e/online_serving/test_z_image.py + - tests/e2e/online_serving/test_zimage_expansion.py + + diffusion_stable_audio_function: + - *stable_audio + - tests/e2e/offline_inference/test_stable_audio_expansion.py + - tests/e2e/online_serving/test_stable_audio_online_expansion.py + + diffusion_lingbot_function: + - *lingbot_video + - *lingbot_world + - tests/e2e/online_serving/test_lingbot_video.py + - tests/e2e/online_serving/test_lingbot_video_moe.py + - tests/e2e/offline_inference/test_lingbot_world_v2.py + + + # Invalid-param jobs: model business code only, plus the reliability scripts. + reliability_invalid_param_h100: + - *qwen3_omni + - *minicpmo_4_5 + - *qwen_image + - *wan22 + - tests/dfx/reliability/invalid_param_test/ + + reliability_invalid_param_l4: + - *qwen3_tts + - *stable_audio + - tests/dfx/reliability/invalid_param_test/ diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 749f35f5a..c8587d8c3 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -10,8 +10,15 @@ Bootstrap mode (``bootstrap-upload-steps.yml``): - Detect docs-only, pytest skip-mark-only, or combined skip-ci from git diff. - When only CI level YAML changes, enable **L2/L3** upload steps for affected levels only. -Test pipeline mode (e.g. test-merge.yml): - - Drop steps whose ``source_file_dependencies`` do not match changed files. +Test pipeline mode (e.g. test-merge.yml, test-nightly.yml, test-weekly.yml): + - Drop steps whose ``source_file_dependencies`` do not match changed files + (string/list keys from ci_source_file_dependencies.yml, or inline path + prefixes). Filtering applies on **PR label** uploads only; ``main`` + env + schedule (NIGHTLY/WEEKLY/merge push) keeps every job and still strips the field. + If no job-key prefix matches, a change to the pipeline YAML being uploaded + or to a path under the ``source_filter_fallback`` registry key keeps every + job so command, env, and hardware edits can be validated before merge. + When any job-key prefix already matches, normal filtering wins. - Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image`` for NPU) + ``plugins`` (see ci_mirror_hardwares.yml). - Omit ``mirror_hardwares`` to compose ``{chip}_{n}`` from pytest ``-m`` SKU @@ -87,6 +94,10 @@ BOOTSTRAP_UPLOAD_IF_KEYS = { } E2E_GROUP_MARKER = "E2E Test" CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml" +CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_dependencies.yml" +# Registry key whose prefixes bypass source filtering. Paths live in +# ci_source_file_dependencies.yml; do not attach this key to a job. +SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback" CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN" CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi' @@ -95,14 +106,7 @@ CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF # ``*_UPLOAD_IF``: full gate for uploading that child pipeline. NIGHTLY_MAIN_IF = 'build.branch == "main" && build.env("NIGHTLY") == "1"' NIGHTLY_LABEL_IF = ( - f"({NIGHTLY_MAIN_IF}) || " - '(build.branch != "main" && (' - 'build.pull_request.labels includes "nightly-test" || ' - 'build.pull_request.labels includes "omni-test" || ' - 'build.pull_request.labels includes "tts-test" || ' - 'build.pull_request.labels includes "diffusion-x2iat-test" || ' - 'build.pull_request.labels includes "diffusion-x2v-test"' - "))" + f'({NIGHTLY_MAIN_IF}) || (build.branch != "main" && build.pull_request.labels includes "nightly-test")' ) WEEKLY_E2E_IF = 'build.branch == "main" && build.env("WEEKLY") == "1"' WEEKLY_MAIN_IF = 'build.branch == "main" && (build.env("WEEKLY") == "1" || build.env("NON_CRITICAL") == "1")' @@ -248,7 +252,7 @@ def _render_bootstrap_pipeline( return yaml.safe_dump(doc, sort_keys=False) -# --- Test pipeline (test-ready.yml, test-merge.yml) --- +# --- Test pipeline (test-ready.yml, test-merge.yml, test-nightly.yml) --- @lru_cache(maxsize=1) @@ -264,6 +268,117 @@ def _load_mirror_hardwares() -> dict[str, dict[str, Any]]: return presets +@lru_cache(maxsize=1) +def _load_source_file_dependencies() -> dict[str, list[str]]: + def flatten(value: Any, *, key: str) -> list[str]: + """Flatten YAML-anchor nested lists into a de-duplicated prefix list.""" + if isinstance(value, str): + if not value.strip(): + raise ValueError(f"empty path in source_file_dependencies[{key!r}]") + return [value] + if isinstance(value, list): + flattened: list[str] = [] + seen: set[str] = set() + for item in value: + for path in flatten(item, key=key): + if path not in seen: + seen.add(path) + flattened.append(path) + return flattened + raise ValueError( + f"source_file_dependencies[{key!r}] must be a list of path prefixes, got {type(value).__name__}", + ) + + if not CI_SOURCE_FILE_DEPENDENCIES_PATH.is_file(): + raise FileNotFoundError( + f"missing CI source_file_dependencies registry: {CI_SOURCE_FILE_DEPENDENCIES_PATH}", + ) + doc = yaml.safe_load(CI_SOURCE_FILE_DEPENDENCIES_PATH.read_text(encoding="utf-8")) + if not isinstance(doc, dict): + raise ValueError(f"invalid CI source_file_dependencies registry: {CI_SOURCE_FILE_DEPENDENCIES_PATH}") + presets = doc.get("source_file_dependencies") + if not isinstance(presets, dict): + raise ValueError(f"source_file_dependencies must be a mapping in {CI_SOURCE_FILE_DEPENDENCIES_PATH}") + loaded: dict[str, list[str]] = {} + for key, value in presets.items(): + if not isinstance(key, str) or not key.strip(): + raise ValueError( + f"source_file_dependencies keys must be non-empty strings in {CI_SOURCE_FILE_DEPENDENCIES_PATH}", + ) + loaded[key] = flatten(value, key=key) + return loaded + + +_REGISTRY_KEY_RE = re.compile(r"^[a-z][a-z0-9_]*$") + + +def _parse_command_text(step: dict[str, Any]) -> str: + """Parse step ``commands`` / ``command`` into a single string.""" + raw = step.get("commands", step.get("command")) + if raw is None: + return "" + if isinstance(raw, str): + return raw + if isinstance(raw, list): + return "\n".join(str(item) for item in raw if item is not None) + return str(raw) + + +def _resolve_source_file_dependencies(step: dict[str, Any]) -> list[str] | None: + """Expand a registry key (or list of keys / inline path prefixes).""" + deps = step.get("source_file_dependencies") + if deps is None: + return None + + def dedupe(*groups: list[str]) -> list[str]: + merged: list[str] = [] + seen: set[str] = set() + for group in groups: + for path in group: + if path not in seen: + seen.add(path) + merged.append(path) + return merged + + def lookup(key: str) -> list[str]: + if key == SOURCE_FILTER_FALLBACK_KEY: + raise ValueError( + f"source_file_dependencies {key!r} bypasses filtering and cannot be attached to a step " + f"({_get_step_label(step)!r})", + ) + registry = _load_source_file_dependencies() + paths = registry.get(key) + if paths is None: + known = ", ".join(sorted(registry)) + raise ValueError( + f"unknown source_file_dependencies {key!r} in step {_get_step_label(step)!r}; known: {known}", + ) + return list(paths) + + prefixes: list[str] + if isinstance(deps, str): + prefixes = lookup(deps) if _REGISTRY_KEY_RE.fullmatch(deps) else [deps] + elif isinstance(deps, list): + if not all(isinstance(item, str) for item in deps): + raise ValueError( + f"source_file_dependencies must be a string key or list of strings in step {_get_step_label(step)!r}", + ) + keyish = [bool(_REGISTRY_KEY_RE.fullmatch(item)) for item in deps] + if deps and all(keyish): + prefixes = dedupe(*(lookup(key) for key in deps)) + elif any(keyish): + raise ValueError( + f"source_file_dependencies in step {_get_step_label(step)!r} mixes registry keys and path prefixes", + ) + else: + prefixes = list(deps) + else: + raise ValueError( + f"source_file_dependencies must be a string key or list in step {_get_step_label(step)!r}", + ) + return prefixes + + @lru_cache(maxsize=1) def _cuda_mirror_chips() -> tuple[str, ...]: """Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset. @@ -278,7 +393,7 @@ def _cuda_mirror_chips() -> tuple[str, ...]: return tuple(sorted(chips, key=lambda chip: (-len(chip), chip))) -def _read_cards_marks(expr: str) -> tuple[set[int], bool]: +def _parse_cards_marks(expr: str) -> tuple[set[int], bool]: """Parse registered ``cards_n`` / ``not cards_n`` from a pytest ``-m`` expr. Returns positive card counts and whether any ``not cards_*`` is present. @@ -295,8 +410,8 @@ def _read_cards_marks(expr: str) -> tuple[set[int], bool]: return positives, has_not_cards -def _read_hardware_marks(expr: str) -> set[str]: - """Return lowercase positive CUDA SKUs that have a mirror preset. +def _parse_hardware_marks(expr: str) -> set[str]: + """Parse lowercase positive CUDA SKUs that have a mirror preset. ``not H100`` is ignored. SKUs with no ``{chip}_*`` preset (e.g. H200) are dropped. """ @@ -338,36 +453,22 @@ def _get_mirror_hw_selector() -> str: return selector -def _read_pytest_marks(commands: Any) -> tuple[set[str], int | Literal["max"] | None]: - """Read pytest ``-m`` from step commands: CUDA SKU chips and cards count. +def _parse_pytest_marks(step: dict[str, Any]) -> tuple[set[str], int | Literal["max"] | None]: + """Parse pytest ``-m`` from step commands: CUDA SKU chips and cards count. *chips* are lowercase SKUs from positive markers (``not H100`` is ignored). *cards* is the max positive ``cards_n``, ``"max"`` when only ``not cards_*`` is present, or ``None``. """ - chunks: list[str] = [] - - def _collect(value: Any) -> None: - if value is None: - return - if isinstance(value, str): - chunks.append(value) - elif isinstance(value, list): - for part in value: - _collect(part) - else: - chunks.append(str(value)) - - _collect(commands) - text = "\n".join(chunks) + text = _parse_command_text(step) chips: set[str] = set() positives: set[int] = set() has_not_cards = False for match in _PYTEST_MARKER_ARG.finditer(text): expr = match.group(1) or match.group(2) or match.group(3) or "" - chips |= _read_hardware_marks(expr) - card_counts, not_cards = _read_cards_marks(expr) + chips |= _parse_hardware_marks(expr) + card_counts, not_cards = _parse_cards_marks(expr) positives.update(card_counts) has_not_cards = has_not_cards or not_cards @@ -484,7 +585,7 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None: if "mirror_hardwares" not in step: if has_pool: return step - chips, cards = _read_pytest_marks(step.get("commands")) + chips, cards = _parse_pytest_marks(step) if not chips and cards is None: return step preset_name = _compose_mirror_hardware_name(chips, cards, step_label=step_label) @@ -552,11 +653,7 @@ def _process_test_steps( processed.append(step) continue - deps = step.get("source_file_dependencies") - if deps is not None and not isinstance(deps, list): - raise ValueError( - f"source_file_dependencies must be a list in step {_get_step_label(step)!r}", - ) + deps = _resolve_source_file_dependencies(step) if changed_files is not None and deps is not None and not _match_source_file(changed_files, deps): _log(f"skip {_get_step_label(step)!r} (no changes under {deps})") continue @@ -595,10 +692,25 @@ def _select_e2e_group_steps(steps: list[Any]) -> list[Any]: return selected +def _any_source_dependency_match(steps: list[Any], changed_files: list[str]) -> bool: + """True when any step with ``source_file_dependencies`` matches *changed_files*.""" + for step in steps: + if not isinstance(step, dict): + continue + deps = _resolve_source_file_dependencies(step) + if deps is not None and _match_source_file(changed_files, deps): + return True + nested = step.get("steps") + if isinstance(nested, list) and _any_source_dependency_match(nested, changed_files): + return True + return False + + def _render_test_pipeline( doc: dict[str, Any], changed_files: list[str] | None, *, + pipeline_path: Path | None = None, e2e_only: bool = False, ) -> dict[str, Any]: """Filter steps by PR diff and strip uploader-only ``source_file_dependencies`` metadata.""" @@ -611,6 +723,13 @@ def _render_test_pipeline( return doc if e2e_only: steps = _select_e2e_group_steps(steps) + # Bypass is a fallback: only when no job-key prefix matched. + if changed_files is not None and pipeline_path is not None: + if not _any_source_dependency_match(steps, changed_files): + bypass = _source_filter_fallback_reason(changed_files, pipeline_path) + if bypass is not None: + _log(f"keep all jobs (no source_file_dependencies match; bypassed by {bypass})") + changed_files = None steps = _process_test_steps(steps, changed_files) return {**doc, "steps": steps} @@ -618,6 +737,63 @@ def _render_test_pipeline( # --- Entry (read file → bootstrap or test render → YAML string) --- +def _repo_relative_posix(path: Path) -> str: + resolved = path.resolve() + try: + return resolved.relative_to(ROOT).as_posix() + except ValueError: + return path.as_posix().replace("\\", "/") + + +def _source_filter_fallback_prefixes() -> list[str]: + """Shared uploader/registry prefixes from ``source_filter_fallback``. Not a job key.""" + registry = _load_source_file_dependencies() + prefixes = registry.get(SOURCE_FILTER_FALLBACK_KEY) + if not prefixes: + raise ValueError( + f"source_file_dependencies[{SOURCE_FILTER_FALLBACK_KEY!r}] must list shared uploader paths " + f"in {CI_SOURCE_FILE_DEPENDENCIES_PATH}", + ) + return prefixes + + +def _source_filter_fallback_reason(changed_files: list[str] | None, pipeline_path: Path) -> str | None: + """Return the changed path that should disable source filtering, if any.""" + if not changed_files: + return None + changed = set(changed_files) + pipeline_rel = _repo_relative_posix(pipeline_path) + if pipeline_rel in changed: + return pipeline_rel + prefixes = _source_filter_fallback_prefixes() + for path in changed_files: + if _match_source_file([path], prefixes): + return path + return None + + +def _changed_files_for_source_filter( + ctx, + *, + force_all: bool, + e2e_only: bool, +) -> list[str] | None: + """Return the diff to filter against, or None to keep every step. + + ``source_file_dependencies`` is label-only: scheduled ``main`` uploads + (NIGHTLY/WEEKLY/post-merge) run the full pipeline. ``--all`` / ``--e2e`` + also disable filtering. Pipeline YAML / ``source_filter_fallback`` matches + are applied later in ``_render_test_pipeline`` only when no job-key prefix + already matched. + """ + if force_all or e2e_only: + return None + if os.environ.get("BUILDKITE_BRANCH", "") == "main": + _log("main branch: keep all jobs (source_file_dependencies is label-only)") + return None + return ctx.changed_files + + def _render_pipeline( path: Path, *, @@ -635,13 +811,18 @@ def _render_pipeline( text = path.read_text(encoding="utf-8") ctx = resolve_ci_context_from_git() - changed_files = None if force_all or e2e_only else ctx.changed_files + changed_files = _changed_files_for_source_filter(ctx, force_all=force_all, e2e_only=e2e_only) doc = yaml.safe_load(text) if not isinstance(doc, dict): raise ValueError(f"invalid pipeline YAML: {path}") - doc = _render_test_pipeline(doc, changed_files, e2e_only=e2e_only) + doc = _render_test_pipeline( + doc, + changed_files, + pipeline_path=path, + e2e_only=e2e_only, + ) return yaml.safe_dump(doc, sort_keys=False) diff --git a/.buildkite/cuda/bootstrap-upload-steps.yml b/.buildkite/cuda/bootstrap-upload-steps.yml index ab4254feb..502a506e6 100644 --- a/.buildkite/cuda/bootstrap-upload-steps.yml +++ b/.buildkite/cuda/bootstrap-upload-steps.yml @@ -22,34 +22,48 @@ steps: agents: queue: "cpu_queue_premerge" - # L3 Test — PR merge-test label (diff-aware) or main+WEEKLY=1 (--e2e, full E2E group) + # L3 Test — PR merge-test label (diff-aware) or main (full pipeline / WEEKLY --e2e) - label: "Upload Merge Pipeline" depends_on: image-build key: upload-merge-pipeline commands: - | EXTRA="" - if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then - EXTRA="--e2e" + if [ "${BUILDKITE_BRANCH}" = "main" ]; then + if [ "${WEEKLY:-}" = "1" ]; then + EXTRA="--e2e" + else + EXTRA="--all" + fi fi python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml agents: queue: "cpu_queue_premerge" - # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) + # L4 Test — main+NIGHTLY=1 (full pipeline), or PR nightly-test label (diff-aware) - label: "Upload Nightly Pipeline" depends_on: image-build key: upload-nightly-pipeline commands: - - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/cuda/test-nightly.yml + - | + EXTRA="" + if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--all" + fi + python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-nightly.yml agents: queue: "cpu_queue_premerge" - # L5 Test — main+WEEKLY=1 / NON_CRITICAL=1 (scheduled), or PR with weekly-test label + # L5 Test — main+WEEKLY=1 / NON_CRITICAL=1 (full pipeline), or PR weekly-test label (diff-aware) - label: "Upload Weekly Pipeline" depends_on: image-build key: upload-weekly-pipeline commands: - - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/cuda/test-weekly.yml + - | + EXTRA="" + if [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--all" + fi + python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-weekly.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 5aa49ea6c..1a458bccf 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -78,21 +78,7 @@ steps: depends_on: upload-merge-pipeline steps: - label: "TTS · Qwen3-TTS CustomVoice Test" - source_file_dependencies: - - tests/e2e/features/prefill_phase/ - - vllm_omni/worker/gpu_model_runner.py - - tests/helpers/media.py - - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - - vllm_omni/model_executor/models/common/snake_activation.py - - vllm_omni/model_executor/models/common/alias_free_activation.py - - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - - vllm_omni/deploy/qwen3_tts.yaml - - vllm_omni/entrypoints/openai/serving_speech.py - - vllm_omni/entrypoints/openai/tts_adapters/ - - vllm_omni/engine/stage_pool.py + source_file_dependencies: tts_qwen3_tts_function commands: - | timeout 40m bash -c " @@ -103,21 +89,7 @@ steps: mirror_hardwares: l4_1 - label: "TTS · Qwen3-TTS Base Test" - source_file_dependencies: - - tests/helpers/media.py - - tests/e2e/online_serving/test_qwen3_tts_base.py - - tests/e2e/offline_inference/test_qwen3_tts_base.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - - vllm_omni/model_executor/models/common/snake_activation.py - - vllm_omni/model_executor/models/common/alias_free_activation.py - - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - - vllm_omni/deploy/qwen3_tts.yaml - - vllm_omni/entrypoints/openai/serving_speech.py - - vllm_omni/entrypoints/openai/tts_adapters/ - - vllm_omni/engine/stage_pool.py - - .buildkite/common/scripts/run_cov_split.sh - - pyproject.toml + source_file_dependencies: tts_qwen3_tts_cov # 40 -> 60: the coverage split runs the model twice. timeout_in_minutes: 60 commands: @@ -132,36 +104,14 @@ steps: mirror_hardwares: l4_1 - label: "Omni · Qwen3-Omni Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_qwen3_omni.py - - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py - - tests/e2e/online_serving/test_qwen3_omni.py - - tests/helpers/runtime.py - - vllm_omni/entrypoints/openai/serving_chat.py - - vllm_omni/model_executor/models/qwen3_omni/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - - vllm_omni/model_executor/models/common/snake_activation.py - - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py - - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - - vllm_omni/deploy/qwen3_omni_moe.yaml - - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml + source_file_dependencies: omni_qwen3_omni_function timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_minicpmo_4_5.py - - tests/e2e/online_serving/test_minicpmo_4_5.py - - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/deploy/minicpmo_4_5.yaml + source_file_dependencies: omni_minicpmo_4_5_function 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' @@ -174,59 +124,14 @@ steps: mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Duplex Test" - source_file_dependencies: - - tests/assets/minicpmo_4_5/response_required_16k.wav - - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav - - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py - - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py - - vllm_omni/deploy/minicpmo_4_5.yaml - - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - - vllm_omni/engine/duplex/ - - vllm_omni/entrypoints/duplex/ - # Public duplex client library + its unit suites: client-side - # changes must re-run the duplex e2e group. The live client probe - # (test_duplex_client_live.py, single GPU, core tier) runs in the - # ready pipeline's duplex group. - - vllm_omni/clients/ - - tests/clients/ - - vllm_omni/entrypoints/duplex_request_client.py - - vllm_omni/outputs/duplex.py - - vllm_omni/model_executor/duplex_sampling.py - - tests/engine/duplex/ - - tests/entrypoints/duplex/ - - tests/model_executor/models/minicpmo_4_5/duplex/ - # Legacy Realtime probe client + video stacking used by the e2e drivers. - - vllm_omni/experimental/fullduplex/ - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/core/sched/ - - vllm_omni/outputs/multimodal_accumulation.py - - vllm_omni/engine/stage_init_utils.py - - vllm_omni/engine/async_omni_engine.py - - vllm_omni/engine/orchestrator.py - - vllm_omni/entrypoints/cli/serve.py - - vllm_omni/entrypoints/openai/api_server.py - - vllm_omni/entrypoints/openai/duplex_capability.py - - vllm_omni/config/stage_config.py + source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_1 - label: "Diffusion · Qwen Image Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image.py - - vllm_omni/diffusion/models/qwen_image/ - - vllm_omni/diffusion/models/dmd2/ - - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py + source_file_dependencies: diffusion_qwen_image_function commands: - | timeout 40m bash -c ' @@ -236,11 +141,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · Wan22 Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_wan22_t2v.py - - tests/e2e/online_serving/test_wan22_t2v.py - - vllm_omni/diffusion/models/wan2_2/ - - vllm_omni/diffusion/models/dmd2/ + source_file_dependencies: diffusion_wan22_function commands: - | timeout 40m bash -c " @@ -250,9 +151,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · HunyuanVideo-1.5 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py - - vllm_omni/diffusion/models/hunyuan_video/ + source_file_dependencies: diffusion_hunyuan_video_function commands: - | timeout 40m bash -c " @@ -262,17 +161,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · MiniMax-H3 L3 DLO DP2 and FastH3" - source_file_dependencies: - - tests/e2e/online_serving/minimax_h3/ - - tests/helpers/assertions.py - - tests/helpers/mark.py - - tests/helpers/media.py - - tests/helpers/runtime.py - - vllm_omni/diffusion/models/minimax_h3/ - - vllm_omni/diffusion/lora/ - - vllm_omni/diffusion/offloader/ - - vllm_omni/diffusion/distributed/ - - vllm_omni/entrypoints/openai/ + source_file_dependencies: diffusion_minimax_h3_function # Keep four server lifecycles in one Buildkite job, but never in one # pytest process. DLO uses multiprocessing shared-memory queues; the # parent's resource_tracker can retain a stale /psm_* entry during diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index c85e9abc1..ee4fc6f3e 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -12,26 +12,34 @@ steps: - group: ":card_index_dividers: Omni Model Test" key: nightly-omni-test-group depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Omni · Function Test with H100 · Single-GPU" + source_file_dependencies: + - omni_qwen3_omni_function + - omni_minicpmo_4_5_function timeout_in_minutes: 120 commands: - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - label: ":full_moon: Omni · Function Test with H100 · 2-GPU" + source_file_dependencies: + - omni_qwen3_omni_function + - omni_minicpmo_4_5_function timeout_in_minutes: 90 commands: - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Omni · Doc Test with L4 · 4-GPU" + source_file_dependencies: omni_qwen2_5_omni_doc timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" @@ -39,12 +47,14 @@ steps: mirror_hardwares: l4_4 - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU" + source_file_dependencies: omni_qwen3_omni_doc timeout_in_minutes: 90 commands: - pytest -s -v tests/examples/ -m "full_model and omni and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Omni · Accuracy Test" + source_file_dependencies: omni_qwen3_omni_accuracy timeout_in_minutes: 180 commands: - export SEED_TTS_WER_EVAL=1 @@ -58,6 +68,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test" + source_file_dependencies: omni_minicpmo_4_5_accuracy timeout_in_minutes: 180 commands: - export SEED_TTS_WER_EVAL=1 @@ -71,6 +82,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · No Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-no-async-chunk timeout_in_minutes: 300 commands: @@ -84,6 +96,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-async-chunk timeout_in_minutes: 180 commands: @@ -97,19 +110,21 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" + source_file_dependencies: omni_minicpmo_4_5_perf key: nightly-omni-performance-minicpmo-4-5 timeout_in_minutes: 180 commands: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json -m "H100 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT mirror_hardwares: h100_1 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_perf key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts timeout_in_minutes: 180 commands: @@ -123,6 +138,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100" + source_file_dependencies: omni_qwen3_omni_function timeout_in_minutes: 45 commands: - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model" --run-level "core_model" @@ -131,10 +147,11 @@ steps: - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: TTS · Function Test with L4" + source_file_dependencies: tts_qwen3_tts_function timeout_in_minutes: 120 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" @@ -144,6 +161,7 @@ steps: mirror_hardwares: l4_1 - label: ":full_moon: TTS · Perf Test · Single-GPU" + source_file_dependencies: tts_qwen3_tts_perf key: nightly-tts-performance-single timeout_in_minutes: 180 commands: @@ -158,6 +176,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: TTS · Perf Test · 2-GPU" + source_file_dependencies: tts_qwen3_tts_perf key: nightly-tts-performance-2gpu timeout_in_minutes: 180 commands: @@ -180,6 +199,7 @@ steps: steps: - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 2-GPU" + source_file_dependencies: diffusion_tiny_model timeout_in_minutes: 30 # NOTE: Single-GPU (core_model) tiny model tests run on the ready label, # so nightly only adds the multi-GPU parallelism configurations. @@ -193,6 +213,7 @@ steps: mirror_hardwares: l4_2 - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 4-GPU" + source_file_dependencies: diffusion_tiny_model timeout_in_minutes: 30 # CFG+TP extra_test_groups request 2×2=4 devices (see get_required_device_count). commands: @@ -204,11 +225,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" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU" + source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 commands: - >- @@ -219,6 +240,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU" + source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 commands: - >- @@ -228,30 +250,26 @@ steps: --run-level "full_model" mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2I(&A&T) · Boogu Edit CFG · 3-GPU" + - label: ":full_moon: Diffusion X2I(&A&T) · Z-Image Function Test · L4 · 4-GPU" + source_file_dependencies: diffusion_z_image_function timeout_in_minutes: 120 commands: - >- pytest -sv - tests/e2e/online_serving/test_boogu_image_edit_expansion.py - -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full" - -m "slow and diffusion and cuda and H100 and cards_3" + tests/e2e/online_serving/test_zimage_expansion.py + -m "full_model and diffusion and L4 and cards_4" --run-level "full_model" - mirror_hardwares: h100_3 - - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · 4-GPU" - timeout_in_minutes: 120 - commands: - - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and cards_4" --run-level "full_model" --ignore=tests/e2e/accuracy mirror_hardwares: l4_4 - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test" + source_file_dependencies: diffusion_text_to_image_doc timeout_in_minutes: 60 commands: - pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test" + source_file_dependencies: diffusion_qwen_image_accuracy timeout_in_minutes: 180 commands: - export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120 @@ -260,6 +278,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test" + source_file_dependencies: diffusion_hunyuan_image3_accuracy timeout_in_minutes: 180 commands: - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct" @@ -267,13 +286,8 @@ steps: - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model mirror_hardwares: h100_4 - - label: ":full_moon: HiDream-O1-Image · Accuracy Test" - timeout_in_minutes: 180 - commands: - - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model - mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU" + source_file_dependencies: diffusion_qwen_image_perf key: nightly-diffusion-x2iat-performance-qwen-image-single timeout_in_minutes: 180 commands: @@ -289,6 +303,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-GPU" + source_file_dependencies: diffusion_qwen_image_perf key: nightly-diffusion-x2iat-performance-qwen-image-4gpu timeout_in_minutes: 180 commands: @@ -309,59 +324,73 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2v-test" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · Single-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · 2-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · Single-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · 2-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test" + source_file_dependencies: diffusion_hunyuan_video_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2V · MAGI-2 Preview Serving Smoke" - timeout_in_minutes: 180 + - label: ":full_moon: Diffusion X2V · LingBot Function Test" + source_file_dependencies: diffusion_lingbot_function + timeout_in_minutes: 120 commands: - - pytest -s -v tests/e2e/online_serving/test_magi2.py -m "advanced_model and diffusion and cuda and H100 and cards_4" --run-level "advanced_model" - mirror_hardwares: h100_4 + - >- + pytest -s -v + tests/e2e/online_serving/test_lingbot_video.py + tests/e2e/online_serving/test_lingbot_video_moe.py + tests/e2e/offline_inference/test_lingbot_world_v2.py + -m "full_model and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Doc Test" + source_file_dependencies: diffusion_image_to_video_doc timeout_in_minutes: 60 commands: - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model" mirror_hardwares: h100_1 - - label: ":full_moon: Diffusion X2V · Other Function Test" + - label: ":full_moon: Diffusion X2V · Wan2.2 Autoround Function Test" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" + source_file_dependencies: diffusion_wan22_accuracy timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model @@ -369,12 +398,14 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · MiniMax H3 I2VA/Ref2VA Accuracy Test" + source_file_dependencies: diffusion_minimax_h3_accuracy timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py -m full_model --run-level full_model mirror_hardwares: h100_4 - - label: ":full_moon: Diffusion X2V · Other Accuracy Test" + - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Accuracy Test" + source_file_dependencies: diffusion_hunyuan_video_accuracy timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model @@ -382,6 +413,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · Single-GPU" + source_file_dependencies: diffusion_wan22_perf key: nightly-diffusion-x2v-performance-single timeout_in_minutes: 180 commands: @@ -396,6 +428,9 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Perf Test · 2-GPU" + source_file_dependencies: + - diffusion_wan22_perf + - diffusion_cosmos3_perf key: nightly-diffusion-x2v-performance-2gpu timeout_in_minutes: 180 commands: @@ -414,6 +449,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5" + source_file_dependencies: diffusion_hunyuan_video_perf key: nightly-diffusion-x2v-performance-hunyuanvideo15 timeout_in_minutes: 180 commands: @@ -431,6 +467,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3" + source_file_dependencies: diffusion_minimax_h3_perf key: nightly-diffusion-x2v-performance-minimax-h3 timeout_in_minutes: 180 commands: @@ -450,54 +487,63 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Diffusion · Distributed Test with L4 · 2-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_2' --run-level "full_model" mirror_hardwares: l4_2 - label: ":full_moon: Diffusion · Distributed Test with L4 · 3-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_3' --run-level "full_model" mirror_hardwares: l4_3 - label: ":full_moon: Diffusion · Distributed Test with L4 · 4-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_4' --run-level "full_model" mirror_hardwares: l4_4 - label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_2' --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion · Distributed Test with H100 · 4-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_4' --run-level "full_model" mirror_hardwares: h100_4 - label: ":full_moon: Diffusion · Offloader Test" + source_file_dependencies: diffusion_offloader timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/offloader -m "full_model and cuda" mirror_hardwares: l4_1 - label: ":full_moon: Diffusion · Quantization Test with H100 · Single-GPU" + source_file_dependencies: diffusion_quantization timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion · MiniMax H3 FP8 Quality Test with H100 · 2-GPU" + source_file_dependencies: diffusion_minimax_h3_function timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100' --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion · Quantization Test with L4" + source_file_dependencies: diffusion_quantization timeout_in_minutes: 60 commands: - pip install "vllm-gguf-plugin==0.0.4" diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 6d8d1e8b9..67f03dbb6 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -54,16 +54,19 @@ steps: mirror_hardwares: l4_1 - label: "Diffusion · Distributed Test · Single-GPU" + source_file_dependencies: diffusion_distributed_attention commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_1' --run-level "core_model" mirror_hardwares: l4_1 - label: "Diffusion · Distributed Test · 2-GPU" + source_file_dependencies: diffusion_distributed_attention commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model" mirror_hardwares: l4_2 - label: "Diffusion · Distributed & Attention Test · 4-GPU" + source_file_dependencies: diffusion_distributed_attention commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_4' --run-level "core_model" - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda and L4 and cards_4' --run-level "core_model" @@ -120,29 +123,14 @@ steps: 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/ + source_file_dependencies: custom_pipeline commands: - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" mirror_hardwares: l4_1 - label: "RLHF · VeRL-Omni E2E Test" - source_file_dependencies: - - 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/ - - vllm_omni/inputs/ - - vllm_omni/outputs/ + source_file_dependencies: rlhf_verl_omni commands: - | timeout 20m bash -c " @@ -151,17 +139,7 @@ steps: mirror_hardwares: l4_1 - label: "Omni · Qwen3-Omni Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen3_omni.py - - tests/helpers/runtime.py - - vllm_omni/entrypoints/openai/serving_chat.py - - vllm_omni/model_executor/models/qwen3_omni/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - - vllm_omni/model_executor/models/common/snake_activation.py - - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py - - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - - vllm_omni/deploy/qwen3_omni_moe.yaml + source_file_dependencies: omni_qwen3_omni_function commands: - | timeout 20m bash -c " @@ -170,14 +148,7 @@ steps: mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Offline Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_minicpmo_4_5.py - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/deploy/minicpmo_4_5.yaml + source_file_dependencies: omni_minicpmo_4_5_function commands: - | timeout 20m bash -c " @@ -186,15 +157,7 @@ steps: mirror_hardwares: h100_1 - label: "Omni · MiniCPM-o 4.5 Online Test" - source_file_dependencies: - - tests/e2e/online_serving/test_minicpmo_4_5.py - - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/deploy/minicpmo_4_5.yaml + source_file_dependencies: omni_minicpmo_4_5_function commands: - | timeout 20m bash -c " @@ -203,48 +166,7 @@ steps: mirror_hardwares: h100_1 - label: "Omni · MiniCPM-o 4.5 Duplex Test" - source_file_dependencies: - - tests/assets/minicpmo_4_5/response_required_16k.wav - - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav - - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py - - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py - - vllm_omni/deploy/minicpmo_4_5.yaml - - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - - vllm_omni/engine/duplex/ - - vllm_omni/entrypoints/duplex/ - # Public duplex client library + its unit suites: client-side - # changes must re-run the duplex e2e group, which includes the - # live client probe (test_duplex_client_live.py, single GPU). - - vllm_omni/clients/ - - tests/clients/ - - tests/e2e/online_serving/test_duplex_client_live.py - - vllm_omni/entrypoints/duplex_request_client.py - - vllm_omni/outputs/duplex.py - - vllm_omni/model_executor/duplex_sampling.py - - tests/engine/duplex/ - - tests/entrypoints/duplex/ - - tests/model_executor/models/minicpmo_4_5/duplex/ - # Legacy Realtime probe client + video stacking used by the e2e drivers. - - vllm_omni/experimental/fullduplex/ - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/core/sched/ - - vllm_omni/outputs/multimodal_accumulation.py - - vllm_omni/engine/stage_init_utils.py - - vllm_omni/engine/async_omni_engine.py - - vllm_omni/engine/orchestrator.py - - vllm_omni/entrypoints/cli/serve.py - - vllm_omni/entrypoints/openai/api_server.py - - vllm_omni/entrypoints/openai/duplex_capability.py - - vllm_omni/config/stage_config.py + source_file_dependencies: omni_minicpmo_4_5_duplex_function commands: - | timeout 35m bash -c " @@ -254,29 +176,14 @@ steps: - label: "Omni · MiniCPM-o 4.5 Perf Test" timeout_in_minutes: 60 - source_file_dependencies: - - tests/dfx/perf/tests/test_minicpmo_4_5_ready.json - - vllm_omni/deploy/minicpmo_4_5.yaml - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py + source_file_dependencies: omni_minicpmo_4_5_perf commands: - export BENCHMARK_DIR=tests/dfx/perf/results - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda' mirror_hardwares: h100_1 - label: "TTS · Qwen3-TTS CustomVoice Test" - source_file_dependencies: - - tests/helpers/media.py - - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - - vllm_omni/model_executor/models/common/snake_activation.py - - vllm_omni/model_executor/models/common/alias_free_activation.py - - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - - vllm_omni/deploy/qwen3_tts.yaml - - vllm_omni/entrypoints/openai/serving_speech.py - - vllm_omni/entrypoints/openai/tts_adapters/ - - vllm_omni/engine/stage_pool.py + source_file_dependencies: tts_qwen3_tts_function commands: - | timeout 20m bash -c " @@ -286,10 +193,7 @@ steps: mirror_hardwares: l4_1 - label: "Diffusion · Wan22 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_wan22_t2v.py - - vllm_omni/diffusion/models/wan2_2/ - - vllm_omni/diffusion/models/dmd2/ + source_file_dependencies: diffusion_wan22_function commands: - | timeout 40m bash -c " @@ -299,12 +203,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · Cosmos3 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_cosmos3.py - - vllm_omni/diffusion/models/cosmos3/ - - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py - - vllm_omni/experimental/world_models/ - - vllm_omni/model_extras/cosmos3.py + source_file_dependencies: diffusion_cosmos3_function commands: - | timeout 40m bash -c " @@ -315,14 +214,7 @@ steps: - label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test" timeout_in_minutes: 90 - source_file_dependencies: - - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py - - tests/helpers/assertions.py - - tests/helpers/runtime.py - - vllm_omni/diffusion/models/minimax_h3/ - - vllm_omni/diffusion/offloader/ - - vllm_omni/diffusion/distributed/ - - vllm_omni/entrypoints/openai/ + source_file_dependencies: diffusion_minimax_h3_function commands: - | timeout 80m bash -c " @@ -332,11 +224,7 @@ steps: mirror_hardwares: h100_2 - label: "Diffusion · Qwen Image Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image.py - - vllm_omni/diffusion/models/qwen_image/ - - vllm_omni/diffusion/models/dmd2/ - - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py + source_file_dependencies: diffusion_qwen_image_function commands: - | timeout 40m bash -c " @@ -346,9 +234,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · HunyuanVideo-1.5 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py - - vllm_omni/diffusion/models/hunyuan_video/ + source_file_dependencies: diffusion_hunyuan_video_function commands: - | timeout 40m bash -c " diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index 614af9572..d53bc56e4 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.buildkite/cuda/test-weekly.yml @@ -29,36 +29,42 @@ steps: steps: - label: "Reliability Test · qwen3-omni" + source_file_dependencies: omni_qwen3_omni_reliability timeout_in_minutes: 180 commands: - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" mirror_hardwares: h100_2 - label: "Reliability Test · wan22" + source_file_dependencies: diffusion_wan22_reliability timeout_in_minutes: 180 commands: - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" mirror_hardwares: h100_2 - label: "Reliability Test · hunyuan-image" + source_file_dependencies: diffusion_hunyuan_image3_reliability timeout_in_minutes: 180 commands: - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" mirror_hardwares: h100_4 - label: "Reliability Test · Invalid parameters · H100 · Single-GPU" + source_file_dependencies: reliability_invalid_param_h100 timeout_in_minutes: 90 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_1" mirror_hardwares: h100_1 - label: "Reliability Test · Invalid parameters · H100 · 2-GPU" + source_file_dependencies: reliability_invalid_param_h100 timeout_in_minutes: 90 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_2" mirror_hardwares: h100_2 - label: "Reliability Test · Invalid parameters · L4" + source_file_dependencies: reliability_invalid_param_l4 timeout_in_minutes: 120 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4" @@ -70,6 +76,7 @@ steps: steps: - label: ":full_moon: Omni · Perf Test · vLLM Text" + source_file_dependencies: omni_qwen3_omni_perf key: weekly-omni-performance-vllm-text timeout_in_minutes: 180 commands: @@ -96,6 +103,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · Multi-Replica" + source_file_dependencies: omni_qwen3_omni_perf key: weekly-omni-performance-multi-replicas timeout_in_minutes: 180 commands: @@ -221,21 +229,262 @@ steps: --run-level "full_model" mirror_hardwares: l4_4 - - label: "Diffusion · H100 · Single-GPU" - timeout_in_minutes: 150 + # Split the old shared H100 Single-GPU sweep: one pytest process collected + # every slow+cards_1 diffusion file and exceeded 150 minutes. + - label: "Diffusion · Qwen-Image · H100 · Single-GPU" + timeout_in_minutes: 120 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and H100 and cards_1" + tests/e2e/online_serving/test_qwen_image_edit.py + tests/e2e/online_serving/test_qwen_image_edit_expansion.py + tests/e2e/online_serving/test_qwen_image_layered.py + tests/e2e/online_serving/test_qwen_image_layered_expansion.py + tests/e2e/accuracy/test_qwen_image_edit.py + tests/e2e/accuracy/test_qwen_image_layered.py + -m "slow and diffusion and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - - label: "Diffusion · H100 · 2-GPU" + - label: "Diffusion · Wan2.1 VACE · H100 · Single-GPU" timeout_in_minutes: 90 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and H100 and cards_2" + tests/e2e/online_serving/test_wan_2_1_vace_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · LongCat · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_longcat_image_expansion.py + tests/e2e/online_serving/test_longcat_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Bagel · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_bagel.py + tests/e2e/online_serving/test_bagel_expansion.py + tests/e2e/offline_inference/test_bagel.py + tests/e2e/offline_inference/test_bagel_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Boogu · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image.py + tests/e2e/online_serving/test_boogu_image_edit.py + tests/e2e/online_serving/test_boogu_image_expansion.py + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Flux · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_flux_2_dev_expansion.py + tests/e2e/offline_inference/test_flux1_schnell_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · LTX · H100 · Single-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_ltx.py + tests/e2e/online_serving/test_ltx25.py + tests/e2e/accuracy/ltx/test_ltx_official_similarity.py + tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Krea2 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_krea2_expansion.py + tests/e2e/offline_inference/test_krea2_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · HiDream-I1 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_hidream_i1_full_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · SANA-Video · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_sana_video_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · SenseNova-U1 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/offline_inference/test_sensenova_u1_expansion.py + tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py + tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · MammothModa2 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/offline_inference/test_mammoth_moda2_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Helios · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/accuracy/test_video_streaming_output_similarity.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + # Split the old shared H100 2-GPU sweep: one pytest process collected every + # slow+cards_2 diffusion file and exceeded 90 minutes. + - label: "Diffusion · Qwen-Image · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_qwen_image_edit_expansion.py + tests/e2e/online_serving/test_qwen_image_layered_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Wan2.1 VACE · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_wan_2_1_vace_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · LongCat · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_longcat_image_expansion.py + tests/e2e/online_serving/test_longcat_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Bagel · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_bagel_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Boogu · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image_expansion.py + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Flux · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_flux_2_dev_expansion.py + tests/e2e/offline_inference/test_flux1_kontext_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · GLM-Image · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_glm_image_expansion.py + tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · LTX · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_ltx.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Krea2 · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/offline_inference/test_krea2_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · DreamZero · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_dreamzero_expansion.py + tests/e2e/accuracy/test_dreamzero.py + -m "slow and diffusion and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 @@ -246,4 +495,33 @@ steps: pytest -sv tests/e2e/ -m "slow and diffusion and H100 and cards_4" --run-level "full_model" + --ignore=tests/e2e/online_serving/test_magi2.py + mirror_hardwares: h100_4 + + # Dedicated: HiDream-O1 accuracy needs a long TP2 run; the shared H100 + # 2-GPU sweep is only 90 minutes. + - label: "Diffusion · HiDream-O1-Image Accuracy · H100 · 2-GPU" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m "slow and diffusion and cuda and H100 and cards_2" --run-level "full_model" + mirror_hardwares: h100_2 + + # Dedicated: Boogu CFG-parallel size 3 needs 3 GPUs; weekly has no H100 3-GPU sweep. + - label: "Diffusion · Boogu Edit CFG · H100 · 3-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full" + -m "slow and diffusion and cuda and H100 and cards_3" + --run-level "full_model" + mirror_hardwares: h100_3 + + # Dedicated: MAGI-2 Preview serving smoke needs ~30 min server init; the + # shared H100 4-GPU sweep is only 90 minutes. + - label: "Diffusion · MAGI-2 Preview Serving Smoke · H100 · 4-GPU" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/online_serving/test_magi2.py -m "slow and diffusion and cuda and H100 and cards_4" --run-level "full_model" mirror_hardwares: h100_4 diff --git a/.buildkite/npu/bootstrap-upload-steps.yml b/.buildkite/npu/bootstrap-upload-steps.yml index 398fbc6dc..d88fcc70b 100644 --- a/.buildkite/npu/bootstrap-upload-steps.yml +++ b/.buildkite/npu/bootstrap-upload-steps.yml @@ -93,7 +93,7 @@ steps: command: | python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/npu/test-npu-ready.yml - # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) + # L4 Test — main+NIGHTLY=1 (scheduled), or PR nightly-test label - label: "Upload Nightly Pipeline" depends_on: - image-build-a2 @@ -102,6 +102,11 @@ steps: - image-build-310p key: upload-nightly-pipeline commands: - - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/npu/test-npu-nightly.yml + - | + EXTRA="" + if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--all" + fi + python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/npu/test-npu-nightly.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/npu/pipeline-npu-a3.yml b/.buildkite/npu/pipeline-npu-a3.yml index a57f5f8e5..3544b384d 100644 --- a/.buildkite/npu/pipeline-npu-a3.yml +++ b/.buildkite/npu/pipeline-npu-a3.yml @@ -96,13 +96,7 @@ # key: upload-nightly-pipeline # if: >- # (build.branch == "main" && build.env("NIGHTLY") == "1") || - # (build.branch != "main" && ( - # build.pull_request.labels includes "nightly-test" || - # build.pull_request.labels includes "omni-test" || - # build.pull_request.labels includes "tts-test" || - # build.pull_request.labels includes "diffusion-x2iat-test" || - # build.pull_request.labels includes "diffusion-x2v-test" - # )) + # (build.branch != "main" && build.pull_request.labels includes "nightly-test") # commands: # - buildkite-agent pipeline upload .buildkite/npu/test-npu-nightly.yml # agents: diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml index e049783ce..f64fa711f 100644 --- a/.buildkite/npu/test-npu-nightly.yml +++ b/.buildkite/npu/test-npu-nightly.yml @@ -4,10 +4,10 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "omni-test" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Omni · Perf Test · No Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-no-async-chunk timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 @@ -24,6 +24,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · Perf Test · Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-async-chunk timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 @@ -40,6 +41,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 Test" + source_file_dependencies: omni_minicpmo_4_5_function mirror_hardwares: a3_npu_2 env: VLLM_WORKER_MULTIPROC_METHOD: spawn @@ -48,6 +50,7 @@ steps: commands: - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model' - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 env: @@ -62,6 +65,7 @@ steps: -m "npu and A3 and omni" --run-level full_model - label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test" + source_file_dependencies: omni_minicpmo_4_5_accuracy timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 env: @@ -82,6 +86,7 @@ steps: buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" + source_file_dependencies: omni_minicpmo_4_5_perf key: nightly-omni-performance-minicpmo-4-5-npu timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 @@ -98,6 +103,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test (A2)" + source_file_dependencies: omni_minicpmo_4_5_perf key: nightly-omni-performance-minicpmo-4-5-npu-a2 timeout_in_minutes: 180 mirror_hardwares: a2b3_npu_1 @@ -115,6 +121,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_perf key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts-npu timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 @@ -133,9 +140,10 @@ steps: - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: TTS · Perf Test" + source_file_dependencies: tts_qwen3_tts_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 env: @@ -151,15 +159,15 @@ steps: EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - - group: ":card_index_dividers: Diffusion I2I Model Test" - key: nightly-HunyuanImage3-x2i-test-group + - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" + key: nightly-diffusion-x2iat-group 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" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test" + source_file_dependencies: diffusion_hunyuan_image3_accuracy timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -169,15 +177,8 @@ steps: DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA commands: - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model - - group: ":card_index_dividers: Diffusion Benchmark Perf Test" - key: nightly-diffusion-benchmark-group - 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: - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 CFGP=2" + source_file_dependencies: diffusion_hunyuan_image3_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -194,6 +195,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2" + source_file_dependencies: diffusion_hunyuan_image3_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -210,6 +212,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4" + source_file_dependencies: diffusion_hunyuan_image3_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -230,10 +233,10 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2v-test" + build.pull_request.labels includes "nightly-test" steps: - - label: ":full_moon: Diffusion X2V · Function Test" + - label: ":full_moon: Diffusion X2V · Wan 2.2 Function Test" + source_file_dependencies: diffusion_wan22_function mirror_hardwares: a2b3_npu_4 env: VLLM_WORKER_MULTIPROC_METHOD: spawn @@ -241,7 +244,8 @@ steps: HF_TOKEN: "${HF_TOKEN}" commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model" - - label: ":full_moon: Diffusion X2V · Perf Test" + - label: ":full_moon: Diffusion X2V · Wan 2.2 Perf Test" + source_file_dependencies: diffusion_wan22_perf key: nightly-diffusion-x2v-performance timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 diff --git a/.claude/skills/vllm-omni-test/SKILL.md b/.claude/skills/vllm-omni-test/SKILL.md index 2d4226788..f0fb009d5 100644 --- a/.claude/skills/vllm-omni-test/SKILL.md +++ b/.claude/skills/vllm-omni-test/SKILL.md @@ -468,10 +468,10 @@ When the user asks for **invalid parameter validation**, **invalid request bodie | **Invalid parameters Test · H100** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"` | Diffusion / omni / video invalid-param tests with `H100` hardware mark | | **Invalid parameters Test · L4** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"` | TTS / lighter models marked `L4` | -- **Trigger:** `build.env("WEEKLY") == "1"` or PR label `weekly-test`. +- **Trigger:** `build.env("WEEKLY") == "1"` (all weekly jobs) or PR label `weekly-test` (Reliability / Perf, then filtered by `source_file_dependencies`). - **Default:** extending an existing `invalid_param_test` script needs **no YAML edit** — the weekly steps already sweep the whole directory. - **Edit YAML only** when adding a **new hardware queue**, a **new top-level script** that must run in isolation, or a **model-specific weekly shard** (mirror neighboring reliability steps). -- Weekly steps **do not** use `source_file_dependencies`. +- Weekly Reliability / Perf jobs declare `source_file_dependencies`. Filtering applies only on the `weekly-test` label; `WEEKLY=1` on `main` keeps every job in those groups. CPU coverage (`WEEKLY=1` only) and the E2E Tests group (`NON_CRITICAL=1` only) do **not** use this field. **Example — append to `test_invalid_image_generation.py`:** @@ -644,7 +644,7 @@ If the test is not already collected by an existing pipeline command (for exampl | **L2** | `test-ready.yml` with `core_model` + `source_file_dependencies` | merge / nightly pipelines | | **Invalid param** | `test-weekly.yml` — **Invalid parameters Test** group (`-m "slow and H100"` / `-m "slow and L4"`). Usually **no YAML edit** when appending cases to existing `invalid_param_test` scripts | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml` | -`test-nightly.yml` shards use **explicit pytest file paths** in `commands` (and PR labels like `diffusion-x2iat-test`); they generally **do not** use `source_file_dependencies`. Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI). +`test-nightly.yml` shards use **explicit pytest file paths** in `commands` plus `source_file_dependencies`; PR trigger is `nightly-test` (path filter), or `NIGHTLY=1` on `main` (all jobs). Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI). **`test-nightly.yml` top-level groups** (each uses `full_model` + type marker in pytest `-m`): @@ -653,7 +653,7 @@ If the test is not already collected by an existing pipeline command (for exampl | **Omni Model Test** | `omni` | Function / doc / accuracy / **perf** for Qwen-Omni family | | **TTS Model Test** | `tts` | `-m "full_model and L4 and tts"` (and H100 variants when enabled); **Perf Test** uses `run_benchmark.py` + `test_tts.json` | | **Diffusion X2I(&A&T) Model Test** | `diffusion` | Image/audio/text diffusion; explicit file shards for **Function**; **Perf Test · <Model>** uses `run_diffusion_benchmark.py` + JSON config | -| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); PR label `diffusion-x2v-test`; separate **Perf** steps when in-tree | +| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); separate **Perf** steps when in-tree | Guidelines when editing YAML: diff --git a/.claude/skills/vllm-omni-test/references/test-routing.md b/.claude/skills/vllm-omni-test/references/test-routing.md index 5b90957d6..a61324e6d 100644 --- a/.claude/skills/vllm-omni-test/references/test-routing.md +++ b/.claude/skills/vllm-omni-test/references/test-routing.md @@ -37,7 +37,7 @@ Every model-centric e2e test must declare **exactly one** type marker: | **Diffusion X2I(&A&T) Model Test** | x2**i** / x2**a** / x2**t** — image, audio, text (non-video) | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero | | **Diffusion X2V Model Test** | x2**v** — video only | Wan2.2, HunyuanVideo 1.5, Wan VACE | -PR labels for selective nightly runs: `diffusion-x2iat-test`, `diffusion-x2v-test` (plus `nightly-test` / `NIGHTLY=1`). +PR label for L4 on a branch: `nightly-test` (jobs then filtered by `source_file_dependencies`). Scheduled `main` uses `NIGHTLY=1` and runs every nightly job. L5 Reliability / Perf use `weekly-test` with the same path filter; CPU coverage and the weekly E2E group are env-only (`WEEKLY=1` / `NON_CRITICAL=1`) and are not source-filtered. ## Level and Marker Mapping @@ -238,7 +238,7 @@ pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4" pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py::test_images_generations_invalid_requests -m "slow and H100" ``` -**Trigger:** `WEEKLY=1` or PR label `weekly-test`. Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). No `source_file_dependencies` on weekly steps. +**Trigger:** `WEEKLY=1` (all jobs) or PR label `weekly-test` (Reliability / Perf, filtered by `source_file_dependencies`). Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). ### Platform-targeted examples diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index 885ebd528..fb3af0ed3 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -15,7 +15,8 @@ Canonical layout (prefer these paths for new changes): │ │ ├── skip_ci.py # skip-ci decision (docs / skip-mark / CI YAML paths) │ │ ├── upload_pipeline.py # Bootstrap + test-pipeline uploader (CUDA/NPU) │ │ └── resolve_skip_ci.sh # Shell helpers for AMD/Intel bootstrap -│ └── ci_mirror_hardwares.yml # CUDA uploader presets (referenced by name only) +│ ├── ci_mirror_hardwares.yml # CUDA/NPU hardware presets (referenced by name only) +│ └── ci_source_file_dependencies.yml # L2–L5 path-prefix presets (referenced by name only) ├── cuda/ # Primary NVIDIA CUDA CI │ ├── pipeline.yml # Bootstrap entry (hook upload) │ ├── bootstrap-upload-steps.yml # Bootstrap child steps (upload_pipeline --upload) @@ -108,19 +109,21 @@ Canonical layout (prefer these paths for new changes): Inferred form: unset/`empty` `MIRROR_HW` matches `H100` or `L4` in `-m` (both present → H100); no match skips the step. `MIRROR_HW=b200` must appear in `-m` or the step is skipped. The uploader does **not** rewrite pytest `-m`; B200 collection comes from the YAML expression and from tests that declare `B200` in `hardware_test` / `hardware_marks` (for example `res={"cuda": ["H100", "B200"]}`). `MIRROR_HW` must be empty or `b200` (case-insensitive); unknown values (for example `b20o`) **fail the upload**. A CUDA preset string such as `mirror_hardwares: h100_4` is skipped when `MIRROR_HW=b200`. Unset `MIRROR_HW` keeps string presets. NPU presets ignore `MIRROR_HW`. A `mirror_hardwares` name that is not a key in `ci_mirror_hardwares.yml` **fails the upload**. + **Path filter in YAML:** `source_file_dependencies: ` (string, or a list of preset names)—preset names in [`common/ci_source_file_dependencies.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/ci_source_file_dependencies.yml). Required on **E2E Test** leaf jobs in L2/L3, and on L4/L5 leaf jobs that should be diff-gated on PR labels; see [Step filtering](#step-filtering). + **Conventions** - **`depends_on`:** leaf jobs depend on `upload-ready-pipeline`, `upload-merge-pipeline`, etc. - **`group` / `label`:** `:card_index_dividers:` groups; labels like `Diffusion · Qwen Image Test`. - **`commands`:** `timeout … pytest …` with markers and `--run-level` for the pipeline level. - - **`source_file_dependencies`:** required on **E2E Test** leaf jobs in L2/L3; see [Step filtering](#step-filtering). + - **`source_file_dependencies`:** required on **E2E Test** leaf jobs in L2/L3, and on L4/L5 leaf jobs that should be diff-gated on PR labels; name a job key from `ci_source_file_dependencies.yml` (see [Step filtering](#step-filtering)). The key must list every script the step runs; `commands` is not scanned. **Adding a job** 1. Pick the level file (`test-ready.yml` / `test-merge.yml` / `test-nightly.yml` / `test-weekly.yml`). 2. Add a step under the right **group** (usually **E2E Test** for model pytest). - 3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload-ready-pipeline` (or `upload-merge-pipeline` / `upload-nightly-pipeline` / `upload-weekly-pipeline`). - 4. For L2/L3 E2E, add `source_file_dependencies` (pytest + model + deploy YAML prefixes). + 3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload--pipeline`. + 4. For L2–L5 jobs that should be diff-gated, add a job key in `common/ci_source_file_dependencies.yml` and set `source_file_dependencies: ` on the step. YAML anchors hold model business-code paths only. The job key (`{model}_function` / `_perf` / `_accuracy` / `_reliability` / `_doc` / `_cov`) aliases that anchor and lists the pytest / example / benchmark scripts the step actually runs. The uploader does **not** parse `commands` for path filters—omitting a script from the key means editing that script will not select the job. 5. Dry-run: ```bash @@ -144,6 +147,8 @@ Canonical layout (prefer these paths for new changes): - A5: `a5_npu_2`, `a5_npu_4`, `a5_npu_8`, `a5_npu_16` - 310P: `310p_npu_1`, `310p_npu_2`, `310p_npu_4` + **Path filter in YAML:** same `source_file_dependencies` presets as CUDA, on **L4** leaf jobs in `test-npu-nightly.yml`. Filtering applies on the PR `nightly-test` label; `main` + `NIGHTLY=1` keeps every job. If no job-key prefix matches, a change to the pipeline YAML being uploaded, or to any path under `source_filter_fallback`, keeps every job in that upload. See [Step filtering](#step-filtering). + **Conventions** - **`depends_on: upload-ready-pipeline`** (or `upload-nightly-pipeline`) ties jobs to bootstrap upload keys. @@ -153,7 +158,7 @@ Canonical layout (prefer these paths for new changes): 1. Edit `test-npu-ready.yml` (L2) or `test-npu-nightly.yml` (L4). 2. Add a step with `mirror_hardwares` (add a new preset in `ci_mirror_hardwares.yml` first if needed). - 3. Set `commands` to your pytest file and markers. + 3. Set `commands` to your pytest file and markers. For L4, set `source_file_dependencies` to a job key from `ci_source_file_dependencies.yml` that aliases the model anchor and lists that pytest file. 4. Dry-run: ```bash @@ -213,7 +218,7 @@ PR diffs drive **two independent skip layers**. Both read changed files from git | Layer | Script | When | What is skipped | Where you configure | | ----- | ------ | ---- | --------------- | ------------------- | | **Bootstrap** | [`skip_ci.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/skip_ci.py) + [`upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/upload_pipeline.py) | Before child test pipelines upload (`cuda/pipeline.yml`, `npu/pipeline-npu.yml`, AMD/Intel bootstraps) | Entire default CI, or whole L2/L3 upload for a platform | Whitelists in `skip_ci.py`; bootstrap `if` injected by step `key` in `upload_pipeline.py` | -| **Step filter** | [`upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/upload_pipeline.py) | While uploading CUDA L2/L3 YAML | Individual Buildkite steps inside `test-ready.yml` / `test-merge.yml` | `source_file_dependencies` on each step or group | +| **Step filter** | [`upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/upload_pipeline.py) | While uploading CUDA L2–L5 YAML | Individual Buildkite steps inside `test-ready.yml` / `test-merge.yml` / `test-nightly.yml` / `test-weekly.yml` | `source_file_dependencies` key on each step (presets in `ci_source_file_dependencies.yml`) | **Changed files** (both layers): @@ -306,40 +311,38 @@ Unit coverage: `tests/buildkite/test_skip_ci.py`. #### Step filtering {#step-filtering} -CUDA **L2** (`.buildkite/cuda/test-ready.yml`) and **L3** (`.buildkite/cuda/test-merge.yml`) only. Bootstrap upload entry: `upload_pipeline.py --upload .buildkite/cuda/bootstrap-upload-steps.yml`. +CUDA **L2** (`.buildkite/cuda/test-ready.yml`), **L3** (`.buildkite/cuda/test-merge.yml`), **L4** (`.buildkite/cuda/test-nightly.yml`), **L5** (`.buildkite/cuda/test-weekly.yml`), and NPU **L4** (`.buildkite/npu/test-npu-nightly.yml`). Bootstrap upload entry: `upload_pipeline.py --upload .buildkite/cuda/bootstrap-upload-steps.yml` (or `npu/bootstrap-upload-steps.yml`). + +**When filtering applies:** only **PR label** uploads (`ready`, `merge-test`, `nightly-test`, `weekly-test`). On **`main` + env** (`NIGHTLY=1`, post-merge L3, `WEEKLY=1` / `NON_CRITICAL=1`, `WEEKLY=1 --e2e`) the uploader keeps every job and still strips the key. If **no** job-key prefix matches the diff, a change to the pipeline YAML being uploaded, or any path under the `source_filter_fallback` key in `ci_source_file_dependencies.yml`, keeps every job in that upload. `source_filter_fallback` is not a job key — do not attach it to a step. When any job-key prefix already matches, normal filtering wins (the fallback does not expand the selection). Without this fallback, command / env / hardware-only edits would match nothing (NPU drops every step; CUDA nightly keeps only the email aggregator). A different pipeline YAML does not trigger the fallback for this upload. **Uploader-only keys** — removed before Buildkite sees the YAML; never used at runtime on agents: | Key | Purpose | | --- | ------- | -| `source_file_dependencies` | List of path **prefixes**. If any changed file equals a prefix or starts with `prefix/`, keep the step (or group); otherwise omit it. | +| `source_file_dependencies` | Named preset (or list of presets) from [`ci_source_file_dependencies.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/ci_source_file_dependencies.yml). The uploader expands only that key; it does **not** scan `commands`. If any changed file equals a listed prefix or starts with `prefix/`, keep the step (or group); otherwise omit it. Inline path-prefix lists are still accepted. | | `mirror_hardwares` | Expand to `agents` + `plugins` (+ optional `image`) from `ci_mirror_hardwares.yml`. | ### Policy -- **Always uploaded** (no key): groups outside **E2E Test**—Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3). -- **Diff-gated**: every **E2E Test** leaf job. List the smallest prefix set per step—pytest file(s), model code under `vllm_omni/model_executor/models/` or `vllm_omni/diffusion/models/`, plus `stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable. A **group** may define the key instead; the whole group drops if no prefix matches. +- **Always uploaded** (no key): L2/L3 groups outside **E2E Test**—Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3). +- **Diff-gated**: every **E2E Test** leaf job in L2/L3, L4 nightly leaf jobs (CUDA and NPU), and L5 weekly jobs whose `if` includes the `weekly-test` label (Reliability / Perf). Jobs gated only by `WEEKLY=1` or `NON_CRITICAL=1` have no `source_file_dependencies`. Registry layout: YAML anchors are model business code only (`vllm_omni/model_executor/models/`, `vllm_omni/diffusion/models/`, plus `stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable). Job keys alias the anchor and **also** list the scripts that job runs (pytest files, `run_cov_split.sh --offline` / `--online` targets, example README runners). Feature keys stay cross-cutting. A **group** may define the key instead; the whole group drops if no prefix matches. ### YAML examples ```yaml - label: "Diffusion · Qwen Image Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image.py - - vllm_omni/diffusion/models/qwen_image/ + source_file_dependencies: diffusion_qwen_image_function + # diffusion_qwen_image_function: *qwen_image + tests/e2e/online_serving/test_qwen_image.py commands: - pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' ... mirror_hardwares: h100_1 - label: "TTS · Qwen3-TTS CustomVoice Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - - vllm_omni/deploy/qwen3_tts.yaml + source_file_dependencies: tts_qwen3_tts_function + # tts_qwen3_tts_function: *qwen3_tts + tests/e2e/online_serving/test_qwen3_tts_customvoice.py commands: - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py ... - mirror_hardwares: l4_4 + mirror_hardwares: l4_1 ``` ### Local dry-run @@ -443,15 +446,16 @@ plugin, which provides the agent. `l4_*` jobs run on the EKS `l4-k8s` queue still needs `mount-buildkite-agent: true`. List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's -`source_file_dependencies` — both change what the job measures, so without them a +`source_file_dependencies` preset — both change what the job measures, so without them a change there is filtered out of normal PR builds and only surfaces in a later nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with a synthetic job (it does not pin real merge labels). -Editing only the surrounding CI YAML still does not schedule them, so a PR that -touches just the wiring needs a full E2E run (or the commands run on a GPU host) -to produce artifacts. When checking a new model's -artifacts, compare `lines-covered` between the online and offline XML rather than -just confirming both files exist. +If no job-key prefix matches, editing the pipeline YAML being uploaded, or any +path under the `source_filter_fallback` registry key, keeps every job in that +upload (including coverage jobs) so the wiring can be validated before merge. +When a job-key prefix already matches, bypass does not expand the selection. A +change to a different pipeline YAML still does not schedule them. When checking a new model's artifacts, compare `lines-covered` +between the online and offline XML rather than just confirming both files exist. ### Validation checklist diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 7492a4b82..06e2b975c 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -13,11 +13,15 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "com from skip_ci import resolve_ci_decision # noqa: E402 from upload_pipeline import ( # noqa: E402 CUDA_HF_TOKEN_EXPORT, + NIGHTLY_LABEL_IF, + _changed_files_for_source_filter, _expand_mirror_hardwares, _get_mirror_hw_selector, _load_bootstrap_steps, + _load_source_file_dependencies, _render_bootstrap_pipeline, _render_test_pipeline, + _resolve_source_file_dependencies, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -101,6 +105,26 @@ def test_npu_docs_only_does_not_upload_ready_on_nightly() -> None: assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"] +def test_nightly_label_if_is_only_nightly_test() -> None: + assert 'labels includes "nightly-test"' in NIGHTLY_LABEL_IF + forbidden = ( + 'includes "omni-test"', + 'includes "tts-test"', + 'includes "diffusion-x2iat-test"', + 'includes "diffusion-x2v-test"', + ) + for needle in forbidden: + assert needle not in NIGHTLY_LABEL_IF + for path in ( + Path(".buildkite/cuda/test-nightly.yml"), + Path(".buildkite/npu/test-npu-nightly.yml"), + Path(".buildkite/common/scripts/upload_pipeline.py"), + ): + text = path.read_text(encoding="utf-8") + for needle in forbidden: + assert needle not in text, f"{path} still references {needle}" + + def test_yaml_gated_l45_only_does_not_unconditionally_build_image() -> None: rendered = _render([".buildkite/cuda/test-nightly.yml"]) assert "if: true" not in rendered @@ -366,8 +390,17 @@ def test_mirror_hw_typo_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: ) -def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]: - rendered = _render_test_pipeline(doc, changed_files=changed_files) +def _surviving_labels( + doc: dict, + changed_files: list[str], + *, + pipeline_path: Path | None = None, +) -> set[str]: + rendered = _render_test_pipeline( + doc, + changed_files=changed_files, + pipeline_path=pipeline_path, + ) labels: set[str] = set() def walk(steps: list | None) -> None: @@ -382,6 +415,17 @@ def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]: return labels +def _iter_steps(doc: dict): + def walk(steps: list | None): + for step in steps or []: + if not isinstance(step, dict): + continue + yield step + yield from walk(step.get("steps")) + + yield from walk(doc.get("steps")) + + # Synthetic coverage-style job: shared inputs that change what the split measures. _COVERAGE_SHARED_INPUTS_DOC = { "steps": [ @@ -426,3 +470,289 @@ def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None: ) assert "Coverage Pilot" not in labels assert "Unrelated Model Test" not in labels + + +def _pipeline_dep_keys(path: Path) -> set[str]: + doc = yaml.safe_load(path.read_text(encoding="utf-8")) + keys: set[str] = set() + + def walk(steps: list | None) -> None: + for step in steps or []: + if not isinstance(step, dict): + continue + deps = step.get("source_file_dependencies") + if isinstance(deps, str) and "/" not in deps: + keys.add(deps) + elif isinstance(deps, list) and deps and all("/" not in item for item in deps): + keys.update(deps) + walk(step.get("steps")) + + walk((doc or {}).get("steps")) + return keys + + +def test_pipeline_source_file_dependency_keys_are_registered() -> None: + _load_source_file_dependencies.cache_clear() + registry = _load_source_file_dependencies() + used = ( + _pipeline_dep_keys(Path(".buildkite/cuda/test-ready.yml")) + | _pipeline_dep_keys(Path(".buildkite/cuda/test-merge.yml")) + | _pipeline_dep_keys(Path(".buildkite/cuda/test-nightly.yml")) + | _pipeline_dep_keys(Path(".buildkite/cuda/test-weekly.yml")) + | _pipeline_dep_keys(Path(".buildkite/npu/test-npu-nightly.yml")) + ) + missing = used - set(registry) + assert not missing, f"unregistered source_file_dependencies keys: {sorted(missing)}" + + +def test_source_file_dependencies_key_expands_from_registry() -> None: + doc = { + "steps": [ + { + "label": "Diffusion · Qwen Image Test", + "source_file_dependencies": "diffusion_qwen_image_function", + "commands": ["pytest"], + }, + ], + } + assert "Diffusion · Qwen Image Test" in _surviving_labels( + doc, + ["vllm_omni/diffusion/models/qwen_image/transformer.py"], + ) + assert "Diffusion · Qwen Image Test" not in _surviving_labels(doc, ["vllm_omni/unrelated.py"]) + + +def test_registry_lists_pytest_targets() -> None: + _load_source_file_dependencies.cache_clear() + resolved = _resolve_source_file_dependencies( + { + "label": "Diffusion · Wan22 Test", + "source_file_dependencies": "diffusion_wan22_function", + "commands": [ + "pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py " + "tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model'", + ], + }, + ) + assert resolved is not None + assert "tests/e2e/offline_inference/test_wan22_t2v.py" in resolved + assert "tests/e2e/online_serving/test_wan22_t2v.py" in resolved + assert "vllm_omni/diffusion/models/wan2_2/" in resolved + + +def test_coverage_key_lists_offline_online_scripts() -> None: + resolved = _resolve_source_file_dependencies( + { + "label": "TTS · Qwen3-TTS Base Test", + "source_file_dependencies": "tts_qwen3_tts_cov", + "commands": [ + ".buildkite/common/scripts/run_cov_split.sh \\\n" + " --offline tests/e2e/offline_inference/test_qwen3_tts_base.py \\\n" + " --online tests/e2e/online_serving/test_qwen3_tts_base.py", + ], + }, + ) + assert resolved is not None + assert "tests/e2e/offline_inference/test_qwen3_tts_base.py" in resolved + assert "tests/e2e/online_serving/test_qwen3_tts_base.py" in resolved + assert ".buildkite/common/scripts/run_cov_split.sh" in resolved + assert "pyproject.toml" in resolved + + +def test_source_file_dependencies_list_of_keys_concatenates() -> None: + resolved = _resolve_source_file_dependencies( + { + "label": "composed", + "source_file_dependencies": ["omni_qwen3_omni_function", "tts_qwen3_tts_function"], + }, + ) + assert resolved is not None + assert "vllm_omni/model_executor/models/qwen3_omni/" in resolved + assert "vllm_omni/model_executor/models/qwen3_tts/" in resolved + assert resolved.count("vllm_omni/model_executor/models/common/snake_activation.py") == 1 + + +def test_unknown_source_file_dependencies_key() -> None: + with pytest.raises(ValueError, match="unknown source_file_dependencies"): + _render_test_pipeline( + {"steps": [{"label": "bad", "source_file_dependencies": "not_a_real_key"}]}, + changed_files=None, + ) + + +def test_source_file_dependencies_rejects_mixed_keys_and_paths() -> None: + with pytest.raises(ValueError, match="mixes registry keys and path prefixes"): + _resolve_source_file_dependencies( + { + "label": "bad", + "source_file_dependencies": ["omni_qwen3_omni_function", "tests/e2e/online_serving/test_qwen3_omni.py"], + }, + ) + + +# Synthetic pipeline: selection depends only on listed deps, not on live job names. +_SOURCE_FILTER_DOC = { + "steps": [ + { + "group": "E2E Tests", + "if": 'build.env("NON_CRITICAL") == "1"', + "steps": [ + { + "label": "Dedicated E2E", + "commands": ["pytest -sv tests/e2e/online_serving/test_magi2.py"], + }, + ], + }, + { + "label": "Omni Sweep", + "source_file_dependencies": ["tests/e2e/online_serving/test_qwen3_omni.py"], + "commands": ["pytest -sv tests/e2e/ -m omni"], + }, + { + "label": "Z-Image Function", + "source_file_dependencies": [ + "tests/e2e/online_serving/test_zimage_expansion.py", + "vllm_omni/diffusion/models/z_image/", + ], + "mirror_hardwares": "l4_4", + "commands": ["pytest -sv tests/e2e/online_serving/test_zimage_expansion.py"], + }, + { + "label": "Tiny Model", + "source_file_dependencies": ["tests/e2e/online_serving/test_tiny.py"], + "commands": ["pytest -sv tests/e2e/ -m tiny"], + }, + { + "label": "Wan Function", + "source_file_dependencies": [ + "tests/e2e/offline_inference/test_wan22_t2v.py", + "tests/e2e/online_serving/test_wan22_t2v.py", + "vllm_omni/diffusion/models/wan2_2/", + ], + "commands": ["pytest -sv tests/e2e/offline_inference/test_wan22_t2v.py"], + }, + { + "label": "Wan Perf", + "source_file_dependencies": ["vllm_omni/diffusion/models/wan2_2/"], + "commands": ["pytest -sv tests/dfx/perf/scripts/run_benchmark.py"], + }, + { + "label": "Doc Test", + "source_file_dependencies": [ + "tests/examples/offline_inference/test_text_to_image.py", + "tests/examples/online_serving/test_text_to_image.py", + "vllm_omni/diffusion/models/qwen_image/", + "vllm_omni/diffusion/models/z_image/", + ], + "commands": ["pytest -sv tests/examples/*/test_text_to_image.py"], + }, + ], +} + + +def test_source_filter_ignores_unrelated_e2e_file() -> None: + """A sweep command is not a dependency; only listed paths select a keyed job.""" + labels = _surviving_labels(_SOURCE_FILTER_DOC, ["tests/e2e/online_serving/test_magi2.py"]) + assert labels == {"Dedicated E2E"} + + +def test_source_filter_selects_job_by_listed_script_not_sibling() -> None: + labels = _surviving_labels(_SOURCE_FILTER_DOC, ["tests/e2e/online_serving/test_zimage_expansion.py"]) + assert labels == {"Dedicated E2E", "Z-Image Function"} + + +def test_source_filter_selects_every_job_sharing_a_path() -> None: + labels = _surviving_labels(_SOURCE_FILTER_DOC, ["vllm_omni/diffusion/models/wan2_2/transformer.py"]) + assert labels == {"Dedicated E2E", "Wan Function", "Wan Perf"} + + +def test_source_filter_selects_composed_doc_job_by_example_or_model() -> None: + for changed in ( + "tests/examples/online_serving/test_text_to_image.py", + "vllm_omni/diffusion/models/z_image/transformer.py", + "vllm_omni/diffusion/models/qwen_image/foo.py", + ): + labels = _surviving_labels(_SOURCE_FILTER_DOC, [changed]) + assert "Doc Test" in labels, changed + assert "Tiny Model" not in labels + assert "Omni Sweep" not in labels + assert "Doc Test" not in _surviving_labels(_SOURCE_FILTER_DOC, ["tests/e2e/online_serving/test_magi2.py"]) + + +def test_source_filter_strips_deps_and_expands_hardware(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "") + rendered = _render_test_pipeline( + _SOURCE_FILTER_DOC, + changed_files=["tests/e2e/online_serving/test_zimage_expansion.py"], + ) + dumped = yaml.safe_dump(rendered) + assert "source_file_dependencies" not in dumped + assert "mirror_hardwares" not in dumped + z_image = next(step for step in _iter_steps(rendered) if step.get("label") == "Z-Image Function") + assert z_image["agents"]["queue"] == "l4-k8s" + + +def test_source_filter_disabled_on_main_branch(monkeypatch: pytest.MonkeyPatch) -> None: + class _Ctx: + changed_files = ["vllm_omni/unrelated.py"] + + monkeypatch.setenv("BUILDKITE_BRANCH", "main") + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) is None + + monkeypatch.setenv("BUILDKITE_BRANCH", "feat/source-filter") + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) == [ + "vllm_omni/unrelated.py", + ] + monkeypatch.setenv("BUILDKITE_BRANCH", "main") + assert _changed_files_for_source_filter(_Ctx(), force_all=True, e2e_only=False) is None + + +def test_source_filter_fallback_is_fallback_when_no_job_key_matches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Fallback only when no listed prefix matched; a matching prefix still wins.""" + monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "") + monkeypatch.setenv("BUILDKITE_BRANCH", "feat/nightly-yaml") + pipeline_yaml = Path(".buildkite/npu/test-npu-nightly.yml") + shared_paths = _load_source_file_dependencies()["source_filter_fallback"] + assert ".buildkite/common/scripts/upload_pipeline.py" in shared_paths + + # Synthetic steps only — do not pin live Buildkite job labels. + doc = { + "steps": [ + {"key": "ungated", "commands": ["true"]}, + { + "key": "gated_a", + "source_file_dependencies": ["pkg/model_a/"], + "commands": ["true"], + }, + { + "key": "gated_b", + "source_file_dependencies": ["pkg/model_b/"], + "commands": ["true"], + }, + ], + } + + def surviving_keys(changed_files: list[str]) -> set[str]: + rendered = _render_test_pipeline( + doc, + changed_files=changed_files, + pipeline_path=pipeline_yaml, + ) + return {step["key"] for step in _iter_steps(rendered) if isinstance(step.get("key"), str)} + + # Only pipeline YAML / shared uploader → no listed prefix match → keep every step. + for changed in [pipeline_yaml.as_posix(), *shared_paths]: + assert surviving_keys([changed]) == {"ungated", "gated_a", "gated_b"}, changed + + # A different pipeline YAML is not a fallback for this upload. + assert surviving_keys([".buildkite/cuda/test-nightly.yml"]) == {"ungated"} + + # A matching source prefix wins over fallback files in the same diff. + assert surviving_keys( + [ + "pkg/model_a/transformer.py", + ".buildkite/common/scripts/upload_pipeline.py", + ], + ) == {"ungated", "gated_a"} diff --git a/tests/e2e/accuracy/test_hidream_o1_image.py b/tests/e2e/accuracy/test_hidream_o1_image.py index 923e53a47..440050be1 100644 --- a/tests/e2e/accuracy/test_hidream_o1_image.py +++ b/tests/e2e/accuracy/test_hidream_o1_image.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -16,7 +16,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs from vllm_omni.inputs.data import OmniDiffusionSamplingParams -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] MODEL_ID = "HiDream-ai/HiDream-O1-Image" PROMPT = "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py index 64578af01..88e5979dc 100644 --- a/tests/e2e/offline_inference/test_lingbot_world_v2.py +++ b/tests/e2e/offline_inference/test_lingbot_world_v2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end smoke test for LingBot-World v2 generation.""" from __future__ import annotations @@ -37,7 +37,7 @@ _OMNI_RUNNER_PARAM = ( ) pytestmark = [ - pytest.mark.slow, + pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.skipif( not _HAS_INPUT_ASSETS, diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py index 24cc2f7e9..7f14e403d 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py @@ -12,7 +12,7 @@ from tests.helpers.runtime import OmniServer, OpenAIClientHandler from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_fl2va -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py index f988a4431..2c61d2a3c 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py @@ -12,7 +12,7 @@ from tests.helpers.runtime import OmniServer, OpenAIClientHandler from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_ref2va -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py index 9fc3e8e31..1fd80454b 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py @@ -21,7 +21,7 @@ from ._common import ( post_sync, ) -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py index 0190598d5..8a1a6542a 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py @@ -20,7 +20,7 @@ from ._common import ( turbo_params, ) -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/test_lingbot_video.py b/tests/e2e/online_serving/test_lingbot_video.py index 8f0f45639..6ada319f5 100644 --- a/tests/e2e/online_serving/test_lingbot_video.py +++ b/tests/e2e/online_serving/test_lingbot_video.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Online serving smoke coverage for the dense LingBot-Video checkpoint.""" @@ -41,7 +41,7 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: @@ -63,7 +63,7 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH ) -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: diff --git a/tests/e2e/online_serving/test_lingbot_video_moe.py b/tests/e2e/online_serving/test_lingbot_video_moe.py index 7e4d16568..a75fa78e0 100644 --- a/tests/e2e/online_serving/test_lingbot_video_moe.py +++ b/tests/e2e/online_serving/test_lingbot_video_moe.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Single-GPU serving smoke for ``robbyant/lingbot-video-moe-30b-a3b``.""" @@ -41,7 +41,7 @@ def _get_server_cases(model: str): ] -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True) def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: @@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH ) -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True) def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: diff --git a/tests/e2e/online_serving/test_magi2.py b/tests/e2e/online_serving/test_magi2.py index 22a30b495..d713abab9 100644 --- a/tests/e2e/online_serving/test_magi2.py +++ b/tests/e2e/online_serving/test_magi2.py @@ -8,7 +8,7 @@ It exercises the recommended resident SP4 topology through ``/v1/videos``. From ``tests/``:: - pytest -s -v e2e/online_serving/test_magi2.py -m "advanced_model and diffusion" --run-level=advanced_model + pytest -s -v e2e/online_serving/test_magi2.py -m "slow and diffusion" --run-level=full_model """ import json @@ -27,8 +27,7 @@ PROMPT = "A red fox walks through fresh snow while wind moves the pine branches. FOUR_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=4) -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", diff --git a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py index 0bf546683..4fa71aead 100644 --- a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py +++ b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py @@ -111,7 +111,6 @@ def _run_t2va_request(client: OpenAIClientHandler, seed: int) -> bytes: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion -@pytest.mark.slow @pytest.mark.parametrize( "omni_server", [ diff --git a/tools/pre_commit/check_buildkite.py b/tools/pre_commit/check_buildkite.py index 362253494..f470f09f9 100644 --- a/tools/pre_commit/check_buildkite.py +++ b/tools/pre_commit/check_buildkite.py @@ -27,6 +27,7 @@ from upload_pipeline import _render_test_pipeline # noqa: E402 SKIP_FILES = { ".buildkite/common/ci_mirror_hardwares.yml", + ".buildkite/common/ci_source_file_dependencies.yml", ".buildkite/npu/pipeline-npu-a3.yml", } From a2486dd00bcc4f6206ebfe4c99f40e8a784e8d06 Mon Sep 17 00:00:00 2001 From: NancyFyong <88076188+NancyFyong@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:50:54 +0800 Subject: [PATCH 002/209] [Core][Diffusion] Add a typed pre-D2H video media contract (#6615) Signed-off-by: NancyFyong Signed-off-by: Samit <285365963@qq.com> Co-authored-by: NancyFyong Co-authored-by: Samit <285365963@qq.com> --- .buildkite/cuda/test-nightly.yml | 8 + docs/.nav.yml | 1 + docs/mkdocs/hooks/generate_argparse.py | 2 + .../device_side_video_postprocess.md | 72 +++ tests/config/test_omni_config.py | 55 ++ .../wan2_2/test_wan22_pipeline_diffuse.py | 111 ++++ tests/diffusion/test_async_output_worker.py | 136 ++++- .../test_device_output_preparation.py | 460 +++++++++++++++ tests/diffusion/test_device_reduction.py | 125 ++++ .../diffusion/test_diffusion_model_runner.py | 114 +++- tests/diffusion/test_ipc_async.py | 30 +- tests/diffusion/test_media_output_contract.py | 546 ++++++++++++++++++ ...test_video_transport_defaults_are_inert.py | 61 ++ tests/diffusion/test_wan_device_reduction.py | 348 +++++++++++ tests/docs/test_generate_argparse.py | 24 + .../test_device_postprocess_equivalence.py | 105 ++++ tests/entrypoints/test_serve.py | 24 + vllm_omni/config/omni_config.py | 9 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 37 +- vllm_omni/diffusion/diffusion_engine.py | 46 +- vllm_omni/diffusion/ipc.py | 257 ++++++++- vllm_omni/diffusion/media.py | 187 ++++++ .../models/wan2_2/pipeline_wan2_2.py | 32 +- .../diffusion/postprocess/device_reduction.py | 152 +++++ vllm_omni/diffusion/postprocess/media.py | 92 +++ .../worker/diffusion_model_runner.py | 22 + .../diffusion/worker/diffusion_worker.py | 14 +- vllm_omni/diffusion/worker/request_batch.py | 24 +- vllm_omni/engine/arg_utils.py | 1 + vllm_omni/entrypoints/cli/serve.py | 18 + 31 files changed, 3054 insertions(+), 60 deletions(-) create mode 100644 docs/user_guide/diffusion/device_side_video_postprocess.md create mode 100644 tests/diffusion/test_device_output_preparation.py create mode 100644 tests/diffusion/test_device_reduction.py create mode 100644 tests/diffusion/test_media_output_contract.py create mode 100644 tests/diffusion/test_video_transport_defaults_are_inert.py create mode 100644 tests/diffusion/test_wan_device_reduction.py create mode 100644 tests/docs/test_generate_argparse.py create mode 100644 tests/e2e/offline_inference/test_device_postprocess_equivalence.py create mode 100644 vllm_omni/diffusion/media.py create mode 100644 vllm_omni/diffusion/postprocess/device_reduction.py create mode 100644 vllm_omni/diffusion/postprocess/media.py diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index ee4fc6f3e..583f935dc 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -389,6 +389,14 @@ steps: - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2V · Wan2.2 Device-Postprocess Equivalence Test" + timeout_in_minutes: 90 + commands: + - | + VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + pytest -s -v tests/e2e/offline_inference/test_device_postprocess_equivalence.py -m "full_model and cuda" --run-level "full_model" + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" source_file_dependencies: diffusion_wan22_accuracy timeout_in_minutes: 180 diff --git a/docs/.nav.yml b/docs/.nav.yml index aa84f2819..8dac28d22 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -71,6 +71,7 @@ nav: - Overview and Compatibility: - Overview: user_guide/diffusion_features.md - Feature Compatibility: user_guide/feature_compatibility.md + - Device-side Video Postprocessing: user_guide/diffusion/device_side_video_postprocess.md - CPU Offloading: - Overview: user_guide/diffusion/cpu_offload.md - Model-Level: user_guide/diffusion/offloader/module_offload.md diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 4c3ad3373..ea0fc9e0c 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import ast +import json import logging import re import sys @@ -124,6 +125,7 @@ def extract_omni_serve_subparser_init(): "make_arg_parser": lambda parser: parser, # no-op for doc "_ensure_vllm_platform": lambda: None, # no-op for doc "_nonneg_finite_float": float, + "_json_object": json.loads, "_parse_stage_overrides": __import__("json").loads, "VLLM_SUBCMD_PARSER_EPILOG": "", "logger": logger, diff --git a/docs/user_guide/diffusion/device_side_video_postprocess.md b/docs/user_guide/diffusion/device_side_video_postprocess.md new file mode 100644 index 000000000..910e64457 --- /dev/null +++ b/docs/user_guide/diffusion/device_side_video_postprocess.md @@ -0,0 +1,72 @@ +# Device-side video postprocessing + +vLLM-Omni can convert decoded WAN2.2 video from normalized floating-point +`[B, C, T, H, W]` tensors to contiguous `uint8` `[B, T, H, W, C]` frames before +the worker-to-engine device-to-host transfer. This reduces the transported video +payload by four times when the original IPC path widens bfloat16 to float32. + +The optimization is disabled by default. Enable it with a structured diffusion +configuration: + +```python +from vllm_omni import Omni + +engine = Omni( + model="Wan-AI/Wan2.2-TI2V-5B-Diffusers", + video_output_transport={"enable_device_postprocess": True}, +) +``` + +For the server CLI, pass the same object as JSON: + +```bash +vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni \ + --video-output-transport '{"enable_device_postprocess": true}' +``` + +For a deploy configuration, place it on the diffusion stage: + +```yaml +stages: + - stage_id: 0 + video_output_transport: + enable_device_postprocess: true +``` + +## Runtime contract + +WAN2.2 emits decoded video through the typed `DiffusionOutput.media` field. The +model runner validates and splits the batch into request-local tensors before +preparing each video for transport. The worker then performs the D2H copy and +shared-memory IPC using the prepared representation. The engine uses the generic +media finalizer instead of a model-specific postprocessor. + +The runtime leaves the video in normalized floating-point form when: + +- `enable_device_postprocess` is false; +- the requested output type is not `np`; or +- frame interpolation still requires the floating-point tensor. + +These are policy fallbacks, not malformed-contract errors. If the device cannot +allocate the temporary conversion buffer, the runner logs a warning and prepares +the request-local float representation instead. Invalid tensor layout, encoding, +value range, lifecycle state, and non-memory runtime failures still fail before +IPC. + +## Precision + +Device preparation converts to float32 before denormalization and quantization. +For bfloat16 WAN output this is more precise than the historical host path, +which denormalizes in bfloat16. Relative to that path, values may differ by at +most one level out of 255. Float32 inputs remain byte-identical after +quantization. + +## Scope + +WAN2.2 is the reference migration for the first contract version. Other video +pipelines continue using their existing output paths until they can declare the +same typed media contract without bypassing model-specific float consumers such +as safety checks or audio/video packaging. + +See [RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) for the +contract, lifecycle, batching rules, and migration plan. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 40229f521..da8ffb452 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1542,6 +1542,61 @@ def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch): assert cfg.quantization_config.get_name() == "int8" +def test_video_output_transport_mapping_is_normalized() -> None: + from vllm_omni.diffusion.data import VideoOutputTransportConfig + + cfg = omni_config_module._DiffusionConfigProjection( + video_output_transport={"enable_device_postprocess": True}, + ) + + assert isinstance(cfg.video_output_transport, VideoOutputTransportConfig) + assert cfg.video_output_transport.enable_device_postprocess is True + + +def test_omni_diffusion_config_normalizes_video_output_transport_mapping() -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig, VideoOutputTransportConfig + + cfg = OmniDiffusionConfig( + model=None, + video_output_transport={"enable_device_postprocess": True}, + ) + + assert isinstance(cfg.video_output_transport, VideoOutputTransportConfig) + assert cfg.video_output_transport.enable_device_postprocess is True + + +def test_video_output_transport_rejects_non_boolean_flag() -> None: + from vllm_omni.diffusion.data import VideoOutputTransportConfig + + with pytest.raises(TypeError, match="enable_device_postprocess must be a bool"): + VideoOutputTransportConfig(enable_device_postprocess="true") # type: ignore[arg-type] + + +def test_video_output_transport_survives_stage_override_filtering() -> None: + from vllm_omni.config.stage_config import build_stage_runtime_overrides, deploy_runtime_override_keys + + transport = {"enable_device_postprocess": True} + overrides = build_stage_runtime_overrides(0, {"video_output_transport": transport}) + + assert "video_output_transport" in deploy_runtime_override_keys() + assert overrides["video_output_transport"] == transport + + +def test_video_output_transport_reaches_default_diffusion_stage() -> None: + from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + + transport = {"enable_device_postprocess": True} + stages = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "model": "unused", + "model_class_name": "UnknownPipeline", + "video_output_transport": transport, + } + ) + + assert stages[0]["engine_args"]["video_output_transport"] == transport + + def test_compact_offload_config_reaches_terminal_config(monkeypatch): from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index de5abeab0..6f7833ff8 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -9,6 +9,7 @@ import pytest import torch from torch import nn +from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -30,6 +31,22 @@ class _StubTextEncoder(nn.Module): return torch.float32 +class _StubVaeConfig: + latents_mean = [0.0, 0.0, 0.0, 0.0] + latents_std = [1.0, 1.0, 1.0, 1.0] + z_dim = 4 + + +class _StubVae(nn.Module): + dtype = torch.float32 + config = _StubVaeConfig() + + def decode(self, latents, return_dict=False): + del return_dict + batch, _, frames, height, width = latents.shape + return (torch.zeros(batch, 3, frames, height, width),) + + class _StubScheduler: def __init__(self, timesteps: list[int]) -> None: self.timesteps = torch.tensor(timesteps, dtype=torch.int64) @@ -75,6 +92,7 @@ def _make_pipeline() -> Wan22Pipeline: pipeline.transformer = _StubTransformer() pipeline.transformer_2 = None pipeline.text_encoder = _StubTextEncoder() + pipeline.vae = _StubVae() pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4) pipeline.scheduler = _StubScheduler([9, 5]) pipeline.od_config = SimpleNamespace(flow_shift=5.0) @@ -227,6 +245,99 @@ def test_forward_batches_text_generators_latents_and_splits_outputs() -> None: torch.testing.assert_close(outputs[1].output, latents_b) +def test_forward_emits_request_local_typed_media_after_vae_decode() -> None: + pipeline = _make_pipeline() + + def _fake_diffuse( + *, + latents, + timesteps, + prompt_embeds, + negative_prompt_embeds, + guidance_low, + guidance_high, + boundary_timestep, + dtype, + attention_kwargs, + latent_condition, + first_frame_mask, + generator, + ): + del ( + timesteps, + prompt_embeds, + negative_prompt_embeds, + guidance_low, + guidance_high, + boundary_timestep, + dtype, + attention_kwargs, + latent_condition, + first_frame_mask, + generator, + ) + return torch.zeros_like(latents) + + pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="prompt", + request_id="request-0", + sampling_params=OmniDiffusionSamplingParams( + num_frames=1, + num_inference_steps=2, + max_sequence_length=32, + output_type="np", + ), + ) + ] + ) + + outputs = pipeline.forward(batch) + + assert len(outputs) == 1 + assert outputs[0].output is None + assert outputs[0].media is not None + assert outputs[0].media.prepared_for_transport is False + assert outputs[0].media.video.tensor.shape == (1, 3, 1, 8, 8) + assert outputs[0].media.video.spec.layout is VideoTensorLayout.BCTHW + assert outputs[0].media.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert outputs[0].media.video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE + + +def test_forward_keeps_legacy_output_on_non_owner_vae_rank() -> None: + # Distributed VAE decode uses broadcast_result=False, so non-owner ranks get + # an empty placeholder instead of the full video. Wrapping that as typed media + # would fail split_diffusion_output_by_request's batch check on every non-owner + # rank, so the pipeline must keep the placeholder on the legacy output field. + pipeline = _make_pipeline() + pipeline.vae.decode = lambda latents, return_dict=False: (torch.empty(0),) # type: ignore[assignment] + pipeline.diffuse = lambda **kwargs: torch.zeros_like(kwargs["latents"]) # type: ignore[method-assign] + + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="prompt", + request_id="request-0", + sampling_params=OmniDiffusionSamplingParams( + num_frames=1, + num_inference_steps=2, + max_sequence_length=32, + output_type="np", + ), + ) + ] + ) + + outputs = pipeline.forward(batch) + + assert len(outputs) == 1 + assert outputs[0].media is None + assert outputs[0].output is not None + assert outputs[0].output.numel() == 0 + + def test_forward_batches_precomputed_prompt_embeddings() -> None: pipeline = _make_pipeline() diffuse_call = {} diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py index a4bff4625..8a81f30d4 100644 --- a/tests/diffusion/test_async_output_worker.py +++ b/tests/diffusion/test_async_output_worker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for WorkerProc async output logic.""" import queue @@ -8,13 +8,23 @@ import time from unittest.mock import MagicMock import pytest +import torch +import vllm_omni.diffusion.ipc as diffusion_ipc from vllm_omni.diffusion.data import ( AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput, OmniACK, ) +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput @@ -182,6 +192,66 @@ class TestReturnResultSyncPath: mock_pack.assert_called_once_with(output) proc.result_mq.enqueue.assert_called_once_with(output) + def test_media_contract_error_is_not_sent_as_an_unpacked_result(self, mocker): + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=ValueError("media was not prepared"), + ) + + with pytest.raises(ValueError, match="media was not prepared"): + proc._return_result(DiffusionOutput()) + + proc.result_mq.enqueue.assert_not_called() + + def test_media_packing_memory_failure_is_reraised_for_typed_media(self, mocker): + """A non-ValueError packing failure on typed media must re-raise, not + swallow it and enqueue a half-packed payload.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=RuntimeError("SHM pack out of memory"), + ) + + output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput)) + with pytest.raises(RuntimeError, match="SHM pack out of memory"): + proc._return_result(output) + + proc.result_mq.enqueue.assert_not_called() + + def test_media_packing_failure_is_reraised_for_dp_container(self, mocker): + """Container dispatch and typed-media failure detection share a walker.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=RuntimeError("SHM pack out of memory"), + ) + + output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput)) + tagged_output = {"dp_rank": 0, "output": output} + with pytest.raises(RuntimeError, match="SHM pack out of memory"): + proc._return_result(tagged_output) + + proc.result_mq.enqueue.assert_not_called() + + def test_legacy_output_packing_failure_is_still_swallowed(self, mocker): + """Non-typed (legacy output) failures keep the historic swallow-and- + enqueue behavior; only typed media re-raises.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=RuntimeError("SHM pack out of memory"), + ) + + proc._return_result(DiffusionOutput(output="data")) + + # Legacy path: warning logged, output still enqueued once. + proc.result_mq.enqueue.assert_called_once_with(DiffusionOutput(output="data")) + def test_non_diffusion_output_skips_async_in_request_mode(self, mocker): """Non-DiffusionOutput/BatchRunnerOutput in request mode uses sync path.""" proc = _make_worker_proc(step_execution=False) @@ -243,6 +313,70 @@ class TestAsyncOutputLoopLogic: ) assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls" + def test_late_field_failure_sends_error_without_partial_payload_or_shm_leak( + self, + mocker, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from multiprocessing import shared_memory + + proc = _make_worker_proc(step_execution=False) + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 3, 2, 4, 5), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ), + prepared_for_transport=True, + ) + output = DiffusionOutput(media=media, trajectory_latents=torch.ones(30)) + original_latents = output.trajectory_latents + + real_array_to_shm = diffusion_ipc._array_to_shm + created_handles: list[dict[str, object]] = [] + + def _fail_on_second(array): + if created_handles: + raise RuntimeError("late-field packing failure") + handle = real_array_to_shm(array) + created_handles.append(handle) + return handle + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second) + + def _pack_without_stream(value, d2h_stream=None): + # The worker loop behavior is under test here. CPU tensors avoid + # introducing a device dependency into this unit test. + del d2h_stream + return diffusion_ipc.pack_diffusion_output_shm(value) + + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=_pack_without_stream, + ) + mock_torch = mocker.MagicMock() + mock_torch.accelerator.current_accelerator.return_value.type = "cpu" + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch) + + proc._async_output_queue.put((output, "late-failure", None)) + proc._async_output_queue.put(None) + proc._async_output_loop() + + messages = [call.args[0] for call in proc.result_mq.enqueue.call_args_list] + assert len(messages) == 1 + assert messages[0].kind is AsyncOutputKind.OUTPUT_READY + assert messages[0].async_output_id == "late-failure" + assert messages[0].error == "Background D2H/SHM packing failed" + assert messages[0].output is None + assert output.media is media + assert output.trajectory_latents is original_latents + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=str(created_handles[0]["name"])) + class TestDrainAsyncOutputs: """Async outputs must finish packing before device memory is released.""" diff --git a/tests/diffusion/test_device_output_preparation.py b/tests/diffusion/test_device_output_preparation.py new file mode 100644 index 000000000..0a6f1ab90 --- /dev/null +++ b/tests/diffusion/test_device_output_preparation.py @@ -0,0 +1,460 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import numpy as np +import pytest +import torch + +import vllm_omni.diffusion.ipc as diffusion_ipc +import vllm_omni.diffusion.postprocess.device_reduction as device_reduction +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, + ensure_request_owned_tensor, + slice_diffusion_media_output, +) +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _config(*, enabled: bool) -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model=None, + video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled), + ) + + +def _sampling(*, output_type: str | None = "np", interpolate: bool = False) -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + output_type=output_type, + enable_frame_interpolation=interpolate, + ) + + +def _media( + tensor: torch.Tensor, + *, + value_range: VideoValueRange = VideoValueRange.NEGATIVE_ONE_TO_ONE, + consumers: frozenset[FloatVideoConsumer] = frozenset(), +) -> DiffusionMediaOutput: + return DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=value_range, + ), + constraints=VideoTransportConstraints(pending_float_consumers=consumers), + ) + ) + + +def test_runtime_prepares_request_local_video_as_uint8() -> None: + video = torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5) + + prepared = prepare_diffusion_media_for_transport( + _media(video), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.tensor.dtype == torch.uint8 + assert prepared.video.tensor.shape == (2, 2, 4, 5, 3) + assert prepared.video.spec == VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + assert prepared.video.constraints.pending_float_consumers == frozenset() + + +def test_disabled_runtime_preserves_float_but_marks_media_prepared() -> None: + video = torch.randn(1, 3, 2, 4, 5) + + prepared = prepare_diffusion_media_for_transport( + _media(video), + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.tensor is video + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + + +def test_device_oom_falls_back_to_request_owned_float_media(monkeypatch: pytest.MonkeyPatch) -> None: + source = torch.randn(2, 3, 2, 4, 5) + sliced = slice_diffusion_media_output(_media(source), 0, 1) + + def raise_oom(video: torch.Tensor, *, do_denormalize: bool) -> torch.Tensor: + raise torch.OutOfMemoryError("injected") + + monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", raise_oom) + + prepared = prepare_diffusion_media_for_transport( + sliced, + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype == source.dtype + assert prepared.video.tensor.shape == (1, 3, 2, 4, 5) + assert prepared.video.tensor.is_contiguous() + assert prepared.video.tensor._base is None + + +def test_device_non_oom_error_is_not_swallowed(monkeypatch: pytest.MonkeyPatch) -> None: + def raise_runtime_error(video: torch.Tensor, *, do_denormalize: bool) -> torch.Tensor: + raise RuntimeError("injected") + + monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", raise_runtime_error) + + with pytest.raises(RuntimeError, match="injected"): + prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + +def test_declared_and_request_float_consumers_are_unioned() -> None: + video = torch.randn(1, 3, 2, 4, 5) + media = _media( + video, + consumers=frozenset({FloatVideoConsumer.VIDEO_GUARDRAILS}), + ) + + prepared = prepare_diffusion_media_for_transport( + media, + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + + assert prepared.video.tensor is video + assert prepared.video.constraints.pending_float_consumers == frozenset( + { + FloatVideoConsumer.FRAME_INTERPOLATION, + FloatVideoConsumer.VIDEO_GUARDRAILS, + } + ) + + +def test_non_numpy_presentation_preserves_float() -> None: + prepared = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=True), + sampling_params=_sampling(output_type="pil"), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + + +def test_zero_to_one_video_is_not_denormalized() -> None: + video = torch.tensor([[[[[0.0, 0.5, 1.0]]]]]).expand(1, 3, 1, 1, 3) + + prepared = prepare_diffusion_media_for_transport( + _media(video, value_range=VideoValueRange.ZERO_TO_ONE), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + expected = np.rint(video.permute(0, 2, 3, 4, 1).numpy() * 255.0).astype(np.uint8) + np.testing.assert_array_equal(prepared.video.tensor.numpy(), expected) + + +def test_invalid_contract_raises_even_when_feature_is_off() -> None: + media = _media(torch.randn(3, 2, 4, 5)) + + with pytest.raises(ValueError, match="rank 5"): + prepare_diffusion_media_for_transport( + media, + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + + +@pytest.mark.parametrize( + ("tensor", "spec", "message"), + [ + ( + torch.randn(1, 2, 4, 5, 3), + VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + "torch.uint8", + ), + ( + torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_ONE, + ), + "incompatible value range", + ), + ( + torch.randn(1, 4, 2, 4, 5), + VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + "3 channels", + ), + ], +) +def test_invalid_spec_tensor_combinations_are_contract_errors( + tensor: torch.Tensor, + spec: VideoTensorSpec, + message: str, +) -> None: + media = DiffusionMediaOutput(video=VideoMediaOutput(tensor=tensor, spec=spec)) + + with pytest.raises(ValueError, match=message): + media.validate() + + +def test_pipeline_cannot_emit_unprepared_uint8_media() -> None: + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + ) + ) + + with pytest.raises(ValueError, match="must emit unprepared media as NORMALIZED_FLOAT"): + prepare_diffusion_media_for_transport( + media, + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + +def test_descriptors_are_immutable() -> None: + media = _media(torch.randn(1, 3, 2, 4, 5)) + + with pytest.raises(FrozenInstanceError): + media.prepared_for_transport = True # type: ignore[misc] + with pytest.raises(FrozenInstanceError): + media.video.spec.layout = VideoTensorLayout.BTHWC # type: ignore[misc] + + +def test_float_fallback_detaches_a_request_view_before_transport() -> None: + source = torch.randn(4, 3, 2, 4, 5) + sliced = slice_diffusion_media_output(_media(source), 1, 3) + + prepared = prepare_diffusion_media_for_transport( + sliced, + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + + tensor = prepared.video.tensor + assert tensor.shape == (2, 3, 2, 4, 5) + assert tensor.is_contiguous() + assert tensor._base is None + assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size() + assert tensor.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + + +def test_full_batch_slice_does_not_duplicate_owned_storage() -> None: + source = torch.randn(2, 3, 2, 4, 5) + + sliced = slice_diffusion_media_output(_media(source), 0, 2) + + assert sliced.video.tensor.untyped_storage().data_ptr() == source.untyped_storage().data_ptr() + + +def test_diffusion_output_to_cpu_moves_prepared_typed_media() -> None: + prepared = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + output = DiffusionOutput(media=prepared, to_cpu=True) + + assert output.media is not None + assert output.media.video.tensor.device.type == "cpu" + + +def test_diffusion_output_to_cpu_rejects_unprepared_media() -> None: + with pytest.raises(ValueError, match="prepared before to_cpu"): + DiffusionOutput(media=_media(torch.randn(1, 3, 2, 4, 5)), to_cpu=True) + + +def test_prepared_media_round_trips_through_typed_ipc(monkeypatch: pytest.MonkeyPatch) -> None: + expected = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 256, 256)), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + output = DiffusionOutput(media=expected) + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + pack_diffusion_output_shm(output) + + assert isinstance(output.media, dict) + assert output.media["__type__"] == "diffusion_media_v1" + assert output.media["video"]["tensor"]["__tensor_shm__"] is True + + unpack_diffusion_output_shm(output) + + assert isinstance(output.media, DiffusionMediaOutput) + assert output.media.prepared_for_transport is True + assert output.media.video.spec == expected.video.spec + torch.testing.assert_close(output.media.video.tensor, expected.video.tensor) + + +def test_float_media_constraints_round_trip_through_typed_ipc(monkeypatch: pytest.MonkeyPatch) -> None: + expected = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 256, 256, dtype=torch.bfloat16)), + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + output = DiffusionOutput(media=expected) + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + pack_diffusion_output_shm(output) + unpack_diffusion_output_shm(output) + + assert isinstance(output.media, DiffusionMediaOutput) + assert output.media.video.tensor.dtype is torch.bfloat16 + assert output.media.video.constraints.pending_float_consumers == frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + torch.testing.assert_close(output.media.video.tensor, expected.video.tensor) + + +def test_ipc_rejects_media_that_skipped_runner_preparation() -> None: + output = DiffusionOutput(media=_media(torch.randn(1, 3, 2, 4, 5))) + + with pytest.raises(ValueError, match="must be prepared"): + pack_diffusion_output_shm(output) + + +def test_ipc_rejects_unknown_media_schema_version() -> None: + prepared = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + output = DiffusionOutput(media=prepared) + pack_diffusion_output_shm(output) + assert isinstance(output.media, dict) + output.media["schema_version"] = 2 + + with pytest.raises(ValueError, match="schema version"): + unpack_diffusion_output_shm(output) + + +def test_prepare_rejects_media_a_pipeline_already_prepared(monkeypatch: pytest.MonkeyPatch) -> None: + # A pipeline must emit unprepared media so the runner applies request policy + # (e.g. frame interpolation). Accepting a forged prepared uint8 payload would + # skip the interpolation constraint; interpolator never runs and gets dropped. + forged = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + ), + prepared_for_transport=True, + ) + + with pytest.raises(ValueError, match="already prepared"): + prepare_diffusion_media_for_transport( + forged, + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + + +def test_prepare_rejects_prepared_media_despite_pending_interpolation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The forged payload is the exact case the reviewer described: device + # postprocessing enabled plus frame_interpolation, but the media already being + # prepared. It must be rejected, not silently returned with interpolation dropped. + forged = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 3, 2, 2, 3, dtype=torch.float32), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ), + prepared_for_transport=True, + ) + + with pytest.raises(ValueError, match="already prepared"): + prepare_diffusion_media_for_transport( + forged, + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + + +def test_request_owned_tensor_rejects_a_detached_storage_view() -> None: + # ``tensor._base is None`` is not proof of ownership: a detached slice keeps a + # nonzero storage offset and the whole source storage. Serializing it would + # leak every other request's pixels, so it must be cloned before transport. + source = torch.randn(4, 3, 2, 4, 5) + detached_slice = source[1:2].detach() + assert detached_slice.is_contiguous() + assert detached_slice._base is None + assert detached_slice.storage_offset() != 0 + + owned = ensure_request_owned_tensor(detached_slice) + + assert owned.shape == detached_slice.shape + assert owned.storage_offset() == 0 + assert owned.untyped_storage().nbytes() == owned.numel() * owned.element_size() + assert owned.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + torch.testing.assert_close(owned, detached_slice) + + +def test_prepared_media_rejects_non_compact_request_storage() -> None: + # A zero-offset view of a bigger buffer is contiguous but not compact, so + # prepared transport must refuse it just as it refuses a storage view. + source = torch.randn(2, 3, 2, 4, 5) + zero_offset_view = source[0:1] + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=zero_offset_view, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ), + prepared_for_transport=True, + ) + + with pytest.raises(ValueError, match="request-local storage"): + media.validate() diff --git a/tests/diffusion/test_device_reduction.py b/tests/diffusion/test_device_reduction.py new file mode 100644 index 000000000..f69bb5eb5 --- /dev/null +++ b/tests/diffusion/test_device_reduction.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""GPU correctness/size tests for device-side video reduction. + +The reduction must be byte-identical to the current production output (SHM +widens bf16->f32, then diffusers postprocess, then the API server's ``*255`` +rounding), otherwise enabling it would silently change pixels. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.postprocess.device_reduction import reduce_video_to_uint8_frames + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] + +_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available() +requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path") + + +def _device() -> str: + return "cuda:0" + + +def _reference_uint8(video: torch.Tensor, *, compute_dtype: torch.dtype | None = torch.float32) -> np.ndarray: + """The diffusers postprocess plus the API server's rounding, for the same input. + + ``compute_dtype`` selects the precision the denormalize runs at. + ``torch.float32`` is the precision the device reduction uses; ``None`` keeps + the input's own dtype, which is what the engine's postprocess actually does. + """ + from diffusers.video_processor import VideoProcessor + + source = video if compute_dtype is None else video.to(compute_dtype) + processed = VideoProcessor(vae_scale_factor=8).postprocess_video(source, output_type="np") + scaled = np.clip(processed.astype(np.float32), 0.0, 1.0) * 255.0 + np.rint(scaled, out=scaled) + return scaled.astype(np.uint8) + + +@requires_gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_device_reduction_matches_a_float32_reference(dtype: torch.dtype) -> None: + """The reduction must be the float32 computation, exactly, for every input dtype.""" + torch.manual_seed(0) + # Values slightly outside [-1, 1] so the clamp is actually exercised. + video = (torch.rand(1, 3, 8, 64, 96, device=_device()) * 2.4 - 1.2).to(dtype) + + reference = _reference_uint8(video) + produced = reduce_video_to_uint8_frames(video).cpu().numpy() + + assert produced.shape == reference.shape == (1, 8, 64, 96, 3) + assert produced.dtype == np.uint8 + np.testing.assert_array_equal(produced, reference) + + +@requires_gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_device_reduction_deviates_from_a_narrow_postprocess_by_at_most_one(dtype: torch.dtype) -> None: + """Pin the cost of computing in float32: at most one 255th. + + A postprocess that denormalizes in the VAE's own dtype (WAN, HunyuanVideo, + LTX-2, Cosmos3) does not match this reduction bit for bit; ~20% of pixels + differ by one on a real 480x832x81 WAN generation. Keep that a bound rather + than a surprise, so a regression widening it fails here. + """ + torch.manual_seed(0) + video = (torch.rand(1, 3, 8, 64, 96, device=_device()) * 2.4 - 1.2).to(dtype) + + narrow = _reference_uint8(video, compute_dtype=None) + produced = reduce_video_to_uint8_frames(video).cpu().numpy() + + deviation = np.abs(produced.astype(np.int16) - narrow.astype(np.int16)) + assert deviation.max() <= 1 + # The narrow path really is different, so the bound above is not vacuous. + assert deviation.any() + + +@requires_gpu +def test_device_reduction_shrinks_hop1_payload_4x() -> None: + """The whole point: the D2H payload drops from float32 to uint8 (4x).""" + video = torch.rand(1, 3, 8, 704, 1280, device=_device(), dtype=torch.bfloat16) * 2 - 1 + + # Production transports bf16 widened to float32 over SHM. + widened_f32_bytes = video.numel() * 4 + reduced = reduce_video_to_uint8_frames(video) + reduced_bytes = reduced.numel() * reduced.element_size() + + assert reduced.dtype == torch.uint8 + assert widened_f32_bytes / reduced_bytes == pytest.approx(4.0, rel=1e-6) + + +@requires_gpu +def test_device_reduction_stays_on_device() -> None: + """Reduction must happen before D2H, so the result is still on the GPU.""" + video = torch.rand(1, 3, 4, 32, 48, device=_device(), dtype=torch.bfloat16) * 2 - 1 + reduced = reduce_video_to_uint8_frames(video) + assert reduced.device.type == "cuda" + assert reduced.is_contiguous() + + +@requires_gpu +def test_device_reduction_without_denormalize_matches_already_unit_range() -> None: + """Pipelines with do_normalize=False emit [0,1]; skip the denorm step.""" + video = torch.rand(1, 3, 4, 16, 16, device=_device(), dtype=torch.float32) # already [0,1] + + produced = reduce_video_to_uint8_frames(video, do_denormalize=False).cpu().numpy() + + expected = video.permute(0, 2, 3, 4, 1).cpu().numpy() + expected = np.rint(np.clip(expected, 0, 1) * 255.0).astype(np.uint8) + np.testing.assert_array_equal(produced, expected) + + +def test_device_reduction_rejects_non_5d() -> None: + """A wrong-rank tensor is a programming error, not a silent reshape.""" + with pytest.raises(ValueError, match=r"\[B, C, F, H, W\]"): + reduce_video_to_uint8_frames(torch.rand(3, 8, 64, 64)) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 9f8d4dec6..69a9e4395 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -10,9 +10,19 @@ import torch import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.data import DiffusionOutput, VideoOutputTransportConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput, NewRequestData from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -144,6 +154,19 @@ def _make_request(): ) +def _raw_video_media(batch_size: int = 1, tensor: torch.Tensor | None = None) -> DiffusionMediaOutput: + return DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor if tensor is not None else torch.randn(batch_size, 3, 2, 4, 5), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + + def _make_request_with_params(req_id: str, sampling_params): return SimpleNamespace( request_id=req_id, @@ -214,6 +237,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar enable_cache_dit_summary=enable_cache_dit_summary, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, + video_output_transport=VideoOutputTransportConfig(), ) runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend( vllm_config=runner.vllm_config, @@ -546,6 +570,48 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): assert fourth.get_request_output("req").finished is True +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_stepwise_prepares_typed_media(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.pipeline = _ChunkStepPipeline([DiffusionOutput(media=_raw_video_media())]) + runner.od_config.streaming_output = True + runner.od_config.step_execution = True + runner.od_config.video_output_transport = VideoOutputTransportConfig(enable_device_postprocess=True) + req = OmniDiffusionRequest( + prompt="a prompt", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=4), + request_id="req", + ) + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + scheduler_output = DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=[NewRequestData(request_id="req", req=req)], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=1, + num_waiting_reqs=0, + ) + + DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + scheduler_output = DiffusionSchedulerOutput( + step_id=1, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData(request_ids=["req"]), + finished_req_ids=set(), + num_running_reqs=1, + num_waiting_reqs=0, + ) + result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output).get_request_output("req").result + + assert result.media.prepared_for_transport is True + assert result.media.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES + + @pytest.mark.core_model @pytest.mark.cpu def test_execute_stepwise_streaming_decodes_final_only_pipeline(monkeypatch): @@ -819,6 +885,7 @@ def _make_batch_runner(pipeline): cache_backend="none", enable_cache_dit_summary=False, parallel_config=SimpleNamespace(use_hsdp=False), + video_output_transport=VideoOutputTransportConfig(), ) runner.kv_transfer_manager = SimpleNamespace( receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None, @@ -938,6 +1005,21 @@ def test_execute_model_batch_uses_runner_output_helper(monkeypatch): assert helper_calls == [(["req-0", "req-1"], ["a", "b"])] +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_prepares_typed_media_before_return(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(media=_raw_video_media())])) + runner.od_config.video_output_transport = VideoOutputTransportConfig(enable_device_postprocess=True) + + result = DiffusionModelRunner.execute_model_batch(runner, _make_scheduler_output(num_reqs=1), runner.od_config) + + prepared = result.runner_outputs[0].result.media + assert prepared.prepared_for_transport is True + assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES + + @pytest.mark.core_model @pytest.mark.cpu def test_split_diffusion_output_by_request_slices_single_and_multi_request_outputs(): @@ -959,6 +1041,36 @@ def test_split_diffusion_output_by_request_slices_single_and_multi_request_outpu assert single[0].output == ["img-0a", "img-0b"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_diffusion_output_by_request_preserves_typed_media(): + reqs = [_make_request(), _make_request()] + batch = DiffusionRequestBatch(requests=reqs) + source = torch.randn(4, 3, 2, 4, 5) + result = DiffusionOutput(media=_raw_video_media(tensor=source)) + + outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2) + + assert [output.media.video.tensor.shape for output in outputs] == [ + torch.Size([2, 3, 2, 4, 5]), + torch.Size([2, 3, 2, 4, 5]), + ] + assert all(output.media.prepared_for_transport is False for output in outputs) + torch.testing.assert_close(outputs[0].media.video.tensor, source[:2]) + torch.testing.assert_close(outputs[1].media.video.tensor, source[2:]) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_diffusion_output_by_request_validates_media_batch_dimension(): + reqs = [_make_request(), _make_request()] + batch = DiffusionRequestBatch(requests=reqs) + result = DiffusionOutput(media=_raw_video_media(batch_size=3)) + + with pytest.raises(ValueError, match=r"request_count \* num_outputs_per_prompt \(4\), got 3"): + split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2) + + @pytest.mark.core_model @pytest.mark.cpu def test_split_diffusion_output_by_request_slices_tuple_outputs(): diff --git a/tests/diffusion/test_ipc_async.py b/tests/diffusion/test_ipc_async.py index 03919095b..43650c581 100644 --- a/tests/diffusion/test_ipc_async.py +++ b/tests/diffusion/test_ipc_async.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for IPC async D2H stream packing.""" +from unittest.mock import ANY + import pytest import torch @@ -43,7 +45,7 @@ class TestPackTensorIfLargeWithD2hStream: tensor = torch.zeros(1) d2h = mocker.MagicMock() _pack_value_if_large(tensor, d2h_stream=d2h) - mock_pack.assert_called_once_with(tensor, d2h_stream=d2h) + mock_pack.assert_called_once_with(tensor, d2h_stream=d2h, created=None) def test_pack_value_if_large_passes_stream_through_dict(self, mocker): """d2h_stream propagates through dict container packing.""" @@ -93,7 +95,7 @@ class TestPackDiffusionFieldsWithD2hStream: d2h = mocker.MagicMock() out = DiffusionOutput(output=torch.zeros(1)) _pack_diffusion_fields(out, d2h_stream=d2h) - mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h) + mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h, created=None) def test_trajectory_fields_packed_with_stream(self, mocker): mock_pack_tensor = mocker.patch( @@ -120,7 +122,12 @@ class TestPackDiffusionOutputShmWithD2hStream: d2h = mocker.MagicMock() out = DiffusionOutput() pack_diffusion_output_shm(out, d2h_stream=d2h) - mock_pack_fields.assert_called_once_with(out, d2h_stream=d2h) + mock_pack_fields.assert_called_once_with( + out, + d2h_stream=d2h, + created=ANY, + pending_updates=ANY, + ) def test_rpc_envelope_path_with_stream(self, mocker): mock_pack_fields = mocker.patch( @@ -130,11 +137,11 @@ class TestPackDiffusionOutputShmWithD2hStream: result = DiffusionOutput() envelope = {"type": "diffusion_rpc_result", "result": result} pack_diffusion_output_shm(envelope, d2h_stream=d2h) - # envelope["result"] is replaced in-place by pack_diffusion_output_shm; - # assert with the original reference, not the replaced one. mock_pack_fields.assert_called_once_with( result, d2h_stream=d2h, + created=ANY, + pending_updates=ANY, ) @pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True]) @@ -147,7 +154,12 @@ class TestPackDiffusionOutputShmWithD2hStream: tagged = {"dp_rank": 1, "output": result} value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged pack_diffusion_output_shm(value, d2h_stream=d2h) - mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h) + mock_pack_fields.assert_called_once_with( + result, + d2h_stream=d2h, + created=ANY, + pending_updates=ANY, + ) def test_runner_output_path_with_stream(self, mocker): mock_pack_fields = mocker.patch( @@ -163,9 +175,9 @@ class TestPackDiffusionOutputShmWithD2hStream: result=result, ) pack_diffusion_output_shm(runner_out, d2h_stream=d2h) - # runner_out.result is replaced in-place by pack_diffusion_output_shm; - # assert with the original reference, not the replaced one. mock_pack_fields.assert_called_once_with( result, d2h_stream=d2h, + created=ANY, + pending_updates=ANY, ) diff --git a/tests/diffusion/test_media_output_contract.py b/tests/diffusion/test_media_output_contract.py new file mode 100644 index 000000000..ae5c2fdbf --- /dev/null +++ b/tests/diffusion/test_media_output_contract.py @@ -0,0 +1,546 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module +import vllm_omni.diffusion.ipc as diffusion_ipc +import vllm_omni.diffusion.postprocess.media as media_postprocess +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, +) +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport +from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _video( + tensor: torch.Tensor, + *, + encoding: VideoTensorEncoding, + value_range: VideoValueRange, + consumers: frozenset[FloatVideoConsumer] = frozenset(), +) -> VideoMediaOutput: + layout = VideoTensorLayout.BTHWC if encoding is VideoTensorEncoding.UINT8_FRAMES else VideoTensorLayout.BCTHW + return VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=layout, + encoding=encoding, + value_range=value_range, + ), + constraints=VideoTransportConstraints(pending_float_consumers=consumers), + ) + + +def _prepared_media(video: VideoMediaOutput) -> DiffusionMediaOutput: + return DiffusionMediaOutput(video=video, prepared_for_transport=True) + + +def _od_config(*, enabled: bool = False) -> OmniDiffusionConfig: + config = OmniDiffusionConfig( + model=None, + video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled), + ) + config.model_class_name = "WanPipeline" + return config + + +def test_finalizer_rejects_media_that_skipped_runner_preparation() -> None: + media = DiffusionMediaOutput( + video=_video( + torch.randn(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + + with pytest.raises(ValueError, match="before transport preparation"): + finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + +def test_uint8_media_finalization_preserves_frames() -> None: + frames = torch.randint(0, 256, (1, 4, 8, 8, 3), dtype=torch.uint8) + + output = finalize_diffusion_media( + _prepared_media( + _video( + frames, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + np.testing.assert_array_equal(output["payload"]["video"], frames.numpy()) + assert output["metadata"] == {} + + +def test_float_media_finalization_uses_the_declared_range() -> None: + video = torch.tensor([[[[[-1.0, 0.0, 1.0]]]]]).expand(1, 3, 1, 1, 3).contiguous() + + output = finalize_diffusion_media( + _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + expected = (video.permute(0, 2, 3, 4, 1).numpy() * 0.5 + 0.5).clip(0.0, 1.0) + np.testing.assert_array_equal(output["payload"]["video"], expected) + + +def test_zero_to_one_float_media_is_not_denormalized() -> None: + video = torch.tensor([[[[[0.0, 0.5, 1.0]]]]]).expand(1, 3, 1, 1, 3).contiguous() + media = _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.ZERO_TO_ONE, + ) + ) + + output = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + expected = video.permute(0, 2, 3, 4, 1).numpy() + np.testing.assert_array_equal(output["payload"]["video"], expected) + + +def test_presentation_is_resolved_outside_the_tensor_descriptor() -> None: + video = torch.zeros(1, 3, 2, 4, 5) + media = _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + + output = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="pil")) + + assert "presentation" not in media.video.__dataclass_fields__ + assert len(output["payload"]["video"]) == 1 + assert len(output["payload"]["video"][0]) == 2 + + +def test_frame_interpolation_consumes_its_constraint(monkeypatch: pytest.MonkeyPatch) -> None: + video = torch.randn(1, 3, 2, 4, 5) + interpolated = torch.randn(1, 3, 3, 4, 5) + monkeypatch.setattr( + media_postprocess, + "interpolate_video_tensor", + lambda *args, **kwargs: (interpolated, 2.0), + ) + + output = finalize_diffusion_media( + _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}), + ) + ), + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + assert output["payload"]["video"].shape == (1, 3, 4, 5, 3) + assert output["metadata"] == {"video": {"video_fps_multiplier": 2.0}} + + +def test_uint8_media_rejects_a_pending_float_consumer() -> None: + media = _prepared_media( + _video( + torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}), + ) + ) + + with pytest.raises(ValueError, match="cannot have pending float consumers"): + media.validate() + + +def test_unconsumed_float_constraint_fails_before_public_formatting() -> None: + media = _prepared_media( + _video( + torch.randn(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + consumers=frozenset({FloatVideoConsumer.VIDEO_GUARDRAILS}), + ) + ) + + with pytest.raises(ValueError, match="pending float consumers"): + finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + +def test_engine_routes_typed_media_exclusively(monkeypatch: pytest.MonkeyPatch) -> None: + media = _prepared_media( + _video( + torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ) + request = OmniDiffusionRequest( + prompt="test", + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + request_id="request-0", + ) + engine = object.__new__(DiffusionEngine) + engine.od_config = _od_config() + engine.post_process_func = lambda _: pytest.fail("legacy postprocessor was called") + engine._post_process_accepts_sampling_params = False + + finalized = {"payload": {"video": np.zeros((1, 1, 2, 2, 3), dtype=np.uint8)}, "metadata": {}} + + def _format_outputs( + *, + request: object, + od_config: object, + diffusion_output: object, + output_data: object, + postprocess_output: object, + ) -> list[object]: + del request, od_config, diffusion_output, output_data + return [postprocess_output] + + monkeypatch.setattr(diffusion_engine_module, "finalize_diffusion_media", lambda *args, **kwargs: finalized) + monkeypatch.setattr(diffusion_engine_module, "normalize_diffusion_postprocess_output", lambda value: value) + monkeypatch.setattr(diffusion_engine_module, "format_diffusion_outputs", _format_outputs) + + result = engine.postprocess_output(request, DiffusionOutput(media=media)) + + assert result == [finalized] + + +def test_typed_media_crosses_runner_ipc_and_engine_without_legacy_postprocess( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sampling_params = OmniDiffusionSamplingParams(output_type="np", num_outputs_per_prompt=1) + raw_media = DiffusionMediaOutput( + video=_video( + torch.linspace(-1.0, 1.0, 1 * 3 * 2 * 4 * 5).reshape(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + prepared = prepare_diffusion_media_for_transport( + raw_media, + od_config=_od_config(enabled=True), + sampling_params=sampling_params, + ) + diffusion_output = DiffusionOutput(media=prepared) + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + pack_diffusion_output_shm(diffusion_output) + unpack_diffusion_output_shm(diffusion_output) + + engine = object.__new__(DiffusionEngine) + engine.od_config = _od_config() + engine.post_process_func = lambda _: pytest.fail("legacy postprocessor was called") + engine._post_process_accepts_sampling_params = False + request = OmniDiffusionRequest(prompt="test", sampling_params=sampling_params, request_id="request-0") + + outputs = engine.postprocess_output(request, diffusion_output) + + assert len(outputs) == 1 + assert outputs[0].images[0].dtype == np.uint8 + assert outputs[0].images[0].shape == (1, 2, 4, 5, 3) + + +def test_diffusion_output_rejects_typed_and_legacy_output_together() -> None: + media = _prepared_media( + _video( + torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ) + with pytest.raises(ValueError, match="both media and legacy output"): + DiffusionOutput(output=torch.zeros(1), media=media) + + +def test_diffusion_output_rejects_postprocessor_on_typed_media() -> None: + media = _prepared_media( + _video( + torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ) + + with pytest.raises(ValueError, match="model-specific post_process_func"): + DiffusionOutput(media=media, post_process_func=lambda value: value) + + +def test_diffusion_output_preserves_the_legacy_positional_constructor() -> None: + # ``media`` was originally declared right after ``output``, which broke + # ``DiffusionOutput(output, trajectory_timesteps)`` by binding the second + # argument to media. It must stay last so out-of-tree positional callers keep + # the pre-existing field order. + output_tensor = torch.randn(1) + timesteps = torch.tensor([1, 2, 3]) + + output = DiffusionOutput(output_tensor, timesteps) + + assert output.output is output_tensor + assert output.trajectory_timesteps is timesteps + assert output.media is None + + +def test_frame_interpolation_uses_the_declared_range_negative_one_to_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # RIFE infers [0,1] vs [-1,1] from the tensor min/max, so an all-nonnegative + # [-1,1] clip is misread as [0,1] and interpolated in the wrong space. The + # interpolator must receive the unit-range form per the declared spec. + video = torch.full((1, 3, 2, 3, 3), 0.25) + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=video, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + ), + ), + prepared_for_transport=True, + ) + + seen: list[torch.Tensor] = [] + + def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]: + seen.append(tensor) + return tensor * tensor, 2.0 + + monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator) + + output = finalize_diffusion_media( + media, + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + # The interpolator is fed the unit-range form 0.25→0.625, not raw 0.25. + assert len(seen) == 1 + assert seen[0].shape == video.shape + assert float(seen[0][0, 0, 0, 0, 0]) == pytest.approx(0.625, abs=1e-6) + # Squaring gives 0.390625; restored to [-1,1] (-0.21875) by the fix, then the + # existing denormalize maps it back to [0,1] as 0.390625 in the public frames. + frames = output["payload"]["video"] + assert np.allclose(np.asarray(frames), 0.390625, atol=1e-6) + + +def test_frame_interpolation_uses_the_declared_range_zero_to_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # With a declared [0,1] range, the all-0.25 clip is already in unit space, so + # squaring gives 0.390625 and there is no restore step. + video = torch.full((1, 3, 2, 3, 3), 0.25) + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=video, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.ZERO_TO_ONE, + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + ), + ), + prepared_for_transport=True, + ) + + def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]: + return tensor * tensor, 2.0 + + monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator) + + output = finalize_diffusion_media( + media, + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + # 0.25 is already [0,1]; squaring gives 0.0625, and no denormalize is applied. + frames = output["payload"]["video"] + assert np.allclose(np.asarray(frames), 0.0625, atol=1e-6) + + +def test_frame_interpolation_clamps_zero_to_one_overshoot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A declared [0,1] clip whose sample is almost entirely inside [0,1] — but + # with one bf16 value overshooting 1 — is still in unit space, so it must be + # clamped to [0,1] before RIFE and interpolated as [0,1]. Without the clamp + # the interpolator classifies it as [-1,1] (amax > 1) and returns the wrong + # space, which is then written back with no *2-1 restore. + video = torch.full((1, 3, 2, 3, 3), 0.25, dtype=torch.bfloat16) + video[0, 0, 0, 0, 0] = 1.25 # single overshoot > 1 -> amax > 1 + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=video, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.ZERO_TO_ONE, + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + ), + ), + prepared_for_transport=True, + ) + + seen: list[torch.Tensor] = [] + + def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]: + seen.append(tensor) + return tensor.float().square(), 2.0 + + monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator) + + output = finalize_diffusion_media( + media, + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + # The interpolator is fed the clamped bf16 unit-range tensor (max now 1.0), + # not the raw 1.25 overshoot that would be misread as [-1,1]. + assert len(seen) == 1 + assert seen[0].dtype is torch.bfloat16 + assert float(seen[0].min()) == pytest.approx(0.25, abs=1e-6) + assert float(seen[0].max()) == pytest.approx(1.0, abs=1e-6) + # Squaring the clamped values gives 0.0625 at the 0.25 samples and 1.0 at the + # overshoot; the overall shape is unchanged (no *2-1 for ZERO_TO_ONE). + frames = output["payload"]["video"] + assert np.allclose(np.asarray(frames).reshape(-1)[1:], 0.0625, atol=1e-6) + assert np.isclose(float(np.asarray(frames).reshape(-1)[0]), 1.0, atol=1e-6) + + +def test_packing_failure_unlinks_created_segments_and_leaves_output_untouched( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Packing a late field can fail after an SHM segment is already created; + # the segment must be unlinked and ``output`` left unmutated so a partial + # object is never enqueued. + prepared = prepare_diffusion_media_for_transport( + DiffusionMediaOutput( + video=_video( + torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ), + od_config=_od_config(enabled=False), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + output = DiffusionOutput(media=prepared, trajectory_latents=torch.randn(30)) + original_media = prepared + original_latents = output.trajectory_latents + + real_array_to_shm = diffusion_ipc._array_to_shm + created_handles: list[dict[str, object]] = [] + + def _fail_on_second(array: np.ndarray) -> dict[str, object]: + if created_handles: + raise RuntimeError("injected packing failure") + handle = real_array_to_shm(array) + created_handles.append(handle) + return handle + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second) + + with pytest.raises(RuntimeError, match="injected packing failure"): + pack_diffusion_output_shm(output) + + assert created_handles, "no SHM segment was created before the injected failure" + # Output was not mutated despite the media segment being created first. + assert output.media is original_media + assert isinstance(output.media, DiffusionMediaOutput) + assert output.trajectory_latents is original_latents + + from multiprocessing import shared_memory + + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=str(created_handles[0]["name"])) + + +def test_second_batch_entry_failure_leaves_the_whole_batch_untouched( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from multiprocessing import shared_memory + + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + first_media = _prepared_media( + _video( + torch.zeros(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + second_media = _prepared_media( + _video( + torch.ones(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + first = DiffusionOutput(media=first_media) + second = DiffusionOutput(media=second_media) + batch = BatchRunnerOutput( + runner_outputs=[ + RunnerOutput(request_id="first", result=first), + RunnerOutput(request_id="second", result=second), + ] + ) + + real_array_to_shm = diffusion_ipc._array_to_shm + created_handles: list[dict[str, object]] = [] + + def _fail_on_second(array: np.ndarray) -> dict[str, object]: + if created_handles: + raise RuntimeError("second batch entry failed") + handle = real_array_to_shm(array) + created_handles.append(handle) + return handle + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second) + + with pytest.raises(RuntimeError, match="second batch entry failed"): + pack_diffusion_output_shm(batch) + + # No per-entry assignment is committed until every entry has packed. + assert first.media is first_media + assert second.media is second_media + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=str(created_handles[0]["name"])) diff --git a/tests/diffusion/test_video_transport_defaults_are_inert.py b/tests/diffusion/test_video_transport_defaults_are_inert.py new file mode 100644 index 000000000..fd7346e76 --- /dev/null +++ b/tests/diffusion/test_video_transport_defaults_are_inert.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pin the default behavior and the WAN reference integration boundary.""" + +from __future__ import annotations + +import inspect +from pathlib import Path + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_transport_defaults() -> None: + transport = VideoOutputTransportConfig() + + assert transport.enable_device_postprocess is False + + +def test_a_default_diffusion_config_carries_a_default_transport() -> None: + config = OmniDiffusionConfig(model=None) + + assert isinstance(config.video_output_transport, VideoOutputTransportConfig) + assert config.video_output_transport.enable_device_postprocess is False + + +def test_wan_uses_the_typed_media_contract() -> None: + import vllm_omni.diffusion.models as models_pkg + + source = (Path(inspect.getfile(models_pkg)).parent / "wan2_2/pipeline_wan2_2.py").read_text() + assert "DiffusionMediaOutput(" in source + assert "reduce_video_to_uint8_frames(" not in source + # The migrated built-in WAN pipeline emits typed media, so it never invokes + # the legacy postprocess hook. The mappings are intentionally retained so + # legacy/out-of-tree WAN replacements (registered with + # post_process_func_name=None) keep the built-in hook instead of exposing raw + # BCTHW tensors. + assert _DIFFUSION_POST_PROCESS_FUNCS["WanPipeline"] == "get_wan22_post_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS["WanDMDPipeline"] == "get_wan22_post_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS["WanT2VDMD2Pipeline"] == "get_wan22_post_process_func" + + +def test_non_reference_models_do_not_use_the_runtime_reducer() -> None: + import vllm_omni.diffusion.models as models_pkg + + root = Path(inspect.getfile(models_pkg)).parent + paths = ( + "cosmos3/pipeline_cosmos3.py", + "hunyuan_video/pipeline_hunyuan_video_1_5.py", + "lingbot_video/pipeline_lingbot_video.py", + "ltx2/ltx2_runtime.py", + "minimax_h3/pipeline_minimax_h3.py", + ) + for relative in paths: + source = (root / relative).read_text() + assert "prepare_diffusion_media_for_transport(" not in source + assert "reduce_video_to_uint8_frames(" not in source diff --git a/tests/diffusion/test_wan_device_reduction.py b/tests/diffusion/test_wan_device_reduction.py new file mode 100644 index 000000000..492aec7a5 --- /dev/null +++ b/tests/diffusion/test_wan_device_reduction.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""WAN reference coverage for the typed pre-D2H media contract.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import get_wan22_post_process_func +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport +from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media +from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_uint8_frames +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] + +_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available() +requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path") + + +def _decoded_wan_video(tensor: torch.Tensor) -> DiffusionMediaOutput: + return DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + + +def _config(*, enabled: bool) -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model=None, + video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled), + ) + + +def _encoder_frames(video_payload: np.ndarray) -> list[np.ndarray]: + return [_coerce_video_to_uint8_frames(video_payload[i]) for i in range(video_payload.shape[0])] + + +@requires_gpu +def test_wan_media_reduction_matches_float32_reference_and_bounds_native_path() -> None: + torch.manual_seed(0) + vae_out = torch.rand(2, 3, 8, 64, 96, device="cuda:0", dtype=torch.bfloat16) * 2.4 - 1.2 + sampling = OmniDiffusionSamplingParams(output_type=None, enable_frame_interpolation=False) + + legacy_postprocess = get_wan22_post_process_func(_config(enabled=False)) + widened_out = legacy_postprocess(vae_out.float(), output_type="np", sampling_params=sampling) + widened_frames = _encoder_frames(widened_out["payload"]["video"]) + native_out = legacy_postprocess(vae_out, output_type="np", sampling_params=sampling) + native_frames = _encoder_frames(native_out["payload"]["video"]) + + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(vae_out), + od_config=_config(enabled=True), + sampling_params=sampling, + ) + reduced_out = finalize_diffusion_media(prepared, sampling_params=sampling) + assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES + assert reduced_out["payload"]["video"].dtype == np.uint8 + reduced_frames = _encoder_frames(reduced_out["payload"]["video"]) + + assert len(widened_frames) == len(reduced_frames) == 2 + for expected, native, produced in zip(widened_frames, native_frames, reduced_frames, strict=True): + assert produced.shape == expected.shape == (8, 64, 96, 3) + np.testing.assert_array_equal(produced, expected) + deviation = np.abs(produced.astype(np.int16) - native.astype(np.int16)) + assert deviation.max() <= 1 + + +@requires_gpu +def test_wan_uint8_media_finalization_does_not_rescale() -> None: + frames = torch.randint(0, 256, (1, 4, 16, 16, 3), dtype=torch.uint8, device="cuda:0") + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=frames, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + ), + prepared_for_transport=True, + ) + + out = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + np.testing.assert_array_equal(out["payload"]["video"], frames.cpu().numpy()) + + +@requires_gpu +def test_diffusion_output_to_cpu_moves_wan_media_off_device() -> None: + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")), + od_config=_config(enabled=False), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + output = DiffusionOutput(media=prepared, to_cpu=True) + + assert output.media is not None + assert output.media.video.tensor.device.type == "cpu" + + +@requires_gpu +def test_small_typed_media_is_on_cpu_before_ipc() -> None: + sampling = OmniDiffusionSamplingParams(output_type="np") + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")), + od_config=_config(enabled=True), + sampling_params=sampling, + ) + output = DiffusionOutput(media=prepared) + + pack_diffusion_output_shm(output) + + assert isinstance(output.media, dict) + packed_tensor = output.media["video"]["tensor"] + assert isinstance(packed_tensor, torch.Tensor) + assert packed_tensor.device.type == "cpu" + unpack_diffusion_output_shm(output) + + +@requires_gpu +def test_wan_float_media_finalization_matches_the_legacy_path() -> None: + video = torch.rand(1, 3, 4, 16, 16, device="cuda:0", dtype=torch.bfloat16) * 2 - 1 + sampling = OmniDiffusionSamplingParams(output_type="np") + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=False), + sampling_params=sampling, + ) + + produced = finalize_diffusion_media(prepared, sampling_params=sampling) + expected = get_wan22_post_process_func(_config(enabled=False))( + video, + output_type="np", + sampling_params=sampling, + ) + + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + np.testing.assert_array_equal(produced["payload"]["video"], expected["payload"]["video"]) + + +@requires_gpu +def test_oom_fallback_media_packs_without_device_widening(monkeypatch: pytest.MonkeyPatch) -> None: + # Regression: when the fp32 reduction OOMs, the fallback keeps the bf16 float + # payload. Async IPC must copy it to the host in bf16 and widen on the CPU; + # a second device-side fp32 conversion would OOM again and drop the payload. + from vllm_omni.diffusion.postprocess import device_reduction + + # Large enough (> 1 MB in bf16) that packing routes the tensor through the + # shared-memory D2H path rather than the small-tensor CPU copy. + video = torch.rand(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) * 2 - 1 + + def _raise_oom(*_args: object, **_kwargs: object) -> torch.Tensor: + raise torch.OutOfMemoryError("forced fp32 reduction OOM") + + monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", _raise_oom) + + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=True), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype == torch.bfloat16 + + output = DiffusionOutput(media=prepared) + + pinned_dtypes: list[torch.dtype] = [] + real_empty = torch.empty + + def _spy_empty(*args: object, **kwargs: object) -> torch.Tensor: + tensor = real_empty(*args, **kwargs) + if kwargs.get("pin_memory"): + pinned_dtypes.append(tensor.dtype) + return tensor + + monkeypatch.setattr(torch, "empty", _spy_empty) + + d2h_stream = torch.Stream(device=torch.device("cuda", 0)) + pack_diffusion_output_shm(output, d2h_stream=d2h_stream) + d2h_stream.synchronize() + + assert pinned_dtypes == [torch.bfloat16], ( + "async packing must D2H in the original dtype and widen on the CPU, " + "never widen bf16 to fp32 on the accelerator before the copy" + ) + + unpack_diffusion_output_shm(output) + assert isinstance(output.media, DiffusionMediaOutput) + assert output.media.video.tensor.dtype == torch.bfloat16 + torch.testing.assert_close(output.media.video.tensor, video.cpu()) + + +@requires_gpu +def test_oom_fallback_recovers_the_allocator(monkeypatch: pytest.MonkeyPatch) -> None: + # A real float32 OOM leaves the CUDA caching allocator in a failed state; + # the fallback must empty the cache and drain the device before marking the + # float representation transport-ready, so the worker D2H of the same tensor + # can actually run. Use a genuine allocation failure (per-process cap) rather + # than injecting the exception, which never poisons the allocator. + from vllm_omni.platforms import current_omni_platform + + video = torch.empty(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) + video.uniform_(-1.0, 1.0) + expected = video.cpu() + d2h_stream = torch.Stream(device=video.device) + base_bytes = video.numel() * video.element_size() + _free, total = current_omni_platform.get_device_memory() + + # Cap the process so the fp32 conversion (2x base bytes) cannot fit, while + # retaining the original bf16 tensor. This forces a real allocator OOM at + # the .to(fp32), rather than injecting an exception from a monkeypatch. + torch.cuda.set_per_process_memory_fraction(min(0.99, (1.2 * base_bytes) / total)) + real_empty_cache = current_omni_platform.empty_cache + real_synchronize = current_omni_platform.synchronize + try: + empty_calls: list[int] = [] + sync_calls: list[bool] = [] + + def _empty_cache() -> None: + real_empty_cache() + empty_calls.append(torch.accelerator.current_device_index()) + + def _synchronize() -> None: + real_synchronize() + sync_calls.append(True) + + monkeypatch.setattr(current_omni_platform, "empty_cache", _empty_cache) + monkeypatch.setattr(current_omni_platform, "synchronize", _synchronize) + + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=True), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + # The fallback ran past a genuine OOM and kept the bf16 float payload. + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype == torch.bfloat16 + assert empty_calls, "fallback must empty the CUDA allocator cache after a real OOM" + assert sync_calls, "fallback must drain the device after a real OOM" + + # Exercise the next production step: async D2H/SHM packing of the same + # tensor must succeed after allocator recovery and preserve its values. + output = DiffusionOutput(media=prepared) + pack_diffusion_output_shm(output, d2h_stream=d2h_stream) + d2h_stream.synchronize() + unpack_diffusion_output_shm(output) + assert isinstance(output.media, DiffusionMediaOutput) + torch.testing.assert_close(output.media.video.tensor, expected) + finally: + torch.cuda.set_per_process_memory_fraction(1.0) + real_empty_cache() + real_synchronize() + + +@requires_gpu +def test_oom_fallback_releases_late_uint8_traceback_before_clone( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A final uint8 allocation OOM must release float intermediates first.""" + from vllm_omni.diffusion.postprocess import device_reduction + from vllm_omni.platforms import current_omni_platform + + # A split request is a non-compact view, so float fallback must allocate a + # request-owned clone after reduction fails. + backing = torch.empty(2, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) + backing.uniform_(-1.0, 1.0) + video = backing[1:2].detach() + expected = video.cpu() + _, total = current_omni_platform.get_device_memory() + real_to = torch.Tensor.to + real_empty_cache = current_omni_platform.empty_cache + real_synchronize = current_omni_platform.synchronize + real_ensure_request_owned = device_reduction.ensure_request_owned_tensor + uint8_attempts: list[int] = [] + fallback_allocations: list[int] = [] + + def _record_fallback_allocation(tensor: torch.Tensor) -> torch.Tensor: + fallback_allocations.append(torch.accelerator.memory_allocated(tensor.device)) + return real_ensure_request_owned(tensor) + + def _fail_real_uint8_allocation( + tensor: torch.Tensor, + *args: object, + **kwargs: object, + ) -> torch.Tensor: + dtype = kwargs.get("dtype") + if dtype is None and args and isinstance(args[0], torch.dtype): + dtype = args[0] + if tensor.is_cuda and dtype is torch.uint8: + # Earlier float32 work runs without a cap. Immediately before the + # real final uint8 allocation, cap the allocator below that request + # but above the post-unwind request-owned bf16 clone requirement. + real_empty_cache() + allocated = torch.accelerator.memory_allocated(tensor.device) + uint8_bytes = tensor.numel() + limit_bytes = allocated + max(1, uint8_bytes // 2) + torch.cuda.set_per_process_memory_fraction(min(0.99, limit_bytes / total)) + uint8_attempts.append(allocated) + return real_to(tensor, *args, **kwargs) + + monkeypatch.setattr(device_reduction, "ensure_request_owned_tensor", _record_fallback_allocation) + monkeypatch.setattr(torch.Tensor, "to", _fail_real_uint8_allocation) + try: + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=True), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + assert uint8_attempts, "the real final .to(uint8) allocation was not attempted" + assert fallback_allocations, "float fallback did not request compact storage" + assert fallback_allocations[0] < uint8_attempts[0], ( + "the failed reduction's float intermediates remained live during fallback" + ) + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype is torch.bfloat16 + assert prepared.video.tensor.storage_offset() == 0 + assert prepared.video.tensor.untyped_storage().data_ptr() != video.untyped_storage().data_ptr() + torch.testing.assert_close(prepared.video.tensor.cpu(), expected) + finally: + torch.cuda.set_per_process_memory_fraction(1.0) + real_empty_cache() + real_synchronize() diff --git a/tests/docs/test_generate_argparse.py b/tests/docs/test_generate_argparse.py new file mode 100644 index 000000000..16b92365f --- /dev/null +++ b/tests/docs/test_generate_argparse.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +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_argparse.py" +SPEC = importlib.util.spec_from_file_location("generate_argparse", HOOK_PATH) +assert SPEC and SPEC.loader +generate_argparse = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(generate_argparse) + + +def test_static_serve_parser_supports_video_output_transport() -> None: + parser = generate_argparse.create_parser_subparser_init(generate_argparse.OmniServeCommand) + + args = parser.parse_args(["--omni", "--video-output-transport", '{"enable_device_postprocess": true}']) + + assert args.video_output_transport == {"enable_device_postprocess": True} diff --git a/tests/e2e/offline_inference/test_device_postprocess_equivalence.py b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py new file mode 100644 index 000000000..a7fbcafd9 --- /dev/null +++ b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""End-to-end WAN check for the typed pre-D2H media contract. + +Runs the whole engine twice with the same seed, once per flag state, and compares +the frames a client receives. WAN denormalizes historical bfloat16 output in the +input dtype while device preparation computes in float32, so the paths may differ +by one 255th. + +Run with a WAN2.2 checkpoint:: + + VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + pytest -s tests/e2e/offline_inference/test_device_postprocess_equivalence.py +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from tests.helpers.mark import hardware_test +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_MODEL") + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.gpu, + pytest.mark.skipif( + not MODEL, + reason="Set VLLM_OMNI_DEVICE_POSTPROCESS_MODEL to a WAN2.2 checkpoint to run this.", + ), +] + +_SEED = 1234 +_STEPS = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_STEPS", "2")) +_SIZE = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_SIZE", "256")) +_FRAMES = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_FRAMES", "17")) + + +def _generate(enable_device_postprocess: bool) -> np.ndarray: + from vllm_omni.entrypoints.omni import Omni + + engine = Omni( + model=MODEL, + num_gpus=1, + video_output_transport={"enable_device_postprocess": enable_device_postprocess}, + ) + try: + # output_type must travel with the request: the pipelines read it from the + # sampling params, and the reduction only covers the "np" path. + sampling_params = OmniDiffusionSamplingParams( + output_type="np", + seed=_SEED, + num_inference_steps=_STEPS, + height=_SIZE, + width=_SIZE, + num_frames=_FRAMES, + ) + outputs = engine.generate({"prompt": "a robot waving"}, sampling_params) + finally: + engine.close() + + video = outputs[0].images[0] + return video.numpy() if hasattr(video, "numpy") else np.asarray(video) + + +def _as_uint8_frames(video: np.ndarray) -> np.ndarray: + """Reproduce what the API server does before encoding. + + A float payload is scaled and rounded there, so that is the only fair place to + compare it against a payload the worker already reduced. + """ + if video.dtype == np.uint8: + return video + return np.rint(np.clip(video, 0.0, 1.0) * 255.0).astype(np.uint8) + + +# A dtype difference moves ~20% of pixels by one; dropping the round moves ~50%. +_MAX_DIFFERING_FRACTION = 0.35 + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_device_postprocess_produces_equivalent_frames() -> None: + float_path = _generate(False) + device_path = _generate(True) + + # Guard against the run silently taking the float path twice, which would make + # the comparison meaningless. + assert float_path.dtype != np.uint8, "the float path returned uint8; the gate may be stuck open" + assert device_path.dtype == np.uint8, "the WAN typed-media path did not prepare uint8 frames" + assert float_path.shape == device_path.shape + + expected = _as_uint8_frames(float_path) + deviation = np.abs(device_path.astype(np.int16) - expected.astype(np.int16)) + differing = float((deviation > 0).mean()) + print(f"maxdiff={deviation.max()} differing={differing:.4f} ({int((deviation > 0).sum())}/{deviation.size})") + + assert deviation.max() <= 1, "the reduction moved a pixel by more than one 255th" + assert differing < _MAX_DIFFERING_FRACTION, ( + f"{differing:.1%} of pixels differ, more than a denormalize-precision difference explains" + ) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index b64758e5d..0089330b2 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -69,6 +69,30 @@ def test_serve_parser_accepts_deploy_config() -> None: assert args.get_explicit_kwargs_dict()["deploy_config"] == "/tmp/deploy.yaml" +def test_serve_parser_accepts_video_output_transport() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + ["serve", "fake-model", "--omni", "--video-output-transport", '{"enable_device_postprocess": true}'] + ) + + expected = {"enable_device_postprocess": True} + assert args.video_output_transport == expected + assert args.get_explicit_kwargs_dict()["video_output_transport"] == expected + + +@pytest.mark.parametrize("value", ["{not json", "[]"]) +def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + with pytest.raises(SystemExit): + parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value]) + + def test_serve_parser_rejects_stage_configs_path() -> None: parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="subcommand") diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 81eb76f0b..2b016be0d 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -718,6 +718,7 @@ class _DiffusionConfigProjection: cache_strategy: str = "none" cache_backend: str = "none" cache_config: Any = field(default_factory=dict) + video_output_transport: object = field(default_factory=dict) enable_cache_dit_summary: bool = False diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True) @@ -812,6 +813,7 @@ class _DiffusionConfigProjection: AttentionConfig, DiffusionCacheConfig, TransformerConfig, + VideoOutputTransportConfig, build_attention_config, parse_kv_cache_skip_selector, validate_dlo_host_registration_options, @@ -854,6 +856,13 @@ class _DiffusionConfigProjection: elif not isinstance(self.cache_config, DiffusionCacheConfig): self.cache_config = DiffusionCacheConfig() + if self.video_output_transport is None: + self.video_output_transport = VideoOutputTransportConfig() + elif isinstance(self.video_output_transport, Mapping): + self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport)) + elif not isinstance(self.video_output_transport, VideoOutputTransportConfig): + raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping") + self._propagate_quantization_from_tf_config(self.tf_model_config) if self.quantization_config is not None: if isinstance(self.quantization_config, QuantizationConfig): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2ac1f99dd..f7d714890 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -458,6 +458,7 @@ class StageDeployConfig: fa_deterministic: bool | None = None cache_backend: str | None = None cache_config: dict[str, Any] | None = None + video_output_transport: dict[str, Any] | None = None enable_cache_dit_summary: bool | None = None step_execution: bool | None = None vae_use_slicing: bool | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c30529520..d27cc36e1 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -29,6 +29,7 @@ from vllm_omni.diffusion.diffusion_kv.config import ( parse_diffusion_kv_cache_mode, ) from vllm_omni.diffusion.lora.manager import LoRABackend +from vllm_omni.diffusion.media import DiffusionMediaOutput from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available from vllm_omni.errors import client_error_metadata @@ -758,6 +759,15 @@ def uses_diffusers_adapter(od_config: object) -> bool: ) +@dataclass +class VideoOutputTransportConfig: + enable_device_postprocess: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.enable_device_postprocess, bool): + raise TypeError("enable_device_postprocess must be a bool") + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -793,6 +803,7 @@ class OmniDiffusionConfig: # Cache backend configuration (NEW) cache_backend: str = "none" # "tea_cache", "deep_cache", etc. cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict) + video_output_transport: VideoOutputTransportConfig = field(default_factory=VideoOutputTransportConfig) enable_cache_dit_summary: bool = False # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the @@ -1258,6 +1269,13 @@ class OmniDiffusionConfig: # If it's neither dict nor DiffusionCacheConfig, convert to empty config self.cache_config = DiffusionCacheConfig() + if self.video_output_transport is None: + self.video_output_transport = VideoOutputTransportConfig() + elif isinstance(self.video_output_transport, Mapping): + self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport)) + elif not isinstance(self.video_output_transport, VideoOutputTransportConfig): + raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping") + # Auto-detect quantization from TransformerConfig if not explicitly set. # This covers the case where tf_model_config is passed at construction # time. For late (post-construction) assignment, callers should use @@ -1618,8 +1636,8 @@ class DiffusionOutput: output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None - # Legacy compatibility fields. New pipeline-specific payloads should be - # carried by output["payload"] instead. + # Legacy compatibility fields. Decoded video uses ``media``; other new + # payloads should be carried by output["payload"] instead. trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None trajectory_latents: torch.Tensor | dict[str, Any] | None = None trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None @@ -1658,7 +1676,18 @@ class DiffusionOutput: # mode) and the receiving side must not initialise a stray CUDA context. to_cpu: bool = False + # Typed video-media contract. Declared last so the pre-existing positional + # constructor order (output, trajectory_timesteps, ...) that out-of-tree + # pipelines rely on is preserved. Mutually exclusive with ``output``. + media: DiffusionMediaOutput | None = None + def __post_init__(self) -> None: + if self.media is not None and not isinstance(self.media, DiffusionMediaOutput): + raise TypeError(f"media must be DiffusionMediaOutput, got {type(self.media).__name__}") + if self.media is not None and self.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + if self.media is not None and self.post_process_func is not None: + raise ValueError("Typed diffusion media cannot carry a model-specific post_process_func") if not self.to_cpu: return @@ -1674,6 +1703,10 @@ class DiffusionOutput: return value self.output = _maybe_to_cpu(self.output) + if self.media is not None: + if not self.media.prepared_for_transport: + raise ValueError("Diffusion media must be prepared before to_cpu=True") + self.media = self.media.to_cpu() self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps) self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents) self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c36ac0459..417090503 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -47,6 +47,7 @@ from vllm_omni.diffusion.output_formatter import ( format_empty_diffusion_outputs, normalize_diffusion_postprocess_output, ) +from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media from vllm_omni.diffusion.registry import ( DiffusionModelRegistry, get_diffusion_post_process_func, @@ -534,26 +535,33 @@ class DiffusionEngine: raise RuntimeError(output.error) logger.debug("Generation completed successfully.") - if output.output is None: - logger.warning("Output is None, returning empty OmniRequestOutput") - return format_empty_diffusion_outputs(request, finished=output.finished) - - # When CPU offload is enabled, move output to CPU before - # post-processing to avoid device OOM — model weights may still - # reside on the device and leave no headroom for intermediates. - output_data = output.output - if self.od_config.enable_cpu_offload: - output_data = _move_tensor_tree_to_cpu(output_data) - - if self.post_process_func is not None: - # Some video pipelines need request-level controls during - # postprocess (for example worker-side frame interpolation). - postprocess_kwargs: dict[str, object] = {} - if self._post_process_accepts_sampling_params: - postprocess_kwargs["sampling_params"] = request.sampling_params - outputs = self.post_process_func(output_data, **postprocess_kwargs) + if output.media is not None: + if output.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + media = output.media.to_cpu() if self.od_config.enable_cpu_offload else output.media + output_data = media.video.tensor + outputs = finalize_diffusion_media(media, sampling_params=request.sampling_params) else: - outputs = output_data + if output.output is None: + logger.warning("Output is None, returning empty OmniRequestOutput") + return format_empty_diffusion_outputs(request, finished=output.finished) + + # When CPU offload is enabled, move output to CPU before + # post-processing to avoid device OOM — model weights may still + # reside on the device and leave no headroom for intermediates. + output_data = output.output + if self.od_config.enable_cpu_offload: + output_data = _move_tensor_tree_to_cpu(output_data) + + if self.post_process_func is not None: + # Some video pipelines need request-level controls during + # postprocess (for example worker-side frame interpolation). + postprocess_kwargs: dict[str, object] = {} + if self._post_process_accepts_sampling_params: + postprocess_kwargs["sampling_params"] = request.sampling_params + outputs = self.post_process_func(output_data, **postprocess_kwargs) + else: + outputs = output_data postprocess_output = normalize_diffusion_postprocess_output(outputs) diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 6eb24c540..f0e88a2c9 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """IPC utilities for transferring large tensors via POSIX shared memory. @@ -11,15 +11,32 @@ serialised through the queue. from __future__ import annotations +from collections.abc import Iterator from typing import Any import numpy as np import torch from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoColorModel, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, +) _SHM_TENSOR_THRESHOLD = 1_000_000 # 1 MB DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result" +_DIFFUSION_MEDIA_WIRE_TYPE = "diffusion_media_v1" + +# Sentinel so compute-then-assign packing can tell "field not packed" apart from +# "field packed to None". +_UNSET = object() def _array_to_shm(array: np.ndarray) -> dict[str, Any]: @@ -61,9 +78,39 @@ def _array_from_shm(handle: dict[str, Any]) -> np.ndarray: return array +def _unlink_shm_handle(handle: dict[str, Any]) -> None: + """Best-effort removal of a shared-memory segment created during packing.""" + from multiprocessing import shared_memory + + name = handle.get("name") + if not name: + return + try: + shm = shared_memory.SharedMemory(name=name) + except FileNotFoundError: + return + try: + shm.close() + finally: + try: + shm.unlink() + except FileNotFoundError: + pass + + +def _unlink_shm_handles(handles: list[dict[str, Any]]) -> None: + for handle in handles: + try: + _unlink_shm_handle(handle) + except Exception: + # Cleanup is best-effort; never mask the original packing failure. + pass + + def _tensor_to_shm( tensor: torch.Tensor, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: """Copy a tensor into POSIX shared memory and return a metadata handle. @@ -82,14 +129,17 @@ def _tensor_to_shm( torch.accelerator.set_stream(d2h_stream) try: t = tensor.detach() - if original_dtype == torch.bfloat16: - t = t.to(torch.float32) cpu = torch.empty(t.shape, dtype=t.dtype, pin_memory=True) cpu.copy_(t, non_blocking=True) finally: torch.accelerator.set_stream(old_stream) d2h_stream.synchronize() tensor = cpu + # numpy has no bfloat16, so widen on the host after D2H. Never widen on + # the accelerator before the copy: a payload that already fell back after + # an fp32 OOM would OOM again on a second device-side fp32 allocation. + if original_dtype == torch.bfloat16: + tensor = tensor.to(torch.float32) else: tensor = tensor.detach().cpu().contiguous() if original_dtype == torch.bfloat16: @@ -101,6 +151,8 @@ def _tensor_to_shm( "torch_dtype": str(original_dtype), } ) + if created is not None: + created.append(handle) return handle @@ -120,6 +172,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: def _pack_tensor_if_large( val: torch.Tensor, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, ) -> torch.Tensor | dict: """Replace a tensor with an SHM handle if it exceeds the threshold. @@ -134,14 +187,16 @@ def _pack_tensor_if_large( except Exception: storage_bytes = view_bytes if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD: - return _tensor_to_shm(val, d2h_stream=d2h_stream) + return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created) return val -def _ndarray_to_shm(array: np.ndarray) -> dict[str, Any]: +def _ndarray_to_shm(array: np.ndarray, created: list[dict[str, Any]] | None = None) -> dict[str, Any]: """Copy a contiguous NumPy array into POSIX shared memory.""" handle = _array_to_shm(array) handle["__ndarray_shm__"] = True + if created is not None: + created.append(handle) return handle @@ -153,6 +208,7 @@ def _ndarray_from_shm(handle: dict[str, Any]) -> np.ndarray: def _pack_value_if_large( val: object, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, ) -> object: """Recursively replace large tensors with SHM handles. @@ -163,15 +219,19 @@ def _pack_value_if_large( the two in sync. """ if isinstance(val, torch.Tensor): - return _pack_tensor_if_large(val, d2h_stream=d2h_stream) + return _pack_tensor_if_large(val, d2h_stream=d2h_stream, created=created) if isinstance(val, np.ndarray): - return _ndarray_to_shm(val) if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD else val + return ( + _ndarray_to_shm(val, created=created) + if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD + else val + ) if isinstance(val, dict): - return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()} + return {key: _pack_value_if_large(value, d2h_stream=d2h_stream, created=created) for key, value in val.items()} if isinstance(val, list): - return [_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val] + return [_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val] if isinstance(val, tuple): - return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val) + return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val) return val @@ -190,18 +250,116 @@ def _unpack_if_shm_handle(val: object) -> object: return val +def _pack_diffusion_media( + media: DiffusionMediaOutput, + *, + d2h_stream: torch.Stream | None, + created: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + media.validate() + if not media.prepared_for_transport: + raise ValueError("Diffusion media must be prepared before IPC packing") + video = media.video + packed_tensor = _pack_tensor_if_large(video.tensor, d2h_stream=d2h_stream, created=created) + if isinstance(packed_tensor, torch.Tensor) and packed_tensor.device.type != "cpu": + packed_tensor = packed_tensor.detach().cpu().contiguous() + return { + "__type__": _DIFFUSION_MEDIA_WIRE_TYPE, + "schema_version": video.schema_version, + "prepared_for_transport": True, + "video": { + "tensor": packed_tensor, + "layout": video.spec.layout.value, + "encoding": video.spec.encoding.value, + "value_range": video.spec.value_range.value, + "color_model": video.spec.color_model.value, + "pending_float_consumers": sorted(consumer.value for consumer in video.constraints.pending_float_consumers), + }, + } + + +def _unpack_diffusion_media(packed: dict[str, Any]) -> DiffusionMediaOutput: + if packed.get("__type__") != _DIFFUSION_MEDIA_WIRE_TYPE: + raise ValueError("Invalid diffusion media wire payload") + video_payload = packed.get("video") + if not isinstance(video_payload, dict): + raise ValueError("Diffusion media wire payload is missing video metadata") + tensor = _unpack_if_shm_handle(video_payload.get("tensor")) + if not isinstance(tensor, torch.Tensor): + raise ValueError("Diffusion media wire payload contains an invalid video tensor") + schema_version = packed.get("schema_version") + if schema_version != 1: + raise ValueError(f"Unsupported video media schema version: {schema_version}") + if packed.get("prepared_for_transport") is not True: + raise ValueError("Diffusion media wire payload was not prepared for transport") + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=VideoTensorLayout(video_payload["layout"]), + encoding=VideoTensorEncoding(video_payload["encoding"]), + value_range=VideoValueRange(video_payload["value_range"]), + color_model=VideoColorModel(video_payload["color_model"]), + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset( + FloatVideoConsumer(value) for value in video_payload.get("pending_float_consumers", []) + ) + ), + schema_version=1, + ), + prepared_for_transport=True, + ) + media.validate() + return media + + def _pack_diffusion_fields( output: DiffusionOutput, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, + pending_updates: list[tuple[DiffusionOutput, str, object]] | None = None, ) -> DiffusionOutput: + if output.media is not None and output.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + if output.media is not None and not isinstance(output.media, DiffusionMediaOutput): + raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(output.media).__name__}") + + # Pack into locals first so a mid-payload failure leaves ``output`` unmutated + # (never enqueue a half-packed object); the caller unlinks any SHM segments + # accumulated in ``created`` on failure. + packed_output = _UNSET if output.output is not None: - output.output = _pack_value_if_large(output.output, d2h_stream=d2h_stream) - if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor): - output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream) - if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor): - output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream) - if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor): - output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream) + packed_output = _pack_value_if_large(output.output, d2h_stream=d2h_stream, created=created) + packed_media = _UNSET + if output.media is not None: + packed_media = _pack_diffusion_media(output.media, d2h_stream=d2h_stream, created=created) + packed_latents = _UNSET + if isinstance(output.trajectory_latents, torch.Tensor): + packed_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream, created=created) + packed_timesteps = _UNSET + if isinstance(output.trajectory_timesteps, torch.Tensor): + packed_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream, created=created) + packed_log_probs = _UNSET + if isinstance(output.trajectory_log_probs, torch.Tensor): + packed_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream, created=created) + + updates = [ + (output, field_name, value) + for field_name, value in ( + ("output", packed_output), + ("media", packed_media), + ("trajectory_latents", packed_latents), + ("trajectory_timesteps", packed_timesteps), + ("trajectory_log_probs", packed_log_probs), + ) + if value is not _UNSET + ] + if pending_updates is None: + for target, field_name, value in updates: + setattr(target, field_name, value) + else: + pending_updates.extend(updates) return output @@ -223,35 +381,84 @@ def pack_diffusion_output_shm( If *d2h_stream* is provided, D2H copies use that stream (non-blocking on the default stream). The caller must synchronize *d2h_stream* afterward. + + Packing is failure-atomic: every SHM segment created during the call is + tracked and unlinked if any part of the payload fails to pack, and field + mutations are committed only after the complete payload succeeds. """ + created: list[dict[str, Any]] = [] + pending_updates: list[tuple[DiffusionOutput, str, object]] = [] + try: + packed = _pack_output_shm( + output, + d2h_stream=d2h_stream, + created=created, + pending_updates=pending_updates, + ) + except BaseException: + _unlink_shm_handles(created) + raise + + # Commit mutations only after every field and batch entry packed. Until + # this point, a failure leaves the complete caller-owned payload untouched. + for target, field_name, value in pending_updates: + setattr(target, field_name, value) + return packed + + +def _iter_diffusion_outputs(output: object) -> Iterator[DiffusionOutput]: + """Yield every packable ``DiffusionOutput`` leaf in a worker payload.""" if isinstance(output, DiffusionOutput): - return _pack_diffusion_fields(output, d2h_stream=d2h_stream) + yield output + return # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput} if isinstance(output, dict) and "dp_rank" in output and "output" in output: inner = output["output"] if isinstance(inner, DiffusionOutput): - output["output"] = _pack_diffusion_fields(inner, d2h_stream=d2h_stream) - return output + yield inner + return if _is_rpc_result_envelope(output): - result = output.get("result") - output["result"] = pack_diffusion_output_shm(result, d2h_stream=d2h_stream) - return output + yield from _iter_diffusion_outputs(output.get("result")) + return result = getattr(output, "result", None) if isinstance(result, DiffusionOutput): - output.result = _pack_diffusion_fields(result, d2h_stream=d2h_stream) + yield result runner_outputs = getattr(output, "runner_outputs", None) if isinstance(runner_outputs, list): for runner_output in runner_outputs: - pack_diffusion_output_shm(runner_output, d2h_stream=d2h_stream) + yield from _iter_diffusion_outputs(runner_output) + + +def payload_carries_typed_media(output: object) -> bool: + """True if *output* carries any typed ``DiffusionOutput.media`` payload.""" + return any(item.media is not None for item in _iter_diffusion_outputs(output)) + + +def _pack_output_shm( + output: object, + *, + d2h_stream: torch.Stream | None, + created: list[dict[str, Any]], + pending_updates: list[tuple[DiffusionOutput, str, object]], +) -> object: + for item in _iter_diffusion_outputs(output): + _pack_diffusion_fields( + item, + d2h_stream=d2h_stream, + created=created, + pending_updates=pending_updates, + ) return output def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: output.output = _unpack_if_shm_handle(output.output) + if isinstance(output.media, dict): + output.media = _unpack_diffusion_media(output.media) output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents) output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps) output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs) diff --git a/vllm_omni/diffusion/media.py b/vllm_omni/diffusion/media.py new file mode 100644 index 000000000..34f4a1224 --- /dev/null +++ b/vllm_omni/diffusion/media.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Literal + +import torch + + +class VideoTensorLayout(str, Enum): + BCTHW = "bcthw" + BTHWC = "bthwc" + + +class VideoTensorEncoding(str, Enum): + NORMALIZED_FLOAT = "normalized_float" + UINT8_FRAMES = "uint8_frames" + + +class VideoValueRange(str, Enum): + NEGATIVE_ONE_TO_ONE = "negative_one_to_one" + ZERO_TO_ONE = "zero_to_one" + ZERO_TO_255 = "zero_to_255" + + +class VideoColorModel(str, Enum): + RGB = "rgb" + + +class FloatVideoConsumer(str, Enum): + FRAME_INTERPOLATION = "frame_interpolation" + VIDEO_GUARDRAILS = "video_guardrails" + + +@dataclass(frozen=True, slots=True) +class VideoTensorSpec: + layout: VideoTensorLayout + encoding: VideoTensorEncoding + value_range: VideoValueRange + color_model: VideoColorModel = VideoColorModel.RGB + + +@dataclass(frozen=True, slots=True) +class VideoTransportConstraints: + pending_float_consumers: frozenset[FloatVideoConsumer] = frozenset() + + +@dataclass(frozen=True, slots=True) +class VideoMediaOutput: + tensor: torch.Tensor + spec: VideoTensorSpec + constraints: VideoTransportConstraints = field(default_factory=VideoTransportConstraints) + schema_version: Literal[1] = 1 + + def validate(self) -> None: + if self.schema_version != 1: + raise ValueError(f"Unsupported video media schema version: {self.schema_version}") + if not isinstance(self.spec, VideoTensorSpec): + raise TypeError(f"Video media spec must be a VideoTensorSpec, got {type(self.spec).__name__}") + if not isinstance(self.constraints, VideoTransportConstraints): + raise TypeError( + f"Video media constraints must be VideoTransportConstraints, got {type(self.constraints).__name__}" + ) + if not isinstance(self.constraints.pending_float_consumers, frozenset): + raise TypeError("pending_float_consumers must be a frozenset") + if not isinstance(self.tensor, torch.Tensor): + raise TypeError(f"Video media tensor must be a torch.Tensor, got {type(self.tensor).__name__}") + if self.tensor.ndim != 5: + raise ValueError(f"Video media tensor must be rank 5, got shape {tuple(self.tensor.shape)}") + if any(size <= 0 for size in self.tensor.shape): + raise ValueError(f"Video media dimensions must be positive, got shape {tuple(self.tensor.shape)}") + if self.spec.color_model is not VideoColorModel.RGB: + raise ValueError(f"Unsupported video color model: {self.spec.color_model}") + if any(not isinstance(consumer, FloatVideoConsumer) for consumer in self.constraints.pending_float_consumers): + raise ValueError("pending_float_consumers must contain only FloatVideoConsumer values") + + if self.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT: + self._validate_normalized_float() + return + if self.spec.encoding is VideoTensorEncoding.UINT8_FRAMES: + self._validate_uint8_frames() + return + raise ValueError(f"Unsupported video tensor encoding: {self.spec.encoding}") + + def _validate_normalized_float(self) -> None: + if self.spec.layout is not VideoTensorLayout.BCTHW: + raise ValueError("NORMALIZED_FLOAT video must use BCTHW layout") + if self.tensor.shape[1] != 3: + raise ValueError(f"RGB BCTHW video must have 3 channels, got {self.tensor.shape[1]}") + if self.tensor.dtype not in {torch.float16, torch.bfloat16, torch.float32}: + raise ValueError(f"NORMALIZED_FLOAT video must use fp16, bf16, or fp32, got {self.tensor.dtype}") + if self.spec.value_range not in { + VideoValueRange.NEGATIVE_ONE_TO_ONE, + VideoValueRange.ZERO_TO_ONE, + }: + raise ValueError(f"NORMALIZED_FLOAT video has incompatible value range: {self.spec.value_range}") + + def _validate_uint8_frames(self) -> None: + if self.spec.layout is not VideoTensorLayout.BTHWC: + raise ValueError("UINT8_FRAMES video must use BTHWC layout") + if self.tensor.shape[-1] != 3: + raise ValueError(f"RGB BTHWC video must have 3 channels, got {self.tensor.shape[-1]}") + if self.tensor.dtype is not torch.uint8: + raise ValueError(f"UINT8_FRAMES video must use torch.uint8, got {self.tensor.dtype}") + if self.spec.value_range is not VideoValueRange.ZERO_TO_255: + raise ValueError(f"UINT8_FRAMES video has incompatible value range: {self.spec.value_range}") + if self.constraints.pending_float_consumers: + raise ValueError("UINT8_FRAMES video cannot have pending float consumers") + if not self.tensor.is_contiguous(): + raise ValueError("UINT8_FRAMES video tensor must be contiguous") + + def with_tensor(self, tensor: torch.Tensor) -> VideoMediaOutput: + return replace(self, tensor=tensor) + + def to_cpu(self) -> VideoMediaOutput: + return self.with_tensor(self.tensor.detach().cpu()) + + +@dataclass(frozen=True, slots=True) +class DiffusionMediaOutput: + video: VideoMediaOutput + # Only the model runner may set this before worker output packing. + prepared_for_transport: bool = False + + def validate(self) -> None: + if not isinstance(self.video, VideoMediaOutput): + raise TypeError(f"Diffusion media video must be a VideoMediaOutput, got {type(self.video).__name__}") + if not isinstance(self.prepared_for_transport, bool): + raise TypeError("prepared_for_transport must be a bool") + self.video.validate() + if not self.prepared_for_transport: + return + if not self.video.tensor.is_contiguous(): + raise ValueError("Prepared video media tensor must be contiguous") + if not _owns_compact_storage(self.video.tensor): + raise ValueError("Prepared video media tensor must own request-local storage") + + def with_video(self, video: VideoMediaOutput) -> DiffusionMediaOutput: + return replace(self, video=video) + + def to_cpu(self) -> DiffusionMediaOutput: + self.validate() + moved = self.with_video(self.video.to_cpu()) + moved.validate() + return moved + + +def _owns_compact_storage(tensor: torch.Tensor) -> bool: + """True if *tensor* is contiguous and backed by exactly its own elements. + + ``tensor._base is None`` is not sufficient: ``batch[i : i + 1].detach()`` is + contiguous with ``_base=None`` yet keeps a nonzero storage offset and the + whole batch storage, so serializing it would leak every other request's + pixels. Require a zero offset and a storage sized to the logical elements. + """ + if not tensor.is_contiguous(): + return False + if tensor.storage_offset() != 0: + return False + try: + storage_bytes = tensor.untyped_storage().nbytes() + except Exception: + return False + return storage_bytes == tensor.numel() * tensor.element_size() + + +def ensure_request_owned_tensor(tensor: torch.Tensor) -> torch.Tensor: + if _owns_compact_storage(tensor): + return tensor + # A contiguous view may still retain another request's full batch storage. + return tensor.clone(memory_format=torch.contiguous_format) + + +def slice_diffusion_media_output( + media: DiffusionMediaOutput, + start: int, + stop: int, +) -> DiffusionMediaOutput: + media.validate() + source = media.video.tensor + tensor = source if start == 0 and stop == source.shape[0] else source[start:stop] + sliced = replace(media, video=media.video.with_tensor(tensor), prepared_for_transport=False) + sliced.validate() + return sliced diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index c9ef325c4..a285e0009 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -26,6 +26,14 @@ from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, Pipe from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import DenoiseProgressMixin from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin @@ -907,6 +915,7 @@ class Wan22Pipeline( if DEBUG_PERF: _t_decode_start = time.perf_counter() + media = None if output_type == "latent": output = latents else: @@ -920,7 +929,27 @@ class Wan22Pipeline( latents.device, latents.dtype ) latents = latents / latents_std + latents_mean - output = self.vae.decode(latents, return_dict=False)[0] + decoded = self.vae.decode(latents, return_dict=False)[0] + # Distributed VAE decode uses broadcast_result=False, so only the + # output-owning rank receives the full [B, C, T, H, W] video; other + # ranks get an empty placeholder. Emit typed media only from the + # owning rank and keep the placeholder on the legacy output field, so + # the media batch-dimension check in split_diffusion_output_by_request + # does not trip on every non-owner rank. + if decoded.dim() == 5: + output = None + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=decoded, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + else: + output = decoded if DEBUG_PERF: current_omni_platform.synchronize() @@ -947,6 +976,7 @@ class Wan22Pipeline( return split_diffusion_output_by_request( DiffusionOutput( output=output, + media=media, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ), req, diff --git a/vllm_omni/diffusion/postprocess/device_reduction.py b/vllm_omni/diffusion/postprocess/device_reduction.py new file mode 100644 index 000000000..ec4ef895f --- /dev/null +++ b/vllm_omni/diffusion/postprocess/device_reduction.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Device-side reduction of decoded video tensors to uint8 frames.""" + +from __future__ import annotations + +from dataclasses import replace + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, + ensure_request_owned_tensor, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + +# VaeImageProcessor.denormalize: (x * 0.5 + 0.5).clamp(0, 1) +_DENORM_SCALE = 0.5 +_DENORM_SHIFT = 0.5 + +logger = init_logger(__name__) + + +def _request_float_consumers( + sampling_params: OmniDiffusionSamplingParams | None, +) -> frozenset[FloatVideoConsumer]: + if sampling_params is not None and sampling_params.enable_frame_interpolation: + return frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + return frozenset() + + +def _request_supports_uint8_frames(sampling_params: OmniDiffusionSamplingParams | None) -> bool: + return sampling_params is None or (sampling_params.output_type or "np") == "np" + + +def _prepare_float_media_for_transport( + media: DiffusionMediaOutput, + video: VideoMediaOutput, +) -> DiffusionMediaOutput: + prepared = replace( + media, + video=video.with_tensor(ensure_request_owned_tensor(video.tensor)), + prepared_for_transport=True, + ) + prepared.validate() + return prepared + + +def prepare_diffusion_media_for_transport( + media: DiffusionMediaOutput, + *, + od_config: OmniDiffusionConfig, + sampling_params: OmniDiffusionSamplingParams | None = None, +) -> DiffusionMediaOutput: + """Validate and prepare one request-local media output before D2H.""" + media.validate() + if media.prepared_for_transport: + # Only the model runner prepares media, exactly once, on media the + # pipeline emits unprepared. Accepting pre-prepared media here would let + # a pipeline forge runner-owned state and skip request policy (e.g. a + # forged uint8 payload silently bypasses frame interpolation). + raise ValueError( + "Media reached transport preparation already prepared; pipelines must " + "emit unprepared media so request policy is applied by the runner" + ) + if media.video.spec.encoding is not VideoTensorEncoding.NORMALIZED_FLOAT: + raise ValueError("Pipelines must emit unprepared media as NORMALIZED_FLOAT video") + + video = media.video + constraints = VideoTransportConstraints( + pending_float_consumers=(video.constraints.pending_float_consumers | _request_float_consumers(sampling_params)) + ) + constrained_video = replace(video, constraints=constraints) + + enabled = od_config.video_output_transport.enable_device_postprocess + output_supported = _request_supports_uint8_frames(sampling_params) + if not enabled or constraints.pending_float_consumers or not output_supported: + if not enabled: + reason = "disabled" + elif constraints.pending_float_consumers: + reason = ",".join(sorted(consumer.value for consumer in constraints.pending_float_consumers)) + else: + reason = "unsupported_presentation" + logger.debug("Device video preparation kept normalized float: reason=%s", reason) + return _prepare_float_media_for_transport(media, constrained_video) + + try: + frames = reduce_video_to_uint8_frames( + constrained_video.tensor, + do_denormalize=constrained_video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + except torch.OutOfMemoryError: + logger.warning("Device video preparation ran out of memory; using normalized float transport") + else: + prepared_video = replace( + constrained_video, + tensor=frames, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + constraints=VideoTransportConstraints(), + ) + prepared = replace(media, video=prepared_video, prepared_for_transport=True) + prepared.validate() + logger.debug( + "Device video preparation converted normalized float to uint8: shape=%s", + tuple(frames.shape), + ) + return prepared + + # Leave the exception handler before allocator cleanup and fallback + # construction. A late conversion OOM otherwise keeps the helper traceback + # (and its live float32 intermediates) reachable while the fallback may need + # to clone a non-compact request view. + current_omni_platform.empty_cache() + current_omni_platform.synchronize() + return _prepare_float_media_for_transport(media, constrained_video) + + +def reduce_video_to_uint8_frames(video: torch.Tensor, *, do_denormalize: bool = True) -> torch.Tensor: + """Reduce a decoded ``[B, C, F, H, W]`` video to uint8 ``[B, F, H, W, C]`` frames. + + Runs denormalize/clamp/permute/round on the input's device so the following + D2H copy carries uint8 instead of float. The result matches + ``VideoProcessor.postprocess_video(output_type="np")`` then the ``*255`` + rounding done in the API server. Pass ``do_denormalize=False`` for VAEs that + already emit ``[0, 1]``. + """ + if video.dim() != 5: + raise ValueError(f"expected a [B, C, F, H, W] video tensor, got shape {tuple(video.shape)}") + + # Match the numpy path, which promotes to float before scaling. + frames = video.to(torch.float32) + if do_denormalize: + frames = frames.mul(_DENORM_SCALE).add(_DENORM_SHIFT).clamp_(0.0, 1.0) + else: + frames = frames.clamp(0.0, 1.0) + frames = frames.permute(0, 2, 3, 4, 1) + frames = frames.mul_(255.0).round_().clamp_(0.0, 255.0).to(torch.uint8) + return frames.contiguous() diff --git a/vllm_omni/diffusion/postprocess/media.py b/vllm_omni/diffusion/postprocess/media.py new file mode 100644 index 000000000..fd2564434 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/media.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import replace + +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoTensorEncoding, + VideoTensorLayout, + VideoTransportConstraints, + VideoValueRange, +) +from vllm_omni.diffusion.postprocess.rife_interpolator import interpolate_video_tensor +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def finalize_diffusion_media( + media: DiffusionMediaOutput, + *, + sampling_params: OmniDiffusionSamplingParams, +) -> dict[str, object]: + """Finalize typed video media without invoking a model postprocessor.""" + media.validate() + if not media.prepared_for_transport: + raise ValueError("Diffusion media reached the engine before transport preparation") + + video = media.video + metadata: dict[str, object] = {} + consumers = video.constraints.pending_float_consumers + if FloatVideoConsumer.FRAME_INTERPOLATION in consumers: + if video.spec.encoding is not VideoTensorEncoding.NORMALIZED_FLOAT: + raise ValueError("Frame interpolation requires NORMALIZED_FLOAT video") + # RIFE guesses [0,1] vs [-1,1] from the tensor's min/max, so relying on + # the sample content misclassifies both an all-nonnegative [-1,1] clip + # (read as [0,1]) and a [0,1] clip whose single bf16 value overshoots 1 + # (read as [-1,1]). Always map to unit range per the declared spec and + # clamp to [0,1] before RIFE, then restore the declared range so + # downstream denormalization is correct regardless of the sample. + to_negative_one = video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE + if to_negative_one: + interp_input = video.tensor.mul(0.5).add(0.5).clamp_(0.0, 1.0) + else: + interp_input = video.tensor.clamp(0.0, 1.0) + interpolated, multiplier = interpolate_video_tensor( + interp_input, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + if to_negative_one: + interpolated = interpolated.mul(2.0).sub(1.0) + consumers = consumers - {FloatVideoConsumer.FRAME_INTERPOLATION} + video = replace( + video, + tensor=interpolated, + constraints=VideoTransportConstraints(pending_float_consumers=frozenset(consumers)), + ) + metadata["video"] = {"video_fps_multiplier": multiplier} + + if consumers: + names = sorted(consumer.value for consumer in consumers) + raise ValueError(f"Video has pending float consumers at finalization: {names}") + + output_type = getattr(sampling_params, "output_type", None) or "np" + if output_type == "latent": + raise ValueError("Latent output cannot be finalized from VideoMediaOutput") + + if video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES: + if video.spec.layout is not VideoTensorLayout.BTHWC: + raise ValueError("UINT8_FRAMES video must use BTHWC layout") + if output_type != "np": + raise ValueError(f"UINT8_FRAMES finalization does not support output_type={output_type!r}") + frames = video.tensor.detach().cpu().numpy() + else: + if video.spec.layout is not VideoTensorLayout.BCTHW: + raise ValueError("NORMALIZED_FLOAT video must use BCTHW layout") + from diffusers.video_processor import VideoProcessor + + do_denormalize = video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE + frames = VideoProcessor(vae_scale_factor=8).postprocess_video( + video.tensor, + output_type=output_type, + do_denormalize=[do_denormalize] * video.tensor.shape[2], + ) + + return { + "payload": {"video": frames}, + "metadata": metadata, + } diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 8804785fa..9dfcd6f3f 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -53,6 +53,7 @@ from vllm_omni.diffusion.models.interface import ( ) from vllm_omni.diffusion.offloader import enable_offload_backend from vllm_omni.diffusion.offloader.config import TEXT_ENCODER_COMPONENT, resolve_offload +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( @@ -680,6 +681,21 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): od_config.cache_backend, ) + def _prepare_output_for_transport( + self, + output: DiffusionOutput, + sampling_params: OmniDiffusionSamplingParams, + ) -> DiffusionOutput: + if output.media is not None: + if output.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + output.media = prepare_diffusion_media_for_transport( + output.media, + od_config=self.od_config, + sampling_params=sampling_params, + ) + return output + def _runner_output_from_outputs( self, reqs: list[OmniDiffusionRequest], @@ -774,6 +790,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): allow_single_output=allow_single_output, pipeline_name=type(self.pipeline).__name__, ) + with record_function("prepare_output_for_transport"): + outputs = [ + self._prepare_output_for_transport(output, req.sampling_params) + for req, output in zip(reqs, outputs, strict=True) + ] if is_primary and outputs and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() @@ -1289,6 +1310,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): clear_pipeline_stage_durations(self.pipeline) result = self.pipeline.post_decode(req) if result is not None: + result = self._prepare_output_for_transport(result, req.sampling) self._attach_stepwise_metadata( req, result, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 4443353a0..b9506f102 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -64,7 +64,11 @@ from vllm_omni.diffusion.distributed.parallel_state import ( model_parallel_is_initialized, ) from vllm_omni.diffusion.forward_context import set_forward_context -from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm +from vllm_omni.diffusion.ipc import ( + DIFFUSION_RPC_RESULT_ENVELOPE, + pack_diffusion_output_shm, + payload_carries_typed_media, +) from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -1131,7 +1135,15 @@ class WorkerProc: # Sync path (original, or async fallback). try: pack_diffusion_output_shm(output) + except (TypeError, ValueError): + raise except Exception as e: + # Typed media that fails to pack is left unprepared/on device (the + # pack is failure-atomic and does not mutate it), so enqueueing it + # would ship a broken payload. Re-raise memory/packing failures here + # instead of swallowing them for typed media. + if payload_carries_typed_media(output): + raise if hasattr(output, "output"): logger.warning("SHM pack failed for model output: %s", e) self._enqueue_result(output) diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index 079a213db..d5f1d5585 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/vllm_omni/diffusion/worker/request_batch.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Request-level batch abstraction for diffusion runner.""" from __future__ import annotations @@ -12,6 +12,7 @@ import numpy as np import torch from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.media import DiffusionMediaOutput, slice_diffusion_media_output from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType @@ -35,6 +36,18 @@ def split_diffusion_output_by_request( """Split a batched DiffusionOutput into one output per request.""" if num_outputs_per_prompt <= 0: raise ValueError(f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}.") + if result.media is not None and result.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + if result.media is not None and not isinstance(result.media, DiffusionMediaOutput): + raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(result.media).__name__}") + if result.media is not None: + expected_batch = req.num_reqs * num_outputs_per_prompt + actual_batch = result.media.video.tensor.shape[0] + if actual_batch != expected_batch: + raise ValueError( + f"Video media batch dimension must equal request_count * num_outputs_per_prompt " + f"({expected_batch}), got {actual_batch}" + ) return [ DiffusionOutput( @@ -43,6 +56,15 @@ def split_diffusion_output_by_request( idx * num_outputs_per_prompt, (idx + 1) * num_outputs_per_prompt, ), + media=( + slice_diffusion_media_output( + result.media, + idx * num_outputs_per_prompt, + (idx + 1) * num_outputs_per_prompt, + ) + if result.media is not None + else None + ), error=result.error, finished=result.finished, stage_durations=result.stage_durations, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 44bd9058c..ee5552751 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -566,6 +566,7 @@ class OrchestratorArgs: diffusion_compile_dynamic: bool | None = None cache_backend: str = "none" cache_config: str | None = None + video_output_transport: dict[str, object] | None = None enable_cache_dit_summary: bool = False step_execution: bool = False vae_use_slicing: bool = False diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 5c0204a33..348e41316 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -71,6 +71,16 @@ def _nonneg_finite_float(value: str) -> float: return parsed +def _json_object(value: str) -> dict[str, object]: + try: + parsed = json.loads(value) + except json.JSONDecodeError as exc: + raise argparse.ArgumentTypeError(f"must be valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise argparse.ArgumentTypeError("must be a JSON object") + return parsed + + def _ensure_vllm_platform(): """Ensure vLLM's current_platform is valid before arg parsing. @@ -669,6 +679,14 @@ class OmniServeCommand(CLISubcommand): 'MagCache: \'{"mag_threshold": 0.24, "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1}\'. ' "Calibration mode: add '\"mag_calibrate\": true'", ) + omni_config_group.add_argument( + "--video-output-transport", + type=_json_object, + default=None, + help=( + "JSON object configuring video output preparation, for example '{\"enable_device_postprocess\": true}'." + ), + ) omni_config_group.add_argument( "--enable-cache-dit-summary", action="store_true", From 7576c4ceb98c2219a3f54b295d27be8b5fed3e31 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 17:56:27 +0800 Subject: [PATCH 003/209] [Bugfix] Bound HWR domain initialization lock waits (#7128) Signed-off-by: Hongsheng Liu --- docs/design/feature/host_weight_runtime.md | 11 +- docs/design/module/host_weight_runtime.md | 12 +- .../test_runtime_resolution.py | 108 ++++++++++++++++++ .../host_weight_runtime/filesystem/store.py | 10 +- vllm_omni/host_weight_runtime/runtime.py | 2 +- 5 files changed, 134 insertions(+), 9 deletions(-) diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index a9bb1b0b8..ac7bc7b3f 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -313,8 +313,15 @@ One process per exact identity owns a build; other workers wait and then acquire leases for the published artifact. Publication is invisible until all payloads and metadata are validated, hashed, fsynced, and atomically renamed. -`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not -cancel synchronous validation, a producer that has already started, or atomic +`coordination_timeout_seconds` bounds domain-initialization and lookup/build +lock acquisition. Store construction and each later resolution or publication +operation have separate budgets from the same wait policy, rather than one +end-to-end startup deadline. A domain-init timeout follows the retryable domain +failure policy below. After contention ends, a fresh runtime construction can +retry initialization; the timed-out runtime retains its original failure. + +The coordination budget does not cancel filesystem I/O, synchronous validation, +a producer that has already started, or atomic publication. A hung in-process producer therefore blocks its owning process and must be handled by external process supervision. Enforceable producer cancellation requires a future process-isolated producer contract. diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md index ac92711fc..c97561667 100644 --- a/docs/design/module/host_weight_runtime.md +++ b/docs/design/module/host_weight_runtime.md @@ -136,8 +136,16 @@ so only when it is marked retryable; unsupported capabilities, identity collisions, producer failures, and publication failures remain visible as startup failures. -`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it -is not a hard wall-clock deadline for in-process validation, production, or +`coordination_timeout_seconds` bounds acquisition of the domain-initialization +lock and lookup/build locks. Store construction and each later resolution or +publication operation receive separate coordination budgets from the same +`WaitPolicy`; they do not share an end-to-end startup deadline. A domain-init +lock timeout is a retryable domain failure: preferred mode may fall back and +required mode fails. Retrying initialization requires constructing a fresh +runtime after contention ends; a waiter never cancels or unlocks the owner. + +The coordination budget is not a hard wall-clock deadline for filesystem I/O, +in-process validation, production, or atomic publication. Once this V1 implementation becomes the producer, the synchronous producer and publication run to completion. A hung producer blocks its owning process and requires external process supervision; enforceable diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py index 21bab37de..cd2de2816 100644 --- a/tests/host_weight_runtime/test_runtime_resolution.py +++ b/tests/host_weight_runtime/test_runtime_resolution.py @@ -4,6 +4,10 @@ from __future__ import annotations +import multiprocessing as mp +import threading +import time +from multiprocessing.connection import Connection from pathlib import Path from typing import cast @@ -15,6 +19,7 @@ from vllm_omni.host_weight_runtime import ( CanonicalJson, ComponentIdentity, FailureCode, + HostWeightError, HostWeightFailure, HostWeightLease, HostWeightLeaseCarrier, @@ -40,12 +45,14 @@ from vllm_omni.host_weight_runtime import ( StoreStatus, TensorWriteSpec, ValidationLevel, + WaitPolicy, WeightArtifactIdentity, WeightProductionSpec, WeightRepresentation, WeightSourceIdentity, ) from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -518,6 +525,107 @@ def test_retryable_store_initialization_failure_obeys_runtime_mode( assert resolution.report.attempts[0].failure.retryable +def _initialize_contended_runtime(config: HostWeightRuntimeConfig, output: Connection) -> None: + try: + output.send("starting") + started = time.monotonic() + runtime = HostWeightRuntime.from_config(config) + resolution = runtime.resolve(_identity()) + try: + output.send((time.monotonic() - started, resolution.report)) + finally: + if resolution.lease is not None: + resolution.lease.close() + finally: + output.close() + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_domain_lock_timeout_and_fresh_runtime_recovery( + tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome +) -> None: + config = HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path), wait=WaitPolicy(0.1)) + initial = HostWeightRuntime.from_config(config) + built = initial.resolve(_identity(), producer=CountingProducer(_identity())) + assert built.lease is not None + built.lease.close() + metadata = {name: (tmp_path / name).read_bytes() for name in ("domain.json", "domain-policy.json")} + lock_path = tmp_path / "locks" / "domain-init.lock" + ctx = mp.get_context("spawn") + reader, writer = ctx.Pipe(duplex=False) + process = ctx.Process(target=_initialize_contended_runtime, args=(config, writer)) + + with FileLock(lock_path, exclusive=True, deadline=None): + process.start() + writer.close() + try: + # Import/bootstrap is outside the short coordination budget. + assert reader.poll(60), "waiter did not start" + assert reader.recv() == "starting" + assert reader.poll(5), "domain initialization did not respect the coordination timeout" + elapsed, report = reader.recv() + assert 0.1 <= elapsed < 5 + assert report.outcome is expected + failure = report.attempts[0].failure + assert failure is not None + assert failure.stage is ResolutionStage.DOMAIN + assert failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failure.retryable + assert "timed out waiting for" in failure.message + assert "domain-init.lock" in failure.message + assert lock_is_active(lock_path), "waiter must not unlock the owner" + process.join(5) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.kill() + process.join(5) + process.close() + reader.close() + + assert {name: (tmp_path / name).read_bytes() for name in metadata} == metadata + recovered = HostWeightRuntime.from_config(config).resolve(_identity()) + assert recovered.report.outcome is ResolutionOutcome.LOCAL_HIT + assert recovered.lease is not None + recovered.lease.close() + + +def test_direct_store_waits_for_domain_lock_released_before_deadline(tmp_path: Path) -> None: + config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path)) + assert config.domain is not None + initial = HostWeightRuntime.from_config(config) + assert initial.store is not None + lock_path = tmp_path / "locks" / "domain-init.lock" + with FileLock(lock_path, exclusive=True, deadline=None) as owner: + release = threading.Timer(0.1, owner.close) + release.start() + try: + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(5)) + finally: + release.join(5) + assert store.root == tmp_path + + +def test_direct_store_domain_timeout_is_typed(tmp_path: Path) -> None: + config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path)) + assert config.domain is not None + HostWeightRuntime.from_config(config) + lock_path = tmp_path / "locks" / "domain-init.lock" + with FileLock(lock_path, exclusive=True, deadline=None): + with pytest.raises(HostWeightError) as caught: + FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(0.01)) + assert caught.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert caught.value.failure.retryable + assert lock_is_active(lock_path) + assert not lock_is_active(lock_path) + + @pytest.mark.parametrize( ("stage", "code"), [ diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index 2b58b77e8..ec6b49f3c 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -23,7 +23,7 @@ import regex as re import torch from safetensors import SafetensorError, safe_open -from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel +from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel, WaitPolicy from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json from ..inspection import ( @@ -429,6 +429,8 @@ class FilesystemHostWeightStore: domain: StorageDomainPolicy, capacity: CapacityPolicy, integrity: IntegrityPolicy, + *, + wait: WaitPolicy = WaitPolicy(), ) -> None: self.domain_policy = domain self.capacity_policy = capacity @@ -464,7 +466,7 @@ class FilesystemHostWeightStore: ) ) try: - self._initialize_domain() + self._initialize_domain(deadline=time.monotonic() + wait.coordination_timeout_seconds) except HostWeightError: raise except OSError as exc: @@ -485,7 +487,7 @@ class FilesystemHostWeightStore: ) ) from exc - def _initialize_domain(self) -> None: + def _initialize_domain(self, *, deadline: float) -> None: try: self.filesystem_type = detect_filesystem_type(self.root) detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type) @@ -545,7 +547,7 @@ class FilesystemHostWeightStore: ) ) - with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None): + with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=deadline): domain_path = self.root / _DOMAIN_FILE if domain_path.exists(): domain_metadata = _read_json_file(domain_path) diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py index fbc726669..17ababce9 100644 --- a/vllm_omni/host_weight_runtime/runtime.py +++ b/vllm_omni/host_weight_runtime/runtime.py @@ -70,7 +70,7 @@ class HostWeightRuntime: from .filesystem import FilesystemHostWeightStore # noqa: PLC0415 try: - store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity) + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=config.wait) except HostWeightError as exc: return cls( config, From ce27ee6038bb188ed874f024ebae431d809d8b12 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 17:56:44 +0800 Subject: [PATCH 004/209] [Bugfix] Escalate diffusion worker shutdown and retain survivors (#7126) Signed-off-by: Hongsheng Liu --- .../test_multiproc_engine_concurrency.py | 133 +++++++++++++++++- .../diffusion/executor/multiproc_executor.py | 71 +++++++--- 2 files changed, 185 insertions(+), 19 deletions(-) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index efe6bf9a3..ade8d8fa7 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -4,8 +4,10 @@ import asyncio import multiprocessing as mp import queue +import signal import threading import time +import weakref from types import SimpleNamespace from unittest.mock import MagicMock, Mock @@ -1330,6 +1332,7 @@ class TestExecutorShutdownCleaner: class FakeProcess: def __init__(self, name): self.name = name + self.pid = 123 self.alive = True self.terminated = False self.join_timeouts = [] @@ -1358,6 +1361,132 @@ class TestExecutorShutdownCleaner: assert first.terminated and second.terminated assert not first.is_alive() and not second.is_alive() + @pytest.mark.parametrize("cooperative", [False, True]) + def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + ctx = mp.get_context("fork") + ready, stop = ctx.Event(), ctx.Event() + + def run_worker(): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + ready.set() + stop.wait() + + dead = ctx.Process(target=lambda: None) + worker = ctx.Process(target=run_worker) + monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2) + monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1) + monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False) + dead.start() + worker.start() + try: + dead.join(5) + assert not dead.is_alive() + assert ready.wait(5), "worker did not install its signal handler" + mq = Mock() + if cooperative: + mq.enqueue.side_effect = lambda *args, **kwargs: stop.set() + else: + mq.enqueue.side_effect = OSError("shutdown queue unavailable") + cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker]) + + cleaner() + + assert not worker.is_alive() + assert worker.exitcode == (0 if cooperative else -signal.SIGKILL) + assert cleaner.processes == [] + assert cleaner.broadcast_mq is None + cleaner() + finally: + for proc in (dead, worker): + if proc.is_alive(): + proc.kill() + proc.join(5) + proc.close() + + @pytest.mark.parametrize("failed_action", ["terminate", "kill"]) + def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + first, second = Mock(pid=101), Mock(pid=102) + first.name, second.name = "first", "second" + for proc in (first, second): + proc.is_alive.side_effect = [True, True, True, False] + getattr(first, failed_action).side_effect = OSError("signal failed") + first.join.side_effect = [OSError("join failed"), None, None] + monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134]) + monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic)) + cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second]) + + cleaner() + + first.kill.assert_called_once() + second.kill.assert_called_once() + assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5] + assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1] + assert cleaner.processes == [] + + def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + survivor = Mock(pid=123) + survivor.name = "surviving-worker" + survivor.is_alive.return_value = True + cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor]) + executor, _, _ = _make_executor() + executor._shutdown_cleaner = cleaner + executor._processes = [survivor] + executor._finalizer = weakref.finalize(executor, cleaner) + executor._pump_stop = threading.Event() + executor._futures_lock = threading.RLock() + executor._rpc_futures = {} + executor._output_futures = {} + executor._batch_split_map = {} + log_error = Mock() + monkeypatch.setattr(executor_module.logger, "error", log_error) + + executor.shutdown() + + assert not executor._finalizer.alive + assert executor._shutdown_cleaner is cleaner + assert executor._processes == [survivor] + log_error.assert_called_once() + assert log_error.call_args.args[1] == [("surviving-worker", 123)] + survivor.is_alive.return_value = False + + executor.shutdown() + executor.shutdown() + + assert executor._shutdown_cleaner is None + assert executor._processes == [] + + def test_concurrent_cleanup_does_not_repeat_process_operations(self): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + joining, release = threading.Event(), threading.Event() + proc = Mock(pid=123) + proc.is_alive.return_value = True + + def join(timeout): + joining.set() + assert release.wait(5) + proc.is_alive.return_value = False + + proc.join.side_effect = join + cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc]) + thread = threading.Thread(target=cleaner) + thread.start() + try: + assert joining.wait(5) + cleaner() + proc.join.assert_called_once() + finally: + release.set() + thread.join(5) + assert not thread.is_alive() + assert cleaner.processes == [] + # ───────── monitor thread & death sentinel integration tests ───────── @@ -1406,7 +1535,7 @@ class TestMultiprocExecutorWorkerMonitor: executor._result_mq = None executor._shutdown_cleaner = None # Use a no-op so shutdown() doesn't crash on None resources. - executor._finalizer = lambda: None + executor._finalizer = weakref.finalize(executor, lambda: None) # ------------------------------------------------------------------ # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates # over them, so they need to exist even when constructed via __new__. @@ -1441,7 +1570,7 @@ class TestMultiprocExecutorWorkerMonitor: executor._broadcast_mq = None executor._result_mq = None executor._shutdown_cleaner = None - executor._finalizer = lambda: None + executor._finalizer = weakref.finalize(executor, lambda: None) proc = _make_short_lived_process() executor._processes = [proc] diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 39e85abcc..57b23656c 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import concurrent.futures @@ -10,7 +13,7 @@ import threading import time import weakref from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from multiprocessing.synchronize import Event from typing import TYPE_CHECKING, Any, cast @@ -42,6 +45,7 @@ _DEQUEUE_TIMEOUT_S = 5.0 _DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0)) _WORKER_SHUTDOWN_GRACE_S = 15.0 _WORKER_TERMINATE_GRACE_S = 5.0 +_WORKER_KILL_GRACE_S = 5.0 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0 @@ -80,33 +84,61 @@ class _ExecutorShutdownCleaner: broadcast_mq: MessageQueue | None = None num_workers: int = 0 processes: list[mp.Process] | None = None + _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) def __call__(self) -> None: """Clean up background resources.""" + # The worker monitor and an explicit shutdown may race. A retry must + # not signal/join the same Process objects while cleanup is in flight. + if not self._lock.acquire(blocking=False): + return + try: + self._cleanup() + finally: + self._lock.release() + + def _cleanup(self) -> None: if self.broadcast_mq is not None: try: for _ in range(self.num_workers): self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0) - self.broadcast_mq = None except Exception as exc: logger.warning("Failed to send shutdown signal: %s", exc) + finally: + self.broadcast_mq = None if self.processes: - join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S - for proc in self.processes: - if not proc.is_alive(): - continue - proc.join(max(0.0, join_deadline - time.monotonic())) - alive = [proc for proc in self.processes if proc.is_alive()] - for proc in alive: - logger.warning("Terminating diffusion worker %s after timeout", proc.name) - proc.terminate() + for action, grace in ( + (None, _WORKER_SHUTDOWN_GRACE_S), + ("terminate", _WORKER_TERMINATE_GRACE_S), + ("kill", _WORKER_KILL_GRACE_S), + ): + if not alive: + break + if action is not None: + for proc in alive: + try: + logger.warning("Calling %s on diffusion worker %s (pid=%s)", action, proc.name, proc.pid) + getattr(proc, action)() + except OSError: + logger.exception("Failed to %s diffusion worker %s (pid=%s)", action, proc.name, proc.pid) - terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S - for proc in alive: - proc.join(max(0.0, terminate_deadline - time.monotonic())) + deadline = time.monotonic() + grace + for proc in alive: + try: + proc.join(max(0.0, deadline - time.monotonic())) + except OSError: + logger.exception("Failed to join diffusion worker %s (pid=%s)", proc.name, proc.pid) + alive = [proc for proc in alive if proc.is_alive()] + + self.processes = alive + if alive: + logger.error( + "Diffusion worker cleanup incomplete after kill: %s; retaining processes for shutdown retry", + [(proc.name, proc.pid) for proc in alive], + ) class MultiprocDiffusionExecutor(DiffusionExecutor): @@ -1006,8 +1038,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): def shutdown(self) -> None: self._closed = True self._pump_stop.set() + cleaner = self._shutdown_cleaner try: - self._finalizer() + if self._finalizer.alive: + self._finalizer() + elif cleaner is not None: + cleaner() finally: pump_threads = getattr(self, "_result_pump_threads", []) for thread in pump_threads: @@ -1031,5 +1067,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): self._rpc_futures.clear() self._output_futures.clear() self._batch_split_map.clear() - self._shutdown_cleaner = None - self._processes = [] + self._processes = (cleaner.processes or []) if cleaner is not None else [] + if not self._processes: + self._shutdown_cleaner = None From f11bf1a05d8dff5d9b03d0012fb9d3dca3bc6b05 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 17:57:20 +0800 Subject: [PATCH 005/209] [Misc] Add standalone safetensors retention diagnostic (#7145) Signed-off-by: Hongsheng Liu --- benchmarks/host_weight_runtime/README.md | 57 +++++++++ .../safetensors_retention.py | 120 ++++++++++++++++++ docs/design/feature/host_weight_runtime.md | 5 + 3 files changed, 182 insertions(+) create mode 100644 benchmarks/host_weight_runtime/README.md create mode 100644 benchmarks/host_weight_runtime/safetensors_retention.py diff --git a/benchmarks/host_weight_runtime/README.md b/benchmarks/host_weight_runtime/README.md new file mode 100644 index 000000000..899811151 --- /dev/null +++ b/benchmarks/host_weight_runtime/README.md @@ -0,0 +1,57 @@ +# Host weight dependency memory diagnostic + +`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from +access to one cached tensor view. It imports neither vLLM nor HWR and uses a +synthetic float32 payload controlled by `--tensor-elements` (default: 64). +Tensor size can affect retention: a negative result for 16 elements does not +rule out growth for 64 elements. The selected size is recorded in the JSON +arguments and used for both file creation and the correctness check. +HWR itself calls `get_tensor()` when +acquiring a lease, then exposes cached views; this probe does not measure HWR +requests or repeated lease acquisition. + +Run one mode per fresh process in the same Python environment and CPU affinity: + +```bash +git rev-parse HEAD > revision.txt +git status --short > working-tree.txt +git diff > working-tree.patch +# Select an allowed CPU from your process affinity; 0 is only an example. +CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json +``` + +For a size comparison, run 16 and 64 elements twice each in fresh processes, +with cached reuse as a control at both sizes: + +```bash +for repetition in 1 2; do + for elements in 16 64; do + for mode in get_tensor reuse; do + CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode "$mode" --tensor-elements "$elements" \ + > "$mode-$elements-$repetition.json" + done + done +done +``` + +The diagnostic sets Torch to one CPU thread, warms up before the baseline, +and reports preparation separately from timed loop work. Checkpoints run GC +and report Linux private/anonymous memory, RSS, and open descriptors. It checks +that transient tensor objects and the final cached view are released. Linux +`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc +support fail explicitly. Keep raw JSON and the repository snapshot with results. + +Compare the final loop sample with the warmed baseline, then inspect the closed +sample. Private-memory growth after Python tensor objects disappear is evidence +of retention in the dependency/allocator path, not proof of a leak's root cause. +RSS can include file-backed pages; it is not a private-memory metric. Samples +also include small diagnostic bookkeeping allocations. Do not assert a portable +memory threshold or turn this synthetic loop into a per-request service estimate. + +Do not drop shared caches, trim allocators, or add serving-time GC to make the +numbers look smaller. A dependency replacement or version change requires its +own reproducible comparison and correctness/lifecycle validation. diff --git a/benchmarks/host_weight_runtime/safetensors_retention.py b/benchmarks/host_weight_runtime/safetensors_retention.py new file mode 100644 index 000000000..55795e7ef --- /dev/null +++ b/benchmarks/host_weight_runtime/safetensors_retention.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import platform +import tempfile +import time +import weakref +from pathlib import Path + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be positive") + return parsed + + +def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]: + gc.collect() + counters = {} + for line in Path("/proc/self/smaps_rollup").read_text().splitlines(): + fields = line.split() + if len(fields) == 3 and fields[2] == "kB": + counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024 + return { + "phase": phase, + "iterations": iterations, + "loop_seconds": loop_seconds, + "rss_bytes": counters["Rss"], + "private_bytes": counters["Private_Clean"] + counters["Private_Dirty"], + "anonymous_bytes": counters["Anonymous"], + "fd_count": len(tuple(Path("/proc/self/fd").iterdir())), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True) + parser.add_argument( + "--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor" + ) + parser.add_argument("--iterations", type=positive_int, default=1_000_000) + parser.add_argument("--sample-every", type=positive_int, default=100_000) + parser.add_argument("--warmup", type=positive_int, default=1000) + args = parser.parse_args() + if not Path("/proc/self/smaps_rollup").is_file(): + parser.error("requires Linux /proc/self/smaps_rollup") + + preparation_start = time.monotonic() + import safetensors + import torch + from safetensors import safe_open + from safetensors.torch import save_file + + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + samples = [sample("before_file", 0, 0.0)] + with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory: + path = Path(directory) / "synthetic.safetensors" + save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path)) + with safe_open(path, framework="pt", device="cpu") as reader: + cached = reader.get_tensor("weight") + assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32)) + for _ in range(args.warmup): + tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached + del tensor + samples.append(sample("baseline", 0, 0.0)) + preparation_seconds = time.monotonic() - preparation_start + completed = 0 + loop_seconds = 0.0 + while completed < args.iterations: + count = min(args.sample_every, args.iterations - completed) + start = time.monotonic() + for _ in range(count): + tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached + last_tensor = weakref.ref(tensor) + del tensor + loop_seconds += time.monotonic() - start + completed += count + snapshot = sample("loop", completed, loop_seconds) + snapshot["last_tensor_alive"] = last_tensor() is not None + if args.mode == "get_tensor": + assert last_tensor() is None, "diagnostic retained a transient tensor" + samples.append(snapshot) + del cached + del reader + samples.append(sample("closed", completed, loop_seconds)) + assert last_tensor() is None, "last tensor survived reader/view teardown" + + print( + json.dumps( + { + "schema_version": 1, + "arguments": vars(args), + "python": platform.python_version(), + "torch": torch.__version__, + "safetensors": safetensors.__version__, + "kernel": platform.release(), + "hostname": platform.node(), + "pid": os.getpid(), + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "torch_threads": torch.get_num_threads(), + "pythonmalloc": os.environ.get("PYTHONMALLOC", "default"), + "preparation_seconds": preparation_seconds, + "samples": samples, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index ac7bc7b3f..15fce07b8 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -42,6 +42,11 @@ V1 does not include: ## Motivation and use cases +For investigating CPU memory retained by dependency tensor materialization, +see the [standalone safetensors diagnostic](https://github.com/vllm-project/vllm-omni/blob/main/benchmarks/host_weight_runtime/README.md). +It distinguishes repeated dependency calls from reuse of cached views and does +not establish per-request HWR leakage. + Model loading can create the same final host representation repeatedly. This is especially expensive when loading performs checkpoint decoding, tensor renaming, TP slicing, quantization, packing, or scale construction before GPU From d8d4062547899d71b9073649f404e61658136a68 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Thu, 10 Sep 2026 02:58:24 -0700 Subject: [PATCH 006/209] [CI] Isolate layerwise offload memory measurements (#6938) Signed-off-by: andyluo7 --- .../test_diffusion_layerwise_offload.py | 269 +++++++++++++++--- 1 file changed, 231 insertions(+), 38 deletions(-) diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index 8e083605f..c386871d3 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import concurrent.futures import gc +import multiprocessing as mp +import os from typing import Any import numpy as np @@ -45,32 +48,98 @@ IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = { "sampler_params": {"height": 480, "width": 640, "num_frames": 5}, } +_OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_test" +# A second CI build can reverse the same exact-head measurement without a +# source edit; the default keeps the existing baseline-first workload. +_MEASUREMENT_ORDER_ENV = "VLLM_OMNI_OFFLOAD_TEST_ORDER" +_MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = { + "baseline-first": (False, True), + "offload-first": (True, False), +} -def check_audio_determinism(audio1, audio2, atol=1e-2): - device = current_omni_platform.device_type - if isinstance(audio1, np.ndarray): - audio1 = torch.from_numpy(audio1).to(device) - if isinstance(audio2, np.ndarray): - audio2 = torch.from_numpy(audio2).to(device) - if not torch.allclose(audio1, audio2, atol=atol): - diff = torch.abs(audio1 - audio2) - print(f"Max difference: {diff.max().item()}") - print(f"Mean difference: {diff.mean().item()}") +class OffloadStateProbe: + """Test-only worker extension exposing the configured offloader state.""" + + model_runner: Any + + def get_offload_state_for_test(self) -> dict[str, Any]: + backend = getattr(self.model_runner, "offload_backend", None) + block_groups = getattr(backend, "_blocks", ()) + group_sizes = [len(group) for group in block_groups] if isinstance(block_groups, (list, tuple)) else [] + + return { + "probe": _OFFLOAD_STATE_PROBE, + "backend_type": ( + None if backend is None else f"{backend.__class__.__module__}.{backend.__class__.__qualname__}" + ), + "requested": bool(getattr(self.model_runner.od_config, "enable_layerwise_offload", False)), + "enabled": bool(backend is not None and backend.is_enabled()), + "block_group_count": len(group_sizes), + "block_group_sizes": group_sizes, + "block_count": sum(group_sizes), + } + + +def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool: + if not np.allclose(audio1, audio2, atol=atol): + diff = np.abs(audio1 - audio2) + print(f"Max difference: {diff.max()}") + print(f"Mean difference: {diff.mean()}") raise AssertionError(f"Audio outputs differ beyond tolerance atol={atol}") return True +def _device_used_mb(device_index: int) -> float: + current_omni_platform.synchronize() + with current_omni_platform.device(device_index): + free_bytes, total_bytes = current_omni_platform.mem_get_info() + return (total_bytes - free_bytes) / (1024**2) + + +def _extract_audio(output: Any) -> np.ndarray | None: + if not output: + return None + multimodal_output = getattr(output[0], "multimodal_output", None) + if not isinstance(multimodal_output, dict): + return None + audio = multimodal_output.get("audio") + if audio is None: + return None + if isinstance(audio, torch.Tensor): + return audio.detach().cpu().numpy() + return np.asarray(audio) + + +def _collect_offload_states(value: Any) -> list[dict[str, Any]]: + if isinstance(value, dict): + return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else [] + if isinstance(value, (list, tuple)): + states: list[dict[str, Any]] = [] + for item in value: + states.extend(_collect_offload_states(item)) + return states + return [] + + +def _measurement_order() -> tuple[bool, bool]: + """Return the requested order without changing the default CI workload.""" + configured = os.environ.get(_MEASUREMENT_ORDER_ENV, "baseline-first") + try: + return _MEASUREMENT_ORDERS[configured] + except KeyError as exc: + allowed = ", ".join(_MEASUREMENT_ORDERS) + raise ValueError(f"{_MEASUREMENT_ORDER_ENV} must be one of: {allowed}; got {configured!r}") from exc + + def run_inference( model_name: str, layerwise_offload: bool = False, num_inference_steps: int = 3, -) -> tuple[float, Any]: +) -> dict[str, Any]: current_omni_platform.empty_cache() device_index = current_omni_platform.current_device() - with current_omni_platform.device(device_index): - free_bytes, total_bytes = current_omni_platform.mem_get_info() - initial_used_mb = (total_bytes - free_bytes) / (1024**2) + initial_used_mb = _device_used_mb(device_index) if model_name in AUDIO_MODEL: params = AUDIO_MODEL_PARAMS @@ -80,10 +149,17 @@ def run_inference( with OmniRunner( model_name, enable_layerwise_offload=layerwise_offload, + worker_extension_cls=f"{OffloadStateProbe.__module__}.{OffloadStateProbe.__qualname__}", # TODO: we might want to add overlapped feature e2e tests # cache_backend="cache_dit", **params["runner_params"], ) as runner: + offload_states = _collect_offload_states( + runner.omni.engine.collective_rpc(method="get_offload_state_for_test", timeout=60) + ) + if not offload_states: + raise AssertionError("The offload-state worker probe returned no diffusion worker results") + # Measure steady-state inference memory, not model construction. Enabling # layerwise offload first loads the model and then replaces each block's # device storage with CPU-backed weights. Monitoring that transition @@ -108,20 +184,136 @@ def run_inference( finally: monitor.stop() + audio = _extract_audio(output) + del output + # DeviceMemoryMonitor reports absolute device usage. Subtract this run's - # starting usage so process-wide compiler/workspace caches retained from a - # previous run do not make the second measurement order-dependent. - peak = max(0.0, monitor.peak_used_mb - initial_used_mb) + # starting usage. Each mode runs in its own spawned process below, so model + # objects and process-local allocator/compiler/backend workspace state cannot + # carry into the other measurement. Shared filesystem caches can persist; + # _MEASUREMENT_ORDER_ENV enables a reversed-order validation run for that. + peak_used_mb = monitor.peak_used_mb + incremental_peak_mb = max(0.0, peak_used_mb - initial_used_mb) + del runner gc.collect() + cleanup_dist_env_and_memory() current_omni_platform.empty_cache() + post_cleanup_used_mb = _device_used_mb(device_index) - return peak, output + return { + "initial_used_mb": initial_used_mb, + "peak_used_mb": peak_used_mb, + "incremental_peak_mb": incremental_peak_mb, + "post_cleanup_used_mb": post_cleanup_used_mb, + "audio": audio, + "offload_states": offload_states, + } + + +def run_inference_isolated( + model_name: str, + layerwise_offload: bool = False, + num_inference_steps: int = 3, +) -> dict[str, Any]: + """Run one measurement in a fresh interpreter and accelerator context.""" + with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn")) as executor: + return executor.submit(run_inference, model_name, layerwise_offload, num_inference_steps).result() + + +def _assert_offload_state(measurement: dict[str, Any], *, expected_enabled: bool) -> None: + states = measurement["offload_states"] + assert states, "Expected at least one offload-state result" + for state in states: + assert state["requested"] is expected_enabled, f"Unexpected offload request state: {state}" + assert state["enabled"] is expected_enabled, f"Unexpected offload state: {state}" + if expected_enabled: + assert state["backend_type"], f"Enabled offloader has no backend type: {state}" + assert state["block_group_count"] > 0, f"Enabled offloader has no block groups: {state}" + assert state["block_count"] > 0, f"Enabled offloader has no managed blocks: {state}" + else: + assert state["block_group_count"] == 0, f"Disabled offloader retained block groups: {state}" + assert state["block_count"] == 0, f"Disabled offloader retained managed blocks: {state}" + + +def _assert_audio_outputs( + model_name: str, + audio_no_offload: np.ndarray | None, + audio_offload: np.ndarray | None, +) -> None: + if model_name not in AUDIO_MODEL: + return + + assert audio_no_offload is not None, "Baseline Stable Audio inference returned no audio" + assert audio_offload is not None, "Layerwise-offloaded Stable Audio inference returned no audio" + # Match the sibling cpu-offload test's tolerance: layerwise offload moves + # blocks across the PCIe bus on a side stream, which can perturb cuBLAS + # algorithm selection and produce ~ULP-level drift larger than 1e-3. + check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2) + + +def _print_measurement(label: str, measurement: dict[str, Any]) -> None: + print( + f"{label}: initial={measurement['initial_used_mb']:.1f} MB, " + f"peak={measurement['peak_used_mb']:.1f} MB, " + f"incremental_peak={measurement['incremental_peak_mb']:.1f} MB, " + f"post_cleanup={measurement['post_cleanup_used_mb']:.1f} MB" + ) + print(f"{label} offload state: {measurement['offload_states']}") + + +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("configured", "expected"), + [ + (None, (False, True)), + ("baseline-first", (False, True)), + ("offload-first", (True, False)), + ], +) +def test_measurement_order(monkeypatch, configured: str | None, expected: tuple[bool, bool]) -> None: + if configured is None: + monkeypatch.delenv(_MEASUREMENT_ORDER_ENV, raising=False) + else: + monkeypatch.setenv(_MEASUREMENT_ORDER_ENV, configured) + + assert _measurement_order() == expected + + +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +def test_measurement_order_rejects_unknown_value(monkeypatch) -> None: + monkeypatch.setenv(_MEASUREMENT_ORDER_ENV, "unknown") + + with pytest.raises(ValueError, match=_MEASUREMENT_ORDER_ENV): + _measurement_order() + + +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("audio_no_offload", "audio_offload", "message"), + [ + (None, None, "Baseline Stable Audio inference returned no audio"), + (None, np.zeros(1), "Baseline Stable Audio inference returned no audio"), + (np.zeros(1), None, "Layerwise-offloaded Stable Audio inference returned no audio"), + ], +) +def test_stable_audio_requires_both_outputs(audio_no_offload, audio_offload, message) -> None: + with pytest.raises(AssertionError, match=message): + _assert_audio_outputs("stabilityai/stable-audio-open-1.0", audio_no_offload, audio_offload) @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) -@pytest.mark.parametrize("model_name", list(MODELS.keys())) +@pytest.mark.parametrize( + "model_name", + [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS], +) def test_layerwise_offload_diffusion_model(model_name: str): """Test that layerwise offloading reduces GPU memory usage. @@ -130,14 +322,15 @@ 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. """ + measurements: dict[bool, dict[str, Any]] = {} try: - # Run without layerwise offloading (baseline) - no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False) - cleanup_dist_env_and_memory() - - # Run with layerwise offloading (1 layer on device) - layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True) - cleanup_dist_env_and_memory() + measurement_order = _measurement_order() + print( + "Measurement order: " + + " -> ".join("layerwise-offload" if enabled else "baseline" for enabled in measurement_order) + ) + for enabled in measurement_order: + measurements[enabled] = run_inference_isolated(model_name, layerwise_offload=enabled) except ValueError as exc: # 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): @@ -145,20 +338,18 @@ def test_layerwise_offload_diffusion_model(model_name: str): f"Skipping: gated HF repo {model_name!r} inaccessible " f"({exc}). See docs/contributing/ci/hf_credentials.md." ) - pytest.fail(f"Inference failed: {exc}") - except Exception: - pytest.fail("Inference failed") + raise - print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB") - print(f"No offload peak memory: {no_offload_peak_memory} MB") + no_offload = measurements[False] + layerwise_offload = measurements[True] + _assert_offload_state(no_offload, expected_enabled=False) + _assert_offload_state(layerwise_offload, expected_enabled=True) + _print_measurement("No offload", no_offload) + _print_measurement("Layerwise offload", layerwise_offload) - if model_name == "stabilityai/stable-audio-open-1.0": - audio_offload = output_offload[0].multimodal_output.get("audio") - audio_no_offload = output_no_offload[0].multimodal_output.get("audio") - # Match the sibling cpu-offload test's tolerance: layerwise offload moves - # blocks across the PCIe bus on a side stream, which can perturb cuBLAS - # algorithm selection and produce ~ULP-level drift larger than 1e-3. - check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2) + audio_no_offload = no_offload["audio"] + audio_offload = layerwise_offload["audio"] + _assert_audio_outputs(model_name, audio_no_offload, audio_offload) is_rocm = torch.version.hip is not None platform = "rocm" if is_rocm else "cuda" @@ -170,6 +361,8 @@ def test_layerwise_offload_diffusion_model(model_name: str): # Verify that layerwise offloading significantly reduces memory usage # Passes only if the actual savings meets the expected savings + no_offload_peak_memory = no_offload["incremental_peak_mb"] + layerwise_offload_peak_memory = layerwise_offload["incremental_peak_mb"] actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, ( f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB " From 43a476c1e5817cedfdf706a6d62c867a43cd5065 Mon Sep 17 00:00:00 2001 From: wkutak Date: Thu, 10 Sep 2026 13:51:09 +0200 Subject: [PATCH 007/209] [Model] Add Cosmos3 mixed W8A8/W8A16 and W4A4/W4A16 denoising (#6560) Signed-off-by: Rahul Steiger Signed-off-by: Wojciech Kutak Co-authored-by: Rahul Steiger --- docs/user_guide/quantization/modelopt.md | 85 +++ recipes/cosmos3/Cosmos3-Nano.md | 32 + .../test_cosmos3_mixed_precision.py | 690 ++++++++++++++++++ .../cosmos3/mixed_precision/__init__.py | 14 + .../cosmos3/mixed_precision/checkpoint.py | 176 +++++ .../models/cosmos3/mixed_precision/config.py | 106 +++ .../models/cosmos3/mixed_precision/runtime.py | 116 +++ .../cosmos3/mixed_precision/strategy.py | 218 ++++++ .../models/cosmos3/pipeline_cosmos3.py | 24 +- .../models/cosmos3/transformer_cosmos3.py | 52 ++ 10 files changed, 1509 insertions(+), 4 deletions(-) create mode 100644 tests/diffusion/quantization/test_cosmos3_mixed_precision.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md index d8da4e1a9..07d05730e 100644 --- a/docs/user_guide/quantization/modelopt.md +++ b/docs/user_guide/quantization/modelopt.md @@ -204,6 +204,91 @@ omni = Omni( | `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages | | `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints | +## Cosmos3 Mixed-Precision Schedule + +The experimental Cosmos3 schedule can use the checkpoint-native quantized GEMM for middle denoising +steps and dense 16-bit activation execution for selected first and last steps. +The schedule supports serialized ModelOpt FP8 and NVFP4 checkpoints: + +| Detected weight format | Native middle steps | A16 boundary steps | +|------------------------|---------------------|--------------------| +| FP8 | W8A8 | W8A16 | +| NVFP4 | W4A4 | W4A16 | + +The native and A16 paths share one quantized weight representation. The A16 +path dequantizes the live backend weight and calls `F.linear`; it does not keep +a second checkpoint-weight snapshot. Scheduled FP8 currently requires +serialized tensorwise scales and a backend that retains canonical FP8 weights. +Scheduled NVFP4 requires a supported CUTLASS-compatible native or FlashInfer +layout. SmoothQuant, Marlin-repacked FP8, and per-channel/per-token FP8 +schedules fail closed. + +The reasoner uses dense A16 execution by default when its weights are FP8 or +NVFP4. Set the nested `reasoner` field to `native` to retain the checkpoint-native +W8A8/W4A4 path. A BF16 reasoner remains BF16 in either mode. The schedule +currently requires tensor parallel size 1 and one active request per worker, +does not support HSDP, and does not support block-scaled FP8. Its live-weight +design is compatible with model-level and standard layer-wise offload, but +those paths still require GPU end-to-end validation. Distributed layer-wise +offload is rejected because its direct loader bypasses ModelOpt post-load +transformations. + +A checkpoint can opt in by adding a versioned policy to its authoritative +`transformer/config.json`. The same policy schema is used for FP8 and NVFP4: + +```json +{ + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": { + "diffusion_step_policy": { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": ["transformer"], + "default_mode": "native", + "first_steps": {"count": 3, "mode": "a16"}, + "last_steps": {"count": 3, "mode": "a16"}, + "overlap": "a16", + "reasoner": "a16" + } + } + } +} +``` + +The example preserves the checkpoint's existing `quant_method` and +`quant_algo`; only `runtime` is added. + +When present and valid, the checkpoint policy is the default. An explicit +runtime object overrides it. To retain ordinary checkpoint-native execution +for every step, disable only the schedule: + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config '{"cosmos3_mixed_precision":{"enabled":false}}' +``` + +This does not disable checkpoint quantization; it prevents installation of the +mixed-step runtime. + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config \ + '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}' +``` + +The presence of `cosmos3_mixed_precision` supplies an explicit runtime +override; an empty object uses the defaults shown above. +The schedule uses vLLM's selected native backend and fails during loading if +its live weight layout cannot support A16 dequantization. +The runtime selects FP8 or NVFP4 from the checkpoint's ModelOpt configuration. +BF16 linears are untouched. Mixed FP8/NVFP4 checkpoints are not supported by +this schedule. + ## Validation and Notes 1. Compare the ModelOpt checkpoint against the BF16 baseline with the same diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 490f5b1f4..f81e5e887 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -115,6 +115,38 @@ shift at the same seed). The pipeline auto-resolves from `model_index.json`; pass `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly. +For a serialized ModelOpt FP8 or NVFP4 checkpoint, an experimental +mixed-precision schedule can use native W8A8/W4A4 in middle denoising steps +and dense W8A16/W4A16 in the first and last steps: + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config \ + '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}' +``` + +The nested object's presence supplies an explicit runtime override; use an +empty object for these defaults. A compatible checkpoint can instead declare +the versioned `runtime.diffusion_step_policy` in its authoritative +`transformer/config.json`. The checkpoint policy is used only when no runtime +override is present. Use +`--additional-config '{"cosmos3_mixed_precision":{"enabled":false}}'` to +disable the checkpoint schedule without disabling checkpoint quantization. +This path currently requires tensor parallel size 1 and does not support HSDP +or block-scaled FP8. It keeps one live quantized weight representation: +scheduled FP8 requires serialized tensorwise scales and a canonical backend, +while scheduled NVFP4 requires a supported CUTLASS-compatible native or +FlashInfer layout. Its dequantization is a correctness baseline; no speedup is +implied. Quantized reasoner weights use A16 by default; set +`"reasoner":"native"` in the nested object to keep W8A8/W4A4. +The checkpoint's ModelOpt configuration selects FP8 or NVFP4; mixed FP8/NVFP4 +checkpoints are not supported by this schedule. BF16 linears are left unchanged. +The schedule currently supports one active request per worker. Model-level and +layer-wise offload are designed to work with this live-weight path but still +require GPU validation. Distributed layer-wise offload and incompatible native +weight layouts fail during loading. + #### Verification Best quality uses the JSON-upsampled prompts from `assets/` (download with diff --git a/tests/diffusion/quantization/test_cosmos3_mixed_precision.py b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py new file mode 100644 index 000000000..026090fe9 --- /dev/null +++ b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.cosmos3.mixed_precision import ( + Cosmos3MixedPrecisionConfig, + Cosmos3MixedPrecisionRuntime, + resolve_mixed_precision_config, +) +from vllm_omni.diffusion.models.cosmos3.mixed_precision import runtime as runtime_impl +from vllm_omni.diffusion.models.cosmos3.mixed_precision.checkpoint import read_checkpoint_policy +from vllm_omni.diffusion.models.cosmos3.mixed_precision.runtime import ( + Cosmos3MixedPrecisionLinearMethod, +) +from vllm_omni.diffusion.models.cosmos3.mixed_precision.strategy import ( + Fp8W8A8W8A16Strategy, + Nvfp4W4A4W4A16Strategy, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _policy( + *, + first_steps: int = 3, + last_steps: int = 3, + reasoner: str = "a16", + scope: list[str] | None = None, +) -> dict[str, object]: + return { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": scope or ["transformer"], + "default_mode": "native", + "first_steps": {"count": first_steps, "mode": "a16"}, + "last_steps": {"count": last_steps, "mode": "a16"}, + "overlap": "a16", + "reasoner": reasoner, + } + + +def _fake_quant_config(name: str, *, serialized: bool = True): + if name == "modelopt": + return SimpleNamespace( + get_name=lambda: name, + is_checkpoint_fp8_serialized=serialized, + ) + if name == "modelopt_fp4": + return SimpleNamespace( + get_name=lambda: name, + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=serialized, + ) + return SimpleNamespace(get_name=lambda: name) + + +def _checkpoint_od_config( + quant_config, + *, + policy: object | None = None, + additional_config: dict[str, object] | None = None, +): + quant_config_name = quant_config.get_name() + quant_algo = { + "modelopt": "FP8", + "modelopt_fp4": "NVFP4", + }.get(quant_config_name, "FP8") + disk_quant_config: dict[str, object] = { + "quant_method": quant_config_name, + "quant_algo": quant_algo, + } + if policy is not None: + disk_quant_config["runtime"] = {"diffusion_step_policy": policy} + return SimpleNamespace( + quantization_config=quant_config, + tf_model_config=SimpleNamespace( + params={"quantization_config": disk_quant_config}, + quant_config=quant_config, + ), + additional_config=additional_config, + ) + + +def _swizzle_blockscale_cpu(scale: torch.Tensor) -> torch.Tensor: + rows, cols = scale.shape + padded_rows = ((rows + 127) // 128) * 128 + padded_cols = ((cols + 3) // 4) * 4 + padded = torch.zeros( + (1, padded_rows, padded_cols), + dtype=scale.dtype, + device=scale.device, + ) + padded[0, :rows, :cols] = scale + padded = padded.reshape(1, padded_rows // 128, 4, 32, padded_cols // 4, 4) + return ( + padded.permute(0, 1, 4, 3, 2, 5) + .contiguous() + .reshape( + padded_rows, + padded_cols, + ) + ) + + +def test_config_parses_nested_schedule_and_reasoner_policy() -> None: + config = Cosmos3MixedPrecisionConfig.from_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 2, + "last_steps": 4, + "reasoner": "native", + } + } + ) + + assert config is not None + assert config.reasoner == "native" + assert [index for index in range(10) if config.use_high_precision(index, 10)] == [0, 1, 6, 7, 8, 9] + + +@pytest.mark.parametrize( + ("first_steps", "last_steps", "selected"), + [ + (0, 2, [5, 6]), + (2, 0, [0, 1]), + (0, 0, []), + (4, 4, list(range(7))), + ], +) +def test_schedule_boundaries_and_overlap( + first_steps: int, + last_steps: int, + selected: list[int], +) -> None: + config = Cosmos3MixedPrecisionConfig( + first_steps=first_steps, + last_steps=last_steps, + ) + assert [index for index in range(7) if config.use_high_precision(index, 7)] == selected + + +def test_one_step_request_honors_boundary_precision() -> None: + config = Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1) + assert config.use_high_precision(0, 1) + + +@pytest.mark.parametrize( + ("values", "message"), + [ + ({"cosmos3_mixed_precision": True}, "must be a mapping"), + ({"cosmos3_mixed_precision": {"first_steps": -1}}, "non-negative"), + ({"cosmos3_mixed_precision": {"reasoner": "fp16"}}, "must be one of"), + ({"cosmos3_mixed_precision": {"enabled": 0}}, "must be a boolean"), + ( + {"cosmos3_mixed_precision": {"enabled": False, "first_steps": 1}}, + "cannot be combined", + ), + ({"cosmos3_mixed_precision": {"unknown": 1}}, "Unknown"), + ], +) +def test_config_rejects_invalid_values(values: dict, message: str) -> None: + with pytest.raises((TypeError, ValueError), match=message): + Cosmos3MixedPrecisionConfig.from_additional_config(values) + + +def test_config_presence_enables_defaults() -> None: + assert Cosmos3MixedPrecisionConfig.from_additional_config({}) is None + config = Cosmos3MixedPrecisionConfig.from_additional_config({"cosmos3_mixed_precision": {}}) + assert config == Cosmos3MixedPrecisionConfig() + + +def test_noop_config_does_not_install_runtime() -> None: + assert ( + Cosmos3MixedPrecisionConfig.from_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 0, + "last_steps": 0, + "reasoner": "native", + } + } + ) + is None + ) + + +def test_additional_config_distinguishes_absent_and_disabled() -> None: + assert Cosmos3MixedPrecisionConfig.resolve_additional_config({}) == (False, None) + assert Cosmos3MixedPrecisionConfig.resolve_additional_config({"cosmos3_mixed_precision": {"enabled": False}}) == ( + True, + None, + ) + assert Cosmos3MixedPrecisionConfig.resolve_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 0, + "last_steps": 0, + "reasoner": "native", + } + } + ) == (True, None) + + +def test_checkpoint_policy_round_trip_through_transformer_config() -> None: + from vllm_omni.diffusion.data import TransformerConfig + + disk = { + "_class_name": "Cosmos3VFMTransformer", + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": _policy(first_steps=2, last_steps=4)}, + }, + } + + tf_config = TransformerConfig.from_dict(disk) + assert tf_config.quant_config is not None + assert tf_config.quant_config.get_name() == "modelopt" + assert tf_config.to_dict()["quantization_config"]["runtime"] == disk["quantization_config"]["runtime"] + assert read_checkpoint_policy(SimpleNamespace(tf_model_config=tf_config)) == Cosmos3MixedPrecisionConfig( + first_steps=2, + last_steps=4, + reasoner="a16", + ) + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_checkpoint_policy_supports_serialized_modelopt_formats(name: str) -> None: + od_config = _checkpoint_od_config(_fake_quant_config(name), policy=_policy()) + assert read_checkpoint_policy(od_config) == Cosmos3MixedPrecisionConfig() + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_checkpoint_policy_rejects_nonserialized_modelopt_formats(name: str) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name, serialized=False), + policy=_policy(), + ) + with pytest.raises(ValueError, match="serialized ModelOpt"): + read_checkpoint_policy(od_config) + + +def test_checkpoint_policy_rejects_mixed_modelopt_format() -> None: + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt_mixed"), policy=_policy())) + + +def test_checkpoint_policy_missing_metadata_preserves_ordinary_path() -> None: + assert read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"))) is None + + +def test_checkpoint_policy_rejects_malformed_runtime_container() -> None: + od_config = _checkpoint_od_config(_fake_quant_config("modelopt")) + od_config.tf_model_config.params["quantization_config"]["runtime"] = True + with pytest.raises(TypeError, match="runtime must be a mapping"): + read_checkpoint_policy(od_config) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("schema_version", 2, "schema_version"), + ("type", "sigma_ranges", "type"), + ("index_space", "scheduler_timestep", "index_space"), + ("default_mode", "a16", "default_mode"), + ("overlap", "native", "overlap"), + ("reasoner", "fp16", "reasoner"), + ], +) +def test_checkpoint_policy_rejects_invalid_schema(field: str, value: object, message: str) -> None: + policy = _policy() + policy[field] = value + with pytest.raises(ValueError, match=message): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy)) + + +@pytest.mark.parametrize( + ("step_range", "message"), + [ + ({"count": -1, "mode": "a16"}, "non-negative"), + ({"count": True, "mode": "a16"}, "non-negative"), + ({"count": 1, "mode": "native"}, "mode"), + ({"mode": "a16"}, "Missing"), + ({"count": 1, "mode": "a16", "sigma": 1.0}, "Unknown"), + ], +) +def test_checkpoint_policy_rejects_invalid_step_range(step_range: object, message: str) -> None: + policy = _policy() + policy["first_steps"] = step_range + with pytest.raises((TypeError, ValueError), match=message): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy)) + + +def test_checkpoint_policy_rejects_unsupported_quantization_config() -> None: + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("compressed-tensors"), policy=_policy())) + + +def test_checkpoint_policy_for_other_scope_is_not_applied() -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(scope=["vae"]), + ) + assert read_checkpoint_policy(od_config) is None + + +def test_additional_config_overrides_and_disables_checkpoint_policy() -> None: + checkpoint = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(first_steps=1, last_steps=1), + ) + assert resolve_mixed_precision_config(checkpoint) == ( + Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1), + "checkpoint", + ) + + checkpoint.additional_config = { + "cosmos3_mixed_precision": { + "first_steps": 2, + "last_steps": 0, + "reasoner": "native", + } + } + assert resolve_mixed_precision_config(checkpoint) == ( + Cosmos3MixedPrecisionConfig(first_steps=2, last_steps=0, reasoner="native"), + "additional_config", + ) + + checkpoint.additional_config = {"cosmos3_mixed_precision": {"enabled": False}} + assert resolve_mixed_precision_config(checkpoint) == (None, "additional_config_disabled") + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_additional_config_supports_serialized_modelopt_formats_without_checkpoint_policy(name: str) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name), + additional_config={"cosmos3_mixed_precision": {}}, + ) + assert resolve_mixed_precision_config(od_config) == (Cosmos3MixedPrecisionConfig(), "additional_config") + + +@pytest.mark.parametrize("has_checkpoint_policy", [False, True]) +@pytest.mark.parametrize( + ("name", "serialized"), + [ + ("modelopt_mixed", True), + ("mxfp8", True), + ("compressed-tensors", True), + ("modelopt", False), + ("modelopt_fp4", False), + ], +) +def test_additional_config_rejects_incompatible_quantization( + name: str, serialized: bool, has_checkpoint_policy: bool +) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name, serialized=serialized), + policy=_policy() if has_checkpoint_policy else None, + additional_config={"cosmos3_mixed_precision": {}}, + ) + with pytest.raises(ValueError, match="serialized ModelOpt"): + resolve_mixed_precision_config(od_config) + + +@pytest.mark.parametrize("additional_config", [None, {"cosmos3_mixed_precision": {}}]) +def test_policy_validates_effective_quantization_config(additional_config: dict[str, object] | None) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(), + additional_config=additional_config, + ) + od_config.quantization_config = _fake_quant_config("modelopt_mixed") + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + resolve_mixed_precision_config(od_config) + + +def test_additional_config_rejects_missing_quantization_config() -> None: + od_config = SimpleNamespace(additional_config={"cosmos3_mixed_precision": {}}) + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + resolve_mixed_precision_config(od_config) + + +@pytest.mark.parametrize("policy", [_policy(), True]) +def test_additional_config_disable_bypasses_incompatible_checkpoint(policy: object) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt_mixed"), + policy=policy, + additional_config={"cosmos3_mixed_precision": {"enabled": False}}, + ) + od_config.quantization_config = _fake_quant_config("modelopt_mixed") + assert resolve_mixed_precision_config(od_config) == (None, "additional_config_disabled") + + +def test_runtime_allows_standard_offload_and_rejects_distributed(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr( + transformer_cosmos3, + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=True, + enable_distributed_layerwise_offload=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + od_config.enable_distributed_layerwise_offload = True + with pytest.raises(ValueError, match="distributed layer-wise offload"): + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + od_config.enable_distributed_layerwise_offload = False + od_config.max_num_seqs = 2 + with pytest.raises(ValueError, match="one active request"): + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + + +class _Layer(torch.nn.Module): + def __init__(self, output_size: int, input_size: int) -> None: + super().__init__() + self.output_size_per_partition = output_size + self.input_size_per_partition = input_size + + +class _BaseMethod: + def __init__(self, *, nvfp4: bool = False) -> None: + self.processed = False + self.apply_calls = 0 + self.fp8_linear = object() + if nvfp4: + from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, + ) + + self.kernel = object.__new__(CutlassNvFp4LinearKernel) + else: + self.kernel = None + + def create_weights(self, *args, **kwargs) -> None: + pass + + def process_weights_after_loading(self, layer) -> None: + self.processed = True + if layer.weight.dtype == torch.float8_e4m3fn: + layer.weight = layer.weight.t().contiguous() + layer.weight_scale = layer.weight_scale.max().reshape(1) + elif layer.weight.dtype == torch.uint8: + layer.weight_scale = _swizzle_blockscale_cpu(layer.weight_scale) + layer.weight_global_scale = layer.weight_scale_2.max().float().reshape(1) + del layer.weight_scale_2 + + def apply(self, layer, x, bias=None): + del bias + self.apply_calls += 1 + return torch.full( + (*x.shape[:-1], layer.output_size_per_partition), + 17, + dtype=x.dtype, + device=x.device, + ) + + +def _fp8_layer(output_size: int = 2, input_size: int = 4) -> _Layer: + layer = _Layer(output_size, input_size) + layer.weight = torch.ones(output_size, input_size, dtype=torch.float8_e4m3fn) + layer.weight_scale = torch.tensor([0.5], dtype=torch.float32) + return layer + + +def _nvfp4_layer(output_size: int = 128, input_size: int = 64) -> _Layer: + layer = _Layer(output_size, input_size) + layer.weight = torch.full((output_size, input_size // 2), 0x22, dtype=torch.uint8) + layer.weight_scale = torch.ones(output_size, input_size // 16, dtype=torch.float8_e4m3fn) + layer.weight_scale_2 = torch.ones(1) + return layer + + +def _runtime_and_method( + strategy, + layer: _Layer, + *, + path: str = "generation", + reasoner: str = "native", +): + config = Cosmos3MixedPrecisionConfig( + first_steps=1, + last_steps=1, + reasoner=reasoner, + ) + runtime = Cosmos3MixedPrecisionRuntime(config) + base = _BaseMethod(nvfp4=isinstance(strategy, Nvfp4W4A4W4A16Strategy)) + method = Cosmos3MixedPrecisionLinearMethod( + base, + strategy, + runtime, + f"{path}.linear", + path, + ) + method.process_weights_after_loading(layer) + return runtime, base, method + + +@pytest.mark.parametrize( + ("strategy", "layer"), + [ + (Fp8W8A8W8A16Strategy(), _fp8_layer()), + (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()), + ], +) +def test_generation_dispatches_native_middle_and_a16_edges(strategy, layer) -> None: + runtime, base, method = _runtime_and_method(strategy, layer) + x = torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16) + + runtime.set_step(1, 3) + assert torch.equal( + method.apply(layer, x), + torch.full((1, layer.output_size_per_partition), 17, dtype=x.dtype), + ) + runtime.set_step(0, 3) + assert torch.equal( + method.apply(layer, x), + torch.nn.functional.linear(x, strategy.materialize(layer)), + ) + assert base.apply_calls == 1 + + +@pytest.mark.parametrize( + ("strategy", "layer"), + [ + (Fp8W8A8W8A16Strategy(), _fp8_layer()), + (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()), + ], +) +def test_a16_uses_only_live_native_weights(strategy, layer) -> None: + runtime, base, method = _runtime_and_method(strategy, layer) + assert base.processed + assert not hasattr(layer, "_cosmos3_precision_weight") + assert not hasattr(layer, "_cosmos3_precision_weight_scale") + + runtime.set_step(0, 3) + output = method.apply( + layer, + torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16), + ) + assert output.shape == (1, layer.output_size_per_partition) + + +def test_a16_follows_live_weight_storage_after_offload_rebinding() -> None: + layer = _fp8_layer() + strategy = Fp8W8A8W8A16Strategy() + runtime, _, method = _runtime_and_method(strategy, layer) + runtime.set_step(0, 3) + x = torch.ones(1, 4, dtype=torch.bfloat16) + before = method.apply(layer, x) + + layer.weight = torch.zeros_like(layer.weight) + after = method.apply(layer, x) + assert not torch.equal(before, after) + assert torch.count_nonzero(after) == 0 + + +def test_fp8_rejects_smoothquant() -> None: + layer = _fp8_layer() + layer.pre_quant_scale = torch.ones(1) + with pytest.raises(ValueError, match="SmoothQuant"): + Fp8W8A8W8A16Strategy().validate_before_processing( + _BaseMethod(), + layer, + "gen.linear", + ) + + +def test_fp8_rejects_non_tensorwise_scales() -> None: + layer = _fp8_layer(output_size=3, input_size=2) + layer.weight_scale = torch.tensor([[1.0], [2.0], [3.0]]) + with pytest.raises(ValueError, match="tensorwise FP8"): + Fp8W8A8W8A16Strategy().validate_before_processing( + _BaseMethod(), + layer, + "gen.linear", + ) + + +def test_nvfp4_reference_materializes_live_cutlass_weights() -> None: + layer = _nvfp4_layer() + strategy = Nvfp4W4A4W4A16Strategy() + _runtime_and_method(strategy, layer) + output = strategy.materialize(layer) + assert output.shape == (128, 64) + assert torch.equal(output, torch.ones_like(output)) + + +def test_nvfp4_rejects_fused_global_scales() -> None: + layer = _nvfp4_layer() + layer.weight_scale_2 = torch.tensor([1.0, 2.0]) + strategy = Nvfp4W4A4W4A16Strategy() + with pytest.raises(ValueError, match="NVFP4 scales"): + strategy.validate_before_processing( + _BaseMethod(nvfp4=True), + layer, + "gen.qkv", + ) + + +@pytest.mark.parametrize( + ("policy", "uses_native"), + [("a16", False), ("native", True)], +) +def test_reasoner_policy_is_independent_of_generation_step( + policy: str, + uses_native: bool, +) -> None: + layer = _fp8_layer() + runtime, base, method = _runtime_and_method( + Fp8W8A8W8A16Strategy(), + layer, + path="reasoner", + reasoner=policy, + ) + runtime.set_step(1, 3) + method.apply(layer, torch.ones(1, 4, dtype=torch.bfloat16)) + assert (base.apply_calls == 1) is uses_native + + +def test_reset_clears_generation_state() -> None: + runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1)) + runtime.set_step(0, 5) + assert runtime.use_high_precision("generation") + runtime.reset() + assert not runtime.use_high_precision("generation") + + +def test_install_discovers_generation_and_opt_in_reasoner(monkeypatch) -> None: + class _FakeLinear(torch.nn.Module): + def __init__(self, prefix: str) -> None: + super().__init__() + self.prefix = prefix + self.input_size_per_partition = 4 + self.output_size_per_partition = 2 + self.quant_method = SimpleNamespace(name="fp8") + + monkeypatch.setattr(runtime_impl, "LinearBase", _FakeLinear) + strategy = Fp8W8A8W8A16Strategy() + monkeypatch.setattr(strategy, "accepts", lambda method: getattr(method, "name", None) == "fp8") + monkeypatch.setattr(runtime_impl, "_STRATEGIES", (strategy,)) + runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(reasoner="a16")) + reasoner = _FakeLinear("reasoner.q_proj") + generation = [ + _FakeLinear("generation.q_proj"), + _FakeLinear("generation.out_proj"), + ] + transformer = SimpleNamespace( + language_model=SimpleNamespace(layers=torch.nn.Sequential(reasoner)), + gen_layers=torch.nn.Sequential(*generation), + ) + runtime.install(transformer) + assert isinstance(reasoner.quant_method, Cosmos3MixedPrecisionLinearMethod) + assert all(isinstance(layer.quant_method, Cosmos3MixedPrecisionLinearMethod) for layer in generation) + + +def test_pipeline_helpers_forward_and_reset_schedule() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + calls = [] + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + pipeline.transformer = SimpleNamespace( + set_mixed_precision_step=lambda step, count: calls.append((step, count)), + reset_mixed_precision=lambda: calls.append("reset"), + ) + pipeline._set_mixed_precision_step(2, 7) + pipeline._reset_mixed_precision() + assert calls == [(2, 7), "reset"] diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py new file mode 100644 index 000000000..1d5ee972b --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Scheduled A16 execution for quantized Cosmos3 checkpoints.""" + +from .checkpoint import resolve_mixed_precision_config +from .config import Cosmos3MixedPrecisionConfig +from .runtime import Cosmos3MixedPrecisionRuntime + +__all__ = [ + "Cosmos3MixedPrecisionConfig", + "Cosmos3MixedPrecisionRuntime", + "resolve_mixed_precision_config", +] diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py new file mode 100644 index 000000000..c8a62f142 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Checkpoint discovery for the Cosmos3 diffusion-step precision policy.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Literal, cast + +from .config import Cosmos3MixedPrecisionConfig, ReasonerPolicy, _non_negative_int + +PolicySource = Literal[ + "none", + "checkpoint", + "additional_config", + "additional_config_disabled", +] + +_COMPONENT = "transformer" +_POLICY_FIELDS = frozenset( + { + "schema_version", + "type", + "index_space", + "scope", + "default_mode", + "first_steps", + "last_steps", + "overlap", + "reasoner", + } +) +_STEP_RANGE_FIELDS = frozenset({"count", "mode"}) +_SUPPORTED_MODEL_OPT_CONFIGS = frozenset( + { + "modelopt", + "modelopt_fp4", + } +) + + +def resolve_mixed_precision_config( + od_config: object, +) -> tuple[Cosmos3MixedPrecisionConfig | None, PolicySource]: + """Resolve an explicit runtime override before the checkpoint default.""" + override_present, override = Cosmos3MixedPrecisionConfig.resolve_additional_config( + getattr(od_config, "additional_config", None) + ) + if override_present: + if override is not None: + _validate_checkpoint_quantization(od_config) + source: PolicySource = "additional_config" if override is not None else "additional_config_disabled" + return override, source + + checkpoint = read_checkpoint_policy(od_config) + if checkpoint is not None: + return checkpoint, "checkpoint" + return None, "none" + + +def read_checkpoint_policy(od_config: object) -> Cosmos3MixedPrecisionConfig | None: + """Read and validate a policy from ``transformer/config.json``. + + Missing metadata preserves ordinary checkpoint behavior. Metadata that is + present but malformed or incompatible fails closed. + """ + tf_config = getattr(od_config, "tf_model_config", None) + params = getattr(tf_config, "params", None) + if not isinstance(params, Mapping): + return None + + disk_quant_config = params.get("quantization_config") + if not isinstance(disk_quant_config, Mapping): + return None + + if "runtime" not in disk_quant_config: + return None + runtime_config = disk_quant_config["runtime"] + if not isinstance(runtime_config, Mapping): + raise TypeError("quantization_config.runtime must be a mapping") + + if "diffusion_step_policy" not in runtime_config: + return None + raw_policy = runtime_config["diffusion_step_policy"] + if not isinstance(raw_policy, Mapping): + raise TypeError("quantization_config.runtime.diffusion_step_policy must be a mapping") + + policy = _parse_policy(raw_policy) + if policy is None: + return None + _validate_checkpoint_quantization(od_config) + return policy + + +def _parse_policy(policy: Mapping[str, object]) -> Cosmos3MixedPrecisionConfig | None: + fields = set(policy) + unknown = fields - _POLICY_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy fields: {sorted(unknown)}") + missing = _POLICY_FIELDS - fields + if missing: + raise ValueError(f"Missing diffusion_step_policy fields: {sorted(missing)}") + + schema_version = policy["schema_version"] + if not isinstance(schema_version, int) or isinstance(schema_version, bool) or schema_version != 1: + raise ValueError("diffusion_step_policy.schema_version must be the integer 1") + if policy["type"] != "first_last_n": + raise ValueError("diffusion_step_policy.type must be 'first_last_n'") + if policy["index_space"] != "denoising_loop_iteration": + raise ValueError("diffusion_step_policy.index_space must be 'denoising_loop_iteration'") + if policy["default_mode"] != "native": + raise ValueError("diffusion_step_policy.default_mode must be 'native'") + if policy["overlap"] != "a16": + raise ValueError("diffusion_step_policy.overlap must be 'a16'") + + scope = policy["scope"] + if not isinstance(scope, list) or not scope or not all(isinstance(item, str) for item in scope): + raise TypeError("diffusion_step_policy.scope must be a non-empty list of strings") + + first_steps = _parse_step_range(policy["first_steps"], "first_steps") + last_steps = _parse_step_range(policy["last_steps"], "last_steps") + + reasoner = policy["reasoner"] + if reasoner not in {"native", "a16"}: + raise ValueError("diffusion_step_policy.reasoner must be 'native' or 'a16'") + + if _COMPONENT not in scope: + return None + return Cosmos3MixedPrecisionConfig( + first_steps=first_steps, + last_steps=last_steps, + reasoner=cast(ReasonerPolicy, reasoner), + ) + + +def _parse_step_range(value: object, name: str) -> int: + if not isinstance(value, Mapping): + raise TypeError(f"diffusion_step_policy.{name} must be a mapping") + fields = set(value) + unknown = fields - _STEP_RANGE_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy.{name} fields: {sorted(unknown)}") + missing = _STEP_RANGE_FIELDS - fields + if missing: + raise ValueError(f"Missing diffusion_step_policy.{name} fields: {sorted(missing)}") + if value["mode"] != "a16": + raise ValueError(f"diffusion_step_policy.{name}.mode must be 'a16'") + return _non_negative_int(value["count"], f"diffusion_step_policy.{name}.count") + + +def _validate_checkpoint_quantization(od_config: object) -> None: + # Match the effective quantization used to construct the transformer linears. + quant_config = getattr(od_config, "quantization_config", None) + if quant_config is None: + tf_config = getattr(od_config, "tf_model_config", None) + quant_config = getattr(tf_config, "quant_config", None) + + get_name = getattr(quant_config, "get_name", None) + name = get_name() if callable(get_name) else None + if name not in _SUPPORTED_MODEL_OPT_CONFIGS: + raise ValueError(f"diffusion_step_policy requires a serialized ModelOpt FP8 or NVFP4 checkpoint, got {name!r}") + + if name == "modelopt": + _require_serialized(quant_config, "is_checkpoint_fp8_serialized", "FP8") + return + if name == "modelopt_fp4": + if getattr(quant_config, "quant_method", "NVFP4") != "NVFP4": + raise ValueError("diffusion_step_policy requires checkpoint-native NVFP4 W4A4 linears") + _require_serialized(quant_config, "is_checkpoint_nvfp4_serialized", "NVFP4") + return + + +def _require_serialized(quant_config: object | None, flag: str, format_name: str) -> None: + if not bool(getattr(quant_config, flag, False)): + raise ValueError(f"diffusion_step_policy requires serialized ModelOpt {format_name} weights") diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py new file mode 100644 index 000000000..13748c8fd --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration parsing and denoising-step policy for mixed precision.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Literal, cast + +ReasonerPolicy = Literal["native", "a16"] + +_REASONER_POLICIES = frozenset({"native", "a16"}) +_CONFIG_FIELDS = frozenset({"enabled", "first_steps", "last_steps", "reasoner"}) + + +@dataclass(frozen=True) +class Cosmos3MixedPrecisionConfig: + """Validated precision policy shared by every quantization strategy.""" + + first_steps: int = 3 + last_steps: int = 3 + reasoner: ReasonerPolicy = "a16" + + @classmethod + def from_additional_config( + cls, + additional_config: Mapping[str, object] | None, + ) -> Cosmos3MixedPrecisionConfig | None: + """Parse Cosmos3 fields from vLLM-Omni's additional configuration.""" + _, config = cls.resolve_additional_config(additional_config) + return config + + @classmethod + def resolve_additional_config( + cls, + additional_config: Mapping[str, object] | None, + ) -> tuple[bool, Cosmos3MixedPrecisionConfig | None]: + """Return whether an override was supplied and its effective policy. + + The separate presence bit distinguishes an absent runtime override from + an explicit disable. That distinction matters when a checkpoint owns a + default policy. + """ + if additional_config is not None and not isinstance(additional_config, Mapping): + raise TypeError("additional_config must be a mapping") + values = additional_config or {} + if "cosmos3_mixed_precision" not in values: + return False, None + raw_config = values["cosmos3_mixed_precision"] + if not isinstance(raw_config, Mapping): + raise TypeError("cosmos3_mixed_precision must be a mapping") + unknown = set(raw_config) - _CONFIG_FIELDS + if unknown: + raise ValueError(f"Unknown cosmos3_mixed_precision fields: {sorted(unknown)}") + + enabled = raw_config.get("enabled", True) + if not isinstance(enabled, bool): + raise TypeError("cosmos3_mixed_precision.enabled must be a boolean") + if not enabled: + conflicting = set(raw_config) - {"enabled"} + if conflicting: + raise ValueError( + "cosmos3_mixed_precision.enabled=false cannot be combined with " + f"schedule fields: {sorted(conflicting)}" + ) + return True, None + + first_steps = _non_negative_int( + raw_config.get("first_steps", 3), + "cosmos3_mixed_precision.first_steps", + ) + last_steps = _non_negative_int( + raw_config.get("last_steps", 3), + "cosmos3_mixed_precision.last_steps", + ) + + reasoner = str(raw_config.get("reasoner", "a16")).lower() + if reasoner not in _REASONER_POLICIES: + raise ValueError( + f"cosmos3_mixed_precision.reasoner must be one of {sorted(_REASONER_POLICIES)}, got {reasoner!r}" + ) + config = cls( + first_steps=first_steps, + last_steps=last_steps, + reasoner=cast(ReasonerPolicy, reasoner), + ) + if first_steps == 0 and last_steps == 0 and reasoner == "native": + return True, None + return True, config + + def use_high_precision(self, step_index: int, num_steps: int) -> bool: + """Select the high path for an actual scheduler-step index.""" + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}") + if step_index < 0 or step_index >= num_steps: + raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}") + return step_index < self.first_steps or step_index >= num_steps - self.last_steps + + +def _non_negative_int(value: object, name: str) -> int: + """Validate an integer configuration field without accepting booleans.""" + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise TypeError(f"{name} must be a non-negative integer") + return value diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py new file mode 100644 index 000000000..f528262e5 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Format-agnostic schedule state, linear discovery, and dispatch.""" + +from __future__ import annotations + +from typing import Literal + +import torch +from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase + +from .config import Cosmos3MixedPrecisionConfig +from .strategy import ( + Cosmos3PrecisionStrategy, + Fp8W8A8W8A16Strategy, + Nvfp4W4A4W4A16Strategy, +) + +PrecisionPath = Literal["reasoner", "generation"] +_STRATEGIES: tuple[Cosmos3PrecisionStrategy, ...] = ( + Fp8W8A8W8A16Strategy(), + Nvfp4W4A4W4A16Strategy(), +) + + +def _strategy_for(method: object | None) -> Cosmos3PrecisionStrategy | None: + return next((strategy for strategy in _STRATEGIES if strategy.accepts(method)), None) + + +class Cosmos3MixedPrecisionLinearMethod(LinearMethodBase): + """Dispatch between a checkpoint-native method and dense A16.""" + + def __init__( + self, + base_method: LinearMethodBase, + strategy: Cosmos3PrecisionStrategy, + runtime: Cosmos3MixedPrecisionRuntime, + module_name: str, + path: PrecisionPath, + ) -> None: + self.base_method = base_method + self.strategy = strategy + self.runtime = runtime + self.module_name = module_name + self.path = path + + def create_weights(self, *args, **kwargs) -> None: + self.base_method.create_weights(*args, **kwargs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.strategy.validate_before_processing( + self.base_method, + layer, + self.module_name, + ) + self.base_method.process_weights_after_loading(layer) + self.strategy.validate_after_processing(layer, self.module_name) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if not self.runtime.use_high_precision(self.path): + return self.base_method.apply(layer, x, bias) + return self.strategy.apply_high(layer, x, bias) + + +class Cosmos3MixedPrecisionRuntime: + """Own one transformer's schedule state and wrapped linear inventory.""" + + def __init__(self, config: Cosmos3MixedPrecisionConfig) -> None: + self.config = config + # TODO: Make this request-local before Cosmos3 supports interleaved + # denoising requests on one transformer instance. + self._generation_high_precision = False + + def install(self, transformer: torch.nn.Module) -> None: + components: dict[PrecisionPath, torch.nn.Module] = { + "generation": transformer.gen_layers, + } + if self.config.reasoner == "a16": + components["reasoner"] = transformer.language_model.layers + + wrapped_count = 0 + for path, component in components.items(): + for local_name, layer in component.named_modules(): + if not isinstance(layer, LinearBase): + continue + base_method = getattr(layer, "quant_method", None) + strategy = _strategy_for(base_method) + if strategy is None: + continue + module_name = getattr(layer, "prefix", None) or f"{path}.{local_name}" + layer.quant_method = Cosmos3MixedPrecisionLinearMethod( + base_method, + strategy, + self, + module_name, + path, + ) + wrapped_count += 1 + + if not wrapped_count: + raise ValueError("Cosmos3 mixed precision found no compatible FP8 or NVFP4 ModelOpt linears") + + def use_high_precision(self, path: PrecisionPath) -> bool: + return self.config.reasoner == "a16" if path == "reasoner" else self._generation_high_precision + + def set_step(self, step_index: int, num_steps: int) -> None: + self._generation_high_precision = self.config.use_high_precision(step_index, num_steps) + + def reset(self) -> None: + self._generation_high_precision = False diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py new file mode 100644 index 000000000..8208b7622 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Format-specific live-weight validation and reference A16 materialization.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptFp8LinearMethod, + ModelOptNvFp4LinearMethod, +) +from vllm.model_executor.layers.quantization.utils import nvfp4_emulation_utils + +_NVFP4_BLOCK_SIZE = 16 + + +class Cosmos3PrecisionStrategy: + """Quantization-format boundary behind the shared precision schedule.""" + + def accepts(self, method: object | None) -> bool: + """Return whether this strategy owns a native linear method.""" + raise NotImplementedError + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + """Reject native transforms that cannot support live dequantization.""" + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + """Validate the live native weight representation.""" + raise NotImplementedError + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + """Materialize the live quantized weight as a dense BF16 matrix.""" + raise NotImplementedError + + def apply_high( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ) -> torch.Tensor: + """Apply the reference dense A16 path.""" + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + weight = self.materialize(layer) + expected_shape = (output_size, input_size) + if tuple(weight.shape) != expected_shape: + raise RuntimeError(f"Dense weight shape {tuple(weight.shape)} does not match {expected_shape}") + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.shape[1] != input_size: + raise ValueError(f"Expected activation width {input_size}, got {x_2d.shape[1]}") + weight = weight.to(dtype=x.dtype) + if bias is not None and bias.dtype != x.dtype: + bias = bias.to(dtype=x.dtype) + output = F.linear(x_2d, weight, bias) + return output.view(*x.shape[:-1], output_size) + + +class Fp8W8A8W8A16Strategy(Cosmos3PrecisionStrategy): + """Use native ModelOpt W8A8 or reference dense W8A16.""" + + def accepts(self, method: object | None) -> bool: + return isinstance(method, ModelOptFp8LinearMethod) + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + if hasattr(layer, "pre_quant_scale"): + raise ValueError(f"{module_name} uses unsupported SmoothQuant pre_quant_scale") + weight_scale = getattr(layer, "weight_scale", None) + if not isinstance(weight_scale, torch.Tensor) or weight_scale.numel() != 1: + raise ValueError(f"{module_name} requires serialized tensorwise FP8 weights") + fp8_kernel = getattr(method, "fp8_linear", None) + from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( + MarlinFP8ScaledMMLinearKernel, + ) + + if isinstance(fp8_kernel, MarlinFP8ScaledMMLinearKernel): + raise ValueError( + f"{module_name} selected Marlin FP8, which repacks the only weight copy; " + "use a backend that retains canonical FP8 weights" + ) + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + weight = getattr(layer, "weight", None) + scale = getattr(layer, "weight_scale", None) + if not isinstance(weight, torch.Tensor) or weight.dtype != torch.float8_e4m3fn: + dtype = getattr(weight, "dtype", None) + raise TypeError(f"{module_name} requires a canonical FP8 weight, got {dtype}") + if weight.ndim != 2: + raise ValueError(f"{module_name} requires a rank-2 FP8 weight, got {tuple(weight.shape)}") + if not isinstance(scale, torch.Tensor): + raise ValueError(f"{module_name} has an FP8 weight but no weight_scale") + if scale.numel() != 1: + raise ValueError(f"{module_name} requires one live FP8 weight scale") + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + if weight.shape[0] < input_size or weight.shape[1] < output_size: + raise ValueError( + f"{module_name} live FP8 weight shape {tuple(weight.shape)} does not cover " + f"(K, N)=({input_size}, {output_size})" + ) + _validate_positive_finite_scale(scale, module_name) + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + weight = layer.weight[:input_size, :output_size] + scale = layer.weight_scale.reshape(1).to(device=weight.device, dtype=torch.float32) + return (weight.to(torch.float32) * scale).t().to(torch.bfloat16) + + +class Nvfp4W4A4W4A16Strategy(Cosmos3PrecisionStrategy): + """Use native ModelOpt W4A4 or reference dense W4A16.""" + + def accepts(self, method: object | None) -> bool: + return isinstance(method, ModelOptNvFp4LinearMethod) + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + kernel = getattr(method, "kernel", None) + from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, + ) + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ) + + compatible_kernels = ( + CutlassNvFp4LinearKernel, + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ) + if not isinstance(kernel, compatible_kernels): + kernel_name = type(kernel).__name__ if kernel is not None else "none" + raise ValueError( + f"{module_name} selected unsupported NVFP4 backend {kernel_name}; " + "live W4A16 dequantization requires a CUTLASS-compatible layout" + ) + global_scale = _nvfp4_global_scale(layer) + if global_scale is None or global_scale.numel() != 1: + raise ValueError(f"{module_name} is missing ModelOpt NVFP4 scales") + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + packed = getattr(layer, "weight", None) + scale = getattr(layer, "weight_scale", None) + global_scale = getattr(layer, "weight_global_scale", None) + if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8: + dtype = getattr(packed, "dtype", None) + raise TypeError(f"{module_name} requires a live packed NVFP4 weight, got {dtype}") + if not isinstance(scale, torch.Tensor) or not isinstance(global_scale, torch.Tensor): + raise ValueError(f"{module_name} is missing live NVFP4 scales") + output_size = int(layer.output_size_per_partition) + input_size = int(layer.input_size_per_partition) + if input_size % _NVFP4_BLOCK_SIZE != 0: + raise ValueError(f"{module_name} input size {input_size} is not divisible by {_NVFP4_BLOCK_SIZE}") + expected_weight = (output_size, input_size // 2) + if packed.shape[0] < expected_weight[0] or packed.shape[1] < expected_weight[1]: + raise ValueError( + f"{module_name} live NVFP4 weight shape {tuple(packed.shape)} does not cover {expected_weight}" + ) + _validate_positive_finite_scale(global_scale, module_name) + lut_handle = nvfp4_emulation_utils.kE2M1ToFloat_handle + if lut_handle.val.device != packed.device: + lut_handle.val = lut_handle.val.to(packed.device) + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + output_size = int(layer.output_size_per_partition) + input_size = int(layer.input_size_per_partition) + packed = layer.weight[:output_size, : input_size // 2] + return nvfp4_emulation_utils.dequantize_to_dtype( + packed, + layer.weight_scale, + layer.weight_global_scale, + torch.bfloat16, + _NVFP4_BLOCK_SIZE, + True, + ) + + +def _validate_positive_finite_scale(scale: torch.Tensor, module_name: str) -> None: + values = scale.detach().float() + if not torch.isfinite(values).all() or not (values > 0).all(): + raise ValueError(f"{module_name} has a non-finite or non-positive weight scale") + + +def _nvfp4_global_scale(layer: torch.nn.Module) -> torch.Tensor | None: + for name in ("weight_scale_2", "weight_global_scale"): + value = getattr(layer, name, None) + if isinstance(value, torch.Tensor): + return value + return None diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 49a4e4027..c60817afb 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1822,6 +1822,16 @@ class Cosmos3OmniDiffusersPipeline( def num_timesteps(self): return self._num_timesteps + def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: + setter = getattr(self.transformer, "set_mixed_precision_step", None) + if setter is not None: + setter(step_index, num_steps) + + def _reset_mixed_precision(self) -> None: + resetter = getattr(self.transformer, "reset_mixed_precision", None) + if resetter is not None: + resetter() + @staticmethod def _distilled_unsupported_error(detail: str) -> ValueError: return ValueError( @@ -2819,7 +2829,8 @@ class Cosmos3OmniDiffusersPipeline( # Each CFG-parallel rank runs exactly one branch (rank 0 -> cond, # else uncond), so session keying loads/stores only this rank's branch. cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0 - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) # Outside the interval, scale=1 makes the combined output equal # the cond branch. Every rank remains on the same iteration and @@ -2858,7 +2869,8 @@ class Cosmos3OmniDiffusersPipeline( uncond_cache: tuple = (None, None) keep_uncond_for_cache = self._cache_requires_paired_cfg() - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) cfg_active = _cfg_active_at(t) @@ -2912,7 +2924,8 @@ class Cosmos3OmniDiffusersPipeline( else: # No CFG: a single cond branch per step. Bespoke (state None) keeps # using the transformer-instance cache exactly as before. - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) self._kv_load_und(kv_state, is_negative=False) noise_pred = self.transformer( @@ -2928,6 +2941,7 @@ class Cosmos3OmniDiffusersPipeline( self._kv_capture_und(kv_state, is_negative=False) _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) finally: + self._reset_mixed_precision() # Cosmos3 currently receives a unique request_id rather than a # reusable rollout session id. Retaining its state would only pin # K/V buffers on device after this generation finishes. @@ -3112,7 +3126,8 @@ class Cosmos3OmniDiffusersPipeline( self.transformer.reset_cache() self._cosmos3_branch_caches = {} try: - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0 @@ -3224,6 +3239,7 @@ class Cosmos3OmniDiffusersPipeline( )[0] latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents finally: + self._reset_mixed_precision() self._cosmos3_branch_caches = None self.transformer.reset_cache() return latents diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index fe7a34c8b..1ded3db2f 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -40,6 +40,12 @@ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_ from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm from vllm_omni.platforms import current_omni_platform +from .mixed_precision import ( + Cosmos3MixedPrecisionConfig, + Cosmos3MixedPrecisionRuntime, + resolve_mixed_precision_config, +) + if TYPE_CHECKING: from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook @@ -136,6 +142,26 @@ def _od_config_get(od_config: Any, key: str, default: Any = None) -> Any: return default if value is None else value +def _validate_mixed_precision_runtime( + config: Cosmos3MixedPrecisionConfig | None, + od_config: OmniDiffusionConfig, +) -> None: + if config is None: + return + if get_tensor_model_parallel_world_size() != 1: + raise ValueError("Cosmos3 mixed precision currently supports tensor parallel size 1 only") + if int(getattr(od_config, "max_num_seqs", 1)) != 1: + raise ValueError("Cosmos3 mixed precision currently supports one active request per worker") + if bool(getattr(od_config, "enable_distributed_layerwise_offload", False)): + raise ValueError( + "Cosmos3 mixed precision does not support distributed layer-wise offload " + "because its direct loader bypasses ModelOpt post-load transformations" + ) + parallel_config = getattr(od_config, "parallel_config", None) + if bool(getattr(parallel_config, "use_hsdp", False)): + raise ValueError("Cosmos3 mixed precision has not validated live backend weights under HSDP") + + def _as_bool(value: Any) -> bool: if isinstance(value, str): return value.strip().lower() in {"1", "true", "yes", "on"} @@ -1241,6 +1267,19 @@ class Cosmos3VFMTransformer(nn.Module): dtype = od_config.dtype quant_config = getattr(od_config, "quantization_config", None) if od_config else None + mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config) + if mixed_precision_config is None: + if mixed_precision_source == "additional_config_disabled": + logger.info("Cosmos3 checkpoint mixed-precision policy disabled by additional_config") + else: + logger.info( + "Cosmos3 mixed precision active (source=%s): first_steps=%d last_steps=%d reasoner=%s", + mixed_precision_source, + mixed_precision_config.first_steps, + mixed_precision_config.last_steps, + mixed_precision_config.reasoner, + ) + _validate_mixed_precision_runtime(mixed_precision_config, od_config) self.language_model = self._language_model_cls( hidden_size=self.hidden_size, @@ -1300,6 +1339,11 @@ class Cosmos3VFMTransformer(nn.Module): ] ) + self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None + if mixed_precision_config is not None: + self.mixed_precision_runtime = Cosmos3MixedPrecisionRuntime(mixed_precision_config) + self.mixed_precision_runtime.install(self) + self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) self.gen_sp_prepare = Cosmos3GenSPPrepare() self.gen_sp_gather = nn.Identity() @@ -1958,3 +2002,11 @@ class Cosmos3VFMTransformer(nn.Module): def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" self.time_embedder.to(torch.float32) + + def set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: + if self.mixed_precision_runtime is not None: + self.mixed_precision_runtime.set_step(step_index, num_steps) + + def reset_mixed_precision(self) -> None: + if self.mixed_precision_runtime is not None: + self.mixed_precision_runtime.reset() From e540dfb72f0439be08c26fe61912d8df572fdc7c Mon Sep 17 00:00:00 2001 From: tlysanhuo <166924864+tlysanhuo@users.noreply.github.com> Date: Thu, 10 Sep 2026 21:44:28 +0800 Subject: [PATCH 008/209] [Test] Use public render_jinja_template in MiniCPM-o native template test (#7362) Signed-off-by: tly <2200895168@qq.com> --- .../models/minicpmo_4_5/test_chat_template.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_chat_template.py b/tests/model_executor/models/minicpmo_4_5/test_chat_template.py index 8f1068392..2894b4a1c 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_chat_template.py +++ b/tests/model_executor/models/minicpmo_4_5/test_chat_template.py @@ -6,7 +6,7 @@ from pathlib import Path import pytest from jinja2 import TemplateError -from transformers.utils.chat_template_utils import _compile_jinja_template +from transformers.utils.chat_template_utils import render_jinja_template pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -26,12 +26,15 @@ TEMPLATE = Path(__file__).resolve().parents[4] / "vllm_omni/transformers_utils/c ], ) def test_native_content_concatenation(parts, expected): - rendered = _compile_jinja_template(TEMPLATE.read_text()).render( - messages=[{"role": "user", "content": parts}], + # render_jinja_template takes a batch of conversations and returns + # (rendered_conversations, continuation_chunks), so [0][0] is the string. + rendered = render_jinja_template( + [[{"role": "user", "content": parts}]], + chat_template=TEMPLATE.read_text(), add_generation_prompt=True, enable_thinking=False, use_tts_template=True, - ) + )[0][0] assert rendered == ( f"<|im_start|>user\n{expected}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n<|tts_bos|>" ) @@ -39,7 +42,8 @@ def test_native_content_concatenation(parts, expected): def test_unsupported_part_is_not_silently_discarded(): with pytest.raises(TemplateError, match="Unsupported MiniCPM-o content part"): - _compile_jinja_template(TEMPLATE.read_text()).render( - messages=[{"role": "user", "content": [{"type": "unknown"}]}], + render_jinja_template( + [[{"role": "user", "content": [{"type": "unknown"}]}]], + chat_template=TEMPLATE.read_text(), add_generation_prompt=True, ) From c8bef008ebad9913cc7990eb420df5cbd322d20e Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:16:57 +0800 Subject: [PATCH 009/209] [Bugfix] Fix video prewarm cache retention and cancel-restart delay (#7363) Signed-off-by: psv666 <2693925048@qq.com> --- .../openai_api/test_serving_video_stream.py | 156 +++++++++++++++++- .../openai/serving_video_stream.py | 6 +- .../entrypoints/openai/video_stream_base.py | 38 +++-- 3 files changed, 182 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index bc80fe0db..b534d6c67 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -1,14 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the serving-layer streaming video WebSocket handler.""" from __future__ import annotations import asyncio import base64 +import gc import io import json import threading +import weakref from typing import Any import pytest @@ -38,10 +40,10 @@ def _b64(data: bytes) -> str: def _text_result(text: str) -> OmniRequestOutput: class Output: - pass + text: str class RequestOutput: - pass + outputs: list[Output] output = Output() output.text = text @@ -52,10 +54,10 @@ def _text_result(text: str) -> OmniRequestOutput: def _audio_result(audio_data: Any) -> OmniRequestOutput: class Output: - pass + multimodal_output: dict[str, Any] class RequestOutput: - pass + outputs: list[Output] output = Output() output.multimodal_output = {"audio": audio_data} @@ -513,6 +515,88 @@ async def test_new_query_cancels_in_flight_query(): assert "session.done" in ws.sent_types() +@pytest.mark.asyncio +@pytest.mark.parametrize("delay_abort", [False, True], ids=["immediate-ack", "delayed-ack"]) +async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatch, delay_abort): + started: asyncio.Queue[str] = asyncio.Queue() + abort_started: asyncio.Queue[str] = asyncio.Queue() + allow_abort = asyncio.Event() + if not delay_abort: + allow_abort.set() + request_ids: list[str] = [] + closed: list[str] = [] + aborted: list[str] = [] + delays: list[float] = [] + original_sleep = asyncio.sleep + + async def record_sleep(delay, result=None): + # Observe requested delays without a machine-dependent latency limit. + delays.append(delay) + return await original_sleep(0, result) + + monkeypatch.setattr(video_stream_base.asyncio, "sleep", record_sleep) + + class BlockingEngine: + async def generate(self, *, request_id, **kwargs): + if request_ids: + assert aborted[-1] == request_ids[-1] + request_ids.append(request_id) + started.put_nowait(request_id) + try: + if len(request_ids) <= 3: + yield _text_result("partial") + await asyncio.Event().wait() + else: + yield _text_result("final") + finally: + closed.append(request_id) + + async def abort(self, request_id): + assert request_id in closed + abort_started.put_nowait(request_id) + await allow_abort.wait() + aborted.append(request_id) + + class PreprocessedHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt_token_ids": [1]} + + ws = TimedWebSocket() + handler = PreprocessedHandler(chat_service=object(), engine_client=BlockingEngine(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + try: + ws.put({"type": "session.config", "modalities": ["text"], "enable_frame_filter": False}) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + ws.put({"type": "video.query", "text": "first"}) + previous_id = await asyncio.wait_for(started.get(), timeout=2.0) + + for _ in range(3): + ws.put({"type": "video.query", "text": "interrupt"}) + assert await asyncio.wait_for(abort_started.get(), timeout=2.0) == previous_id + if delay_abort: + assert previous_id not in aborted + assert started.empty() + allow_abort.set() + previous_id = await asyncio.wait_for(started.get(), timeout=2.0) + if delay_abort: + allow_abort.clear() + + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert aborted == request_ids[:-1] + assert len(set(request_ids)) == 4 + assert [msg["text"] for msg in ws.sent if msg["type"] == "response.text.done"] == ["final"] + assert "error" not in ws.sent_types() + assert "session.done" in ws.sent_types() + assert not [delay for delay in delays if delay > 0], "Restart must not add a timed grace period after abort" + finally: + allow_abort.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio async def test_video_done_waits_for_in_flight_query(): query_started = asyncio.Event() @@ -588,6 +672,68 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch): assert "session.done" in ws.sent_types() +@pytest.mark.asyncio +@pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"]) +async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames): + frame_a = _make_jpeg(255, 0, 0) + frame_b = _make_jpeg(0, 255, 0) + decode_started = asyncio.Event() + release_decode = asyncio.Event() + frame_b_accepted = asyncio.Event() + decoded_images: dict[bytes, weakref.ReferenceType[Image.Image]] = {} + prewarm_task: asyncio.Task | None = None + original_to_thread = asyncio.to_thread + + async def controlled_to_thread(function, *args, **kwargs): + nonlocal prewarm_task + if function is video_stream_base._decode_frame_bytes: + if args[0] == frame_a: + prewarm_task = asyncio.current_task() + decode_started.set() + await release_decode.wait() + image = await original_to_thread(function, *args, **kwargs) + decoded_images[args[0]] = weakref.ref(image) + return image + return await original_to_thread(function, *args, **kwargs) + + class AckWebSocket(TimedWebSocket): + async def send_json(self, data): + await super().send_json(data) + if data.get("type") == "video.frame.ack" and data.get("frame_id") == "B": + frame_b_accepted.set() + + monkeypatch.setattr(video_stream_base.asyncio, "to_thread", controlled_to_thread) + ws = AckWebSocket() + handler = QwenOmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + session_task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "max_frames": max_frames, "enable_frame_filter": False}) + try: + ws.put({"type": "video.frame", "frame_id": "A", "data": _b64(frame_a)}) + await asyncio.wait_for(decode_started.wait(), timeout=5.0) + ws.put({"type": "video.frame", "frame_id": "B", "data": _b64(frame_b)}) + await asyncio.wait_for(frame_b_accepted.wait(), timeout=5.0) + ack = next(message for message in ws.sent if message.get("frame_id") == "B") + assert ack["accepted"] is True + assert ack.get("dropped_frame_id") == ("A" if max_frames == 1 else None) + + # Finish A's real decode only after B has either evicted A or joined it. + release_decode.set() + assert prewarm_task is not None + await asyncio.wait_for(asyncio.shield(prewarm_task), timeout=5.0) + gc.collect() + assert not session_task.done() + # A finished task must not keep an evicted PIL image alive for the session. + assert (decoded_images[frame_a]() is not None) == (max_frames == 2) + finally: + release_decode.set() + ws.put({"type": "video.done"}) + await asyncio.wait_for(session_task, timeout=5.0) + + gc.collect() + assert decoded_images[frame_a]() is None + assert not any(message.get("type") == "error" for message in ws.sent) + + @pytest.mark.asyncio async def test_client_cannot_send_internal_frame_decode_failed_message(): captured_frames: list[list[str]] = [] diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index af0083fcb..e05a40eef 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen-Omni streaming video WebSocket handler. Accepts video frames incrementally via WebSocket, buffers them, and @@ -26,12 +26,14 @@ Protocol: from __future__ import annotations +from collections.abc import Mapping from typing import Any from vllm_omni.entrypoints.openai.video_stream_base import ( _BAD_FRAME, _DEFAULT_CONFIG_TIMEOUT, _DEFAULT_IDLE_TIMEOUT, + PrewarmedFrame, StreamingVideoSessionConfig, VideoStreamTurnTrigger, ) @@ -59,7 +61,7 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: n_buf = len(frame_buffer) if n_buf <= config.num_frames: diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index f975adbef..ab5d51c33 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Base WebSocket handler for streaming video input understanding. Shared session loop, frame/audio buffering, EVS pre-filter, prewarm, @@ -37,7 +37,8 @@ import uuid import wave from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable +from enum import Enum +from typing import Any, Final, Protocol, TypeAlias, runtime_checkable import torch from fastapi import WebSocket, WebSocketDisconnect @@ -61,7 +62,14 @@ _MAX_BUFFER_FRAMES = 64 _MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 _MAX_MSG_QUEUE = 200 _CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length -_BAD_FRAME = object() + + +class _FrameStatus(Enum): + BAD = "bad" + + +_BAD_FRAME: Final = _FrameStatus.BAD +PrewarmedFrame: TypeAlias = tuple[Any, str] | _FrameStatus def _decode_frame_bytes(raw_bytes: bytes) -> Any: @@ -83,7 +91,7 @@ class VideoStreamPipelineHooks(Protocol): audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: """Build OpenAI-style messages and the current user message.""" ... @@ -169,7 +177,7 @@ class OmniStreamingVideoHandler: audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: raise NotImplementedError @@ -219,7 +227,7 @@ class OmniStreamingVideoHandler: frame_buffer: list[str] = [] # base64-encoded JPEG frames frame_metadata: list[dict[str, Any]] = [] # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. - frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME + frame_pil_cache: dict[str, PrewarmedFrame] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME frame_filter = ( FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None ) @@ -306,7 +314,6 @@ class OmniStreamingVideoHandler: await self._engine_client.abort(prev_request_id) except Exception: pass - await asyncio.sleep(0.1) prev_was_interrupted = False request_id = f"video-{uuid.uuid4().hex[:12]}" @@ -429,8 +436,12 @@ class OmniStreamingVideoHandler: async def _prewarm(b64: str, b: bytes, u: str) -> None: try: pil = await asyncio.to_thread(_decode_frame_bytes, b) - frame_pil_cache[b64] = (pil, u) + # The frame may have been evicted while decoding. + if b64 in frame_buffer: + frame_pil_cache[b64] = (pil, u) except Exception: + if b64 not in frame_buffer: + return frame_pil_cache[b64] = _BAD_FRAME logger.warning("prewarm decode failed for frame (len=%d)", len(b)) try: @@ -577,7 +588,7 @@ class OmniStreamingVideoHandler: query_text: str, request_id: str, interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Build prompt, run inference, stream text + audio response.""" @@ -616,10 +627,13 @@ class OmniStreamingVideoHandler: query_text: str, request_id: str, interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Direct engine_client.generate() path for async_chunk audio.""" + engine_client = self._engine_client + assert engine_client is not None, "_process_query must validate the engine client" + from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) @@ -686,7 +700,7 @@ class OmniStreamingVideoHandler: audio_tail_tensors: list[Any] = [] try: - result_gen = self._engine_client.generate( + result_gen = engine_client.generate( prompt=engine_prompt, request_id=request_id, output_modalities=config.modalities, @@ -989,6 +1003,8 @@ class OmniStreamingVideoHandler: ) mixin = AudioMixin() resp = mixin.create_audio(audio_obj) + # base64_encode=True selects the string result of create_audio(). + assert isinstance(resp.audio_data, str) return resp.audio_data @staticmethod From 05ab5092695169ec711c2a0462e90e1b194968ad Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Thu, 10 Sep 2026 16:32:50 +0200 Subject: [PATCH 010/209] Cosmos3 action policy improvements (#6460) Signed-off-by: Maciej Bala Signed-off-by: MaciejBalaNV Co-authored-by: Hongsheng Liu --- recipes/cosmos3/Cosmos3-Nano.md | 92 +++++++---- tests/config/test_config_factory.py | 71 +++++++++ .../cosmos3/test_cosmos3_framework_compat.py | 144 ++++++++++++++++++ .../models/cosmos3/test_cosmos3_pipeline.py | 78 ++++++++++ .../openai_api/test_openpi_connection.py | 90 +++++++++++ .../openai_api/test_openpi_serving.py | 78 ++++++++++ tests/entrypoints/test_serve.py | 20 +++ tests/model_extras/test_model_extras.py | 1 + vllm_omni/config/pipeline_registry.py | 6 + vllm_omni/config/stage_config.py | 16 +- vllm_omni/deploy/cosmos3_policy_droid.yaml | 49 ++++++ vllm_omni/diffusion/models/cosmos3/action.py | 11 ++ .../models/cosmos3/pipeline_cosmos3.py | 29 +++- vllm_omni/diffusion/models/cosmos3/utils.py | 116 +++++++++----- vllm_omni/entrypoints/cli/serve.py | 11 ++ vllm_omni/entrypoints/openai/api_server.py | 19 ++- vllm_omni/entrypoints/openpi/connection.py | 63 +++++--- vllm_omni/entrypoints/openpi/serving.py | 30 +++- .../model_executor/models/cosmos3/__init__.py | 2 + .../model_executor/models/cosmos3/pipeline.py | 41 +++++ vllm_omni/model_extras/cosmos3.py | 1 + 21 files changed, 863 insertions(+), 105 deletions(-) create mode 100644 tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py create mode 100644 vllm_omni/deploy/cosmos3_policy_droid.yaml create mode 100644 vllm_omni/model_executor/models/cosmos3/__init__.py create mode 100644 vllm_omni/model_executor/models/cosmos3/pipeline.py diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index f81e5e887..6ef24bcf4 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -50,9 +50,10 @@ mode is selected per request: **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way (`domain_name=droid_lerobot`). -- **DROID OpenPI policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and - connect an OpenPI-compatible websocket client to `/v1/realtime/robot/openpi`. - This path returns action chunks directly instead of an mp4. +- **DROID policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and connect + RoboLab or another OpenPI-compatible client to + `/v1/realtime/robot/openpi`. This endpoint returns action chunks directly + instead of an mp4. Action requests can use `input_reference` or `video_reference` for video input. `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics` @@ -314,40 +315,71 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \ curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}' curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4 -# DROID OpenPI policy server (websocket action serving). -# Requires cosmos_framework on PYTHONPATH because the pipeline reuses the -# reference RoboLab action transforms. If your checkpoint config already -# includes policy_server_config, omit the stage_overrides file and flag. -cat > cosmos3_droid_openpi_stage_overrides.json <<'JSON' -{ - "0": { - "model_config": { - "policy_server_config": { - "image_resolution": [540, 640], - "n_external_cameras": 2, - "needs_wrist_camera": true, - "needs_stereo_camera": false, - "needs_session_id": true, - "action_space": "joint_position" - } - } - } -} -JSON +# DROID websocket policy server. Use the tested cosmos-framework revision +# directly from source; installing the package can introduce dependency conflicts. +export COSMOS_FRAMEWORK_ROOT=/path/to/cosmos-framework +git -C "$COSMOS_FRAMEWORK_ROOT" checkout c14617c2bc93dacbf69674fb964eec93182933d9 +export PYTHONPATH="$COSMOS_FRAMEWORK_ROOT" +# Validate every cosmos-framework symbol used by the action-policy pipeline +# before allocating the model. +python - <<'PY' +from vllm_omni.diffusion.models.cosmos3.utils import ( + get_robolab_domain_id, + preflight_cosmos3_action_framework_imports, +) + +preflight_cosmos3_action_framework_imports() +print("cosmos-framework action imports OK; DROID domain:", get_robolab_domain_id("droid_lerobot")) +PY + +# The bundled deploy config (vllm_omni/deploy/cosmos3_policy_droid.yaml) +# selects the registered cosmos3_policy pipeline and carries the DROID +# checkpoint's serving defaults: JSON prompt formatting (a property of this +# checkpoint's training recipe, not a generic server default), the +# model-specific OpenPI binary handshake metadata (policy_server_config), and +# guardrails off (policy serving emits robot actions; the guardrail stack is +# not part of it, so --no-guardrails is implied). +# Policy checkpoints cannot be auto-detected — they share their HF metadata +# with the T2I/video Cosmos3 checkpoints — so --deploy-config is required. +# Per-checkpoint model_config tweaks can be layered on top with +# --stage-overrides; dict overrides deep-merge with the deploy yaml. +export VLLM_OMNI_ROOT=/absolute/path/to/vllm-omni vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ --omni \ --host 0.0.0.0 --port 8000 \ - --model-class-name Cosmos3OmniDiffusersPipeline \ - --no-guardrails \ - --stage-overrides "$(cat cosmos3_droid_openpi_stage_overrides.json)" + --deploy-config "$VLLM_OMNI_ROOT/vllm_omni/deploy/cosmos3_policy_droid.yaml" \ + --robot-openpi-idle-timeout 0 -# Point an OpenPI websocket client at: -# ws://localhost:8000/v1/realtime/robot/openpi -# The first server message is policy_server_config. Each infer request sends a -# msgpack-numpy observation dict and receives a writable float32 action array. +# From a RoboLab checkout: +python policies/cosmos3/run.py \ + --remote-uri ws://localhost:8000/v1/realtime/robot/openpi \ + --task BananaInBowlTask + +# For a server started with "--api-key $POLICY_API_KEY", pass the same value as +# a Bearer token. --remote-token takes precedence over COSMOS3_API_TOKEN. +COSMOS3_API_TOKEN="$POLICY_API_KEY" python policies/cosmos3/run.py \ + --remote-uri wss://policy.example/v1/realtime/robot/openpi \ + --task BananaInBowlTask ``` +The endpoint sends `policy_server_config` as its initial binary MsgPack +message and returns the action array directly. RoboLab accepts that direct +array as well as proxy responses shaped as `{"action": ...}` or +`{"actions": ...}`, and turns structured server errors into client exceptions. + +RoboLab supplies `session_id=robolab-episode--env-`, so +parallel environments have independent policy state even though they share +one WebSocket. The server tracks interleaved session IDs independently. +The default receive-idle timeout is 30 seconds; set +`--robot-openpi-idle-timeout 0` when simulator steps between replans can take +longer, or set another non-negative timeout in seconds. + +If vLLM-Omni is started with `--api-key` or `VLLM_API_KEY`, the standard +OpenPI route uses the normal API authentication middleware. RoboLab's +`--remote-token` and `COSMOS3_API_TOKEN` send the required +`Authorization: Bearer ...` header. + #### Notes - **Measured latency (1x B300, bf16, guardrails off):** diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 51c7d549e..af898ae50 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -197,6 +197,25 @@ class TestStageConfig: omega_config = config.to_omegaconf() assert omega_config.engine_args.max_num_seqs == 32 + def test_to_omegaconf_dict_override_deep_merges_yaml_dict(self): + """A partial dict CLI override (e.g. --no-guardrails riding on + model_config) must layer onto the YAML dict, not clobber siblings.""" + config = StageConfig( + stage_id=0, + model_stage="diffusion", + stage_type=StageType.DIFFUSION, + yaml_engine_args={ + "model_config": { + "guardrails": True, + "policy_server_config": {"action_space": "joint_position"}, + } + }, + runtime_overrides={"model_config": {"guardrails": False}}, + ) + omega_config = config.to_omegaconf() + assert omega_config.engine_args.model_config.guardrails is False + assert omega_config.engine_args.model_config.policy_server_config.action_space == "joint_position" + def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self): config = StageConfig( stage_id=1, @@ -623,6 +642,58 @@ class TestPipelineDiscovery: assert p.hf_architectures == ("SomeCollidingArch",) +class TestCosmos3PolicyPipeline: + """Cosmos3 policy serving resolves via deploy yaml selection (π0 precedent).""" + + def test_registered_without_capturing_base_cosmos3_checkpoints(self): + assert "cosmos3_policy" in OMNI_PIPELINES + # T2I/video Cosmos3 checkpoints report model_type=cosmos3_omni and the + # same model_index.json _class_name as policy checkpoints; they must + # keep resolving through the single-stage diffusion fallback, so the + # policy pipeline must not be reachable by auto-detection. + assert "cosmos3_omni" not in OMNI_PIPELINES + pipeline = OMNI_PIPELINES["cosmos3_policy"] + assert pipeline.hf_architectures == () + assert pipeline.diffusers_class_name is None + + def test_droid_deploy_yaml_carries_policy_server_config(self): + deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) + assert deploy.pipeline == "cosmos3_policy" + + stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_policy"], deploy) + assert len(stages) == 1 + stage = stages[0].to_omegaconf() + + assert stage.stage_type == "diffusion" + assert stage.final_output_type == "action" + assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline" + # OpenPI websocket handshake metadata must reach the serving layer via + # engine_args.model_config (ServingRealtimeRobotOpenPI reads it there). + psc = stage.engine_args.model_config.policy_server_config + assert list(psc.image_resolution) == [540, 640] + assert psc.n_external_cameras == 2 + assert psc.needs_wrist_camera is True + assert psc.needs_stereo_camera is False + assert psc.needs_session_id is True + assert psc.action_space == "joint_position" + # The DROID checkpoint's training recipe expects JSON-formatted prompts. + assert stage.default_sampling_params.extra_args.format_prompt_as_json is True + + def test_no_guardrails_cli_flag_keeps_policy_server_config(self): + """--no-guardrails becomes a partial model_config CLI override; it must + deep-merge with the deploy yaml's model_config, not clobber the + policy_server_config the OpenPI handshake depends on.""" + deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) + stages, _ = StageConfigFactory._create_legacy_from_registry( + OMNI_PIPELINES["cosmos3_policy"], + {"model_config": {"guardrails": False}}, + user_deploy_config=deploy, + ) + stage = stages[0].to_omegaconf() + assert stage.engine_args.model_config.guardrails is False + assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position" + + class TestStagePipelineConfig: def test_frozen(self): s = StagePipelineConfig(stage_id=0, model_stage="a") diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py b/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py new file mode 100644 index 000000000..25f06b7a9 --- /dev/null +++ b/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from importlib.util import find_spec + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_FRAMEWORK_AVAILABLE = find_spec("cosmos_framework") is not None +requires_cosmos_framework = pytest.mark.skipif( + not _FRAMEWORK_AVAILABLE, + reason="cosmos_framework source is not present on PYTHONPATH", +) + + +def test_framework_preflight_checks_every_policy_symbol(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils + + imported = [] + + def record_import(module_name: str, symbol_name: str, error_message: str): + imported.append((module_name, symbol_name, error_message)) + return object() + + monkeypatch.setattr(cosmos3_utils, "lazy_import", record_import) + + cosmos3_utils.preflight_cosmos3_action_framework_imports() + + expected_symbols = { + "ActionTransformPipeline", + "FlowUniPCMultistepScheduler", + "build_abs_pose_from_components", + "convert_rotation", + "get_domain_id", + "pose_abs_to_rel", + "pose_rel_to_abs", + } + assert {symbol_name for _, symbol_name, _ in imported} == expected_symbols + assert [(module_name, symbol_name) for module_name, symbol_name, _ in imported] == list( + cosmos3_utils._ROBOLAB_FRAMEWORK_IMPORTS.values() + ) + + +def test_lazy_import_reports_framework_symbol_skew(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils + + monkeypatch.setattr(cosmos3_utils, "import_module", lambda _module_name: object()) + + with pytest.raises(ImportError, match="MissingSymbol") as exc_info: + cosmos3_utils.lazy_import( + "cosmos_framework.fake_module", + "MissingSymbol", + "Cosmos Framework compatibility error.", + ) + + assert isinstance(exc_info.value.__cause__, AttributeError) + + +@requires_cosmos_framework +def test_action_domain_table_matches_current_framework_and_legacy_aliases() -> None: + from cosmos_framework.data.generator.action.utils.domain_utils import ( + EMBODIMENT_TO_DOMAIN_ID as FRAMEWORK_DOMAIN_IDS, + ) + + from vllm_omni.diffusion.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID as VLLM_DOMAIN_IDS + from vllm_omni.diffusion.models.cosmos3.utils import _LEGACY_EMBODIMENT_ALIASES + + canonical_vllm_domain_ids = { + name: domain_id for name, domain_id in VLLM_DOMAIN_IDS.items() if name not in _LEGACY_EMBODIMENT_ALIASES + } + + assert canonical_vllm_domain_ids == FRAMEWORK_DOMAIN_IDS + assert {alias: VLLM_DOMAIN_IDS[alias] for alias in _LEGACY_EMBODIMENT_ALIASES} == { + alias: FRAMEWORK_DOMAIN_IDS[canonical_name] for alias, canonical_name in _LEGACY_EMBODIMENT_ALIASES.items() + } + + +@requires_cosmos_framework +def test_current_cosmos_framework_policy_modules_import_and_instantiate() -> None: + from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id + from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + convert_rotation, + pose_abs_to_rel, + pose_rel_to_abs, + ) + from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline + from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import ( + FlowUniPCMultistepScheduler, + ) + + transform = ActionTransformPipeline( + max_action_dim=64, + cfg_dropout_rate=0.0, + format_prompt_as_json=True, + ) + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, + shift=1.0, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(4, device=torch.device("cpu"), shift=5.0) + + assert transform.prompt_json_formatter is not None + assert len(scheduler.timesteps) == 4 + assert get_domain_id("droid_lerobot") == 8 + assert convert_rotation([0.0, 0.0, 0.0, 1.0], "quat_xyzw", "rot6d").shape == (6,) + assert callable(pose_abs_to_rel) + assert callable(pose_rel_to_abs) + assert callable(build_abs_pose_from_components) + + +@requires_cosmos_framework +def test_real_action_transform_pipeline_matches_droid_policy_contract() -> None: + from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline + + transform = ActionTransformPipeline( + max_action_dim=64, + cfg_dropout_rate=0.0, + format_prompt_as_json=True, + ) + sample = { + "ai_caption": "Pick up the cube.", + "video": torch.zeros((3, 3, 192, 320), dtype=torch.uint8), + "action": torch.zeros((3, 8), dtype=torch.float32), + "conditioning_fps": torch.tensor(15, dtype=torch.long), + "mode": "wam", + "domain_id": torch.tensor(8, dtype=torch.long), + "viewpoint": "concat_view", + "additional_view_description": "Wrist view above two exterior views.", + } + + transformed = transform(sample, "256") + + assert isinstance(transformed["ai_caption"], dict) + assert transformed["action"].shape == (3, 64) + assert int(transformed["raw_action_dim"].item()) == 8 + assert transformed["sequence_plan"].has_action is True + assert transformed["mode"] == "wam" + assert int(transformed["domain_id"].item()) == 8 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 72e4ab962..302409aeb 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -3,6 +3,7 @@ from __future__ import annotations +import json import sys import types from dataclasses import dataclass @@ -21,6 +22,20 @@ from vllm_omni.experimental.world_models.session_state import SessionStateManage pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.mark.parametrize( + ("alias", "canonical_name"), + [ + ("galbot", "embodiment_b"), + ("agibot_gear_gripper", "embodiment_c_gripper"), + ("agibot_gear_gripper_ext", "embodiment_c_gripper_ext"), + ], +) +def test_action_domain_table_preserves_legacy_aliases(alias: str, canonical_name: str) -> None: + from vllm_omni.diffusion.models.cosmos3.action import resolve_domain_id + + assert resolve_domain_id(domain_name=alias) == resolve_domain_id(domain_name=canonical_name) + + def test_pipeline_declares_layerwise_offload_components() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline @@ -479,6 +494,69 @@ def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None: assert captured["session_id"] == "robolab-request-7" +@pytest.mark.parametrize("format_prompt_as_json", [False, True]) +def test_robolab_input_builder_threads_prompt_format_and_uses_wam( + make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, + format_prompt_as_json: bool, +) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + pipeline.transformer = StubCosmos3Transformer(action_gen=True, action_dim=64) + captured: dict[str, Any] = {} + + def fake_transform(sample, resolution): + captured["sample_mode"] = sample["mode"] + captured["resolution"] = resolution + sample["sequence_plan"] = SimpleNamespace( + condition_frame_indexes_action=[0], + action_start_frame_offset=1, + ) + sample["raw_action_dim"] = torch.tensor(8) + sample["image_size"] = torch.tensor([16, 16, 16, 16]) + if format_prompt_as_json: + sample["ai_caption"] = {"actions": {"instruction": sample["ai_caption"]}} + return sample + + def fake_get_transform(*, format_prompt_as_json: bool): + captured["format_prompt_as_json"] = format_prompt_as_json + return fake_transform + + pipeline._get_robolab_transform = fake_get_transform + monkeypatch.setattr(pipeline_cosmos3, "get_robolab_domain_id", lambda name: 8) + obs = { + "prompt": "Pick up the cube.", + "observation/image": np.zeros((16, 16, 3), dtype=np.uint8), + "observation/joint_position": np.zeros(7, dtype=np.float32), + "observation/gripper_position": np.zeros(1, dtype=np.float32), + } + sampling_params = make_sampling_params( + extra_args={ + "robot_obs": obs, + "action_chunk_size": 2, + "image_height": 16, + "image_width": 16, + "format_prompt_as_json": format_prompt_as_json, + } + ) + + inputs = pipeline._build_robolab_policy_inputs(sampling_params, request_id="request-1") + + assert inputs is not None + assert captured == { + "sample_mode": "wam", + "resolution": "480", + "format_prompt_as_json": format_prompt_as_json, + } + assert inputs.domain_id == 8 + assert inputs.raw_action_dim == 8 + if format_prompt_as_json: + assert json.loads(inputs.prompt) == {"actions": {"instruction": "Pick up the cube."}} + else: + assert inputs.prompt == "Pick up the cube." + + @pytest.mark.parametrize( ("prompt", "sampling_params", "message"), [ diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 19ec23b6d..a6f12f661 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -310,6 +310,29 @@ def test_handle_connection_closes_websocket_on_idle_timeout(monkeypatch): serving.infer.assert_not_called() +def test_handle_connection_can_disable_idle_timeout(monkeypatch): + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + + async def fail_wait_for(*_args, **_kwargs): + pytest.fail("asyncio.wait_for should not be used when the idle timeout is disabled") + + monkeypatch.setattr(openpi_connection.asyncio, "wait_for", fail_wait_for) + websocket = FakeWebSocket([{"type": "websocket.disconnect"}]) + serving = _serving_mock() + + asyncio.run( + openpi_connection.RobotRealtimeConnection( + websocket, + serving, + idle_timeout=None, + ).handle_connection() + ) + + assert websocket.accepted is True + assert websocket.closed is False + serving.infer.assert_not_called() + + def test_handle_connection_keeps_session_state_per_websocket(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) requests = { @@ -343,6 +366,73 @@ def test_handle_connection_keeps_session_state_per_websocket(monkeypatch): assert calls[2].kwargs == {"session_id": "session-b", "reset": True} +def test_handle_connection_keeps_interleaved_parallel_session_state(monkeypatch): + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + requests = { + b"a1": {"prompt": "first-a", "session_id": "session-a"}, + b"b1": {"prompt": "first-b", "session_id": "session-b"}, + b"a2": {"prompt": "second-a", "session_id": "session-a"}, + b"b2": {"prompt": "second-b", "session_id": "session-b"}, + } + monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data])) + serving = _serving_mock() + websocket = FakeWebSocket( + [{"type": "websocket.receive", "bytes": key} for key in (b"a1", b"b1", b"a2", b"b2")] + + [{"type": "websocket.disconnect"}] + ) + + asyncio.run(openpi_connection.RobotRealtimeConnection(websocket, serving).handle_connection()) + + assert [call.kwargs for call in serving.infer.await_args_list] == [ + {"session_id": "session-a", "reset": True}, + {"session_id": "session-b", "reset": True}, + {"session_id": "session-a", "reset": False}, + {"session_id": "session-b", "reset": False}, + ] + + +def test_handle_connection_bounds_tracked_sessions_without_resetting_live_ones(monkeypatch): + """A client may hold one socket across many episodes, minting a fresh + session id per episode and never sending ``reset`` (RoboLab does exactly + this). Session tracking must stay bounded, and evicting finished episodes + must not disturb the sessions that are still running. + """ + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + monkeypatch.setattr(openpi_connection, "MAX_TRACKED_SESSIONS", 8) + + # One long-lived session interleaved with far more short episodes than the cap. + episodes = 40 + keys = [] + requests = {} + for episode in range(episodes): + for session_id in (f"episode-{episode}", "live-env"): + key = f"{session_id}@{episode}".encode() + keys.append(key) + requests[key] = {"prompt": "pick", "session_id": session_id} + + monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data])) + serving = _serving_mock() + websocket = FakeWebSocket( + [{"type": "websocket.receive", "bytes": key} for key in keys] + [{"type": "websocket.disconnect"}] + ) + connection = openpi_connection.RobotRealtimeConnection(websocket, serving) + + asyncio.run(connection.handle_connection()) + + assert len(connection._seen_sessions) <= 8 + + resets_by_session: dict[str, list[bool]] = {} + for call in serving.infer.await_args_list: + resets_by_session.setdefault(call.kwargs["session_id"], []).append(call.kwargs["reset"]) + + # Each new episode resets exactly once, on its first observation. + for episode in range(episodes): + assert resets_by_session[f"episode-{episode}"] == [True] + # The continuously active session is never evicted, so it resets only once + # despite ``episodes`` other sessions churning past the cap. + assert resets_by_session["live-env"] == [True] + [False] * (episodes - 1) + + def test_handle_connection_reset_endpoint_resets_next_infer(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) requests = { diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 072e734e4..82ee6e7dd 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -15,6 +15,7 @@ from starlette.testclient import TestClient from vllm_omni.entrypoints.openpi import connection as openpi_connection from vllm_omni.entrypoints.openpi import serving as openpi_serving +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -177,6 +178,22 @@ def test_create_policy_server_returns_none_without_policy_config(): assert serving is None +def test_legacy_policy_request_defaults_key_no_longer_controls_serving(): + od_config = SimpleNamespace( + model_config={ + "policy_server_config": {}, + "policy_request_defaults": "malformed legacy value", + } + ) + + serving = openpi_serving.ServingRealtimeRobotOpenPI.create_policy_server( + engine_client=SimpleNamespace(get_diffusion_od_config=lambda: od_config), + model_name="generic-policy", + ) + + assert serving is not None + + def test_policy_server_config_allows_explicit_empty_config(): serving = openpi_serving.ServingRealtimeRobotOpenPI( engine_client=_engine_with_policy_config(policy_config={}), @@ -221,6 +238,44 @@ def test_build_request_uses_unique_engine_request_id_per_inference(): assert request_a.request_id != request_b.request_id +def test_build_request_clones_stage_defaults_before_protocol_fields(): + default_params = OmniDiffusionSamplingParams( + extra_args={ + "format_prompt_as_json": True, + "session_id": "configured-session-must-not-win", + "reset": False, + "nested": {"values": []}, + } + ) + od_config = SimpleNamespace(model_config={"policy_server_config": {}}) + engine_client = SimpleNamespace( + get_diffusion_od_config=lambda: od_config, + default_sampling_params_list=[default_params], + ) + serving = openpi_serving.ServingRealtimeRobotOpenPI( + engine_client=engine_client, + model_name="custom-policy", + ) + + request_a = serving._build_request( + {"prompt": "pick up the object"}, + session_id="wire-session", + reset=True, + ) + request_b = serving._build_request( + {"prompt": "pick up the object"}, + session_id="wire-session", + reset=False, + ) + + assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True + assert request_a.sampling_params.extra_args["session_id"] == "wire-session" + assert request_a.sampling_params.extra_args["reset"] is True + request_a.sampling_params.extra_args["nested"]["values"].append("request-a") + assert request_b.sampling_params.extra_args["nested"] == {"values": []} + assert default_params.extra_args["nested"] == {"values": []} + + def test_build_request_forwards_seed_to_sampling_params(): """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds.""" serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) @@ -257,6 +312,29 @@ def test_infer_keeps_session_state_but_uses_unique_engine_request_ids(): assert sampling_params_b.extra_args["reset"] is False +def test_infer_yields_event_loop_while_engine_is_running(): + class DelayedEngine(RecordingEngine): + def generate(self, *, prompt, request_id, sampling_params_list): + async def _generate(): + await asyncio.sleep(0.03) + yield SimpleNamespace(multimodal_output={"actions": [0.0]}) + + return _generate() + + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=DelayedEngine()) + + async def run_test(): + infer_task = asyncio.create_task(serving.infer({"prompt": "pick"}, session_id="session-a", reset=True)) + ticks = 0 + while not infer_task.done(): + await asyncio.sleep(0.005) + ticks += 1 + await infer_task + return ticks + + assert asyncio.run(run_test()) >= 2 + + def test_two_websocket_clients_without_session_id_do_not_conflict(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", _json_pack) monkeypatch.setattr(openpi_connection, "_unpack", _json_unpack) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 0089330b2..083f96d9f 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -112,6 +112,26 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: assert args.cfg_parallel_size == 4 +def test_serve_parser_accepts_robot_openpi_idle_timeout() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "0"]) + + assert args.robot_openpi_idle_timeout == 0 + assert args.get_explicit_kwargs_dict()["robot_openpi_idle_timeout"] == 0 + + +def test_serve_parser_rejects_negative_robot_openpi_idle_timeout() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + with pytest.raises(SystemExit): + parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "-1"]) + + def test_serve_parser_accepts_ulysses_a2a_permute() -> None: parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="subcommand") diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 6bf823339..e55691e10 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -177,6 +177,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na "resolution", "image_size", "use_state", + "format_prompt_as_json", "observation", "robot_obs", "deterministic_seed", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index b711f0cfb..a69fb08c6 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -46,6 +46,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_SINGLE_STAGE_PIPELINE, BAGEL_THINK_PIPELINE, ) +from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE @@ -149,6 +150,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "lingbot_world": LINGBOT_WORLD_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, + # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3 + # checkpoints (which stay on the single-stage diffusion fallback), so this + # entry is only reachable through a deploy yaml's ``pipeline:`` key + # (see deploy/cosmos3_policy_droid.yaml). + "cosmos3_policy": COSMOS3_POLICY_PIPELINE, "gepard": GEPARD_PIPELINE, "glm_image": GLM_IMAGE_PIPELINE, "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f7d714890..0e962b4cc 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1130,10 +1130,22 @@ class StageConfig: _apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides) reconcile_diffusion_attention_overrides(engine_args, runtime_overrides) - # CLI overrides take precedence over YAML defaults + # CLI overrides take precedence over YAML defaults. Most dict-valued + # overrides are deep-merged so a partial CLI dict (e.g. --no-guardrails + # riding on ``model_config``) layers onto the deploy YAML instead of + # clobbering sibling keys such as ``policy_server_config`` — the same + # rationale as the platform-overlay deep-merge. Legacy atomic mappings + # are handled explicitly below. for key, value in runtime_overrides.items(): if value is not None and key not in ("devices", "max_batch_size", "num_replicas"): - engine_args[key] = value + existing = engine_args.get(key) + # ``omni_kv_config`` is an atomic legacy override: callers use + # a partial mapping to replace the topology-provided transfer + # role, rather than to add fields to it. + if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, dict): + engine_args[key] = _get_recursively_merged_dict(existing, value) + else: + engine_args[key] = value # Build runtime config from YAML defaults + CLI overrides runtime: dict[str, Any] = dict(self.yaml_runtime) diff --git a/vllm_omni/deploy/cosmos3_policy_droid.yaml b/vllm_omni/deploy/cosmos3_policy_droid.yaml new file mode 100644 index 000000000..cbcbcb42c --- /dev/null +++ b/vllm_omni/deploy/cosmos3_policy_droid.yaml @@ -0,0 +1,49 @@ +# Cosmos3 policy (DROID) deploy: single diffusion stage serving the OpenPI +# realtime API (/v1/realtime/robot/openpi). +# +# Topology (COSMOS3_POLICY_PIPELINE) is declared in +# vllm_omni/model_executor/models/cosmos3/pipeline.py. Policy checkpoints share +# their HF metadata (model_type=cosmos3_omni, model_index.json _class_name= +# Cosmos3OmniDiffusersPipeline) with the T2I/video Cosmos3 checkpoints, which +# must keep using the default single-stage diffusion fallback — so this +# pipeline cannot be auto-detected and must be selected explicitly: +# +# vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni \ +# --deploy-config /absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml +# +# Per-checkpoint model_config tweaks (e.g. policy_server_config values) can +# still be layered on top with --stage-overrides; dict overrides deep-merge +# with this file. + +pipeline: cosmos3_policy +# async_chunk defaults to True, which the deploy validator rejects for a +# single-stage pipeline (no next-stage processor). Must be false here. +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + # OpenPI policy serving is one observation -> one action reply. + max_num_seqs: 1 + model_class_name: Cosmos3OmniDiffusersPipeline + # JSON prompt formatting is a property of this DROID checkpoint's training + # recipe, not a generic server default. + default_sampling_params: + extra_args: + format_prompt_as_json: true + model_config: + # Policy serving emits robot actions, not text/video; the Cosmos3 + # guardrail stack (gated nvidia/Cosmos-1.0-Guardrail repo) is not part + # of it. Equivalent to launching with --no-guardrails. + guardrails: false + # OpenPI handshake metadata advertised to the robot client. Sent verbatim + # (the pipeline never reads these); they describe the camera inputs and + # action space the client must produce for this checkpoint. + policy_server_config: + image_resolution: [540, 640] + n_external_cameras: 2 + needs_wrist_camera: true + needs_stereo_camera: false + needs_session_id: true + action_space: joint_position diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index dc755b39c..affba4195 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -38,13 +38,24 @@ EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { "bridge_orig_lerobot": 7, "droid_lerobot": 8, "robomind-franka": 8, + "embodiment_b": 9, + # Backward-compatible alias for the pre-generator Cosmos Framework name. "galbot": 9, "robomind-franka-dual": 12, "robomind-ur": 13, "agibotworld": 15, + "embodiment_c_gripper": 15, + "embodiment_c_gripper_ext": 15, + # Backward-compatible aliases for the pre-generator Cosmos Framework names. "agibot_gear_gripper": 15, "agibot_gear_gripper_ext": 15, + "xdof_yam": 16, + "molmoact2_yam": 16, + "abc_yam": 16, "fractal": 20, + "drawanything": 21, + "behavior1k_lerobot": 22, + "maniparena": 23, } diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index c60817afb..241eae805 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -138,6 +138,7 @@ from .utils import ( ensure_gripper_array, extract_robolab_image, extract_robolab_prompt_image, + get_robolab_domain_id, lazy_action_transform_pipeline, make_robolab_action_postprocess_inputs, next_robolab_seed, @@ -1075,7 +1076,7 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = None self._num_timesteps = None self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None - self._robolab_transform = None + self._robolab_transforms: dict[bool, Any] = {} # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on # this pipeline. Tells the sequential-CFG loop to keep paired @@ -1394,11 +1395,18 @@ class Cosmos3OmniDiffusersPipeline( return val return default - def _get_robolab_transform(self): - if self._robolab_transform is None: + def _get_robolab_transform(self, *, format_prompt_as_json: bool = False): + transforms = getattr(self, "_robolab_transforms", None) + if transforms is None: + transforms = {} + self._robolab_transforms = transforms + if format_prompt_as_json not in transforms: action_dim = int(getattr(self.transformer, "action_dim", 64)) - self._robolab_transform = lazy_action_transform_pipeline(action_dim) - return self._robolab_transform + transforms[format_prompt_as_json] = lazy_action_transform_pipeline( + action_dim, + format_prompt_as_json=format_prompt_as_json, + ) + return transforms[format_prompt_as_json] def _build_robolab_policy_inputs( self, @@ -1443,7 +1451,8 @@ class Cosmos3OmniDiffusersPipeline( resolution = str(extra_param("resolution", ROBOLAB_DEFAULT_RESOLUTION)) fps = float(extra_param("conditioning_fps", ROBOLAB_DEFAULT_CONDITIONING_FPS)) domain_name = str(extra_param("domain_name", ROBOLAB_DEFAULT_DOMAIN_NAME)) - domain_id = resolve_domain_id(domain_name=domain_name, require_explicit=True) + domain_id = get_robolab_domain_id(domain_name) + format_prompt_as_json = self._truthy(extra_param("format_prompt_as_json", False)) if use_state and history_length < 1: raise ValueError("RoboLab history_length must be >= 1 when use_state is true.") @@ -1506,7 +1515,9 @@ class Cosmos3OmniDiffusersPipeline( "action": action, # Cosmos Framework consumes this as an integer conditioning bucket. "conditioning_fps": torch.tensor(fps, dtype=torch.long), - "mode": ACTION_MODE_POLICY, + # Cosmos Framework 1.2 renamed the internal policy transform mode to + # ``wam``. The public vLLM action_mode remains ``policy``. + "mode": "wam", "domain_id": torch.tensor(domain_id, dtype=torch.long), "viewpoint": "concat_view", "additional_view_description": ROBOLAB_CONCAT_VIEW_DESCRIPTION, @@ -1514,7 +1525,9 @@ class Cosmos3OmniDiffusersPipeline( if history_action is not None: sample["history_action"] = history_action - sample = self._get_robolab_transform()(sample, resolution) + sample = self._get_robolab_transform(format_prompt_as_json=format_prompt_as_json)(sample, resolution) + if isinstance(sample.get("ai_caption"), dict): + sample["ai_caption"] = json.dumps(sample["ai_caption"]) sequence_plan = sample["sequence_plan"] video_tensor = sample["video"].float() / 127.5 - 1.0 raw_action_dim_tensor = sample.get("raw_action_dim") diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py index 76df64b3b..8d104e987 100644 --- a/vllm_omni/diffusion/models/cosmos3/utils.py +++ b/vllm_omni/diffusion/models/cosmos3/utils.py @@ -78,6 +78,12 @@ ROBOLAB_CONCAT_VIEW_DESCRIPTION = ( "sides, with the robot visible." ) +_LEGACY_EMBODIMENT_ALIASES = { + "galbot": "embodiment_b", + "agibot_gear_gripper": "embodiment_c_gripper", + "agibot_gear_gripper_ext": "embodiment_c_gripper_ext", +} + @dataclass(frozen=True) class RoboLabPolicyInputs: @@ -252,33 +258,83 @@ def extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None: return ensure_rgb_uint8_image(image, "multi_modal_data.image") +_ROBOLAB_FRAMEWORK_IMPORT_ERROR = ( + "Cosmos3 RoboLab policy serving requires a compatible cosmos_framework source on PYTHONPATH." +) +_ROBOLAB_FRAMEWORK_IMPORTS: dict[str, tuple[str, str]] = { + "ActionTransformPipeline": ( + "cosmos_framework.data.generator.action.utils.transforms", + "ActionTransformPipeline", + ), + "convert_rotation": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "convert_rotation", + ), + "pose_abs_to_rel": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "pose_abs_to_rel", + ), + "pose_rel_to_abs": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "pose_rel_to_abs", + ), + "build_abs_pose_from_components": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "build_abs_pose_from_components", + ), + "FlowUniPCMultistepScheduler": ( + "cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc", + "FlowUniPCMultistepScheduler", + ), + "get_domain_id": ( + "cosmos_framework.data.generator.action.utils.domain_utils", + "get_domain_id", + ), +} + + def lazy_import(module_name: str, symbol_name: str, error_message: str): try: module = import_module(module_name) - except ModuleNotFoundError as exc: - raise ModuleNotFoundError(error_message) from exc - return getattr(module, symbol_name) + return getattr(module, symbol_name) + except (ImportError, AttributeError) as exc: + raise ImportError( + f"{error_message} Could not import required symbol {symbol_name!r} from {module_name!r}." + ) from exc -def lazy_action_transform_pipeline(max_action_dim: int): - ActionTransformPipeline = lazy_import( - "cosmos_framework.data.vfm.action.transforms", - "ActionTransformPipeline", - "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the " - "golden ActionTransformPipeline can be reused.", +def _load_robolab_framework_symbol(symbol_name: str): + module_name, exported_name = _ROBOLAB_FRAMEWORK_IMPORTS[symbol_name] + return lazy_import( + module_name, + exported_name, + _ROBOLAB_FRAMEWORK_IMPORT_ERROR, ) - return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0) + + +def preflight_cosmos3_action_framework_imports() -> None: + """Validate every Cosmos Framework symbol used by action policy serving.""" + for symbol_name in _ROBOLAB_FRAMEWORK_IMPORTS: + _load_robolab_framework_symbol(symbol_name) + + +def lazy_action_transform_pipeline(max_action_dim: int, *, format_prompt_as_json: bool = False): + ActionTransformPipeline = _load_robolab_framework_symbol("ActionTransformPipeline") + return ActionTransformPipeline( + max_action_dim=max_action_dim, + cfg_dropout_rate=0.0, + format_prompt_as_json=format_prompt_as_json, + ) + + +def get_robolab_domain_id(domain_name: str) -> int: + get_domain_id = _load_robolab_framework_symbol("get_domain_id") + canonical_name = _LEGACY_EMBODIMENT_ALIASES.get(domain_name.lower().strip(), domain_name) + return int(get_domain_id(canonical_name)) def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device): - FlowUniPCMultistepScheduler = lazy_import( - "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc", - "FlowUniPCMultistepScheduler", - ( - "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the " - "golden FlowUniPCMultistepScheduler can be reused." - ), - ) + FlowUniPCMultistepScheduler = _load_robolab_framework_symbol("FlowUniPCMultistepScheduler") scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, @@ -290,38 +346,22 @@ def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.de def convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray: - convert_rotation = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "convert_rotation", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + convert_rotation = _load_robolab_framework_symbol("convert_rotation") return convert_rotation(value, src, dst) def pose_abs_to_rel(*args, **kwargs) -> np.ndarray: - pose_abs_to_rel_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "pose_abs_to_rel", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + pose_abs_to_rel_func = _load_robolab_framework_symbol("pose_abs_to_rel") return pose_abs_to_rel_func(*args, **kwargs) def pose_rel_to_abs(*args, **kwargs) -> np.ndarray: - pose_rel_to_abs_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "pose_rel_to_abs", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + pose_rel_to_abs_func = _load_robolab_framework_symbol("pose_rel_to_abs") return pose_rel_to_abs_func(*args, **kwargs) def build_abs_pose_from_components(*args, **kwargs) -> np.ndarray: - build_abs_pose_from_components_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "build_abs_pose_from_components", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + build_abs_pose_from_components_func = _load_robolab_framework_symbol("build_abs_pose_from_components") return build_abs_pose_from_components_func(*args, **kwargs) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 348e41316..598cbe893 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -930,6 +930,15 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Disable Cosmos3 text/video safety guardrails for this server.", ) + omni_config_group.add_argument( + "--robot-openpi-idle-timeout", + type=_nonneg_finite_float, + default=30.0, + help=( + "Seconds the /v1/realtime/robot/openpi endpoint waits for the next request " + "before closing an idle WebSocket (default: 30). Set to 0 to disable the timeout." + ), + ) # Enable diffusion pipeline profiling omni_config_group.add_argument( @@ -1011,6 +1020,8 @@ def run_headless(args: TrackingNamespace) -> None: # Filter down to a dict of things explicitly requested by the user args_dict = args.get_explicit_kwargs_dict() + # This controls only the API-process WebSocket and is not a stage-engine option. + args_dict.pop("robot_openpi_idle_timeout", None) # This CLI-only negative alias is consumed below when selecting the # launcher log_stats value; it is not a per-stage config override. args_dict.pop("disable_log_stats", None) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index adb7d6e9d..565362312 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -485,6 +485,8 @@ async def build_async_omni_from_stage_config( try: kwargs = args.get_explicit_kwargs_dict() + # This controls only the API-process WebSocket and is not an AsyncOmni option. + kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) async_omni = AsyncOmni(model=model, **kwargs) @@ -1494,11 +1496,24 @@ async def realtime_robot_openpi(websocket: WebSocket): serving = getattr(websocket.app.state, "openai_serving_realtime_robot", None) if serving is None: + logger.warning( + "Rejecting robot OpenPI WebSocket: policy serving is disabled. " + "The diffusion stage must provide model_config.policy_server_config." + ) await websocket.accept() - await websocket.send_json({"type": "error", "error": "Robot policy not available", "code": "unsupported"}) + await websocket.send_json( + { + "type": "error", + "error": "Robot policy not available", + "code": "unsupported", + } + ) await websocket.close() return - connection = RobotRealtimeConnection(websocket, serving) + state_args = getattr(websocket.app.state, "args", None) + configured_timeout = getattr(state_args, "robot_openpi_idle_timeout", 30.0) + idle_timeout = None if configured_timeout == 0 else configured_timeout + connection = RobotRealtimeConnection(websocket, serving, idle_timeout=idle_timeout) await connection.handle_connection() diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index 9b8bfe0d8..112254fbe 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -29,6 +29,7 @@ mappings. from __future__ import annotations import asyncio +from collections import OrderedDict from typing import Any import msgspec @@ -45,10 +46,12 @@ logger = init_logger(__name__) _DEFAULT_IDLE_TIMEOUT = 30.0 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024 _MISSING = object() +# Upper bound on the per-connection set of seen session ids. +MAX_TRACKED_SESSIONS = 1024 def _pack_numpy(obj: Any) -> Any: - if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"): + if isinstance(obj, np.ndarray | np.generic) and obj.dtype.kind in ("V", "O", "c"): raise ValueError(f"Unsupported dtype: {obj.dtype}") if isinstance(obj, np.ndarray): if not obj.flags.c_contiguous: @@ -165,17 +168,38 @@ class RobotRealtimeConnection: self, websocket: WebSocket, serving: ServingRealtimeRobotOpenPI, - idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + idle_timeout: float | None = _DEFAULT_IDLE_TIMEOUT, ) -> None: self.websocket = websocket self.serving = serving self._idle_timeout = idle_timeout self._current_session_id: str | None = None - self._call_count = 0 + # Session ids seen on this connection, most-recently-used last. + self._seen_sessions: OrderedDict[str, None] = OrderedDict() def reset(self) -> None: self._current_session_id = None - self._call_count = 0 + self._seen_sessions.clear() + + def _mark_session_seen(self, session_id: str) -> bool: + """Record ``session_id``; return True when it is the first sighting. + + A session's first observation is the one that carries ``reset`` to the + policy, so this needs per-session memory. + """ + if session_id in self._seen_sessions: + self._seen_sessions.move_to_end(session_id) + return False + + self._seen_sessions[session_id] = None + if len(self._seen_sessions) > MAX_TRACKED_SESSIONS: + evicted, _ = self._seen_sessions.popitem(last=False) + logger.debug( + "Robot OpenPI connection is tracking more than %d sessions; dropped least-recently-used %s", + MAX_TRACKED_SESSIONS, + evicted, + ) + return True async def _send_error(self, message: str) -> None: await self.websocket.send_bytes(_pack({"type": "error", "message": message})) @@ -199,18 +223,22 @@ class RobotRealtimeConnection: await self.websocket.send_bytes(_pack(metadata)) while True: - try: - msg = await asyncio.wait_for( - self.websocket.receive(), - timeout=self._idle_timeout, - ) - except asyncio.TimeoutError: - logger.info("Robot OpenPI connection idle timeout after %.1f seconds", self._idle_timeout) + idle_timeout = self._idle_timeout + if idle_timeout is None: + msg = await self.websocket.receive() + else: try: - await self.websocket.close() - except Exception: - logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True) - return + msg = await asyncio.wait_for( + self.websocket.receive(), + timeout=idle_timeout, + ) + except asyncio.TimeoutError: + logger.info("Robot OpenPI connection idle timeout after %.1f seconds", idle_timeout) + try: + await self.websocket.close() + except Exception: + logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True) + return if msg.get("type") == "websocket.disconnect": break @@ -245,13 +273,12 @@ class RobotRealtimeConnection: session_id, ) self._current_session_id = session_id - self._call_count = 0 - self._call_count += 1 + reset = self._mark_session_seen(session_id) actions = await self.serving.infer( obs, session_id=session_id, - reset=self._call_count <= 1, + reset=reset, ) await self.websocket.send_bytes(_pack(actions)) except Exception: diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index f2448dfa1..80a3839ab 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -148,16 +148,32 @@ class ServingRealtimeRobotOpenPI: `AsyncOmni.generate()` and routed to the diffusion stage. """ from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.entrypoints.openai.stage_params import ( + clone_sampling_params, + get_default_sampling_params_list, + ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams - # An optional integer ``seed`` in the inference message becomes the engine request's - # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one. + # The engine applies stage default_sampling_params only to requests + # that carry no explicit params; this endpoint always passes explicit + # params, so start from a clone of the diffusion stage's defaults + # (e.g. a policy deploy yaml's ``extra_args``) and layer the OpenPI + # protocol fields on top. seed = obs.pop("seed", None) - extra_args = { - "reset": reset, - "session_id": session_id, - "robot_obs": obs, - } + sampling_params = OmniDiffusionSamplingParams() + for default_params in get_default_sampling_params_list(self.engine_client): + if isinstance(default_params, OmniDiffusionSamplingParams): + sampling_params = clone_sampling_params(default_params) + break + + extra_args = sampling_params.extra_args or {} + extra_args.update( + { + "reset": reset, + "session_id": session_id, + "robot_obs": obs, + } + ) prompt = obs.get("prompt", "") sampling_params = OmniDiffusionSamplingParams( diff --git a/vllm_omni/model_executor/models/cosmos3/__init__.py b/vllm_omni/model_executor/models/cosmos3/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/cosmos3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/cosmos3/pipeline.py b/vllm_omni/model_executor/models/cosmos3/pipeline.py new file mode 100644 index 000000000..d26b4f46c --- /dev/null +++ b/vllm_omni/model_executor/models/cosmos3/pipeline.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 policy single-stage topology for online OpenPI serving. + +Cosmos3 policy checkpoints (e.g. ``nvidia/Cosmos3-Nano-Policy-DROID``) are one +diffusion stage: robot observation -> action chunk. They are +registered in ``OMNI_PIPELINES`` because online OpenPI serving +(``/v1/realtime/robot/openpi``) requires a registered pipeline so the deploy +yaml's ``model_config.policy_server_config`` reaches the websocket handshake. + +This pipeline declares neither ``hf_architectures`` nor +``diffusers_class_name``: policy checkpoints share their HF metadata +(``model_type=cosmos3_omni``, ``model_index.json`` ``_class_name= +Cosmos3OmniDiffusersPipeline``) with the T2I/video Cosmos3 checkpoints, which +must keep resolving through the default single-stage diffusion fallback. +Select this pipeline explicitly via a deploy yaml ``pipeline:`` key, e.g. +``vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni --deploy-config +/absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml``. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +COSMOS3_POLICY_PIPELINE = PipelineConfig( + model_type="cosmos3_policy", + model_arch="Cosmos3OmniDiffusersPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Cosmos3OmniDiffusersPipeline", + ), + ), +) diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py index 01cecf94f..452946b08 100644 --- a/vllm_omni/model_extras/cosmos3.py +++ b/vllm_omni/model_extras/cosmos3.py @@ -35,6 +35,7 @@ COSMOS3_EXTRA_BODY_PARAMS = frozenset( "resolution", "image_size", "use_state", + "format_prompt_as_json", "observation", "robot_obs", "deterministic_seed", From 627630d7e4725d94c84b64e9996f97bdec2044fe Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:11:27 +0800 Subject: [PATCH 011/209] [BugFix][CI] Restore diff-aware source filtering for post-merge L3 (#7371) Signed-off-by: wangyu <410167048@qq.com> --- .../common/ci_source_file_dependencies.yml | 16 +++++++---- .buildkite/common/scripts/upload_pipeline.py | 27 +++++++++---------- .buildkite/cuda/bootstrap-upload-steps.yml | 11 +++----- .buildkite/cuda/test-merge.yml | 3 ++- .buildkite/cuda/test-nightly.yml | 1 + docs/contributing/ci/ci_settings.md | 2 +- tests/buildkite/test_upload_pipeline.py | 24 +++++++++-------- 7 files changed, 45 insertions(+), 39 deletions(-) diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 0455922c8..724794c95 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -11,6 +11,8 @@ # YAML anchors are nested lists; the loader flattens them. x-qwen3-tts: &qwen3_tts + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/worker/gpu_model_runner.py - vllm_omni/model_executor/models/qwen3_tts/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -23,6 +25,8 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/engine/stage_pool.py x-minicpmo-4-5: &minicpmo_4_5 + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/model_executor/models/minicpmo_4_5/ - vllm_omni/model_executor/models/step_audio2/ - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ @@ -43,7 +47,6 @@ x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/core/sched/ - vllm_omni/outputs/multimodal_accumulation.py - vllm_omni/engine/stage_init_utils.py - - vllm_omni/engine/async_omni_engine.py - vllm_omni/engine/orchestrator.py - vllm_omni/entrypoints/cli/serve.py - vllm_omni/entrypoints/openai/api_server.py @@ -51,11 +54,15 @@ x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/config/stage_config.py x-qwen2-5-omni: &qwen2_5_omni + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/model_executor/models/qwen2_5_omni/ - vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py - vllm_omni/deploy/qwen2_5_omni.yaml x-qwen3-omni: &qwen3_omni + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -116,11 +123,9 @@ x-lingbot-world: &lingbot_world source_file_dependencies: - # Shared uploader/registry paths. Not a job key — do not attach this to a step. - # Used only when no job-key prefix matches: then a hit keeps every job. + # Hardware presets and shared test helpers. Not a job key — do not attach this + # to a step. Used only when no job-key prefix matches: then a hit keeps every job. source_filter_fallback: - - .buildkite/common/scripts/upload_pipeline.py - - .buildkite/common/ci_source_file_dependencies.yml - .buildkite/common/ci_mirror_hardwares.yml - tests/helpers/ @@ -280,6 +285,7 @@ source_file_dependencies: - *wan22 - tests/e2e/offline_inference/test_wan22_t2v.py - tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py + - tests/e2e/offline_inference/test_device_postprocess_equivalence.py - tests/e2e/online_serving/test_wan22_t2v.py - tests/e2e/online_serving/test_wan22_expansion.py diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index c8587d8c3..71ed75119 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -13,12 +13,14 @@ Bootstrap mode (``bootstrap-upload-steps.yml``): Test pipeline mode (e.g. test-merge.yml, test-nightly.yml, test-weekly.yml): - Drop steps whose ``source_file_dependencies`` do not match changed files (string/list keys from ci_source_file_dependencies.yml, or inline path - prefixes). Filtering applies on **PR label** uploads only; ``main`` + env - schedule (NIGHTLY/WEEKLY/merge push) keeps every job and still strips the field. - If no job-key prefix matches, a change to the pipeline YAML being uploaded - or to a path under the ``source_filter_fallback`` registry key keeps every - job so command, env, and hardware edits can be validated before merge. - When any job-key prefix already matches, normal filtering wins. + prefixes). Filtering applies on **PR label** uploads and on post-merge + ``main`` L3 uploads. Scheduled ``main`` uploads that pass ``--all`` + (NIGHTLY/WEEKLY) or ``--e2e`` keep every selected job and still strip + the field. If no job-key prefix matches, a change to the pipeline YAML + being uploaded or to a path under the ``source_filter_fallback`` registry + key keeps every job so command, env, and hardware edits can be validated + before merge. When any job-key prefix already matches, normal filtering + wins. - Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image`` for NPU) + ``plugins`` (see ci_mirror_hardwares.yml). - Omit ``mirror_hardwares`` to compose ``{chip}_{n}`` from pytest ``-m`` SKU @@ -780,17 +782,14 @@ def _changed_files_for_source_filter( ) -> list[str] | None: """Return the diff to filter against, or None to keep every step. - ``source_file_dependencies`` is label-only: scheduled ``main`` uploads - (NIGHTLY/WEEKLY/post-merge) run the full pipeline. ``--all`` / ``--e2e`` - also disable filtering. Pipeline YAML / ``source_filter_fallback`` matches - are applied later in ``_render_test_pipeline`` only when no job-key prefix - already matched. + ``--all`` / ``--e2e`` disable filtering (scheduled NIGHTLY/WEEKLY full + uploads and weekly E2E sweeps). Post-merge ``main`` L3 and PR-label + uploads filter against the commit/PR diff. Pipeline YAML / + ``source_filter_fallback`` matches are applied later in + ``_render_test_pipeline`` only when no job-key prefix already matched. """ if force_all or e2e_only: return None - if os.environ.get("BUILDKITE_BRANCH", "") == "main": - _log("main branch: keep all jobs (source_file_dependencies is label-only)") - return None return ctx.changed_files diff --git a/.buildkite/cuda/bootstrap-upload-steps.yml b/.buildkite/cuda/bootstrap-upload-steps.yml index 502a506e6..3f5aa1cb1 100644 --- a/.buildkite/cuda/bootstrap-upload-steps.yml +++ b/.buildkite/cuda/bootstrap-upload-steps.yml @@ -22,19 +22,16 @@ steps: agents: queue: "cpu_queue_premerge" - # L3 Test — PR merge-test label (diff-aware) or main (full pipeline / WEEKLY --e2e) + # L3 Test — PR merge-test label or post-merge on main (both diff-aware); + # main+WEEKLY=1 uploads the E2E group only (--e2e). - label: "Upload Merge Pipeline" depends_on: image-build key: upload-merge-pipeline commands: - | EXTRA="" - if [ "${BUILDKITE_BRANCH}" = "main" ]; then - if [ "${WEEKLY:-}" = "1" ]; then - EXTRA="--e2e" - else - EXTRA="--all" - fi + if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--e2e" fi python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml agents: diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 1a458bccf..d600e8fbe 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -1,5 +1,6 @@ # CUDA CI L3 (`advanced_model` e2e) plus L1 CPU Simple Test (`core_model and cpu`). -# Uploaded by Upload Merge Pipeline: PR label `merge-test`, or post-merge on main. +# Uploaded by Upload Merge Pipeline: PR label `merge-test`, or post-merge on main +# (both diff-aware via source_file_dependencies). env: VLLM_WORKER_MULTIPROC_METHOD: spawn diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 583f935dc..cfaca90fe 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -390,6 +390,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 Device-Postprocess Equivalence Test" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - | diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index fb3af0ed3..9931a418d 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -313,7 +313,7 @@ Unit coverage: `tests/buildkite/test_skip_ci.py`. CUDA **L2** (`.buildkite/cuda/test-ready.yml`), **L3** (`.buildkite/cuda/test-merge.yml`), **L4** (`.buildkite/cuda/test-nightly.yml`), **L5** (`.buildkite/cuda/test-weekly.yml`), and NPU **L4** (`.buildkite/npu/test-npu-nightly.yml`). Bootstrap upload entry: `upload_pipeline.py --upload .buildkite/cuda/bootstrap-upload-steps.yml` (or `npu/bootstrap-upload-steps.yml`). -**When filtering applies:** only **PR label** uploads (`ready`, `merge-test`, `nightly-test`, `weekly-test`). On **`main` + env** (`NIGHTLY=1`, post-merge L3, `WEEKLY=1` / `NON_CRITICAL=1`, `WEEKLY=1 --e2e`) the uploader keeps every job and still strips the key. If **no** job-key prefix matches the diff, a change to the pipeline YAML being uploaded, or any path under the `source_filter_fallback` key in `ci_source_file_dependencies.yml`, keeps every job in that upload. `source_filter_fallback` is not a job key — do not attach it to a step. When any job-key prefix already matches, normal filtering wins (the fallback does not expand the selection). Without this fallback, command / env / hardware-only edits would match nothing (NPU drops every step; CUDA nightly keeps only the email aggregator). A different pipeline YAML does not trigger the fallback for this upload. +**When filtering applies:** **PR label** uploads (`ready`, `merge-test`, `nightly-test`, `weekly-test`) and **post-merge `main` L3** (`test-merge.yml`). Scheduled **`main` + env** uploads that pass `--all` (`NIGHTLY=1`, `WEEKLY=1` / `NON_CRITICAL=1`) or `--e2e` (`WEEKLY=1` on L2/L3) keep every selected job and still strip the key. If **no** job-key prefix matches the diff, a change to the pipeline YAML being uploaded, or any path under the `source_filter_fallback` key in `ci_source_file_dependencies.yml`, keeps every job in that upload. `source_filter_fallback` is not a job key — do not attach it to a step. When any job-key prefix already matches, normal filtering wins (the fallback does not expand the selection). **Uploader-only keys** — removed before Buildkite sees the YAML; never used at runtime on agents: diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 06e2b975c..a3192dada 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -692,19 +692,24 @@ def test_source_filter_strips_deps_and_expands_hardware(monkeypatch: pytest.Monk assert z_image["agents"]["queue"] == "l4-k8s" -def test_source_filter_disabled_on_main_branch(monkeypatch: pytest.MonkeyPatch) -> None: +def test_source_filter_respects_force_all_and_uses_diff_on_main( + monkeypatch: pytest.MonkeyPatch, +) -> None: class _Ctx: changed_files = ["vllm_omni/unrelated.py"] + # Post-merge main L3 must still filter by the commit diff. monkeypatch.setenv("BUILDKITE_BRANCH", "main") - assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) is None + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) == [ + "vllm_omni/unrelated.py", + ] monkeypatch.setenv("BUILDKITE_BRANCH", "feat/source-filter") assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) == [ "vllm_omni/unrelated.py", ] - monkeypatch.setenv("BUILDKITE_BRANCH", "main") assert _changed_files_for_source_filter(_Ctx(), force_all=True, e2e_only=False) is None + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=True) is None def test_source_filter_fallback_is_fallback_when_no_job_key_matches( @@ -715,7 +720,6 @@ def test_source_filter_fallback_is_fallback_when_no_job_key_matches( monkeypatch.setenv("BUILDKITE_BRANCH", "feat/nightly-yaml") pipeline_yaml = Path(".buildkite/npu/test-npu-nightly.yml") shared_paths = _load_source_file_dependencies()["source_filter_fallback"] - assert ".buildkite/common/scripts/upload_pipeline.py" in shared_paths # Synthetic steps only — do not pin live Buildkite job labels. doc = { @@ -742,7 +746,7 @@ def test_source_filter_fallback_is_fallback_when_no_job_key_matches( ) return {step["key"] for step in _iter_steps(rendered) if isinstance(step.get("key"), str)} - # Only pipeline YAML / shared uploader → no listed prefix match → keep every step. + # Only pipeline YAML / fallback paths → no listed prefix match → keep every step. for changed in [pipeline_yaml.as_posix(), *shared_paths]: assert surviving_keys([changed]) == {"ungated", "gated_a", "gated_b"}, changed @@ -750,9 +754,7 @@ def test_source_filter_fallback_is_fallback_when_no_job_key_matches( assert surviving_keys([".buildkite/cuda/test-nightly.yml"]) == {"ungated"} # A matching source prefix wins over fallback files in the same diff. - assert surviving_keys( - [ - "pkg/model_a/transformer.py", - ".buildkite/common/scripts/upload_pipeline.py", - ], - ) == {"ungated", "gated_a"} + assert surviving_keys(["pkg/model_a/transformer.py", *shared_paths[:1]]) == { + "ungated", + "gated_a", + } From ffcaaa9439823d9b927d201397bb9d4db7218f14 Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:45:28 +0800 Subject: [PATCH 012/209] [Bugfix] Fail when a diffusion LoRA adapter binds no layer (#7349) Signed-off-by: Guangjian --- tests/diffusion/lora/test_lora_manager.py | 44 +++++++++++++++++++++++ vllm_omni/diffusion/lora/manager.py | 13 +++++-- 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 0814bf15f..10b023bf0 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -411,6 +411,50 @@ def test_lora_manager_rolls_back_all_layers_when_activation_fails(): assert len(first.set_calls) == first_calls_after_success + 2 +def test_lora_manager_rejects_adapter_that_binds_no_layer(): + """An adapter whose target modules match nothing must fail, not silently no-op.""" + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + layer = _DummyLoRALayer(n_slices=1, output_slices=(2,)) + manager._lora_modules = {"transformer.blocks.0.attn.to_q": layer} + + # Adapter-side name the engine does not expose, e.g. a diffusers-style + # checkpoint against a differently named engine layout. + mismatched = LoRALayerWeights( + module_name="unet.down_blocks.0.attn.to_q", + rank=2, + lora_alpha=2, + lora_a=torch.ones((2, 2)), + lora_b=torch.ones((2, 2)), + ) + manager._registered_adapters = { + 3: type( + "LM", + (), + { + "id": 3, + "loras": {"unet.down_blocks.0.attn.to_q": mismatched}, + "get_lora": lambda self, key: self.loras.get(key), + }, + )() + } + + with pytest.raises(ValueError, match="applies to no layer") as excinfo: + manager._activate_adapter(3, scale=1.0) + + # The message must name what was received so the mismatch is diagnosable. + assert "unet.down_blocks.0.attn.to_q" in str(excinfo.value) + + # Nothing was bound and the adapter must not be left marked active. + assert manager._active_adapter_id is None + assert layer.set_calls == [] + assert layer.reset_calls >= 1 + + def _dummy_lora_request(adapter_id: int) -> LoRARequest: return LoRARequest( lora_name=f"adapter_{adapter_id}", diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 2f6d47780..11d193531 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -556,9 +556,9 @@ class DiffusionLoRAManager: def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None: binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None) - lora_names_by_id = ( - {id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {} - ) + # Track bindings unconditionally. The zero-binding guard below needs this + # bookkeeping on every pipeline, not only the ones that supply a validator. + lora_names_by_id = {id(weights): name for name, weights in lora_model.loras.items()} bound_lora_names: set[str] = set() def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None: @@ -677,6 +677,13 @@ class DiffusionLoRAManager: scale, ) + if not bound_lora_names: + raise ValueError( + f"LoRA adapter {lora_model.id} applies to no layer: expected target modules in " + f"{sorted(self._expected_lora_modules)} but received {sorted(lora_model.loras)}. " + "Activating it would leave the base model unchanged." + ) + if callable(binding_validator): binding_validator( lora_model=lora_model, From 5a386d789f7646f42e7c97fe873848de9dd400aa Mon Sep 17 00:00:00 2001 From: summer <128961079+zhang-keliang@users.noreply.github.com> Date: Fri, 11 Sep 2026 08:04:08 +0800 Subject: [PATCH 013/209] [Bugfix] Fix host-memory leak on aborted /v1/images/generations (#6462) (#6561) Signed-off-by: summer <128961079+zhang-keliang@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/entrypoints/test_async_omni.py | 89 +++++++++++++++++++ tests/metrics/test_emit_calls.py | 3 +- tests/metrics/test_prometheus.py | 2 - vllm_omni/entrypoints/client_request_state.py | 6 ++ vllm_omni/entrypoints/omni_base.py | 18 +--- 5 files changed, 100 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 2549889b3..a69bcce09 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -329,6 +329,95 @@ def test_abort_enqueues_synthetic_finished_when_engine_returns_empty(): asyncio.run(run()) +@pytest.mark.cpu +def test_abort_drops_consumed_metric_message_ids_with_state(): + """Exercise the production leak path from #6462: a metric-bearing + ``OutputMessage`` goes through ``_handle_output_message`` (which populates + the per-request de-dup set), the request is aborted (which, since #6367, + keeps ``request_states`` registered and enqueues a terminal abort output), + and the set must be released with the state by ``generate()``'s normal + cleanup — with no independent request-keyed metric state surviving on the + instance at any point. A reintroduced + class-level ``_consumed_metric_messages``-style map populated by the + production handler fails the sweep below. + """ + import time as _time + from collections.abc import Mapping + + from vllm_omni.engine.messages import OutputMessage + from vllm_omni.entrypoints.client_request_state import ClientRequestState + from vllm_omni.metrics.stats import OrchestratorAggregator + + async def run(): + omni = get_async_omni_instance() + omni.engine.get_stage_metadata = lambda stage_id: SimpleNamespace( + stage_type="llm", + final_output=True, + final_output_type="text", + ) + + rid = "req-1-cccc" + state = ClientRequestState(request_id=rid, external_request_id="req-1") + state.metrics = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=_time.time(), + final_stage_id_for_e2e=0, + ) + state.input_stream_task = None + omni.request_states[rid] = state + + msg = OutputMessage( + request_id=rid, + stage_id=0, + engine_outputs=SimpleNamespace(final_output_type="text"), + metrics=SimpleNamespace( + num_tokens_in=3, + num_tokens_out=5, + stage_gen_time_ms=1.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + ), + finished=False, + ) + handled, out_rid, out_stage, out_state = omni._handle_output_message(msg) + assert handled is False and out_state is state + # The production handler recorded the de-dup entry on the state... + assert id(msg) in state.consumed_metric_message_ids + # ...and de-duplicates a replay of the same message object. + omni._handle_output_message(msg) + assert len(state.consumed_metric_message_ids) == 1 + + assert not hasattr(omni, "_consumed_metric_messages") + + # #6367 contract: abort() keeps the state registered and enqueues a + # terminal abort output; generate()'s ``finally`` owns the cleanup. + await omni.abort("req-1") + assert rid in omni.request_states + terminal = state.queue.get_nowait() + assert terminal.finished is True + assert terminal.engine_outputs.outputs[0].finish_reason == "abort" + # The de-dup set still lives only on the (still-registered) state. + for name, value in vars(omni).items(): + if name == "request_states": + continue + assert not (isinstance(value, Mapping) and rid in value), ( + f"request-keyed metric state kept outside request_states in {name!r}" + ) + + # generate()'s normal cleanup releases the state and the set with it. + omni._log_summary_and_cleanup(rid) + assert rid not in omni.request_states + for name, value in vars(omni).items(): + assert not (isinstance(value, Mapping) and rid in value), ( + f"request-keyed state survived cleanup in {name!r}" + ) + assert not hasattr(omni, "_consumed_metric_messages") + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index b46f4080f..02a98fc00 100644 --- a/tests/metrics/test_emit_calls.py +++ b/tests/metrics/test_emit_calls.py @@ -41,7 +41,6 @@ def _make_omni_base_with_mock_prom(mocker): obj = object.__new__(OmniBase) obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = True return obj, obj.prom_metrics @@ -270,7 +269,7 @@ class TestStageWorkloadMetricScope: obj = object.__new__(OmniBase) obj._enable_ar_profiler = False - obj._consumed_metric_messages = {} + obj.request_states = {"req-replay": SimpleNamespace(consumed_metric_message_ids=set())} obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) obj.mod_metrics = mocker.Mock() obj.engine = SimpleNamespace( diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index e45365cf1..ca8dd8d5e 100644 --- a/tests/metrics/test_prometheus.py +++ b/tests/metrics/test_prometheus.py @@ -154,7 +154,6 @@ class TestRequestLifecycleGauges: obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test") obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = False # Simulate request lifecycle: start (counter 0→1, dict {} → {req}), @@ -183,7 +182,6 @@ class TestRequestLifecycleGauges: obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2") obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = False # Two in flight, one finalizes — running should report 1, waiting 0. diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 0aa0c7c5a..22f1c6d5e 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -39,3 +39,9 @@ class ClientRequestState: # without re-querying stage_pools. self.audio_emit_stage_id: int | None = None self.audio_emit_replica_id: int | None = None + # De-dup set for metric messages: OmniBase populates this in + # ``_handle_output_message`` / ``_process_single_result`` so the same + # ``id(msg)`` isn't counted twice into per-request metrics. Kept on + # the request state (not a class-level dict) so it is released with + # the state — see #6462 / #6561. + self.consumed_metric_message_ids: set[int] = set() diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 264b37fb5..899ff116d 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -231,7 +231,6 @@ class OmniBase(PDDisaggregationMixin): self.async_chunk = bool(getattr(self.engine, "async_chunk", False)) self.request_states: dict[str, ClientRequestState] = {} - self._consumed_metric_messages: dict[str, set[int]] = {} self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list @@ -298,13 +297,6 @@ class OmniBase(PDDisaggregationMixin): """True when a non-empty stage pool has lost all of its replicas.""" return len(pool.clients) > 0 and self._live_replica_count(pool) == 0 - def _consumed_metric_message_ids(self, request_id: str) -> set[int]: - consumed_by_request = getattr(self, "_consumed_metric_messages", None) - if consumed_by_request is None: - consumed_by_request = {} - self._consumed_metric_messages = consumed_by_request - return consumed_by_request.setdefault(request_id, set()) - @property def is_running(self) -> bool: return self.engine.is_alive() @@ -449,9 +441,6 @@ class OmniBase(PDDisaggregationMixin): ) finally: self.request_states.pop(request_id, None) - consumed_by_request = getattr(self, "_consumed_metric_messages", None) - if consumed_by_request is not None: - consumed_by_request.pop(request_id, None) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. @@ -532,7 +521,7 @@ class OmniBase(PDDisaggregationMixin): stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type) msg_id = id(msg) - consumed = self._consumed_metric_message_ids(req_id) + consumed = req_state.consumed_metric_message_ids if msg_id not in consumed: req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type) submit_ts = msg.stage_submit_ts @@ -642,8 +631,9 @@ class OmniBase(PDDisaggregationMixin): output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) if finished and _m is not None: msg_id = id(result) - consumed = self._consumed_metric_message_ids(req_id) - if msg_id not in consumed: + req_state = self.request_states.get(req_id) + consumed = req_state.consumed_metric_message_ids if req_state is not None else None + if consumed is not None and msg_id not in consumed: metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id) From 5c69fc2a8a6bca8f87505d1e0cad112daa5941e4 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:44:32 -0700 Subject: [PATCH 014/209] [Refactor] Declare model-local KV held outside the paged manager (#6171) Signed-off-by: Yueqian Lin --- docs/design/feature/model_local_kv_caches.md | 132 ++++++ docs/design/index.md | 1 + .../models/test_model_local_kv.py | 404 ++++++++++++++++++ tests/tools/test_check_tts_adapter.py | 12 +- .../models/mimo_audio/mimo_audio_llm.py | 67 +++ .../ming_flash_omni/ming_flash_omni_talker.py | 11 + .../models/ming_flash_omni/talker_module.py | 58 ++- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 52 +++ .../model_executor/models/model_local_kv.py | 347 +++++++++++++++ .../qwen3_tts/segmented_graph_wrapper.py | 86 ++++ .../modeling_qwen3_tts_tokenizer_v2.py | 69 +++ vllm_omni/worker/gpu_model_runner.py | 47 ++ 12 files changed, 1283 insertions(+), 3 deletions(-) create mode 100644 docs/design/feature/model_local_kv_caches.md create mode 100644 tests/model_executor/models/test_model_local_kv.py create mode 100644 vllm_omni/model_executor/models/model_local_kv.py diff --git a/docs/design/feature/model_local_kv_caches.md b/docs/design/feature/model_local_kv_caches.md new file mode 100644 index 000000000..371b989ff --- /dev/null +++ b/docs/design/feature/model_local_kv_caches.md @@ -0,0 +1,132 @@ +# Model-local KV caches + +Some models keep their attention KV as HuggingFace `transformers` cache objects +instead of using the engine's paged KV manager. Each such model declares what it +holds by implementing `model_local_kv_specs()`, and +`OmniGPUModelRunner.load_model` sums and logs the result after load. + +## What is and is not already accounted for + +This distinction decides what the protocol is for, and an earlier revision of +this page got it wrong by claiming none of this memory was accounted for. + +CUDA-graph capture happens *during* `load_model` -- MiMo in its constructor, +Qwen3-TTS in `load_weights`. `determine_available_memory()` then samples +per-process GPU memory through NVML and subtracts it from the requested budget +(`vllm_omni/worker/base.py`). So the **graph-resident** portion is already +charged. Subtracting the declared total from the KV budget would double-count +it. + +What is *not* charged is the **per-request** portion, which is allocated after +profiling. That half is not small: Qwen3-TTS's codec decoder is 4.44 MiB per +request, so a server at `max_num_seqs=256` reaches roughly 1.1 GiB that no +per-stage number predicted. + +The consumer therefore reports and does not subtract. + +## What the four models declare + +| Model | Cache | Bounded by | Per row | Rows | +|---|---|---|---|---| +| Qwen3-TTS codec decoder (async-chunk only) | sliding `DynamicCache` | `sliding_window - 1 = 71` | 4.44 MiB | `max_num_seqs`, x2 for the working copy | +| Qwen3-TTS graph pool | retained captures | same | 4.44 MiB | fixed, from captured shapes | +| MiMo-Audio local transformer | `DynamicCache` | `group_size + max(delay_pattern) = 11` | 704 KiB | `max_num_seqs` | +| MiMo-Audio graph pool | captured | same | 704 KiB | fixed 261 (bucket sum) | +| MiniCPM-o Whisper encoder (streaming only) | `EncoderDecoderCache` | `embed_positions` rows = 1500 | 140.62 MiB | fixed 1, per session | +| ming_flash_omni talker | `StaticCache` | hardcoded `max_cache_len = 2048` | 24.00 MiB | fixed 1 | + +**No cache is bounded by `max_model_len`.** Each has its own mechanism: a +sliding window, a decode-loop trip count, an encoder frame limit, a constant. + +**ming's geometry does not exist in this repo.** `self._llm_config` is a +`Qwen2Config` from the checkpoint, so layers, kv-heads, head-dim and dtype are +only knowable after load. That is what makes this a runtime query rather than a +table. + +## Lifetime is not multiplicity + +`scope` records when memory appears and goes. It does **not** decide how wide +the allocation is, and an earlier revision that derived scaling from it was +wrong in two different directions at once: + +- ming's talker is per-call but serialized -- `forward` takes + `runtime_additional_information[0]` and runs the whole AR loop inline -- so + scaling it by `max_num_seqs` over-reported it by that factor. +- MiniCPM-o's encoder is per *session*, not per sequence, so `max_num_seqs` is + simply the wrong number for it. + +So width is declared explicitly, by naming the driver. `RowDriver` has two +values, because two are all the known caches need: + +- `FIXED` -- a count the model controls, with `rows_fixed` and a required + `rows_reason`. `rows_fixed=1` on a cache that looks per-request is a claim a + reviewer must be able to check. +- `MAX_NUM_SEQS` -- one row per in-flight sequence; only the engine knows the + value. + +A third value belongs here when a model actually has a cache the engine widens +by some other number. An earlier revision added one speculatively, for a single +declarer, and paid for it with a class hierarchy, an engine-capacity object and +a conditional-annotation field. + +`rows` counts row-equivalents at peak and deliberately does not distinguish +"one allocation N rows wide" from "N allocations of one row". Those cost the +same, and a revision that tried to carry both got the object topology wrong in +three of four declarers: MiniCPM-o keeps one cache object per session while +asserting batch size 1, and both graph pools keep one distinct object per +captured bucket. The layout belongs in `allocation_note`, where it cannot be +mistaken for arithmetic. + +## Declaring nothing is a valid declaration + +A declaration must be inert when its path is. Qwen3-TTS holds KV only in +async-chunk mode -- stateless decoding runs `_forward_exact`, which calls the +transformer without `use_cache` -- so the stateless path declares nothing. +Getting this wrong is easy in both directions, and both have happened here: one +revision hung the declaration off the CUDA-graph wrapper and reported zero under +`enforce_eager`, where the eager async-chunk path does allocate; the next +declared unconditionally and invented roughly a gigabyte for stateless +deployments. + +Where the model cannot see the deciding setting, it declares the honest unit +and says so in `allocation_note`. MiniCPM-o's encoder cache exists only on the +streaming path and is one object per session; the session cap is not visible +from the model, so it declares one session rather than scaling by a number it +would have to guess at. + +## Finding the declarers + +`collect_model_local_kv_specs()` walks `named_modules()`. Two wrinkles: + +- The runner's `self.model` may be a `CUDAGraphWrapper`, which is not an + `nn.Module` but forwards attribute access to its runnable. The walk therefore + works either way, but entering through both the wrapper and its runnable would + count a root-level declarer twice. The collector calls the supported + `unwrap()` accessor first. +- An owner that is not a submodule at all (Qwen3-TTS's graph wrapper, ming's + `MingAudioGenerator`) is reached by one explicit hop from the module that + holds it. + +## Adding a model + +Return one entry per distinct allocation. Prefer `spec_from_hf_config()` and +pass the same config object the cache is built from; it raises rather than +guessing when an attribute is missing. `capacity_source` is diagnostic text -- +never branch on it. + +## Known limits + +Row counts taken from live state after warmup miss caches `transformers` has +not yet materialized, which makes Qwen3-TTS's graph-pool figure a floor for the +xvec path. + +MiMo's graph-pool cache has no reachable Python object -- it is a local inside +the captured frame -- so its width comes from the captured bucket list rather +than from tensors. Its two entries are also mutually exclusive per call: a call +either replays a bucket or runs eagerly. The sum is the right ceiling for the +process and the wrong number for a single request. + +MiMo captures every bucket gated only on `torch.cuda.is_available()` +(`mimo_audio_llm.py:670`), so those 179.44 MiB stay resident even when the +deploy config sets `enforce_eager: true`. Gating it is a one-line change that +needs a model boot to verify, so it is noted rather than bundled here. diff --git a/docs/design/index.md b/docs/design/index.md index cff0b1eb3..6481aa49b 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -26,6 +26,7 @@ implementation contract; it is not, by itself, a general support claim. - [Async Omni Output Materialization](feature/omni_async_output_materialization.md) - [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md) - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md) +- [Model-local KV Caches](feature/model_local_kv_caches.md) - [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md) ### Communication diff --git a/tests/model_executor/models/test_model_local_kv.py b/tests/model_executor/models/test_model_local_kv.py new file mode 100644 index 000000000..5e2c9cdca --- /dev/null +++ b/tests/model_executor/models/test_model_local_kv.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A declaration must match the cache that actually gets allocated. + +Two earlier versions of this file could not fail. The first asserted hand-typed +constants against each other. The second built real `transformers` caches but +compared them to specs written in the test, so no production declaration was +ever executed. What matters here is that a real declarer is called and its +numbers are checked against a real cache built the way the model builds it. +""" + +import ast +import pathlib +from types import SimpleNamespace + +import pytest +import torch +from transformers import Qwen2Config +from transformers.cache_utils import DynamicCache, StaticCache + +from vllm_omni.model_executor.models.model_local_kv import ( + HasModelLocalKV, + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + collect_model_local_kv_specs, + spec_from_hf_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +MODELS_DIR = pathlib.Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "models" + + +def _total_bytes(model: object, max_num_seqs: int = 1) -> int: + """Sum declared peak bytes across a model's declarations.""" + return sum(spec.peak_bytes(max_num_seqs) for _, spec in collect_model_local_kv_specs(model)) + + +def _cache_bytes(cache) -> int: + total = 0 + for layer in cache.layers: + for tensor in (layer.keys, layer.values): + if tensor is not None: + total += tensor.numel() * tensor.element_size() + return total + + +def _qwen2_config(layers=4, kv_heads=2, heads=8, hidden=256) -> Qwen2Config: + return Qwen2Config( + num_hidden_layers=layers, + num_attention_heads=heads, + num_key_value_heads=kv_heads, + hidden_size=hidden, + intermediate_size=hidden * 2, + vocab_size=1000, + ) + + +# -------------------------------------------------------------------------- +# A real declarer, called, checked against a real cache. +# -------------------------------------------------------------------------- + + +def test_ming_declaration_matches_the_cache_it_describes(): + """Calls MingAudioGenerator.model_local_kv_specs and checks the bytes. + + Builds the declarer without its weights -- only `_llm_config` and `_model` + are read -- then builds the `StaticCache` that `_init_kv_cache` builds from + the same config and compares. A wrong layer count, kv-head count, dtype or + capacity in the production declaration fails here. + """ + from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator + + config = _qwen2_config() + generator = object.__new__(MingAudioGenerator) + generator._llm_config = config + generator._model = torch.nn.Linear(1, 1).to(torch.float16) + + (spec,) = generator.model_local_kv_specs() + + cache = StaticCache( + config=config, + max_batch_size=1, + max_cache_len=MingAudioGenerator._STATIC_CACHE_LEN, + device="cpu", + dtype=torch.float16, + ) + head_dim = config.hidden_size // config.num_attention_heads + for layer_idx in range(config.num_hidden_layers): + keys = torch.zeros(1, config.num_key_value_heads, 1, head_dim, dtype=torch.float16) + cache.update(keys, keys.clone(), layer_idx) + + assert spec.peak_bytes(1) == _cache_bytes(cache) + + +def test_ming_does_not_scale_with_max_num_seqs(): + """forward() takes runtime_additional_information[0] and loops inline. + + One worker holds one of these no matter how many sequences the scheduler + admits. A previous revision derived the multiplier from scope and reported + `max_num_seqs` times the real figure. + """ + from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator + + generator = object.__new__(MingAudioGenerator) + generator._llm_config = _qwen2_config() + generator._model = torch.nn.Linear(1, 1).to(torch.float16) + + (spec,) = generator.model_local_kv_specs() + assert spec.peak_bytes(64) == spec.peak_bytes(1) + + +def test_growing_cache_declaration_matches_a_real_dynamic_cache(): + """The declared positions are the extent the tensor reaches when full.""" + layers, kv_heads, head_dim, positions = 4, 2, 32, 11 + config = _qwen2_config(layers=layers, kv_heads=kv_heads, heads=kv_heads, hidden=kv_heads * head_dim) + spec = spec_from_hf_config( + config, + name="local_transformer", + dtype=torch.bfloat16, + physical_capacity_positions=positions, + capacity_source="group_size + max(delay_pattern)", + scope=ModelLocalKVScope.INVOCATION, + rows=RowDriver.FIXED, + rows_reason="test", + ) + + cache = DynamicCache() + for layer_idx in range(layers): + keys = torch.zeros(1, kv_heads, positions, head_dim, dtype=torch.bfloat16) + cache.update(keys, torch.zeros_like(keys), layer_idx) + + assert spec.peak_bytes() == _cache_bytes(cache) + + +def test_multiplicity_does_not_claim_an_object_layout(): + """One allocation B rows wide costs the same as B one-row allocations. + + The protocol reports bytes and deliberately does not distinguish the two. + An earlier revision carried a separate `allocations` count and got the + object topology wrong in three of four declarers, which is worse than not + claiming it: the layout belongs in `allocation_note`, where it cannot be + mistaken for arithmetic. + """ + config = _qwen2_config(layers=4, kv_heads=2, heads=2, hidden=64) + common = dict( + config=config, + name="c", + dtype=torch.bfloat16, + physical_capacity_positions=11, + capacity_source="test", + scope=ModelLocalKVScope.INVOCATION, + ) + batched = spec_from_hf_config(rows=RowDriver.MAX_NUM_SEQS, **common) + + assert batched.row_count(64) == 64 + assert batched.peak_bytes(64) == batched.bytes_per_row * 64 + assert not hasattr(batched, "allocations"), "the object-count axis was removed on purpose" + + +def test_a_fixed_driver_ignores_concurrency(): + """A serialized or capture-pinned cache does not widen with max_num_seqs.""" + spec = spec_from_hf_config( + _qwen2_config(), + name="graph_pool", + dtype=torch.bfloat16, + physical_capacity_positions=11, + capacity_source="test", + scope=ModelLocalKVScope.MODEL, + rows=RowDriver.FIXED, + rows_fixed=261, + rows_reason="sum of captured buckets", + ) + assert spec.row_count(1) == 261 + assert spec.row_count(64) == 261 + assert spec.peak_bytes(64) == spec.bytes_per_row * 261 + + +def test_qwen3_tts_declares_nothing_on_the_stateless_path(): + """Calls the real decoder declarer with no cache-holding mode configured. + + `_forward_exact` runs `pre_transformer` without `use_cache`, so stateless + decoding allocates no KV. A revision that declared a per-request cache + unconditionally invented over a gigabyte at max_num_seqs=256, and the + revision before it declared nothing at all under enforce_eager, where the + eager async-chunk path does allocate. Both are regressions this pins. + """ + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import ( + modeling_qwen3_tts_tokenizer_v2 as mod, + ) + + decoder = object.__new__(mod.Qwen3TTSTokenizerV2Decoder) + # nn.Module.__setattr__ needs the machinery __init__ would have built, and + # this instance deliberately has none of it. + setattr_ = object.__setattr__ + setattr_( + decoder, + "config", + SimpleNamespace( + sliding_window=72, + num_hidden_layers=8, + num_key_value_heads=16, + head_dim=64, + hidden_size=1024, + num_attention_heads=16, + ), + ) + setattr_(decoder, "parameters", lambda: iter([torch.zeros(1, dtype=torch.float32)])) + + setattr_(decoder, "_cudagraph_wrapper", None) + assert decoder.model_local_kv_specs() == [] + + setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=False, model_local_kv_specs=lambda: [])) + assert decoder.model_local_kv_specs() == [] + + setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=True, model_local_kv_specs=lambda: [])) + names = [s.name for s in decoder.model_local_kv_specs()] + assert names == ["codec_decoder_request", "codec_decoder_working_copy"], names + + +def test_fixed_rejects_unreviewable_values(): + """A bare or unexplained constant cannot be reviewed, so it is rejected.""" + common = dict( + name="x", + layers=1, + kv_heads=1, + head_dim=1, + dtype=torch.float16, + physical_capacity_positions=1, + capacity_source="", + scope=ModelLocalKVScope.REQUEST, + ) + with pytest.raises(ValueError, match="rows_reason"): + ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=1, rows_reason=" ", **common) + with pytest.raises(ValueError, match=">= 1"): + ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=0, rows_reason="test", **common) + + +def test_rows_is_required(): + """Defaulting the batch extent is how the previous guarantee died. + + The dataclass declared it required while the helper every declarer uses + defaulted it to 1, so nothing enforced it where it mattered. + """ + with pytest.raises(TypeError): + spec_from_hf_config( # type: ignore[call-arg] + _qwen2_config(), + name="x", + dtype=torch.float16, + physical_capacity_positions=1, + capacity_source="", + scope=ModelLocalKVScope.REQUEST, + ) + + +def test_spec_from_hf_config_raises_rather_than_guessing(): + class Empty: + pass + + with pytest.raises(ValueError, match="layer count"): + spec_from_hf_config( + Empty(), + name="x", + dtype=torch.float16, + physical_capacity_positions=1, + capacity_source="", + scope=ModelLocalKVScope.REQUEST, + rows=RowDriver.FIXED, + rows_reason="test", + ) + + +# -------------------------------------------------------------------------- +# Collection +# -------------------------------------------------------------------------- + + +def _spec(**overrides) -> ModelLocalKVSpec: + kwargs = dict( + name="c", + layers=2, + kv_heads=2, + head_dim=4, + dtype=torch.float16, + physical_capacity_positions=8, + capacity_source="test", + scope=ModelLocalKVScope.REQUEST, + rows=RowDriver.FIXED, + rows_reason="test", + ) + kwargs.update(overrides) + return ModelLocalKVSpec(**kwargs) + + +class _Owner(torch.nn.Module): + def model_local_kv_specs(self): + return [_spec()] + + +class _CUDAGraphWrapperLike: + """Mirrors vllm.compilation.cuda_graph.CUDAGraphWrapper. + + Not an nn.Module, and its __getattr__ forwards everything the runnable + has. That forwarding is why a previous revision's "unwrap fix" was + unnecessary, and why a walk that does not unwrap can enter the tree twice. + """ + + def __init__(self, runnable): + self.runnable = runnable + + def unwrap(self): + return self.runnable + + def __getattr__(self, key): + return getattr(self.runnable, key) + + +def test_collector_finds_owners_nested_in_the_module_tree(): + class Mid(torch.nn.Module): + def __init__(self): + super().__init__() + self.inner = _Owner() + + class Top(torch.nn.Module): + def __init__(self): + super().__init__() + self.mid = Mid() + + model = Top() + assert [path for path, _ in collect_model_local_kv_specs(model)] == ["mid.inner"] + # 2 layers * 2 kv * 4 head_dim * 8 positions * 2 bytes * 2 (K and V) = 512 + assert _total_bytes(model) == 512 + + +def test_a_root_declarer_behind_a_wrapper_is_counted_once(): + """The wrapper forwards attributes, so a naive walk sees it twice.""" + + class Root(_Owner): + pass + + wrapped = _CUDAGraphWrapperLike(Root()) + collected = collect_model_local_kv_specs(wrapped) + + assert len(collected) == 1, collected + assert _total_bytes(wrapped) == 512 + + +def test_collector_reaches_owners_through_a_wrapper(): + class Top(torch.nn.Module): + def __init__(self): + super().__init__() + self.owner = _Owner() + + wrapped = _CUDAGraphWrapperLike(Top()) + assert [path for path, _ in collect_model_local_kv_specs(wrapped)] == ["owner"] + + +def test_collector_survives_an_owner_that_raises(): + class Broken(torch.nn.Module): + def model_local_kv_specs(self): + raise RuntimeError("checkpoint field missing") + + class Top(torch.nn.Module): + def __init__(self): + super().__init__() + self.broken = Broken() + self.ok = _Owner() + + assert _total_bytes(Top()) == 512 + + +def test_undeclared_models_report_zero_not_an_error(): + assert _total_bytes(object()) == 0 + assert _total_bytes(torch.nn.Linear(2, 2)) == 0 + + +def test_protocol_is_structural(): + assert isinstance(_Owner(), HasModelLocalKV) + assert not isinstance(torch.nn.Linear(2, 2), HasModelLocalKV) + + +@pytest.mark.parametrize( + ("relative_path", "owner"), + [ + ("qwen3_tts/segmented_graph_wrapper.py", "CUDAGraphDecoderWrapper"), + ("qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py", "Qwen3TTSTokenizerV2Decoder"), + ("mimo_audio/mimo_audio_llm.py", "MiMoAudioLLMForConditionalGeneration"), + ("minicpmo_4_5/minicpmo_4_5_omni_llm.py", "MiniCPMWhisperEncoder"), + ("ming_flash_omni/talker_module.py", "MingAudioGenerator"), + ("ming_flash_omni/ming_flash_omni_talker.py", "MingFlashOmniTalkerForConditionalGeneration"), + ], +) +def test_known_cache_owners_still_declare(relative_path: str, owner: str): + """Parsed rather than imported: these modules need CUDA and real weights. + + Weak by construction -- it only proves the method exists. The ming tests + above are the ones that check a declaration's contents. + """ + tree = ast.parse((MODELS_DIR / relative_path).read_text()) + classes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef) and node.name == owner] + assert classes, f"{owner} not found in {relative_path}" + methods = {child.name for cls in classes for child in cls.body if isinstance(child, ast.FunctionDef)} + assert "model_local_kv_specs" in methods, f"{owner} no longer declares its model-local KV" diff --git a/tests/tools/test_check_tts_adapter.py b/tests/tools/test_check_tts_adapter.py index 47c6ac89f..3107900c0 100644 --- a/tests/tools/test_check_tts_adapter.py +++ b/tests/tools/test_check_tts_adapter.py @@ -63,7 +63,15 @@ def test_gate_passes_on_current_tree(): assert main([]) == 0 -def test_budgets_are_the_real_counts(): +def test_budgets_are_not_exceeded(): + """Budgets are a ceiling, not a target. + + This used to assert equality, which re-imposed through pytest exactly what + #6008 removed from the checker: a PR that deleted branches failed CI unless + it also hand-edited the constant. That is what broke `main` for 5h19m in + #5746. The checker reports an under-budget count as a passing notice + (check_tts_adapter.py:140); the test must agree. + """ root = Path(__file__).resolve().parents[2] serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py" - assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES + assert len(_model_type_branches(ast.parse(serving.read_text()))) <= MAX_MODEL_TYPE_BRANCHES diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index bb01e522c..ec4faadb8 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -58,6 +58,12 @@ from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import MiMoAudioConfig +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -744,6 +750,67 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor def get_language_model(self) -> torch.nn.Module: return self.model + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare the local transformer's decode cache and its graph copies. + + The cache is created per call inside ``base_local_forward`` and is + bounded by the decode loop, not by sequence length: the loop runs + ``group_size + max(delay_pattern)`` steps regardless of how long the + utterance is. + + The second entry is the one that matters. ``base_local_forward`` runs + inside ``torch.cuda.graph`` during capture, so each captured bucket + leaves a cache's worth of tensors pinned in the graph pool for the + process lifetime. There is no Python object left to inspect -- the + cache is a local in the captured frame -- so the width comes from the + buckets that captured successfully, which is why it reads + ``local_forward_cg_by_bs`` rather than MIMO_CUDAGRAPH_BATCH_SIZES: a + bucket whose capture raised is logged and skipped. + + Both entries are declared, and summing them over-states a single step: + a given call either replays a captured bucket or runs eagerly, never + both. The graph pool is resident regardless, so the sum is the right + ceiling for the process and the wrong number for one request. + """ + delay_iters = self.group_size + max(self.delay_pattern) + dtype = next(self.local_transformer.parameters()).dtype + bound = f"group_size({self.group_size}) + max(delay_pattern)({max(self.delay_pattern)})" + specs = [ + spec_from_hf_config( + self.local_config, + name="local_transformer", + dtype=dtype, + physical_capacity_positions=delay_iters, + capacity_source=bound, + scope=ModelLocalKVScope.INVOCATION, + # One cache, B rows wide -- not B caches. base_local_forward + # builds a single DynamicCache whose batch dimension is the + # number of requests in the group. + rows=RowDriver.MAX_NUM_SEQS, + allocation_note="one batched allocation per eager call; graph replay uses the pool entry instead", + ) + ] + captured_rows = sum(self.local_forward_cg_by_bs) + if captured_rows: + specs.append( + spec_from_hf_config( + self.local_config, + name="local_transformer_graph_pool", + dtype=dtype, + physical_capacity_positions=delay_iters, + capacity_source=bound, + scope=ModelLocalKVScope.MODEL, + rows=RowDriver.FIXED, + rows_fixed=captured_rows, + rows_reason=f"sum of captured buckets {sorted(self.local_forward_cg_by_bs)}", + allocation_note=( + "pinned in the CUDA graph pool; capture is gated only on cuda.is_available(), " + "so this stays resident under enforce_eager too" + ), + ) + ) + return specs + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: if kwargs.get("modality_preprocess") is None: mm_dummy_embeddings = [] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index d34cce465..f01e20688 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -158,6 +158,17 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) patch_size=self.patch_size, ) + def model_local_kv_specs(self): + """Forward to the generator, which owns the talker's StaticCache. + + ``MingAudioGenerator`` is a plain object rather than a submodule, so + the module-tree walk cannot reach it on its own. + """ + generator = getattr(self, "audio_generator", None) + if generator is None: + return [] + return generator.model_local_kv_specs() + @property def device(self) -> torch.device: return next(self.model.parameters()).device diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 8e3e55611..293ec2fb3 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -41,6 +41,12 @@ from vllm_omni.model_executor.models.common.ming.cfm_solver import ( integrate_cfm_steps, ) from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) logger = init_logger(__name__) @@ -575,6 +581,13 @@ class MingAudioGenerator: for a single TTS request. The generator is stateless across requests. """ + _STATIC_CACHE_LEN = 2048 + """Positions the talker's StaticCache is preallocated for. + + Read by both ``_init_kv_cache`` and ``model_local_kv_specs`` so the + declaration cannot drift from the allocation. + """ + def __init__( self, config, @@ -609,6 +622,49 @@ class MingAudioGenerator: # trailing latent frames prepended on each decode call self._vae_decode_pad_frames = 32 + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare the talker's preallocated decode cache. + + This is the case that a static table cannot hold. ``self._llm_config`` + is a ``Qwen2Config`` resolved from the checkpoint at load time, so + layers, kv-heads and head-dim do not exist anywhere in this repo, and + the dtype is whatever the weights loaded as. Reading them off the same + objects ``_init_kv_cache`` builds from is the only way to get a real + number. + + Unlike the other three this is not a ceiling that a short utterance + stays under. transformers v5 initializes ``StaticLayer`` lazily, but + the first write to a layer allocates all ``max_cache_len`` positions at + once rather than growing, so a prefill brings the full extent into + existence on step 0 of every generation. + + Width is fixed at one row, not ``max_num_seqs``. ``forward`` takes only + ``runtime_additional_information[0]`` and runs the whole AR loop inline + before returning, so one worker holds one cache no matter how many + sequences the scheduler admits. Declaring this per-sequence -- which an + earlier revision did, by deriving the multiplier from scope -- reported + ``max_num_seqs`` times the real figure. + + The caveat is that nothing enforces the serialization. ``StaticCache`` + is built outside ``_sampler_pool``'s lock, so overlapping calls to + ``generate_latents`` would coexist. If this model ever gains a + concurrent entry point, this row count is the line to revisit. + """ + return [ + spec_from_hf_config( + self._llm_config, + name="talker_llm_static_cache", + dtype=next(self._model.parameters()).dtype, + physical_capacity_positions=self._STATIC_CACHE_LEN, + capacity_source=f"hardcoded max_cache_len={self._STATIC_CACHE_LEN} in _init_kv_cache", + scope=ModelLocalKVScope.INVOCATION, + rows=RowDriver.FIXED, + rows_fixed=1, + rows_reason="forward() handles one request and runs the AR loop inline", + allocation_note="full extent on first write per layer, not grown; rebuilt per text segment", + ) + ] + @cached_property def _sampler_pool(self) -> CFMGraphExecutorPool | None: device = next(self._model.parameters()).device @@ -790,7 +846,7 @@ class MingAudioGenerator: def _init_kv_cache( self, use_static_cache: bool, device: torch.device, dtype: torch.dtype ) -> tuple[StaticCache | None, int]: - max_cache_len = 2048 + max_cache_len = self._STATIC_CACHE_LEN if not use_static_cache: return None, max_cache_len cache = StaticCache( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 315ae159d..d70bc30d0 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -136,6 +136,13 @@ def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]: from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) + logger = init_logger(__name__) hf_logger = logging.get_logger(__name__) @@ -2621,6 +2628,51 @@ class MiniCPMWhisperEncoder(WhisperEncoder): [MiniCPMWhisperEncoderLayer(config, layer_idx=i) for i in range(config.encoder_layers)] ) + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare the streaming encoder's self-attention cache. + + Bounded by learned position embeddings, not by ``max_model_len``: the + cache cannot outgrow ``embed_positions``, and past that point the + forward pass repeats the last position rather than extending. + + Whisper is encoder-only self-attention here, so kv-heads equal + attention heads. Layer count is read off the built ``self.layers`` + rather than the config, so a partially built encoder reports what it + actually has. + + Declared per session, not per sequence. Each streaming session state + owns its own ``audio_past_key_values`` and the encoder runs at batch + size one, so ``max_num_seqs`` is the wrong number here. + + One session is the unit rather than the configured cap, because the cap + is not visible from the model and the setting that carries it cannot + distinguish "duplex off" from "one session" -- so a cap-scaled figure + would report a cache that a non-duplex deployment never allocates. An + earlier revision grew a capacity driver and an engine-side field to + carry that number for this one declarer; the arithmetic is a + multiplication the reader can do when they need the ceiling. + """ + return [ + spec_from_hf_config( + self.config, + name="whisper_encoder_self_attn", + dtype=self.conv1.weight.dtype, + layers=len(self.layers), + kv_heads=self.config.encoder_attention_heads, + head_dim=self.config.d_model // self.config.encoder_attention_heads, + physical_capacity_positions=int(self.embed_positions.weight.shape[0]), + capacity_source="embed_positions rows (max_source_positions)", + scope=ModelLocalKVScope.SESSION, + rows=RowDriver.FIXED, + rows_fixed=1, + rows_reason="one EncoderDecoderCache per streaming session, encoder runs at batch size 1", + allocation_note=( + "only allocated on the streaming path, where use_cache is set; scales with the " + "configured duplex session cap" + ), + ) + ] + def forward( self, input_features, diff --git a/vllm_omni/model_executor/models/model_local_kv.py b/vllm_omni/model_executor/models/model_local_kv.py new file mode 100644 index 000000000..7328e9967 --- /dev/null +++ b/vllm_omni/model_executor/models/model_local_kv.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Declaration protocol for attention KV kept outside the paged manager. + +Several models hold their own attention KV as HuggingFace ``transformers`` cache +objects rather than using the engine's paged KV manager. That memory is +allocated after the profiling run that decides how much KV the engine may claim, +so no per-stage footprint number accounts for it, and nothing catches a cache +that silently stops working. + +This module lets a model *declare* what it holds. It describes; it never +allocates. Allocation for the diffusion path is owned by the DiT KV manager +(RFC #5244 / PR #6094) and this protocol must not grow into a second allocator. + +Why a runtime query rather than a static table: geometry is frequently not +knowable until weights are loaded. ``ming_flash_omni``'s talker builds its cache +from a ``Qwen2Config`` that comes from the checkpoint, so layers, kv-heads, +head-dim and dtype simply do not exist in this repo. A table would have a hole +in it; a post-load query does not. + +Why capacity is a resolved number rather than a formula: the four known caches +are each bounded by a different mechanism -- a sliding window, a decode-loop +trip count, an encoder frame limit, a hardcoded constant -- and not one of them +is ``max_model_len``. Only the model knows its own bound. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Protocol, runtime_checkable + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +__all__ = [ + "HasModelLocalKV", + "ModelLocalKVScope", + "ModelLocalKVSpec", + "RowDriver", + "collect_model_local_kv_specs", + "spec_from_hf_config", +] + + +class ModelLocalKVScope(str, Enum): + """How long one allocation lives. + + Diagnostic only. Lifetime and size are independent: a per-call cache can be + the widest thing a model owns, and a process-lifetime one can be a single + row. ``ModelLocalKVSpec.rows`` carries the size; this carries only the + reader's mental model of when the memory appears and goes. + """ + + INVOCATION = "invocation" + """Dies when the call returns (e.g. a per-step working copy).""" + + REQUEST = "request" + """Retained across steps of one request, released when it finishes.""" + + SESSION = "session" + """Outlives a request; belongs to a duplex/streaming session.""" + + MODEL = "model" + """Lives as long as the model (e.g. captured into a CUDA-graph pool).""" + + +class RowDriver(str, Enum): + """What sets how many rows of a cache are live at peak. + + Two values, because two are all the known caches need. A third belongs + here only when a model actually has a cache the engine widens by some + other number -- adding one speculatively costs a branch everywhere and + buys nothing. + """ + + FIXED = "fixed" + """A count the model controls: a captured-bucket sum, a serialized call path.""" + + MAX_NUM_SEQS = "max_num_seqs" + """One row per in-flight sequence. Only the engine knows the value.""" + + +@dataclass(frozen=True) +class ModelLocalKVSpec: + """One declared KV allocation. + + A model returns one entry per *distinct lifetime*, not per cache object: if + the same geometry exists as both a retained per-request cache and a + short-lived working copy, that is two entries. + """ + + name: str + layers: int + kv_heads: int + head_dim: int + dtype: torch.dtype + + physical_capacity_positions: int + """Positions the tensor can physically hold. + + Not the logical sequence position. A sliding-window cache truncates on every + write, so a stream of thousands of tokens may only ever occupy + ``sliding_window - 1`` slots. + """ + + capacity_source: str + """Free-text note on where the bound comes from, for diagnostics only. + + Never branch on this. It exists so a reader can tell "71 because sliding + window" from "2048 because someone typed 2048", which is otherwise + invisible. + """ + + scope: ModelLocalKVScope + """How long one allocation lives. Diagnostic only. + + Deliberately not the multiplier. An earlier revision derived scaling from + scope, which made every non-``MODEL`` cache multiply by ``max_num_seqs``. + That over-reported ming by that factor (its calls are serialized) and used + the wrong driver entirely for MiniCPM-o (duplex sessions, not sequences). + Lifetime does not imply replication topology; ``rows`` says what does. + """ + + rows: RowDriver + """What sets how many rows of this shape are live at peak. + + A row is one batch entry's worth of the geometry above. Deliberately does + not distinguish "one allocation N rows wide" from "N allocations of one + row": those cost the same, and an earlier revision that tried to carry both + got the object topology wrong in three of four declarers. Bytes are what + this protocol reports; describe the layout in ``allocation_note``. + + Required, because defaulting it silently understates every cache whose + multiplicity is not 1. + """ + + rows_fixed: int = 1 + """The count, when ``rows`` is ``FIXED``. Ignored otherwise.""" + + rows_reason: str = "" + """Why ``rows_fixed`` is that number. Required when ``rows`` is ``FIXED``. + + A bare constant cannot be reviewed. ``rows_fixed=1`` on a cache that looks + per-request is exactly the claim a reader has to be able to check. + """ + + allocation_note: str | None = None + """Optional detail such as "rebuilt per text segment". Diagnostic only.""" + + def __post_init__(self) -> None: + if self.rows is RowDriver.FIXED: + if self.rows_fixed < 1: + raise ValueError(f"rows_fixed must be >= 1, got {self.rows_fixed}") + if not self.rows_reason.strip(): + raise ValueError("a FIXED row count needs rows_reason; an unexplained constant cannot be reviewed") + + @property + def bytes_per_row(self) -> int: + """Bytes for a single batch row of this cache.""" + return ( + 2 # K and V + * self.layers + * self.kv_heads + * self.head_dim + * self.physical_capacity_positions + * torch.empty((), dtype=self.dtype).element_size() + ) + + def row_count(self, max_num_seqs: int = 1) -> int: + """Rows live at peak, given the engine's concurrency.""" + if self.rows is RowDriver.FIXED: + return self.rows_fixed + return max(1, max_num_seqs) + + def peak_bytes(self, max_num_seqs: int = 1) -> int: + """Peak bytes once the engine's concurrency is known. + + The model supplies geometry and names the driver; only the engine knows + the number behind it. + """ + return self.bytes_per_row * self.row_count(max_num_seqs) + + +def _resolve(config: object, candidates: Sequence[str], what: str) -> int: + for attr in candidates: + value = getattr(config, attr, None) + if value is not None: + return int(value) + raise ValueError( + f"Cannot determine {what} from {type(config).__name__}: tried " + f"{', '.join(candidates)}. Pass it explicitly to spec_from_hf_config()." + ) + + +def spec_from_hf_config( + config: object, + *, + name: str, + physical_capacity_positions: int, + capacity_source: str, + scope: ModelLocalKVScope, + dtype: torch.dtype, + rows: RowDriver, + rows_fixed: int = 1, + rows_reason: str = "", + allocation_note: str | None = None, + layers: int | None = None, + kv_heads: int | None = None, + head_dim: int | None = None, +) -> ModelLocalKVSpec: + """Build a spec from the same config object the cache is built from. + + Every known model-local cache is constructed from an HF config, so the + geometry half of a declaration is the same three lookups each time. Pass + ``config`` and this derives them; pass ``layers``/``kv_heads``/``head_dim`` + explicitly to override when the config uses encoder-decoder naming or when + the built module is a more truthful source than the config. + + Raises rather than guessing when an attribute is absent: a silently wrong + geometry would under-report, which is the failure this protocol exists to + prevent. ``rows`` is keyword-required for the same reason -- an earlier + revision defaulted the batch extent to 1 here while documenting it as + required on the dataclass, so every declarer that used this helper got the + default and the guarantee was decorative. + """ + if layers is None: + layers = _resolve(config, ("num_hidden_layers", "encoder_layers", "num_layers"), "layer count") + if kv_heads is None: + kv_heads = _resolve( + config, + ("num_key_value_heads", "num_attention_heads", "encoder_attention_heads"), + "kv head count", + ) + if head_dim is None: + explicit = getattr(config, "head_dim", None) + if explicit is not None: + head_dim = int(explicit) + else: + hidden = _resolve(config, ("hidden_size", "d_model"), "hidden size") + heads = _resolve(config, ("num_attention_heads", "encoder_attention_heads"), "attention head count") + head_dim = hidden // heads + + return ModelLocalKVSpec( + name=name, + layers=layers, + kv_heads=kv_heads, + head_dim=head_dim, + dtype=dtype, + physical_capacity_positions=physical_capacity_positions, + capacity_source=capacity_source, + scope=scope, + rows=rows, + rows_fixed=rows_fixed, + rows_reason=rows_reason, + allocation_note=allocation_note, + ) + + +@runtime_checkable +class HasModelLocalKV(Protocol): + """Implemented by whichever object owns the cache. + + Implement it on the owner, not on the registered model. The owner is + usually several levels down (a codec decoder, a talker backbone), and + ``collect_model_local_kv_specs`` finds it by walking the module tree, so + no intermediate class has to forward anything. + """ + + def model_local_kv_specs(self) -> Sequence[ModelLocalKVSpec]: ... + + +def _unwrap_to_module(model: object) -> object: + """Unwrap runner-side wrappers down to the real module tree. + + ``vllm.compilation.cuda_graph.CUDAGraphWrapper`` is a plain callable, not + an ``nn.Module``, but its ``__getattr__`` forwards anything the runnable + has -- including ``named_modules``. So a check for ``named_modules`` does + not detect the wrapper, and the wrapper and its runnable both answer to the + same attributes. That matters twice over: it means walking the wrapper + happens to work, and it means a naive walk can enter the tree as both the + wrapper and the root module and count a root-level declarer twice. + + Prefer the supported ``unwrap()`` accessor, which returns the runnable + directly. + """ + seen: set[int] = set() + current = model + while current is not None and id(current) not in seen: + seen.add(id(current)) + unwrap = getattr(current, "unwrap", None) + if not callable(unwrap): + return current + try: + nxt = unwrap() + except Exception: + return current + if nxt is None or nxt is current: + return current + current = nxt + return model + + +def collect_model_local_kv_specs(model: object) -> list[tuple[str, ModelLocalKVSpec]]: + """Collect every declaration in a loaded model, with the owner's path. + + Walks ``named_modules()`` because the cache owner is an inner module in all + known cases. Requiring each registered model to forward the call would put + the burden on classes that do not own a cache and would silently report + zero the moment someone forgot -- exactly the failure mode this is meant to + surface. + + A raising declaration is logged and skipped rather than propagated: + reporting memory must not be able to break model load. + """ + root = _unwrap_to_module(model) + + owners: list[tuple[str, object]] = [] + seen_ids: set[int] = set() + if isinstance(root, HasModelLocalKV): + owners.append(("", root)) + seen_ids.add(id(root)) + + named_modules = getattr(root, "named_modules", None) + if callable(named_modules): + for path, module in named_modules(): + if id(module) in seen_ids or not isinstance(module, HasModelLocalKV): + continue + seen_ids.add(id(module)) + owners.append((path, module)) + + collected: list[tuple[str, ModelLocalKVSpec]] = [] + for path, owner in owners: + try: + specs = owner.model_local_kv_specs() + except Exception: + logger.warning( + "model_local_kv_specs() raised on %s; skipping its declaration", + type(owner).__name__, + exc_info=True, + ) + continue + collected.extend((path, spec) for spec in specs) + return collected diff --git a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py index a1eb071ee..3c73f5d13 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py @@ -17,6 +17,13 @@ from transformers.cache_utils import DynamicCache from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) + logger = init_logger(__name__) @@ -166,6 +173,85 @@ class CUDAGraphDecoderWrapper: if self._stats_log_every > 0 and self._stats_total_requests % self._stats_log_every == 0: self.log_decode_stats() + @staticmethod + def _retained_kv_caches(states: dict) -> list[tuple[int, DynamicCache]]: + """Find the KV caches a capture kept alive, with their batch rows. + + Walks the stored ``cache`` payload instead of assuming which capture + path populates it: ``_decode_icl_first_chunk`` fills its dict on the + decoder side, so the only reliable count is the one taken from the + objects themselves after warmup. + """ + found: list[tuple[int, DynamicCache]] = [] + for state in states.values(): + payload = state.get("cache") + if not isinstance(payload, dict): + continue + for value in payload.values(): + if not isinstance(value, DynamicCache): + continue + layers = getattr(value, "layers", None) + keys = getattr(layers[0], "keys", None) if layers else None + if keys is None: + continue + found.append((int(keys.shape[0]), value)) + return found + + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare only the graph-resident copies this wrapper retains. + + The per-decode cache is declared by the decoder itself, because it + exists whether or not graphs were captured. What is specific to this + wrapper is that every captured shape keeps its dummy ``DynamicCache`` + alive for replay (``combined_states``, ``icl_prefix_states`` and + ``xvec_prefix_states`` all store ``cache``), for the process lifetime. + + Rows come from the state dicts rather than the configured capture + lists: a capture that raised is logged and skipped, so the configured + shape count would over-report. The trade is that a capture whose + tensors transformers has not yet materialized contributes nothing, so + this is a floor rather than a contract. + """ + config = self.decoder.config + try: + dtype = next(self.decoder.parameters()).dtype + except StopIteration: # parameterless stub, only reachable in tests + return [] + + # Physical, not logical: the sliding window truncates on every write, + # so a stream of thousands of frames never occupies more than this. + physical = self.prefix_length - 1 + retained = [ + *self._retained_kv_caches(self.combined_states), + *self._retained_kv_caches(self.icl_prefix_states), + *self._retained_kv_caches(self.xvec_prefix_states), + ] + batched_captures = sum(rows for rows, _ in retained) + + if not batched_captures: + return [] + return [ + spec_from_hf_config( + config, + name="codec_decoder_graph_pool", + dtype=dtype, + physical_capacity_positions=physical, + capacity_source=f"sliding_window({self.prefix_length}) - 1", + scope=ModelLocalKVScope.MODEL, + rows=RowDriver.FIXED, + rows_fixed=batched_captures, + rows_reason=( + f"rows across {len(retained)} retained caches in " + f"{len(self.combined_states)} suffix / {len(self.icl_prefix_states)} icl-prefix / " + f"{len(self.xvec_prefix_states)} xvec-prefix captures" + ), + allocation_note=( + "counted from caches transformers has materialized; captures whose tensors are still " + "unallocated report nothing, so this is a floor for the xvec path" + ), + ) + ] + def log_decode_stats(self) -> None: if not getattr(self, "_stats_enabled", False) or self._stats_total_requests == 0: return diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index ea11e64a7..9572d9ee2 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -854,6 +854,75 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): self._cudagraph_enabled = False self._cudagraph_wrapper = None + def model_local_kv_specs(self): + """Declare this decoder's attention KV. + + Only the async-chunk path holds KV. In stateless mode the decoder runs + ``_forward_exact``, which calls ``pre_transformer`` without + ``use_cache`` and allocates no cache at all, so declaring a per-request + cache there invents memory that does not exist. The previous revision + did exactly that and reported over a gigabyte for stateless + deployments; the one before it hung the declaration off the CUDA-graph + wrapper and so reported zero under ``enforce_eager``, where the eager + async-chunk path does allocate. Both mistakes come from tying this to + the wrong condition. + + Two entries when async-chunk is on: the retained per-request cache, and + the working copy ``_decode_icl_first_chunk`` deep-copies from it, which + is live at the same time. + + The wrapper, when present, adds the graph-resident copies. It is a + plain object rather than a submodule, so the module-tree walk cannot + reach it; this is the hop it needs. + """ + from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + RowDriver, + spec_from_hf_config, + ) + + config = self.config + try: + dtype = next(self.parameters()).dtype + except StopIteration: # parameterless stub, only reachable in tests + return [] + + window = int(getattr(config, "sliding_window", 0) or 0) + if window <= 1: + return [] + + specs = [] + wrapper = getattr(self, "_cudagraph_wrapper", None) + # async_chunk is the wrapper's record of the mode; without a wrapper the + # decoder has not been told, and only the stateless path is reachable. + if wrapper is not None and getattr(wrapper, "async_chunk", False): + common = dict( + config=config, + dtype=dtype, + physical_capacity_positions=window - 1, + capacity_source=f"sliding_window({window}) - 1", + rows=RowDriver.MAX_NUM_SEQS, + ) + specs.append( + spec_from_hf_config( + name="codec_decoder_request", + scope=ModelLocalKVScope.REQUEST, + allocation_note="retained across chunks of one request", + **common, + ) + ) + specs.append( + spec_from_hf_config( + name="codec_decoder_working_copy", + scope=ModelLocalKVScope.INVOCATION, + allocation_note="deep-copied from the retained cache each chunk; live alongside it", + **common, + ) + ) + if wrapper is not None: + specs.extend(wrapper.model_local_kv_specs()) + return specs + def precompute_snake_caches(self): """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules.""" count = 0 diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 80d55020f..8dd9f5f58 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -36,6 +36,7 @@ from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_omni.core.prefix_cache import OmniTensorPrefixCache from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding +from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -194,6 +195,52 @@ class OmniGPUModelRunner(GPUModelRunner): self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() self._prewarm_attention_capture_workspaces() + self._report_model_local_kv() + + def _report_model_local_kv(self) -> None: + """Log attention KV this model holds outside the paged manager. + + Reporting only. Two of these caches are captured into CUDA graphs + during ``load_model``, so NVML already charges them to the process + before ``determine_available_memory()`` samples it; subtracting the + declared total from the KV budget would double-count them. What is not + charged is the per-request half, which appears after profiling and + scales with ``max_num_seqs``. + + Wrapped whole: a memory report has no business breaking model load, + and this runs on platforms these declarations have never been + exercised on. + """ + try: + self._log_model_local_kv() + except Exception: + logger.warning("Model-local KV reporting failed; continuing", exc_info=True) + + def _log_model_local_kv(self) -> None: + specs = collect_model_local_kv_specs(getattr(self, "model", None)) + if not specs: + return + max_num_seqs = max(1, int(self.scheduler_config.max_num_seqs)) + total = sum(spec.peak_bytes(max_num_seqs) for _, spec in specs) + logger.info( + "Model-local KV (outside the paged manager): %.2f MiB across %d declaration(s) at max_num_seqs=%d", + total / (1 << 20), + len(specs), + max_num_seqs, + ) + for path, spec in specs: + logger.info( + " %s.%s: %.2f MiB (%d rows from %s x %.2f MiB/row, scope=%s, %d positions from %s)", + path or type(self.model).__name__, + spec.name, + spec.peak_bytes(max_num_seqs) / (1 << 20), + spec.row_count(max_num_seqs), + spec.rows.value if spec.rows.value != "fixed" else f"fixed({spec.rows_fixed}; {spec.rows_reason})", + spec.bytes_per_row / (1 << 20), + spec.scope.value, + spec.physical_capacity_positions, + spec.capacity_source, + ) def _maybe_enable_output_token_ids_for_model_sampler(self) -> None: if getattr(self.model, "logitsprocs_need_output_token_ids", False): From 498c41b8d60270086c9a226c08a55ff6547800a3 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Thu, 10 Sep 2026 20:50:12 -0400 Subject: [PATCH 015/209] [Realtime] Emit current (non-beta) OpenAI audio/transcript event names (#7339) Signed-off-by: Nick Cao Co-authored-by: Claude Sonnet 5 --- docs/design/fullduplex-personaplex.md | 4 +- docs/design/fullduplex.md | 14 +- docs/serving/full_duplex_api.md | 2 +- docs/serving/realtime_duplex_api.md | 40 +++--- examples/online_serving/barge_in_client.py | 2 +- .../minicpmo/realtime_duplex_demo.py | 10 +- .../minicpmo/realtime_web/app/static/app.js | 8 +- tests/benchmarks/patch/test_patch.py | 4 +- tests/benchmarks/test_omniinteract.py | 6 +- tests/clients/test_duplex_client.py | 24 ++-- tests/e2e/features/fullduplex/test_client.py | 12 +- .../fullduplex/test_omni_duplex_eval_cli.py | 2 +- .../minicpmo_realtime_duplex_scenarios.py | 38 +++--- .../nemotron_voicechat_realtime_duplex.py | 8 +- .../personaplex_realtime_duplex.py | 6 +- .../run_minicpmo_realtime_duplex_demo_pair.py | 5 +- ...run_minicpmo_realtime_duplex_server_vad.py | 13 +- ...minicpmo_realtime_duplex_soft_interrupt.py | 7 +- .../test_minicpmo_realtime_duplex_drivers.py | 125 +++++++++++------- tests/engine/test_duplex_import_boundary.py | 4 +- .../duplex/test_websocket_actor.py | 6 +- .../openai/test_duplex_session_attachment.py | 2 +- .../openai_api/test_duplex_handler.py | 40 +++--- ...st_minicpmo_realtime_duplex_simple_demo.py | 8 +- .../test_minicpmo_realtime_web_static.py | 4 +- .../duplex/test_unified_runtime.py | 2 +- .../duplex/omni_duplex_eval_clock.py | 6 +- vllm_omni/benchmarks/omniinteract.py | 10 +- vllm_omni/clients/duplex.py | 6 +- .../entrypoints/duplex/realtime_output.py | 8 +- vllm_omni/entrypoints/duplex/warmup.py | 2 +- vllm_omni/entrypoints/duplex/websocket.py | 2 - .../entrypoints/openai/serving_speech.py | 2 +- vllm_omni/experimental/fullduplex/client.py | 8 +- 34 files changed, 251 insertions(+), 189 deletions(-) diff --git a/docs/design/fullduplex-personaplex.md b/docs/design/fullduplex-personaplex.md index 249554bbb..33d125c97 100644 --- a/docs/design/fullduplex-personaplex.md +++ b/docs/design/fullduplex-personaplex.md @@ -20,7 +20,7 @@ The target is the engine-native path: -> PersonaPlex Talker -> streaming PersonaPlex Code2Wav -> PersonaPlex data-plane projector - -> response.audio.delta + response.audio_transcript.delta + -> response.output_audio.delta + response.output_audio_transcript.delta ``` The standalone `/api/chat` and `/v1/audio/duplex` server that accompanied the @@ -197,7 +197,7 @@ model-neutral native results: It owns per-request audio and text cursors so a cumulative output cannot replay old audio. The generic projector turns each native result into the Realtime -`response.audio.delta` / `response.audio_transcript.delta` pair (and +`response.output_audio.delta` / `response.output_audio_transcript.delta` pair (and `response.output_text.delta` for text) under one `response_id`; the full mapping is the name map in the [Realtime Duplex API](../serving/realtime_duplex_api.md) serving guide. diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 87fb1f79f..7d3d5ec95 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -606,11 +606,11 @@ scoped to one `response_id` unless stated otherwise. | Public event | Trigger and ordering | Cardinality | | --- | --- | --- | | `response.created` | First event for a visible assistant response; precedes its output-item, content-part, speak, text, and audio events. | Exactly once for every visible response. | -| `response.speak` | The model selected speak; emitted after `response.created` and no later than the first `response.audio.delta`. It carries decision metadata, not transcript text. | At most once. | -| `response.audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.audio_transcript.delta` for the same chunk. | Zero or more. | -| `response.audio_transcript.delta` | Append-only transcript contribution paired with an audio delta; it may be empty for a text-less audio unit and must not repeat or overlap earlier text. | One per audio delta. | -| `response.audio.done` | Closes the audio stream after its final delta and before response terminal events. It is omitted for a response with no audio or transcript. | At most once. | -| `response.audio_transcript.done` | Contains the exact concatenation of all non-overlapping transcript deltas for the response. | At most once, and only for a non-empty transcript. | +| `response.speak` | The model selected speak; emitted after `response.created` and no later than the first `response.output_audio.delta`. It carries decision metadata, not transcript text. | At most once. | +| `response.output_audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.output_audio_transcript.delta` for the same chunk. | Zero or more. | +| `response.output_audio_transcript.delta` | Append-only transcript contribution paired with an audio delta; it may be empty for a text-less audio unit and must not repeat or overlap earlier text. | One per audio delta. | +| `response.output_audio.done` | Closes the audio stream after its final delta and before response terminal events. It is omitted for a response with no audio or transcript. | At most once. | +| `response.output_audio_transcript.done` | Contains the exact concatenation of all non-overlapping transcript deltas for the response. | At most once, and only for a non-empty transcript. | | `response.output_item.done` / `conversation.item.done` | Finalize the assistant item after content-part terminal events and before `response.done`. | At most once each. | | `response.done` | Terminal event for a created response; follows audio, transcript, content-part, and item terminal events. | Exactly once for every created response that reaches a terminal state. | | `rate_limits.updated` | Compatibility event emitted immediately after `response.done`; the current payload has an empty rate-limit list. | Once per emitted `response.done`. | @@ -625,8 +625,8 @@ response and never reuse that response ID for later model turns. Compatibility changes must update this table and its protocol contract tests in the same change. The golden transcript test requires joined -`response.audio_transcript.delta` values to equal -`response.audio_transcript.done.transcript`; response tests separately enforce +`response.output_audio_transcript.delta` values to equal +`response.output_audio_transcript.done.transcript`; response tests separately enforce at-most-once `response.speak` and terminal event cardinality. ### Session internal ledgers diff --git a/docs/serving/full_duplex_api.md b/docs/serving/full_duplex_api.md index f648d6fdb..4e8bf9748 100644 --- a/docs/serving/full_duplex_api.md +++ b/docs/serving/full_duplex_api.md @@ -86,7 +86,7 @@ A typical `/v1/realtime?duplex=1` session follows this lifecycle: 3. Send `input_audio_buffer.append` events while microphone audio arrives. 4. Send `input_audio_buffer.commit` at a user-turn boundary when required by the model policy. -5. Consume `response.created`, transcript deltas, `response.audio.delta`, and +5. Consume `response.created`, transcript deltas, `response.output_audio.delta`, and `response.done` or `response.listen` events. 6. Send `playback.ack` after audio has been played when the session advertises playback acknowledgement support. diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index d37d02291..c855c4491 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -311,13 +311,19 @@ work unmodified. The overlap falls into the three tiers used by the event catalo | Direction | Messages | | --- | --- | | Client → server | `session.update`, `input_audio_buffer.append` / `.commit` / `.clear`, `conversation.item.create` / `.retrieve` / `.truncate` / `.delete`, `response.create`, `response.cancel`, `output_audio_buffer.clear` (OpenAI: WebRTC-only; here also a WebSocket event) | -| Server → client | `error` (same `{type, code, message, event_id, param}` envelope and the same classes `invalid_request_error` / `server_error` / `rate_limit_error`), `session.created`, `session.updated`, `input_audio_buffer.committed` / `.cleared` / `.speech_started` / `.speech_stopped`, `conversation.item.created` / `.added` / `.done` / `.deleted` / `.retrieved` / `.truncated`, `conversation.item.input_audio_transcription.completed`, `response.created`, `response.done`, `response.output_item.added` / `.done`, `response.content_part.added` / `.done`, `response.audio.delta` / `.done`, `response.audio_transcript.delta` / `.done`, `response.output_text.delta` / `.done`, `response.function_call_arguments.delta` / `.done`, `rate_limits.updated`, `output_audio_buffer.cleared` | +| Server → client | `error` (same `{type, code, message, event_id, param}` envelope and the same classes `invalid_request_error` / `server_error` / `rate_limit_error`), `session.created`, `session.updated`, `input_audio_buffer.committed` / `.cleared` / `.speech_started` / `.speech_stopped`, `conversation.item.created` / `.added` / `.done` / `.deleted` / `.retrieved` / `.truncated`, `conversation.item.input_audio_transcription.completed`, `response.created`, `response.done`, `response.output_item.added` / `.done`, `response.content_part.added` / `.done`, `response.output_audio.delta` / `.done`, `response.output_audio_transcript.delta` / `.done`, `response.output_text.delta` / `.done`, `response.function_call_arguments.delta` / `.done`, `rate_limits.updated`, `output_audio_buffer.cleared` | | Session fields | `model`, `modalities` / `output_modalities`, `instructions`, `voice`, `input_audio_format` / `output_audio_format` (beta spelling) **and** `audio.input.format` / `audio.output.format` objects (GA spelling), `turn_detection.server_vad {threshold, prefix_padding_ms, silence_duration_ms}`, `input_audio_transcription`, `input_audio_noise_reduction`, `tools`, `tool_choice`, `temperature`, `max_response_output_tokens`, `speed`, `tracing` | | Object shapes | `realtime.session`, `realtime.item`, `realtime.response`; item content parts `input_text` / `input_audio` / `output_text` / `output_audio`; `function_call` and `function_call_output` items; `previous_item_id` chaining; `output_index` / `content_index` addressing | | Sequencing | `response.created → output_item.added → content_part.added → deltas → *.done → content_part.done → output_item.done → response.done → rate_limits.updated`; `speech_started → speech_stopped → committed` | -The server deliberately emits and accepts **both** the beta and the GA -spellings at once (`response.audio.delta` alongside `output_modalities`, +Audio/transcript output uses only the current (non-beta) OpenAI event names +(`response.output_audio.delta`, not the deprecated `response.audio.delta`), +matching `openai.types.realtime` in the official `openai` Python SDK and what +clients built against it (e.g. `livekit-plugins-openai`'s non-Azure code path) +expect; a client still on the old beta SDK types would not recognize these +events. Session fields and `conversation.item.*` events, in contrast, are +deliberately dual: the server emits and accepts **both** the beta and the GA +spellings at once (`output_modalities` alongside `modalities`, `conversation.item.added` alongside `conversation.item.created`) so either generation of OpenAI client parses the stream. @@ -329,7 +335,7 @@ A stock client ignores the extra keys; the extensions are additive. | --- | --- | | `session.created` / `session.updated` | top-level `incarnation`, `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` | | `response.created` / `response.done` / `response.listen` | `response_id` at top level; the raw duplex event under `response.metadata` (`duplex_event`); `status_details.reason` uses vLLM reasons (`barge_in`, `client_cancelled`, `new_response`, …) | -| `response.audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` | +| `response.output_audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` | | `response.speak` (inserted before the first delta) | not an OpenAI event, but rides the OpenAI response envelope (`response_id`, `item_id`, `output_index`, `content_index`) | | `input_audio_buffer.append` | `is_speech`, `video_frames` (+ `max_slice_nums`), `duration_ms`, `audio_end_ms`, per-event `format` / `sample_rate_hz` | | `input_audio_buffer.commit` | `final`, `response_create`, `is_speech:false` (silence declaration) | @@ -443,12 +449,12 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | `response.output_item.added` | 1 | Assistant item (or `function_call` item) attached to the response. | | `response.content_part.added` | 1 | Audio/text content part opened on the assistant item. | | `response.speak` | 3 | Model chose to speak (native lane); at most once per response, before the first audio delta. | -| `response.audio.delta` | 2 | One ordered audio chunk. OpenAI name; adds `format`, `sample_rate_hz`, `metadata{epoch, model_speak, end_of_turn, audio_text_marks, playback}`. | -| `response.audio_transcript.delta` | 1 | Transcript text paired one-to-one with each audio delta. | +| `response.output_audio.delta` | 2 | One ordered audio chunk. OpenAI name; adds `format`, `sample_rate_hz`, `metadata{epoch, model_speak, end_of_turn, audio_text_marks, playback}`. | +| `response.output_audio_transcript.delta` | 1 | Transcript text paired one-to-one with each audio delta. | | `response.output_text.delta` | 1 | Text delta for text-modality responses. | | `response.output_text.done` | 1 | Final text of a text-modality response. | -| `response.audio.done` | 1 | Audio stream closed for the response. | -| `response.audio_transcript.done` | 1 | Full transcript (concatenation of the deltas). | +| `response.output_audio.done` | 1 | Audio stream closed for the response. | +| `response.output_audio_transcript.done` | 1 | Full transcript (concatenation of the deltas). | | `response.content_part.done` | 1 | Content part finalized with its transcript/text. | | `response.output_item.done` | 1 | Assistant item finalized. | | `response.done` | 2 | Terminal event; `status` ∈ `completed` \| `cancelled` \| `failed`. OpenAI name; vLLM-Omni `status_details.reason` values and raw duplex event under `metadata`. | @@ -902,11 +908,11 @@ deferred commit during an active response (`event.response_create_deferred`) } ``` -`S→C response.audio.delta` (one ~1 s unit, 24 kHz PCM16) +`S→C response.output_audio.delta` (one ~1 s unit, 24 kHz PCM16) ```json { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, @@ -926,10 +932,10 @@ deferred commit during an active response (`event.response_create_deferred`) } ``` -`S→C response.audio_transcript.delta` (exactly one per audio delta) +`S→C response.output_audio_transcript.delta` (exactly one per audio delta) ```json -{"type": "response.audio_transcript.delta", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "delta": "It is sunny"} +{"type": "response.output_audio_transcript.delta", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "delta": "It is sunny"} ``` `S→C response.output_text.delta` / `response.output_text.done` (text modality) @@ -942,14 +948,14 @@ deferred commit during an active response (`event.response_create_deferred`) {"type": "response.output_text.done", "response_id": "resp_02", "item_id": "item_resp_02", "output_index": 0, "content_index": 0, "text": "Tomorrow looks cloudy."} ``` -`S→C response.audio.done` / `response.audio_transcript.done` +`S→C response.output_audio.done` / `response.output_audio_transcript.done` ```json -{"type": "response.audio.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0} +{"type": "response.output_audio.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0} ``` ```json -{"type": "response.audio_transcript.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "transcript": "It is sunny, 24 degrees."} +{"type": "response.output_audio_transcript.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "transcript": "It is sunny, 24 degrees."} ``` `S→C response.content_part.done` @@ -1191,8 +1197,8 @@ Realtime projector renames and fans out these events before they reach a | `session.create` (`open_session`, `session.config`) | first `session.update` | | `input.committed` | `conversation.item.added` / `conversation.item.created`, `input_audio_buffer.committed`, `conversation.item.input_audio_transcription.completed`, `conversation.item.done` | | `input.cancelled` | `input_audio_buffer.cleared` | -| `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.audio.delta` and `response.audio_transcript.delta` | -| `response.output_audio.done` / `response.output_text.done` | `response.audio.done` / `response.output_text.done` | +| `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.output_audio.delta` and `response.output_audio_transcript.delta` | +| `response.output_audio.done` / `response.output_text.done` | `response.output_audio.done` / `response.output_text.done` | | `response.text.delta` | `response.output_text.delta` | | `response.message` (chat-fallback raw chunk) | passed through | | `audio.cancelled` | optional `output_audio_buffer.cleared`, then the cancelled terminal events and `response.done` with `status: "cancelled"` | diff --git a/examples/online_serving/barge_in_client.py b/examples/online_serving/barge_in_client.py index 1ac627d35..b144da9f4 100644 --- a/examples/online_serving/barge_in_client.py +++ b/examples/online_serving/barge_in_client.py @@ -65,7 +65,7 @@ from vllm_omni.clients.duplex import ( # noqa: E402 ) _TRANSCRIPT_DELTA_EVENTS = { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py index 9cbc7f29b..9fc75621e 100644 --- a/examples/online_serving/minicpmo/realtime_duplex_demo.py +++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py @@ -63,14 +63,14 @@ class _StreamingOutputWriter: def handle(self, event: dict[str, object]) -> None: event_type = event.get("type") if event_type in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", }: delta = event.get("delta") if isinstance(delta, str) and delta: print(delta, end="", file=sys.stderr, flush=True) return - if event_type != "response.audio.delta": + if event_type != "response.output_audio.delta": return delta = event.get("delta") or event.get("audio") @@ -509,12 +509,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: audio = collector.audio_bytes() first_text_at_s = collector.first_received_at( - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", after_s=commit_sent_at_s, ) first_audio_at_s = collector.first_received_at( - "response.audio.delta", + "response.output_audio.delta", after_s=commit_sent_at_s, ) response_created_at_s = collector.first_received_at( @@ -536,7 +536,7 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: for event in collector.events if event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", } ] diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/app.js b/examples/online_serving/minicpmo/realtime_web/app/static/app.js index 2943eec74..1ea2ba093 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/app.js +++ b/examples/online_serving/minicpmo/realtime_web/app/static/app.js @@ -356,7 +356,7 @@ case 'response.speak': beginAssistant(responseId); break; - case 'response.audio.delta': + case 'response.output_audio.delta': currentResponseId = responseId || currentResponseId; assistantActive = true; setModel('Speaking'); @@ -364,13 +364,13 @@ .then((decoded) => feedPlayback(decoded, responseId)) .catch((error) => appendLog(`audio decode failed: ${error.message || error}`, true)); break; - case 'response.audio.done': + case 'response.output_audio.done': requestPlaybackDrain(responseId); break; - case 'response.audio_transcript.delta': + case 'response.output_audio_transcript.delta': addTranscript('assistant', event.delta || ''); break; - case 'response.audio_transcript.done': + case 'response.output_audio_transcript.done': finishTranscript('assistant', event.transcript || ''); break; case 'conversation.item.input_audio_transcription.delta': diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index accb30100..6dd35e677 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -168,7 +168,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) ) self.events.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": f"turn {self.response_count}", }, @@ -176,7 +176,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) ) self.events.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": base64.b64encode(b"\x00\x00" * 2400).decode(), "sample_rate_hz": 24_000, diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py index f9488b11c..b9dd10870 100644 --- a/tests/benchmarks/test_omniinteract.py +++ b/tests/benchmarks/test_omniinteract.py @@ -65,7 +65,7 @@ def _audio( value: int = 1, ) -> dict[str, object]: event: dict[str, object] = { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "format": "pcm16", "delta": base64.b64encode(bytes((value, 0)) * samples).decode(), @@ -76,7 +76,7 @@ def _audio( def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]: - return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value} + return {"type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": value} def _listen(*, buffering: bool = False) -> dict[str, object]: @@ -427,7 +427,7 @@ def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: P assert result.success and result.eligible_for_official_eval assert summary["scene_type"] == "multi_turn" assert context and context.spans - assert all("delta" not in event for event in context.events if event["type"] == "response.audio.delta") + assert all("delta" not in event for event in context.events if event["type"] == "response.output_audio.delta") oi.publish_deferred_case_artifacts(tmp_path / "out", case, result) directory = oi._output_dir(tmp_path / "out", case) assert all((directory / name).is_file() for name in oi.SUCCESS_ARTIFACTS) diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index 1be3eb943..ce1795e0f 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -279,7 +279,7 @@ async def test_response_handle_flow(): sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) sock.feed( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-1", "delta": _b64(chunk), "sample_rate_hz": 24_000, @@ -288,7 +288,7 @@ async def test_response_handle_flow(): ) sock.feed( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-1", "delta": "hi there", "server_event_seq": 5, @@ -370,7 +370,7 @@ async def test_terminal_listen_keeps_spoken_decision(): sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) sock.feed( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-1", "delta": _b64(chunk), "sample_rate_hz": 24_000, @@ -418,7 +418,7 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): ) for payload in (b"c1", b"c2", b"c3"): # _feed must never block the reader, even against a full queue. - handle._feed(AudioDelta({"type": "response.audio.delta", "delta": _b64(payload)})) + handle._feed(AudioDelta({"type": "response.output_audio.delta", "delta": _b64(payload)})) handle._finish(None) chunks = [chunk async for chunk in handle.audio()] assert chunks == [b"c3"] # oldest chunks were dropped, the sentinel landed @@ -561,7 +561,7 @@ def test_event_collector_accumulates_audio(): collector = EventCollector() collector.add({"type": "response.created", "response": {"id": "r1"}}, received_at_s=1.0) collector.add( - {"type": "response.audio.delta", "response_id": "r1", "delta": _b64(b"ab"), "sample_rate_hz": 16_000}, + {"type": "response.output_audio.delta", "response_id": "r1", "delta": _b64(b"ab"), "sample_rate_hz": 16_000}, received_at_s=1.1, ) collector.add( @@ -632,7 +632,7 @@ def test_event_collector_partitions_audio_by_response(): collector.add({"type": "response.created", "response": {"id": "resp-a"}}) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio-a").decode("ascii"), "sample_rate_hz": 16_000, @@ -643,7 +643,7 @@ def test_event_collector_partitions_audio_by_response(): assert collector.audio_bytes("resp-a") == b"audio-a" assert collector.output_sample_rate_hz == 16_000 assert collector.first_received_at("response.created") is not None - assert collector.last_received_at("response.audio.delta") is not None + assert collector.last_received_at("response.output_audio.delta") is not None def test_event_collector_reports_engine_token_and_audio_intervals(): @@ -664,7 +664,7 @@ def test_event_collector_reports_engine_token_and_audio_intervals(): for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "sample_rate_hz": 16_000, @@ -676,11 +676,11 @@ def test_event_collector_reports_engine_token_and_audio_intervals(): received_at_s=received_at_s, ) collector.add( - {"type": "response.audio_transcript.delta", "response_id": "resp-a", "delta": ""}, + {"type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": ""}, received_at_s=10.1, ) collector.add( - {"type": "response.audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, + {"type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, received_at_s=10.15, ) collector.add( @@ -752,7 +752,7 @@ def test_response_timing_ignores_unowned_session_level_metrics(): ) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "metadata": { @@ -887,7 +887,7 @@ def test_reference_audio_data_url(tmp_path): def test_event_collector_response_text_joins_deltas_per_response(): collector = EventCollector() collector.add({"type": "response.created", "response_id": "r1"}) - collector.add({"type": "response.audio_transcript.delta", "response_id": "r1", "delta": "he"}) + collector.add({"type": "response.output_audio_transcript.delta", "response_id": "r1", "delta": "he"}) collector.add({"type": "response.output_text.delta", "response_id": "r2", "delta": "other"}) collector.add({"type": "response.text.delta", "response_id": "r1", "delta": "llo"}) assert collector.response_text("r1") == "hello" diff --git a/tests/e2e/features/fullduplex/test_client.py b/tests/e2e/features/fullduplex/test_client.py index 266df1fa0..250391fe2 100644 --- a/tests/e2e/features/fullduplex/test_client.py +++ b/tests/e2e/features/fullduplex/test_client.py @@ -241,7 +241,7 @@ def test_realtime_event_collector_partitions_audio_by_response(): collector.add({"type": "response.created", "response": {"id": "resp-a"}}) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio-a").decode("ascii"), "sample_rate_hz": 16_000, @@ -252,7 +252,7 @@ def test_realtime_event_collector_partitions_audio_by_response(): assert collector.audio_bytes("resp-a") == b"audio-a" assert collector.output_sample_rate_hz == 16_000 assert collector.first_received_at("response.created") is not None - assert collector.last_received_at("response.audio.delta") is not None + assert collector.last_received_at("response.output_audio.delta") is not None def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): @@ -273,7 +273,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "sample_rate_hz": 16_000, @@ -286,7 +286,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): ) collector.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "", }, @@ -294,7 +294,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): ) collector.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "hello", }, @@ -369,7 +369,7 @@ def test_response_timing_ignores_unowned_session_level_metrics(): ) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "metadata": { diff --git a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py index 192859de5..a4d7faf9b 100644 --- a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py +++ b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py @@ -119,7 +119,7 @@ async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path await websocket.send( json.dumps( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "r1", "delta": base64.b64encode(b"\0\0" * 240).decode(), "sample_rate_hz": 24_000, diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py index ebd59fc3e..f5bba11e4 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py @@ -213,7 +213,7 @@ class DemoState: item_id = item.get("id") if isinstance(item_id, str) and item_id not in self.assistant_item_ids: self.assistant_item_ids.append(item_id) - elif event_type == "response.audio.delta": + elif event_type == "response.output_audio.delta": delta = event.get("delta") or event.get("audio") if isinstance(delta, str) and delta: try: @@ -386,8 +386,8 @@ class DemoState: "response.output_item.added", "response.content_part.added", "response.speak", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", "response.content_part.done", "response.output_item.done", "response.done", @@ -402,8 +402,8 @@ class DemoState: ) if index is None and event_type not in { "response.speak", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", }: return {} if index is not None: @@ -451,13 +451,13 @@ class DemoState: "response.output_item.added", "response.content_part.added", "response.speak", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", "response.content_part.done", "response.output_item.done", "response.done", ] - if "response.audio.delta" not in indices_by_type: + if "response.output_audio.delta" not in indices_by_type: listen_index = self.first_index( "response.listen", lambda event: ( @@ -515,7 +515,7 @@ class DemoState: def model_speak_before_audio_ok(self) -> bool: speak_index = self.first_index("response.speak") - audio_index = self.first_index("response.audio.delta") + audio_index = self.first_index("response.output_audio.delta") return speak_index is not None and audio_index is not None and speak_index < audio_index def response_done(self, response_id: str | None) -> bool: @@ -532,7 +532,7 @@ class DemoState: return sum( 1 for event in self.events - if event.get("type") == "response.audio.delta" and self._event_response_id(event) == response_id + if event.get("type") == "response.output_audio.delta" and self._event_response_id(event) == response_id ) def response_playback_sent_ms(self, response_id: str | None) -> int: @@ -565,14 +565,16 @@ class DemoState: return "".join( str(event.get("delta", "")) for event in self.events - if event.get("type") == "response.audio_transcript.delta" and self._event_response_id(event) == response_id + if event.get("type") == "response.output_audio_transcript.delta" + and self._event_response_id(event) == response_id ) def response_transcript_done(self, response_id: str) -> list[str]: return [ str(event.get("transcript", "")) for event in self.events - if event.get("type") == "response.audio_transcript.done" and self._event_response_id(event) == response_id + if event.get("type") == "response.output_audio_transcript.done" + and self._event_response_id(event) == response_id ] def completed_response_ids(self) -> list[str]: @@ -603,7 +605,7 @@ class DemoState: return 0 stale = 0 for index, event in enumerate(self.events): - if event.get("type") != "response.audio.delta": + if event.get("type") != "response.output_audio.delta": continue metadata = event.get("metadata") if not isinstance(metadata, dict): @@ -1126,7 +1128,7 @@ async def _send_clean_turn( state, lambda: state.response_audio_delta_count(response_id) > 0, timeout_s=timeout_s, - label=f"{transcript} response.audio.delta", + label=f"{transcript} response.output_audio.delta", ) await _wait_for( state, @@ -1341,7 +1343,7 @@ async def _send_listen_only_overlap_pair( state, lambda: state.response_audio_delta_count(first_response_id) > 0, timeout_s=timeout_s, - label=f"{transcripts[0]} response.audio.delta", + label=f"{transcripts[0]} response.output_audio.delta", ) if state.response_done(first_response_id): raise RuntimeError("first response completed before overlap input could start") @@ -1600,7 +1602,9 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: transcript_hints_enabled=transcript_hints_enabled and not continuous_input, ) model_speak_event_ok = state.model_speak_before_audio_ok() - realtime_audio_lifecycle_ok = state.count("response.audio.delta") > 0 and state.count("response.audio.done") > 0 + realtime_audio_lifecycle_ok = ( + state.count("response.output_audio.delta") > 0 and state.count("response.output_audio.done") > 0 + ) completed_response_ids = state.completed_response_ids() observed_turn_response_ids = [response_id for response_id in turn_response_ids if isinstance(response_id, str)] expected_empty_response_ids = { @@ -1625,7 +1629,7 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: state.count("response.created") == state.count("response.done") and len(state.response_ids) == len(completed_response_ids) and len(completed_response_ids) == len(set(completed_response_ids)) - and state.count("response.audio.done") <= state.count("response.done") + and state.count("response.output_audio.done") <= state.count("response.done") and _response_cardinality_ok( completed_response_ids, expected_turns=expected_turns, diff --git a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py index 11c35cecc..6aba7cfe7 100644 --- a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py +++ b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py @@ -215,7 +215,7 @@ async def run(args: argparse.Namespace) -> dict[str, object]: client.events.count("response.function_call_arguments.done") > 0 if args.expect_function_call else ( - client.events.count("response.audio.delta") >= args.minimum_audio_chunks + client.events.count("response.output_audio.delta") >= args.minimum_audio_chunks and ( args.allow_incomplete_response or client.events.count("response.done") > completed_responses_at_commit @@ -279,10 +279,10 @@ async def run(args: argparse.Namespace) -> dict[str, object]: transcript = "".join( str(event.get("delta", "")) for event in later - if event.get("type") == "response.audio_transcript.delta" + if event.get("type") == "response.output_audio_transcript.delta" ).lower() return ( - any(event.get("type") == "response.audio.delta" for event in later) + any(event.get("type") == "response.output_audio.delta" for event in later) and any(event.get("type") == "response.done" for event in later) and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript) ) @@ -297,7 +297,7 @@ async def run(args: argparse.Namespace) -> dict[str, object]: await asyncio.sleep(args.drain_s) audio = client.events.audio_bytes() - audio_events = _events(client, "response.audio.delta") + audio_events = _events(client, "response.output_audio.delta") rates = {event.get("sample_rate_hz") for event in audio_events} if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}: raise AssertionError(f"unexpected output sample rates: {rates}") diff --git a/tests/e2e/online_serving/personaplex_realtime_duplex.py b/tests/e2e/online_serving/personaplex_realtime_duplex.py index 26ee22037..20da1598d 100644 --- a/tests/e2e/online_serving/personaplex_realtime_duplex.py +++ b/tests/e2e/online_serving/personaplex_realtime_duplex.py @@ -244,7 +244,7 @@ def _audio_frame_stats( def _response_ids(client: RawRealtimeProbe) -> set[str]: return { str(event["response_id"]) - for event in _events(client, "response.audio.delta") + for event in _events(client, "response.output_audio.delta") if isinstance(event.get("response_id"), str) } @@ -268,7 +268,7 @@ def _session_result( raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}") rates = { int(event["sample_rate_hz"]) - for event in _events(client, "response.audio.delta") + for event in _events(client, "response.output_audio.delta") if isinstance(event.get("sample_rate_hz"), int) } if rates != {SAMPLE_RATE_HZ}: @@ -289,7 +289,7 @@ def _session_result( ) runtime_metadata = [ metadata["vllm_omni"] - for event in _events(client, "response.audio.delta") + for event in _events(client, "response.output_audio.delta") if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict) ] if not any( diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py index f5539968b..b5d04097b 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Run two MiniCPM-o Realtime duplex demo processes concurrently. This E2E driver covers the same process boundary users exercise manually: @@ -20,7 +23,7 @@ from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[3] DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py" -AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"} +AUDIO_DELTA_EVENTS = {"response.output_audio.delta"} def _canonical_path(path: str) -> Path: diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py index 4047a0901..0fb1c57b3 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import base64 from pathlib import Path @@ -41,7 +44,10 @@ async def run_server_vad_interrupt(args) -> dict[str, object]: next(e for e in client.events.events if e.get("type") == "response.created") ) assert target_id - await until(lambda: _events(client.events.events, target_id, "response.audio.delta"), "response.audio.delta") + await until( + lambda: _events(client.events.events, target_id, "response.output_audio.delta"), + "response.output_audio.delta", + ) cursor = len(client.events.events) await client.stream_pcm16(interrupt + bytes(16_000 * 2 * 800 // 1000), chunk_ms=args.chunk_ms, realtime=True) await until(lambda: _events(client.events.events, target_id, "response.done"), "cancelled response.done") @@ -64,10 +70,11 @@ async def run_server_vad_interrupt(args) -> dict[str, object]: terminal = done[0] if len(done) == 1 else None trailing = events[events.index(terminal) + 1 :] if terminal is not None else events stale = any( - e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) == target_id for e in trailing + e.get("type") == "response.output_audio.delta" and RealtimeEventCollector.response_id(e) == target_id + for e in trailing ) followup_audio = any( - e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) != target_id + e.get("type") == "response.output_audio.delta" and RealtimeEventCollector.response_id(e) != target_id for e in events[cursor:] ) ok = ( diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py index 3aff0cf8d..5e0cfff3a 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Validate MiniCPM-o Realtime duplex soft-interrupt delta streaming. This E2E driver runs the public ``realtime_duplex_demo.py`` against a live @@ -21,9 +24,9 @@ from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[3] DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py" -AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"} +AUDIO_DELTA_EVENTS = {"response.output_audio.delta"} TRANSCRIPT_DELTA_EVENTS = { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", } diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index 22c8e6b34..d45c535db 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -168,9 +168,9 @@ def test_realtime_duplex_demo_pair_launches_demo_processes_concurrently(tmp_path "events = [", " {'type': 'session.created'},", " {'type': 'response.created', 'response': {'id': response_id}},", - " {'type': 'response.audio.delta', 'response_id': response_id, " + " {'type': 'response.output_audio.delta', 'response_id': response_id, " "'delta': 'AAAA', 'sample_rate_hz': 24000, '_client_received_at_s': 10.1},", - " {'type': 'response.audio.delta', 'response_id': response_id, " + " {'type': 'response.output_audio.delta', 'response_id': response_id, " "'delta': 'AAAA', 'sample_rate_hz': 24000, '_client_received_at_s': 10.6},", " {'type': 'response.done', 'response_id': response_id},", " {'type': 'session.closed'},", @@ -237,8 +237,8 @@ def test_realtime_duplex_demo_pair_rejects_false_green_audio_contract(tmp_path): response_id = "resp-a" events = [ {"type": "response.created", "response": {"id": response_id}}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "sample_rate_hz": 24000}, - {"type": "response.audio.delta", "response_id": "resp-other", "delta": "AAAA", "sample_rate_hz": 24000}, + {"type": "response.output_audio.delta", "response_id": response_id, "delta": "AAAA", "sample_rate_hz": 24000}, + {"type": "response.output_audio.delta", "response_id": "resp-other", "delta": "AAAA", "sample_rate_hz": 24000}, {"type": "response.done", "response_id": response_id}, ] (output / "events.jsonl").write_text( @@ -459,25 +459,25 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.1, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "中国古代四大发明", "_client_received_at_s": 2.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "是造纸术。", "_client_received_at_s": 2.6, @@ -486,25 +486,25 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 4.0}, {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 4.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 4.1, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "一加一等于", "_client_received_at_s": 4.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 4.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "二。", "_client_received_at_s": 4.6, @@ -561,25 +561,25 @@ def test_realtime_duplex_soft_interrupt_accepts_followup_done_after_commit(tmp_p {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.1, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "中国古代四大发明", "_client_received_at_s": 2.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "是造纸术。", "_client_received_at_s": 2.6, @@ -589,25 +589,25 @@ def test_realtime_duplex_soft_interrupt_accepts_followup_done_after_commit(tmp_p {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.5}, {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 3.5}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "一加一等于", "_client_received_at_s": 3.6, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.7, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "二。", "_client_received_at_s": 3.7, @@ -653,8 +653,18 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp events = [ {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.15}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.15, + }, {"type": "response.done", "response_id": response_id, "_client_received_at_s": 2.2}, {"type": "response.listen", "_client_received_at_s": 2.3}, {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, @@ -692,10 +702,25 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_commit_during_speak events = [ {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.2}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.2, + }, {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 4.0}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 4.0, + }, {"type": "response.done", "response_id": response_id, "_client_received_at_s": 5.0}, ] (output / "events.jsonl").write_text( @@ -729,8 +754,18 @@ def test_realtime_duplex_soft_interrupt_response_required_rejects_single_respons events = [ {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.15}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.15, + }, {"type": "response.done", "response_id": response_id, "_client_received_at_s": 2.2}, {"type": "response.listen", "_client_received_at_s": 2.3}, {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, @@ -767,19 +802,19 @@ def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating( {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.2, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "一加一等于二", "_client_received_at_s": 2.2, @@ -788,19 +823,19 @@ def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating( {"type": "response.listen", "_client_received_at_s": 2.4}, {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.2, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "不知道", "_client_received_at_s": 3.2, @@ -834,7 +869,7 @@ def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating( events = [ event for event in events - if event.get("response_id") != second_response_id or event["type"] != "response.audio_transcript.delta" + if event.get("response_id") != second_response_id or event["type"] != "response.output_audio_transcript.delta" ] (output / "events.jsonl").write_text( "".join(demo.json.dumps(event) + "\n" for event in events), @@ -1123,7 +1158,7 @@ def test_realtime_duplex_demo_partitions_timing_by_response_identity(): ) state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": base64.b64encode(b"audio").decode("ascii"), "metadata": { @@ -1142,7 +1177,7 @@ def test_realtime_duplex_demo_partitions_timing_by_response_identity(): ) state.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": "audio", }, @@ -1251,8 +1286,8 @@ def test_realtime_duplex_demo_distinguishes_late_and_missing_commit(): {"type": "response.output_item.added", "response_id": response_id}, {"type": "response.content_part.added", "response_id": response_id}, {"type": "response.speak", "response_id": response_id}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA"}, - {"type": "response.audio.done", "response_id": response_id}, + {"type": "response.output_audio.delta", "response_id": response_id, "delta": "AAAA"}, + {"type": "response.output_audio.done", "response_id": response_id}, {"type": "response.content_part.done", "response_id": response_id}, {"type": "response.output_item.done", "response_id": response_id}, {"type": "response.done", "response_id": response_id}, @@ -1314,7 +1349,7 @@ def test_realtime_duplex_demo_distinct_inputs_compare_audio_content(): def _add_response_transcript(state, response_id, *, transcript, audio=True): state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": "YQ==" if audio else "", } @@ -1322,14 +1357,14 @@ def _add_response_transcript(state, response_id, *, transcript, audio=True): if transcript: state.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": transcript, } ) state.add( { - "type": "response.audio_transcript.done", + "type": "response.output_audio_transcript.done", "response_id": response_id, "transcript": transcript, } @@ -1404,7 +1439,7 @@ def test_realtime_duplex_demo_speak_gate_rejects_text_channel(): { "type": "response.speak", "response_id": "resp-1", - "text": "must be emitted through response.audio_transcript.delta", + "text": "must be emitted through response.output_audio_transcript.delta", } ) @@ -1450,7 +1485,7 @@ def test_realtime_duplex_demo_gate_rejects_delta_done_mismatch(): demo = _load_demo_module() state = demo.DemoState() _add_response_transcript(state, "resp-1", transcript="delta文本") - done = next(event for event in state.events if event.get("type") == "response.audio_transcript.done") + done = next(event for event in state.events if event.get("type") == "response.output_audio_transcript.done") done["transcript"] = "另一个done文本" result = demo._evaluate_transcript_integrity( @@ -2180,7 +2215,7 @@ def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_fi state.add({"type": "response.created", "response": {"id": "resp-first"}}) state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-first", "delta": "YQ==", } @@ -2290,7 +2325,7 @@ def test_realtime_duplex_demo_writes_audio_per_response(tmp_path): state.add({"type": "response.created", "response": {"id": response_id}}) state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": demo.base64.b64encode(payload).decode(), "sample_rate_hz": 24000, diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py index f30266d28..013ab8661 100644 --- a/tests/engine/test_duplex_import_boundary.py +++ b/tests/engine/test_duplex_import_boundary.py @@ -122,8 +122,8 @@ from vllm_omni.clients.duplex import EventCollector collector = EventCollector() chunk = base64.b64encode(b"\\x00\\x00" * 2400).decode("ascii") collector.add({"type": "response.created", "response_id": "r1"}, received_at_s=1.0) -collector.add({"type": "response.audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.1) -collector.add({"type": "response.audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.2) +collector.add({"type": "response.output_audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.1) +collector.add({"type": "response.output_audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.2) summary = collector.timing_summary(after_s=0.0, response_id="r1") if "audio_output" not in summary or "request_metrics" not in summary: raise SystemExit("timing_summary produced no audio_output/request_metrics sections") diff --git a/tests/entrypoints/duplex/test_websocket_actor.py b/tests/entrypoints/duplex/test_websocket_actor.py index 9bdb73667..4e5d151b9 100644 --- a/tests/entrypoints/duplex/test_websocket_actor.py +++ b/tests/entrypoints/duplex/test_websocket_actor.py @@ -122,12 +122,12 @@ async def test_stale_fence_payload_is_dropped_before_websocket_send(): actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2) writer = asyncio.create_task(actor.writer_loop()) - await actor.send_json({"type": "response.audio.delta", "epoch": 1}) - await actor.send_json({"type": "response.audio.delta", "epoch": 2}) + await actor.send_json({"type": "response.output_audio.delta", "epoch": 1}) + await actor.send_json({"type": "response.output_audio.delta", "epoch": 2}) await actor.close_writer() await writer - assert websocket.sent == [{"type": "response.audio.delta", "epoch": 2}] + assert websocket.sent == [{"type": "response.output_audio.delta", "epoch": 2}] assert actor.stale_output_dropped == 1 diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index b4ea9c328..a73159064 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -78,7 +78,7 @@ def test_event_journal_sequences_acknowledges_and_replays_exact_payloads() -> No clock = _Clock() journal = DuplexEventJournal(max_bytes=4096, ttl_s=60.0, clock=clock) - first = journal.record({"type": "response.audio.delta", "delta": "AAAA"}) + first = journal.record({"type": "response.output_audio.delta", "delta": "AAAA"}) second = journal.record({"type": "response.done", "response_id": "resp-1"}) assert first.sequence == 1 diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index d3bbc5d01..0f9698457 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -2197,9 +2197,11 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ) audio_events = [ - payload for payload in payloads if payload["type"] in {"response.output_audio.delta", "response.audio.delta"} + payload + for payload in payloads + if payload["type"] in {"response.output_audio.delta", "response.output_audio.delta"} ] - assert {payload["type"] for payload in audio_events} == {"response.audio.delta"} + assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} assert {payload["format"] for payload in audio_events} == {"pcm16"} assert {payload["sample_rate_hz"] for payload in audio_events} == {24000} @@ -2227,14 +2229,14 @@ def test_native_realtime_protocol_ignores_removed_legacy_event_switches(): ) event_types = {payload["type"] for payload in payloads} - assert "response.audio.delta" in event_types - assert "response.audio.done" in event_types - assert "response.audio_transcript.delta" in event_types - assert "response.audio_transcript.done" in event_types - assert "response.output_audio.delta" not in event_types - assert "response.output_audio.done" not in event_types - assert "response.output_audio_transcript.delta" not in event_types - assert "response.output_audio_transcript.done" not in event_types + assert "response.output_audio.delta" in event_types + assert "response.output_audio.done" in event_types + assert "response.output_audio_transcript.delta" in event_types + assert "response.output_audio_transcript.done" in event_types + assert "response.audio.delta" not in event_types + assert "response.audio.done" not in event_types + assert "response.audio_transcript.delta" not in event_types + assert "response.audio_transcript.done" not in event_types assert "response.output_item.created" not in event_types assert "response.text.delta" not in event_types @@ -2256,9 +2258,9 @@ def test_native_realtime_protocol_emits_terminal_audio_transcript_events(): ) by_type = {payload["type"]: payload for payload in payloads} - assert "response.audio.done" in by_type - assert "response.audio_transcript.delta" in by_type - assert "response.audio_transcript.done" in by_type + assert "response.output_audio.done" in by_type + assert "response.output_audio_transcript.delta" in by_type + assert "response.output_audio_transcript.done" in by_type assert by_type["response.content_part.done"]["part"]["transcript"] == "hello" assert by_type["response.output_item.done"]["item"]["object"] == "realtime.item" assert by_type["response.done"]["response"]["output"][0]["object"] == "realtime.item" @@ -2284,8 +2286,10 @@ def test_native_realtime_protocol_terminal_transcript_equals_joined_deltas(): ) ) - transcript = "".join(str(event["delta"]) for event in events if event["type"] == "response.audio_transcript.delta") - done = [str(event["transcript"]) for event in events if event["type"] == "response.audio_transcript.done"] + transcript = "".join( + str(event["delta"]) for event in events if event["type"] == "response.output_audio_transcript.delta" + ) + done = [str(event["transcript"]) for event in events if event["type"] == "response.output_audio_transcript.done"] assert transcript == "It's Canberra. Next question." assert done == [transcript] @@ -2374,7 +2378,7 @@ def test_native_realtime_protocol_response_done_emits_audio_done_before_terminal payloads = protocol._from_duplex_event({"type": "response.done", "response_id": "resp-done"}) event_types = [payload["type"] for payload in payloads] - assert event_types.index("response.audio.done") < event_types.index("response.content_part.done") + assert event_types.index("response.output_audio.done") < event_types.index("response.content_part.done") assert event_types.index("response.content_part.done") < event_types.index("response.output_item.done") assert event_types.index("conversation.item.done") < event_types.index("response.done") assert event_types.index("response.output_item.done") < event_types.index("response.done") @@ -3181,7 +3185,7 @@ async def test_auto_response_committed_overlap_does_not_precreate_empty_response def commit_overlap_after_first_response(ws: TimedWebSocket, data: dict[str, Any]) -> None: nonlocal overlap_sent - if data.get("type") != "response.audio.delta" or overlap_sent: + if data.get("type") != "response.output_audio.delta" or overlap_sent: return overlap_sent = True chunk = { @@ -8679,7 +8683,7 @@ async def test_minicpmo_native_duplex_uses_segment_text_metadata_for_transcript_ protocol = NativeRealtimeSessionProtocol(TimedWebSocket()) # type: ignore[arg-type] realtime_events = [event for payload in sent for event in protocol._from_duplex_event(payload)] transcript_deltas = [ - m for m in realtime_events if m.get("type") == "response.audio_transcript.delta" and m.get("delta") + m for m in realtime_events if m.get("type") == "response.output_audio_transcript.delta" and m.get("delta") ] assert len(created) == 2 assert len(done) == 2 diff --git a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py index d4a70e1b5..ba17ec656 100644 --- a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py +++ b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py @@ -60,7 +60,7 @@ def test_post_commit_decision_accepts_drained_speak_after_commit_ack(): events = [ {"type": "input_audio_buffer.committed"}, {"type": "response.created", "response": {"id": "resp-a"}}, - {"type": "response.audio.delta", "response_id": "resp-a", "delta": "AAAA"}, + {"type": "response.output_audio.delta", "response_id": "resp-a", "delta": "AAAA"}, {"type": "response.done", "response_id": "resp-a"}, ] @@ -267,7 +267,7 @@ def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, writer.handle( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(first_pcm).decode("ascii"), "sample_rate_hz": 24_000, @@ -275,14 +275,14 @@ def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, ) writer.handle( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "你好", } ) writer.handle( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(second_pcm).decode("ascii"), "sample_rate_hz": 24_000, diff --git a/tests/examples/test_minicpmo_realtime_web_static.py b/tests/examples/test_minicpmo_realtime_web_static.py index 5fdab87d1..926ae36ef 100644 --- a/tests/examples/test_minicpmo_realtime_web_static.py +++ b/tests/examples/test_minicpmo_realtime_web_static.py @@ -42,8 +42,8 @@ def test_client_uses_proxy_relative_realtime_url_and_model_policy_session(): assert "input_audio_buffer.commit" not in source assert "playback.ack" in source assert "event.event || event" in source - assert "response.audio.delta" in source - assert "response.audio_transcript.delta" in source + assert "response.output_audio.delta" in source + assert "response.output_audio_transcript.delta" in source assert "conversation.item.input_audio_transcription" in source assert "force_barge_in" not in source assert "server_vad" not in source diff --git a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py index 8d4c1b589..1f03e894b 100644 --- a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py @@ -61,7 +61,7 @@ def _audio_client( frames[:voiced_frames] = 1000 raw = frames.tobytes() event = { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "response-1", "sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ, "metadata": { diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py index 03997fc8e..494f9a856 100644 --- a/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py @@ -49,7 +49,11 @@ def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "m last: float | None = None for event in events: event_type = event.get("type") - if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}: + if event_type not in { + "response.output_text.delta", + "response.output_audio_transcript.delta", + "response.text.delta", + }: continue delta = event.get("delta") if not isinstance(delta, str) or not delta.strip(): diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py index b821d0235..6b861d339 100644 --- a/vllm_omni/benchmarks/omniinteract.py +++ b/vllm_omni/benchmarks/omniinteract.py @@ -233,7 +233,7 @@ class _Playback: if response_id: self.completed.add(response_id) continue - if event.get("type") != "response.audio.delta": + if event.get("type") != "response.output_audio.delta": continue encoded = event.get("delta") or event.get("audio") if not response_id: @@ -241,13 +241,13 @@ class _Playback: if not isinstance(encoded, str): raise ValueError("response audio payload is missing") if event.get("format") is None: - self._warn_once("response.audio.delta omitted format; assumed pcm16") + self._warn_once("response.output_audio.delta omitted format; assumed pcm16") elif event.get("format") != "pcm16": raise ValueError("OmniInteract output must be pcm16") rate = event.get("sample_rate_hz") if rate is None: rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE - self._warn_once(f"response.audio.delta omitted sample_rate_hz; assumed {rate}") + self._warn_once(f"response.output_audio.delta omitted sample_rate_hz; assumed {rate}") if rate != OUTPUT_SAMPLE_RATE: raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio") try: @@ -683,7 +683,7 @@ def _collect_output( timing = response_times.setdefault(response_id, [start_s, end_s]) timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s) - text_event_types = {"response.audio_transcript.delta", "response.output_text.delta", "response.text.delta"} + text_event_types = {"response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta"} text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types} if None in text_ids or not text_ids <= created: raise ValueError("response transcript has no matching response.created") @@ -726,7 +726,7 @@ def _collect_output( **{key: value for key, value in event.items() if key not in {"delta", "audio"}}, "audio_bytes": audio_sizes[index], } - if event.get("type") == "response.audio.delta" + if event.get("type") == "response.output_audio.delta" else dict(event) for index, event in enumerate(collector.events) ] diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index eb8a6ded0..15a191cc1 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -460,9 +460,7 @@ _EVENT_TYPES: dict[str, type[DuplexEvent]] = { "session.expired": SessionExpired, "response.created": ResponseCreated, "response.done": ResponseDone, - "response.audio.delta": AudioDelta, "response.output_audio.delta": AudioDelta, - "response.audio_transcript.delta": TranscriptDelta, "response.output_audio_transcript.delta": TranscriptDelta, "response.text.delta": TextDelta, "response.output_text.delta": TextDelta, @@ -1359,7 +1357,7 @@ class EventCollector: if self.response_id(event) == response_id and event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } @@ -1406,7 +1404,7 @@ class EventCollector: if ( event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } diff --git a/vllm_omni/entrypoints/duplex/realtime_output.py b/vllm_omni/entrypoints/duplex/realtime_output.py index 82ce9761e..fcbdf87d7 100644 --- a/vllm_omni/entrypoints/duplex/realtime_output.py +++ b/vllm_omni/entrypoints/duplex/realtime_output.py @@ -168,7 +168,7 @@ class RealtimeOutputProjector: text = "" payloads.append( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "item_id": self._response_item_id(response_id), "output_index": 0, @@ -839,7 +839,7 @@ class RealtimeOutputProjector: ) payloads.append( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "item_id": item_id, "output_index": 0, @@ -872,7 +872,7 @@ class RealtimeOutputProjector: state.audio_done_emitted = True payloads.append( { - "type": "response.audio.done", + "type": "response.output_audio.done", "response_id": response_id, "item_id": item_id, "output_index": 0, @@ -882,7 +882,7 @@ class RealtimeOutputProjector: if transcript: payloads.append( { - "type": "response.audio_transcript.done", + "type": "response.output_audio_transcript.done", "response_id": response_id, "item_id": item_id, "output_index": 0, diff --git a/vllm_omni/entrypoints/duplex/warmup.py b/vllm_omni/entrypoints/duplex/warmup.py index c380775ef..5051a1d31 100644 --- a/vllm_omni/entrypoints/duplex/warmup.py +++ b/vllm_omni/entrypoints/duplex/warmup.py @@ -108,7 +108,7 @@ async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None: sent += 1 await asyncio.sleep(0.08) # Wait for the pipeline's first audio output so every stage ran. - saw_audio = await _recv_until(lambda e: e.get("type") == "response.audio.delta", 30) + saw_audio = await _recv_until(lambda e: e.get("type") == "response.output_audio.delta", 30) # Close the session EXPLICITLY: a bare websocket close parks the # session in its disconnect grace window, where it would keep # occupying the max_sessions=1 slot and block the first client. diff --git a/vllm_omni/entrypoints/duplex/websocket.py b/vllm_omni/entrypoints/duplex/websocket.py index df50fd0c3..d51a7ed57 100644 --- a/vllm_omni/entrypoints/duplex/websocket.py +++ b/vllm_omni/entrypoints/duplex/websocket.py @@ -55,9 +55,7 @@ MODEL_OUTPUT_EVENTS = frozenset( "response.output_item.added", "response.content_part.added", "response.output_audio.delta", - "response.audio.delta", "response.output_audio.done", - "response.audio.done", "response.output_text.delta", "response.output_text.done", "response.text.delta", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index a9256d953..2f4e472f8 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1600,7 +1600,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Field naming follows the OpenAI ``speech.audio.delta`` schema, which carries the base64 chunk in ``audio`` (not ``delta`` — that is the - Realtime API ``response.audio.delta`` convention, a different event). + Realtime API ``response.output_audio.delta`` convention, a different event). See https://platform.openai.com/docs/api-reference/audio-streaming. The terminal ``speech.audio.done`` event carries a ``usage`` object diff --git a/vllm_omni/experimental/fullduplex/client.py b/vllm_omni/experimental/fullduplex/client.py index 252fc92df..45cdde482 100644 --- a/vllm_omni/experimental/fullduplex/client.py +++ b/vllm_omni/experimental/fullduplex/client.py @@ -234,7 +234,7 @@ class RealtimeEventCollector: event_type = stored_event.get("type") if event_type == "response.created" and response_id and response_id not in self.response_ids: self.response_ids.append(response_id) - if event_type == "response.audio.delta": + if event_type == "response.output_audio.delta": delta = stored_event.get("delta") or stored_event.get("audio") if isinstance(delta, str) and response_id: try: @@ -263,7 +263,7 @@ class RealtimeEventCollector: if self.response_id(event) == response_id and event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } @@ -322,7 +322,7 @@ class RealtimeEventCollector: if ( event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } @@ -337,7 +337,7 @@ class RealtimeEventCollector: if isinstance(stage0, dict): stage0_metrics = stage0 - if event.get("type") != "response.audio.delta" or ( + if event.get("type") != "response.output_audio.delta" or ( response_id is not None and event_response_id != response_id ): continue From 44fed170321fd8a8e0f117c9229dfcb13023f8c3 Mon Sep 17 00:00:00 2001 From: Jim Ban <77719403+BANANASJIM@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:03:26 -0700 Subject: [PATCH 016/209] [Bugfix][Core] Clean up failed HWR atomic metadata writes (#6956) Signed-off-by: BANANASJIM Co-authored-by: Hongsheng Liu --- .../test_filesystem_store.py | 149 ++++++++++++++++++ .../host_weight_runtime/filesystem/store.py | 35 ++-- 2 files changed, 168 insertions(+), 16 deletions(-) diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py index 7fe716c5d..d149d34bd 100644 --- a/tests/host_weight_runtime/test_filesystem_store.py +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -20,6 +20,7 @@ from collections.abc import Sequence from dataclasses import replace from multiprocessing.process import BaseProcess from pathlib import Path +from typing import IO, Any import pytest import torch @@ -408,6 +409,154 @@ def _make_store( return FilesystemHostWeightStore(domain, capacity or CapacityPolicy(), integrity or IntegrityPolicy()) +class _FaultyTemporaryFile: + """NamedTemporaryFile proxy that injects a write or close failure.""" + + def __init__( + self, + handle: IO[bytes], + *, + write_error: BaseException | None = None, + close_error: BaseException | None = None, + ) -> None: + self._handle = handle + self._write_error = write_error + self._close_error = close_error + + def __getattr__(self, name: str) -> object: + return getattr(self._handle, name) + + def __enter__(self) -> _FaultyTemporaryFile: + return self + + def __exit__(self, *exc_info: object) -> None: + self._handle.close() + if self._close_error is not None: + raise self._close_error + + def write(self, data: bytes) -> int: + if self._write_error is not None: + raise self._write_error + return self._handle.write(data) + + +def _inject_temporary_file_failure( + monkeypatch: pytest.MonkeyPatch, + *, + write_error: BaseException | None = None, + close_error: BaseException | None = None, +) -> None: + original = tempfile.NamedTemporaryFile + + def faulty(*args: Any, **kwargs: Any) -> _FaultyTemporaryFile: + handle: IO[bytes] = original(*args, **kwargs) + return _FaultyTemporaryFile(handle, write_error=write_error, close_error=close_error) + + monkeypatch.setattr(tempfile, "NamedTemporaryFile", faulty) + + +def test_atomic_json_replace_failure_removes_temporary_file( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "domain.json" + path.write_bytes(b"existing metadata") + replace_error = OSError(errno.EIO, "injected replace failure") + + def fail_replace(source: os.PathLike[str], destination: os.PathLike[str]) -> None: + raise replace_error + + monkeypatch.setattr(os, "replace", fail_replace) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert error.value is replace_error + assert path.read_bytes() == b"existing metadata" + assert not list(tmp_path.glob(".domain.json.*.tmp")) + + +def test_atomic_json_cleanup_failure_preserves_replace_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "domain.json" + replace_error = OSError(errno.EIO, "injected replace failure") + original_unlink = Path.unlink + cleanup_attempted = False + + def fail_replace(source: os.PathLike[str], destination: os.PathLike[str]) -> None: + raise replace_error + + def fail_temporary_unlink(target: Path, missing_ok: bool = False) -> None: + nonlocal cleanup_attempted + if target.parent == tmp_path and target.name.startswith(".domain.json."): + cleanup_attempted = True + raise PermissionError(errno.EACCES, "injected cleanup failure") + original_unlink(target, missing_ok=missing_ok) + + monkeypatch.setattr(os, "replace", fail_replace) + monkeypatch.setattr(Path, "unlink", fail_temporary_unlink) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert cleanup_attempted + assert error.value is replace_error + assert len(list(tmp_path.glob(".domain.json.*.tmp"))) == 1 + + +@pytest.mark.parametrize("close_fails", [False, True]) +def test_atomic_json_write_failure_preserves_error_when_cleanup_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + close_fails: bool, +) -> None: + path = tmp_path / "domain.json" + write_error = OSError(errno.EIO, "injected write failure") + original_unlink = Path.unlink + cleanup_attempted = False + + def fail_temporary_unlink(target: Path, missing_ok: bool = False) -> None: + nonlocal cleanup_attempted + if target.parent == tmp_path and target.name.startswith(".domain.json."): + cleanup_attempted = True + raise PermissionError(errno.EACCES, "injected cleanup failure") + original_unlink(target, missing_ok=missing_ok) + + _inject_temporary_file_failure( + monkeypatch, + write_error=write_error, + close_error=OSError(errno.EIO, "injected close failure") if close_fails else None, + ) + monkeypatch.setattr(Path, "unlink", fail_temporary_unlink) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert cleanup_attempted + assert error.value is write_error + assert not path.exists() + assert len(list(tmp_path.glob(".domain.json.*.tmp"))) == 1 + + +def test_atomic_json_close_failure_removes_temporary_file( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "domain.json" + close_error = OSError(errno.EIO, "injected close failure") + + _inject_temporary_file_failure(monkeypatch, close_error=close_error) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert error.value is close_error + assert not path.exists() + assert not list(tmp_path.glob(".domain.json.*.tmp")) + + def _publish_test_artifact( store: FilesystemHostWeightStore, identity: WeightArtifactIdentity | None = None, diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index ec6b49f3c..1542adefc 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -12,6 +12,7 @@ import logging import os import shutil import stat +import tempfile import time import uuid from collections import defaultdict @@ -298,25 +299,27 @@ def _move_artifact_locked(source: Path, destination: Path) -> None: def _write_atomic_json(path: Path, value: object, *, mode: int = 0o444) -> None: data = canonical_json(value) - temporary = path.parent / f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp" - fd = os.open(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600) + handle = tempfile.NamedTemporaryFile(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp", delete=False) + temporary = Path(handle.name) + write_error: BaseException | None = None try: - offset = 0 - while offset < len(data): - count = os.write(fd, data[offset:]) - if count <= 0: - raise OSError(errno.EIO, f"short write for {temporary}") - offset += count - os.fsync(fd) - os.fchmod(fd, mode) - os.fsync(fd) + with handle: + try: + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) + os.fchmod(handle.fileno(), mode) + os.fsync(handle.fileno()) + except BaseException as error: + write_error = error + raise + os.replace(temporary, path) except BaseException: - os.close(fd) - temporary.unlink(missing_ok=True) + with contextlib.suppress(OSError): + temporary.unlink(missing_ok=True) + if write_error is not None: + raise write_error raise - else: - os.close(fd) - os.replace(temporary, path) _fsync_directory(path.parent) From 3d952d1334f78ac6a4adc595974264457a8eb810 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 11 Sep 2026 09:26:35 +0800 Subject: [PATCH 017/209] [Bugfix] Keep MiniMax-H3 reference audio budgets separate (#7281) Signed-off-by: david6666666 <530634352@qq.com> --- .../minimax_h3/test_minimax_h3_contract.py | 50 +++++++++++++++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 5 +- 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index a94f842c5..5470cd9c4 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -2455,6 +2455,56 @@ def test_g4_standalone_audio_duration_and_total_duration_contract(): ) +def test_ref2va_video_soundtrack_does_not_consume_standalone_audio_budget(): + from PIL import Image + + from vllm_omni.diffusion.cache.cachedit.runtime import RequestScopedCacheDiTRuntime + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "ref2va" + pipeline.supported_tasks = frozenset({"ref2va"}) + pipeline.device = torch.device("cpu") + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline._resolve_sigma_positions = Mock(return_value=(None, 50)) + pipeline._quality_policy = MiniMaxH3QualityPolicy(None) + pipeline._cache_dit_runtime = Mock(spec=RequestScopedCacheDiTRuntime) + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._prepare_reference_videos = Mock(return_value=[{"input_has_audio": True}]) + pipeline._encode_visual_conditions = Mock(return_value=(torch.ones(1, 96), [(1, 16, 16), (17, 48, 84)])) + # A 2.35-second video soundtrack and a 15-second standalone reference + # each satisfy their modality's separate 15-second duration budget. + embedded = torch.full((188, 32), 1.0) + standalone = torch.full((1200, 32), 2.0) + pipeline._encode_reference_audio_conditions = Mock(return_value=(embedded, [94], standalone, [600])) + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_inference_steps=50, + extra_args={"task": "ref2va", "duration": 15.0}, + ) + + context = pipeline._prepare_request_inputs( + prompt="A person carrying books down a street", + multi_modal_data={ + "image": Image.new("RGB", (256, 256)), + "video": ["reference.mp4"], + "audio": (torch.zeros(1, 15 * 16000), 16000), + }, + sampling=sampling, + ) + + assert context["audio_condition_lengths"] == [94, 600] + assert [block["kind"] for block in context["ref_blocks"]] == ["image", "video_audio", "audio"] + torch.testing.assert_close(context["audio_condition"], torch.cat([embedded, standalone])) + + def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index ff1f8a3e5..fb9f02b28 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -2458,8 +2458,9 @@ class MiniMaxH3Pipeline( if audio_lengths: if any(length < 80 or length > 600 for length in audio_lengths): raise OmniClientError("MiniMax H3 audio references must each be between 2 and 15 seconds") - if sum(audio_lengths) > 600: - raise OmniClientError("MiniMax H3 audio references must be at most 15 seconds in total") + # Video soundtracks and standalone audio have separate + # 15-second budgets, validated before encoding. Do not merge + # those budgets when concatenating their conditioning rows. if len(audio_lengths) == 1: ref_audio_t = audio_lengths[0] From bad50980c9e2679d4fa6352b0fbb121af9dfbb40 Mon Sep 17 00:00:00 2001 From: Yan Cao <31481315+yancaocn@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:43:34 +0800 Subject: [PATCH 018/209] [Bugfix] Fix Helios USP: per-component split for correct sequence parallelism (#6930) Signed-off-by: yancaocn Co-authored-by: yancaocn --- tests/diffusion/models/helios/__init__.py | 0 .../models/helios/test_helios_usp_2rank.py | 258 ++++++++++++++++++ .../models/helios/test_helios_usp_shard.py | 225 +++++++++++++++ .../models/helios/helios_transformer.py | 90 +++++- 4 files changed, 561 insertions(+), 12 deletions(-) create mode 100644 tests/diffusion/models/helios/__init__.py create mode 100644 tests/diffusion/models/helios/test_helios_usp_2rank.py create mode 100644 tests/diffusion/models/helios/test_helios_usp_shard.py diff --git a/tests/diffusion/models/helios/__init__.py b/tests/diffusion/models/helios/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/helios/test_helios_usp_2rank.py b/tests/diffusion/models/helios/test_helios_usp_2rank.py new file mode 100644 index 000000000..7b16d1d4f --- /dev/null +++ b/tests/diffusion/models/helios/test_helios_usp_2rank.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Two-process USP smoke test for the Helios transformer. + +Pins the multi-rank mechanism that ``test_helios_usp_shard.py`` (unit tests +of ``_sp_split_seq`` alone) cannot see: + + * the manual ``_sp_shard_depth`` bump in ``forward()`` must make attn1 run + the Ulysses all-to-all (not ``NoParallelAttention``); + * attn2 cross-attention must stay rank-local (``skip_sequence_parallel= + True``): if replicated text K/V were fed through Ulysses they would be + ws-fold duplicated and outputs would diverge from the single-rank run; + * end-to-end correctness: ws=2 output equals the ws=1 baseline for the same + weights/inputs (divisible shapes, so the padding path is not involved). + +Runs ``HeliosTransformer3DModel.forward`` with tiny random weights via +``torch.multiprocessing.spawn``: first a baseline (world_size=1), then an +SP run (world_size=2, ulysses_degree=2); rank-0 outputs are compared. + +Requires >= 2 devices; skipped otherwise. +""" + +import os +import tempfile + +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelConfig, + get_sp_plan_from_model, +) +from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context +from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel +from vllm_omni.platforms import current_omni_platform + +_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2) +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, *_L4_TWO_GPU] + +RANDOM_SEED = 42 +DTYPE = torch.float32 # keep both runs on the SDPA fallback for tight tolerance + +# Tiny but architecturally faithful config: +# inner_dim = 4 * 32 = 128, rope_dim sums to head_dim (32), +# current chunk 2x4x6 latent -> (1,2,2) patch -> 2*2*3 = 12 tokens, +# short history 1x4x6 latent -> (1,2,2) patch -> 6 tokens. +# All component lengths are even, so ws=2 needs no padding replicas. +TINY_MODEL_KWARGS = dict( + patch_size=(1, 2, 2), + num_attention_heads=4, + attention_head_dim=32, + in_channels=4, + out_channels=4, + text_dim=64, + freq_dim=32, + ffn_dim=128, + num_layers=2, + rope_dim=(12, 10, 10), + rope_theta=10000.0, + guidance_cross_attn=True, + zero_history_timestep=True, + has_multi_term_memory_patch=True, + is_amplify_history=False, +) + +BATCH = 1 +CUR_FRAMES, LAT_H, LAT_W = 2, 4, 6 # -> 12 current tokens +HIST_FRAMES = 1 # short history -> 6 tokens +TEXT_LEN = 5 + + +def _update_env(envs_dict: dict[str, str]) -> None: + for k, v in envs_dict.items(): + os.environ[k] = v + + +def _make_inputs(device: torch.device) -> dict[str, torch.Tensor]: + torch.manual_seed(RANDOM_SEED + 1) + return { + "hidden_states": torch.randn(BATCH, 4, CUR_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE), + "timestep": torch.full((BATCH,), 500.0, device=device), + "encoder_hidden_states": torch.randn( + BATCH, TEXT_LEN, TINY_MODEL_KWARGS["text_dim"], device=device, dtype=DTYPE + ), + "indices_latents_history_short": torch.tensor([[0]], device=device), + "latents_history_short": torch.randn(BATCH, 4, HIST_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE), + } + + +def _spy_strategies(model: HeliosTransformer3DModel) -> dict[str, list[str]]: + """Record the parallel strategy each Attention resolves to during forward.""" + seen: dict[str, list[str]] = {"attn1": [], "attn2": []} + for name, attn in (("attn1", model.blocks[0].attn1.attn), ("attn2", model.blocks[0].attn2.attn)): + original = attn._get_active_parallel_strategy + + def spy(original=original, name=name): + strategy = original() + seen[name].append(getattr(strategy, "name", type(strategy).__name__)) + return strategy + + attn._get_active_parallel_strategy = spy + return seen + + +def helios_forward_on_model( + local_rank: int, + world_size: int, + ulysses_degree: int, + output_file: str, + model_state_file: str, + is_baseline: bool, +) -> None: + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + torch.set_default_device(device) + torch.set_default_dtype(DTYPE) + + _update_env( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "12361", + } + ) + init_distributed_environment() + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=ulysses_degree, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=1, + cfg_parallel_size=1, + ) + od_config = OmniDiffusionConfig.from_kwargs( + model="test_model", + dtype=DTYPE, + parallel_config=parallel_config, + diffusion_attention_backend="TORCH_SDPA", + ) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=ulysses_degree, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config): + torch.manual_seed(RANDOM_SEED) + model = HeliosTransformer3DModel(**TINY_MODEL_KWARGS).to(device).to(DTYPE) + + if is_baseline and local_rank == 0: + with open(model_state_file, "wb") as f: + torch.save({k: v.cpu() for k, v in model.state_dict().items()}, f) + + if world_size > 1: + torch.distributed.barrier() + + with open(model_state_file, "rb") as f: + model_state = torch.load(f) + model.load_state_dict({k: v.to(device).to(DTYPE) for k, v in model_state.items()}) + + if not is_baseline: + # Mirror registry._apply_sequence_parallel_if_enabled: install the + # proj_out gather hook from the model's own _sp_plan and flip + # sp_plan_hooks_applied so sp_active follows _sp_shard_depth + # (i.e. the manual bump in Helios forward is what enables Ulysses). + plan = get_sp_plan_from_model(model) + assert plan is not None, "HeliosTransformer3DModel must define _sp_plan" + apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=ulysses_degree), plan) + get_forward_context().sp_plan_hooks_applied = True + + seen = _spy_strategies(model) + inputs = _make_inputs(device) + output = model(**inputs).sample + + assert output.shape == (BATCH, 4, CUR_FRAMES, LAT_H, LAT_W), f"unexpected output shape {output.shape}" + + if is_baseline: + assert seen["attn1"] and all(n == "none" for n in seen["attn1"]), ( + f"baseline attn1 must use NoParallelAttention ('none'), saw {seen['attn1']}" + ) + else: + # attn1: SP communication must be active inside the block loop + # (regression here = the manual _sp_shard_depth bump was lost). + assert seen["attn1"] and all(n == "ulysses" for n in seen["attn1"]), ( + f"attn1 must run Ulysses under USP (manual _sp_shard_depth bump), saw {seen['attn1']}" + ) + # attn2: replicated text K/V must stay rank-local + # (regression here = the ws-fold duplicated-keys bug). + assert seen["attn2"] and all(n == "none" for n in seen["attn2"]), ( + f"attn2 must skip sequence parallel (replicated text K/V), saw {seen['attn2']}" + ) + + if local_rank == 0: + with open(output_file, "wb") as f: + torch.save(output.detach().cpu().float(), f) + + destroy_distributed_env() + + +def test_helios_usp2_matches_single_rank(): + """ws=2 (Ulysses) must reproduce the ws=1 output for identical weights/inputs.""" + available_devices = current_omni_platform.get_device_count() + if available_devices < 2: + pytest.skip(f"Test requires 2 devices but only {available_devices} available") + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + baseline_output_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + sp_output_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + model_state_file = f.name + + try: + torch.multiprocessing.spawn( + helios_forward_on_model, + args=(1, 1, baseline_output_file, model_state_file, True), + nprocs=1, + ) + torch.multiprocessing.spawn( + helios_forward_on_model, + args=(2, 2, sp_output_file, model_state_file, False), + nprocs=2, + ) + + with open(baseline_output_file, "rb") as f: + baseline = torch.load(f) + with open(sp_output_file, "rb") as f: + sp = torch.load(f) + + assert baseline.shape == sp.shape + max_abs_diff = (baseline - sp).abs().max().item() + assert torch.allclose(baseline, sp, atol=1e-3, rtol=1e-3), ( + f"USP=2 output diverges from single-rank baseline: max_abs_diff={max_abs_diff:.6e}" + ) + finally: + for path in [baseline_output_file, sp_output_file, model_state_file]: + if os.path.exists(path): + os.remove(path) diff --git a/tests/diffusion/models/helios/test_helios_usp_shard.py b/tests/diffusion/models/helios/test_helios_usp_shard.py new file mode 100644 index 000000000..88301a851 --- /dev/null +++ b/tests/diffusion/models/helios/test_helios_usp_shard.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for Helios per-component sequence splitting. + +These tests verify the manual sharding logic that replaces the broken +rope/blocks.0 hooks in the Helios ``_sp_plan``. They exercise the *real* +``HeliosTransformer3DModel._sp_split_seq`` (the parallel-state getters it +imports lazily inside the method body are monkeypatched), so the model-side +helper cannot drift silently from what is tested. For the multi-rank +mechanism around the helper (manual ``_sp_shard_depth`` bump, attn1/attn2 +parallel strategy, ws=2 vs ws=1 output equivalence) see +``test_helios_usp_2rank.py``. + +Coverage: + 1. USP disabled (ws == 1) returns the tensor unchanged. + 2. Divisible split gives each rank the expected slice. + 3. Non-divisible split pads with replicated tokens (no crash). + 4. Per-component split: history + current tokens are distributed so + that *both* ranks receive some current tokens (no mosaic regression). + 5. Edge cases: 1-D tensor, empty seq, 4-D tensor. +""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch + +import vllm_omni.diffusion.distributed.parallel_state as sp_parallel_state +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@contextmanager +def _fake_sp_state(ws: int, rank: int) -> Iterator[None]: + """Patch the parallel-state getters used by the real helper. + + ``_sp_split_seq`` imports these names from the module *inside* its + function body, so replacing the module attributes takes effect. + """ + prev_ws = sp_parallel_state.get_sequence_parallel_world_size + prev_rank = sp_parallel_state.get_sequence_parallel_rank + sp_parallel_state.get_sequence_parallel_world_size = lambda: ws + sp_parallel_state.get_sequence_parallel_rank = lambda: rank + try: + yield + finally: + sp_parallel_state.get_sequence_parallel_world_size = prev_ws + sp_parallel_state.get_sequence_parallel_rank = prev_rank + + +def _sp_split_seq(x: torch.Tensor, ws: int, rank: int) -> torch.Tensor: + """Call the real model-side helper without constructing a model.""" + with _fake_sp_state(ws, rank): + return HeliosTransformer3DModel._sp_split_seq(None, x) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestSpSplitSeqDisabled: + """When USP is off (ws == 1) the helper must be a no-op.""" + + def test_ws1_returns_unchanged(self): + x = torch.randn(2, 100, 64) + out = _sp_split_seq(x, ws=1, rank=0) + assert torch.equal(out, x) + assert out is x + + +class TestSpSplitSeqDivisible: + """Divisible sequences: each rank gets exactly seq_len / ws tokens.""" + + def test_each_rank_gets_correct_slice(self): + for ws in [2, 4, 8]: + seq_len = ws * 10 + x = torch.arange(seq_len * 4).float().reshape(1, seq_len, 4) + for rank in range(ws): + out = _sp_split_seq(x, ws=ws, rank=rank) + n = seq_len // ws + expected = x[:, rank * n : (rank + 1) * n, :] + assert out.shape == (1, n, 4) + assert torch.equal(out, expected) + + def test_output_is_contiguous(self): + x = torch.randn(1, 20, 8) + out = _sp_split_seq(x, ws=2, rank=0) + assert out.is_contiguous() + + +class TestSpSplitSeqNonDivisible: + """Non-divisible sequences are padded with replicated tokens.""" + + def test_non_divisible_pads_instead_of_crash(self): + """400x272 → 9*25*17 = 3825 tokens, USP=2 → not divisible. + Should pad to 3826 and split, not crash.""" + x = torch.randn(1, 3825, 8) + for rank in range(2): + out = _sp_split_seq(x, ws=2, rank=rank) + # 3825 → pad 1 → 3826 → /2 = 1913 + assert out.shape[1] == 1913 + + def test_non_divisible_pads_with_replicated_last_token(self): + """Verify the padding is the replicated last token, not zeros.""" + x = torch.randn(1, 7, 4) # 7 % 2 = 1, pad 1 + last_token = x[:, -1:, :].clone() + out = _sp_split_seq(x, ws=2, rank=0) + # rank 0 gets tokens [0:4] (first 4 of the padded 8) + assert out.shape == (1, 4, 4) + # rank 1 gets tokens [4:8] — token 7 (the padded one) = last token + out1 = _sp_split_seq(x, ws=2, rank=1) + assert out1.shape == (1, 4, 4) + assert torch.equal(out1[:, -1:, :], last_token) + + def test_non_divisible_7_4(self): + """7 tokens, ws=4 → pad 1 → 8 → /4 = 2 per rank.""" + x = torch.randn(1, 7, 8) + for rank in range(4): + out = _sp_split_seq(x, ws=4, rank=rank) + assert out.shape[1] == 2 + + def test_non_divisible_15_4(self): + """15 tokens, ws=4 → pad 1 → 16 → /4 = 4 per rank.""" + x = torch.randn(1, 15, 8) + for rank in range(4): + out = _sp_split_seq(x, ws=4, rank=rank) + assert out.shape[1] == 4 + + +class TestSpSplitSeqPerComponent: + """Simulate Helios forward: split current + history components so that + *both* ranks receive some current tokens (no mosaic regression).""" + + def test_both_ranks_get_current_tokens(self): + """history=2400, current=540, ws=2. + After per-component split, each rank must have > 0 current tokens.""" + current = torch.randn(1, 540, 64) + history = torch.randn(1, 2400, 64) + + for rank in range(2): + cur_local = _sp_split_seq(current, ws=2, rank=rank) + hist_local = _sp_split_seq(history, ws=2, rank=rank) + + assert cur_local.shape[1] > 0, f"Rank {rank} received 0 current tokens" + assert cur_local.shape[1] == 270 # 540 / 2 + assert hist_local.shape[1] == 1200 # 2400 / 2 + + original_context_length = cur_local.shape[1] + assert original_context_length > 0 # no mosaic regression + + def test_gather_restores_full_current(self): + """After splitting current into 2 ranks, simulating proj_out + gather must restore the full current sequence length so that + the original unpatchify reshape succeeds.""" + current = torch.randn(1, 540, 64) + local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] + gathered = torch.cat(local_chunks, dim=1) + assert gathered.shape[1] == 540 # full current restored + + def test_gather_with_padding_truncates_correctly(self): + """Non-divisible current (3825) → pad to 3826 → split → + gather → 3826. Slice to 3825 for unpatchify.""" + current = torch.randn(1, 3825, 64) + local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] + gathered = torch.cat(local_chunks, dim=1) + assert gathered.shape[1] == 3826 # padded size + # Simulate the slice before unpatchify + expected_seq = 3825 + if gathered.shape[1] > expected_seq: + gathered = gathered[:, :expected_seq, :] + assert gathered.shape[1] == 3825 # correct after slice + + def test_rotary_emb_aligned_with_hidden_states(self): + """rotary_emb and hidden_states must be split at the same + positions so that tokens align inside the attention layer.""" + hidden = torch.randn(1, 540, 5120) + rotary = torch.randn(1, 540, 256) + for rank in range(2): + h_local = _sp_split_seq(hidden, ws=2, rank=rank) + r_local = _sp_split_seq(rotary, ws=2, rank=rank) + assert h_local.shape[1] == r_local.shape[1] + + def test_whole_split_causes_mosaic_but_per_component_does_not(self): + """Demonstrates why per-component split is needed: + with whole split (simulated), rank 0 gets 0 current tokens; + with per-component split, both ranks get current tokens.""" + history = torch.randn(1, 2400, 64) + current = torch.randn(1, 540, 64) + full = torch.cat([history, current], dim=1) # 2940 + + # Whole split (the OLD broken behavior) + whole_n = full.shape[1] // 2 # 1470 + rank0_whole = full[:, :whole_n, :] + rank0_current = rank0_whole.shape[1] - 2400 # negative → 0 current + assert rank0_current <= 0, "Whole split should give rank 0 no current" + + # Per-component split (the NEW correct behavior) + for rank in range(2): + cur_local = _sp_split_seq(current, ws=2, rank=rank) + assert cur_local.shape[1] == 270 # both ranks have current + + +class TestSpSplitSeqEdgeCases: + """Edge cases that should not crash or silently misbehave.""" + + def test_1d_tensor_skipped(self): + """1-D tensors must be returned unchanged (dim < 2 guard).""" + x = torch.randn(100) + out = _sp_split_seq(x, ws=2, rank=0) + assert torch.equal(out, x) + + def test_empty_seq_skipped(self): + """Empty sequence dimension must be returned unchanged.""" + x = torch.randn(1, 0, 64) + out = _sp_split_seq(x, ws=2, rank=0) + assert out.shape == x.shape + + def test_4d_tensor_split_on_dim1(self): + """4-D tensor [B, S, H, D] should split along dim=1.""" + x = torch.randn(1, 20, 8, 32) + out = _sp_split_seq(x, ws=2, rank=1) + assert out.shape == (1, 10, 8, 32) diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index 4e3cf5078..6c6ec223d 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/vllm_omni/diffusion/models/helios/helios_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from Helios (https://github.com/BestWishYsh/Helios) import math @@ -27,10 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.distributed.sp_plan import ( - SequenceParallelInput, - SequenceParallelOutput, -) +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelOutput if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -448,6 +445,12 @@ class HeliosCrossAttention(nn.Module): num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, + # Text K/V is replicated on every SP rank, so Ulysses must not be + # used here: its all-to-all would treat the per-rank replicas as + # sequence shards and ws-fold duplicate the keys. Sharded Q over + # full replicated K/V is correct locally and needs no + # communication (same as Wan2.2 cross-attn). + skip_sequence_parallel=True, ) def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -624,16 +627,50 @@ class HeliosTransformer3DModel(nn.Module): _hsdp_shard_conditions = [_is_transformer_block] _sp_plan = { - "rope": { - 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), - }, - "blocks.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), - }, + # "rope" and "blocks.0" are intentionally omitted. + # The rope hook splits along dim=1 of the 5D output [B, D, T, H, W], + # which is the frequency dimension, NOT the sequence dimension. + # The blocks.0 hook splits the concatenated hidden_states as a whole, + # which can put all history tokens in one rank (original_context_length=0). + # Instead, hidden_states and rotary_emb are split per-component in + # forward() after flatten+transpose, ensuring each rank gets half + # of each component (history and current). "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } + def _sp_split_seq(self, x: torch.Tensor) -> torch.Tensor: + """Split tensor along sequence dim (dim=1) for Ulysses SP. + + If seq_len is not divisible by world_size, replicate the last + token(s) to pad to a divisible length. This avoids crashing + on non-divisible history components (which use different patch + sizes and thus produce different token counts). + + Residual caveat: the replicated padding tokens attend unmasked + in attn1 on the last rank (at most ws-1 replicas per component). + The impact is bounded, but this is the first place to look if + odd resolutions ever show quality drift. + """ + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + ) + + ws = get_sequence_parallel_world_size() + if ws > 1 and x.dim() >= 2 and x.shape[1] > 0: + seq_len = x.shape[1] + remainder = seq_len % ws + if remainder != 0: + pad = ws - remainder + # Replicate the last token to fill the padding + last = x[:, -1:, ...].expand(-1, pad, *([-1] * (x.dim() - 2))) + x = torch.cat([x, last], dim=1) + seq_len = x.shape[1] + r = get_sequence_parallel_rank() + n = seq_len // ws + x = x[:, r * n : (r + 1) * n, ...].contiguous() + return x + def __init__( self, patch_size: tuple[int, ...] = (1, 2, 2), @@ -860,6 +897,9 @@ class HeliosTransformer3DModel(nn.Module): device=hidden_states.device, ) rotary_emb = rotary_emb.flatten(2).transpose(1, 2) + # USP: per-component split (each rank gets half of current frames) + hidden_states = self._sp_split_seq(hidden_states) + rotary_emb = self._sp_split_seq(rotary_emb) original_context_length = hidden_states.shape[1] # 2. Process short history latents @@ -876,6 +916,9 @@ class HeliosTransformer3DModel(nn.Module): device=latents_history_short.device, ) rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_short = self._sp_split_seq(latents_history_short) + rotary_emb_history_short = self._sp_split_seq(rotary_emb_history_short) hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) @@ -896,6 +939,9 @@ class HeliosTransformer3DModel(nn.Module): rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_mid = self._sp_split_seq(latents_history_mid) + rotary_emb_history_mid = self._sp_split_seq(rotary_emb_history_mid) hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) @@ -916,6 +962,9 @@ class HeliosTransformer3DModel(nn.Module): rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_long = self._sp_split_seq(latents_history_long) + rotary_emb_history_long = self._sp_split_seq(rotary_emb_history_long) hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) @@ -958,6 +1007,18 @@ class HeliosTransformer3DModel(nn.Module): timestep_proj = timestep_proj.permute(0, 2, 1, 3) # 6. Transformer blocks + # Manually increment _sp_shard_depth so that attention layers + # know SP is active and perform Ulysses All-to-All communication. + # Normally this is done by split_output=True hooks (e.g. rope), + # but those are removed because they split the wrong dimension. + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + if is_forward_context_available(): + get_forward_context()._sp_shard_depth += 1 + hidden_states = hidden_states.contiguous() encoder_hidden_states = encoder_hidden_states.contiguous() rotary_emb = rotary_emb.contiguous() @@ -979,6 +1040,11 @@ class HeliosTransformer3DModel(nn.Module): hidden_states = self.proj_out(hidden_states) # 8. Unpatchify + # proj_out gather may include padded tokens from non-divisible + # sequences. Slice to the expected size before reshape. + expected_seq = post_patch_num_frames * post_patch_height * post_patch_width + if hidden_states.shape[1] > expected_seq: + hidden_states = hidden_states[:, :expected_seq, :] hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 ) From 96360ccb00d0c4c415eaec4369c827e7396e910a Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Fri, 11 Sep 2026 16:28:42 +0800 Subject: [PATCH 019/209] [Perf][Diffusion] Optimize HSDP startup via Rank-0 shared weight loading and accelerated LoRA delta computation (#7005) Signed-off-by: samithuang <285365963@qq.com> --- docs/mkdocs/hooks/search_index.py | 2 +- tests/config/test_omni_config.py | 1 + tests/diffusion/lora/test_loader.py | 12 + .../model_loader/test_diffusers_loader.py | 449 +++++++++++++++++- tests/engine/test_stage_engine_args.py | 1 + vllm_omni/config/omni_config.py | 3 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 1 + vllm_omni/diffusion/lora/loader.py | 21 +- .../model_loader/diffusers_loader.py | 151 +++++- .../model_loader/host_weight_plan.py | 3 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/entrypoints/cli/serve.py | 7 + 13 files changed, 637 insertions(+), 16 deletions(-) diff --git a/docs/mkdocs/hooks/search_index.py b/docs/mkdocs/hooks/search_index.py index 1b9e268d9..8cf23e869 100644 --- a/docs/mkdocs/hooks/search_index.py +++ b/docs/mkdocs/hooks/search_index.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Reduce browser search-index memory for generated API reference pages.""" from material.plugins.search.plugin import SearchIndex diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index da8ffb452..fdb2cd4b9 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -714,6 +714,7 @@ def test_sub_config_fields_match_structured_scopes(): "max_cudagraph_capture_size", "enable_flashinfer_autotune", "enable_multithread_weight_load", + "enable_broadcast_weight_load", "num_weight_load_threads", "disable_autocast", # Per-stage checkpoint resolution for repos whose stages live in diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py index deb5ec80d..4fd876937 100644 --- a/tests/diffusion/lora/test_loader.py +++ b/tests/diffusion/lora/test_loader.py @@ -509,3 +509,15 @@ def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture): assert pipeline2.lora_is_fused is False pipeline2.lora_is_fused = True assert pipeline2.lora_is_fused is True + + +def test_prepare_lora_delta_with_compute_device(): + base_key = "blocks.0.attn.to_q" + state_dict = { + f"{base_key}.lora_A.weight": torch.ones(RANK, HEAD_DIM), + f"{base_key}.lora_B.weight": torch.ones(HEAD_DIM, RANK), + } + delta, used_keys = _prepare_lora_delta(state_dict, base_key, compute_device=torch.device("cpu")) + assert delta is not None + assert delta.shape == (HEAD_DIM, HEAD_DIM) + assert len(used_keys) == 2 diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 565949ff6..c7a87ac34 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -6,9 +6,11 @@ Tests for the DiffusersPipelineLoader. """ import json +import os +import tempfile from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -260,17 +262,16 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1 -def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present(): +def test_maybe_fuse_distilled_lora_skips_when_lora_path_unset(): cfg = SimpleNamespace( lora_backend="distill", - lora_path="/path/to/lora.safetensors", + lora_path=None, lora_scale=1.0, dtype=torch.bfloat16, quantization_config=None, parallel_config=SimpleNamespace(use_hsdp=False), ) loader = DiffusersPipelineLoader(LoadConfig(), cfg) - loader._hwr_state = {"warm_snapshot": {"transformer.weight": (1234, "hash")}} model = nn.Module() model.load_lora_weights = MagicMock() @@ -278,7 +279,7 @@ def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present(): loader._maybe_fuse_distilled_lora(model) model.load_lora_weights.assert_not_called() - assert getattr(model, "lora_is_fused", False) is True + assert getattr(model, "lora_is_fused", False) is False def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot(): @@ -1053,6 +1054,34 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): assert loader.take_host_weight_plan() is None +def test_dlo_mmap_plan_with_distilled_lora_falls_back_to_ordinary_loader(monkeypatch): + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + lora_backend=LoRABackend.DISTILL, + lora_path="/fake/lora.safetensors", + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.load_lora_weights = MagicMock() + calls: list[str] = [] + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + + assert loader.load_model(load_device="cpu") is model + assert calls == ["load", "process"] + assert model.load_lora_weights.call_count == 1 + assert loader.take_host_weight_plan() is None + + def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch): from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, @@ -1487,3 +1516,413 @@ def test_load_model(prefetch_helios_model, mock_tp_group): ) model = loader.load_model(load_device="cpu") assert isinstance(model, HeliosPipeline) + + +def test_hsdp_broadcast_weight_load_rank0(mocker): + """Ensure that on rank 0 with enable_broadcast_weight_load=True, weights are loaded, fused, and broadcast.""" + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=4) + mocker.patch("torch.distributed.get_rank", return_value=0) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=4, + ), + enable_broadcast_weight_load=True, + lora_backend=LoRABackend.DISTILL, + lora_path="/fake/lora.safetensors", + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._maybe_fuse_distilled_lora = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model: events.append("fuse") + ) + loader._broadcast_model_weights = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, **_kwargs: events.append("broadcast") + ) + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + assert events == ["load", "fuse", "broadcast", "process", "shard"] + + +def test_hsdp_broadcast_weight_load_rank_nonzero(mocker): + """Ensure that on rank > 0 with enable_broadcast_weight_load=True, loading is skipped and broadcast is received.""" + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=4) + mocker.patch("torch.distributed.get_rank", return_value=1) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=4, + ), + enable_broadcast_weight_load=True, + lora_backend=LoRABackend.DISTILL, + lora_path="/fake/lora.safetensors", + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.load_lora_weights = mocker.Mock() # type: ignore[assignment] + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._maybe_fuse_distilled_lora = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model: events.append("fuse") + ) + loader._broadcast_model_weights = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, **_kwargs: events.append("broadcast") + ) + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + # Rank 1 should NOT call load_weights or _maybe_fuse_distilled_lora from disk + assert "load" not in events + assert "fuse" not in events + assert events == ["broadcast", "process", "shard"] + assert getattr(model, "lora_is_fused", False) is True + + +def test_broadcast_model_weights_invokes_dist_broadcast(mocker): + broadcast_calls = [] + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=0) + mocker.patch("torch.distributed.get_backend", return_value="gloo") + mocker.patch("torch.distributed.broadcast", side_effect=lambda tensor, src: broadcast_calls.append((tensor, src))) + mocker.patch("torch.distributed.barrier") + + model = nn.Module() + model.linear = nn.Linear(2, 2) + model.register_buffer("buf", torch.tensor([1.0, 2.0])) + model.register_buffer("int_buf", torch.tensor([1, 2, 3], dtype=torch.long)) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0) + + # float32 tensors (linear.weight, linear.bias, buf) are coalesced into 1 bucket, + # and int64 tensor (int_buf) is in a separate bucket -> 2 coalesced broadcast calls + assert len(broadcast_calls) == 2 + assert all(src == 0 for _, src in broadcast_calls) + + +def test_broadcast_model_weights_receiver_copies_data(mocker): + """Ensure receiver rank copies broadcasted flat tensor data into its model parameters/buffers.""" + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=1) + mocker.patch("torch.distributed.get_backend", return_value="gloo") + mocker.patch("torch.distributed.barrier") + + # Simulate broadcast by filling the flat tensor with source values + def mock_broadcast(tensor, src): + if tensor.dtype == torch.float32: + tensor.fill_(3.14) + elif tensor.dtype == torch.long: + tensor.fill_(42) + + mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast) + + model = nn.Module() + model.linear = nn.Linear(2, 2) + model.register_buffer("buf", torch.zeros(2)) + model.register_buffer("int_buf", torch.zeros(3, dtype=torch.long)) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0) + + assert torch.allclose(model.linear.weight, torch.tensor(3.14)) + assert torch.allclose(model.linear.bias, torch.tensor(3.14)) + assert torch.allclose(model.buf, torch.tensor(3.14)) + assert torch.equal(model.int_buf, torch.tensor([42, 42, 42], dtype=torch.long)) + + +def test_broadcast_model_weights_mixed_cpu_and_accelerator_devices(mocker): + """Ensure that when some pipeline modules are on CPU (transformer) and others on accelerator (VAE/encoder), + broadcasting normalizes tensors to CPU before concatenation and succeeds without device mismatch. + """ + dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=0) + mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo") + mocker.patch("torch.distributed.broadcast") + mocker.patch("torch.distributed.barrier") + + model = nn.Module() + # Transformer parameters on CPU + model.transformer = nn.Linear(4, 4) + # VAE / Encoder parameters on accelerator + model.vae = nn.Linear(4, 4).to(dev) + model.text_encoder = nn.Linear(4, 4).to(dev) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + # Should not raise RuntimeError: Tensors must have same device + loader._broadcast_model_weights(model, target_device=dev, src_rank=0) + + +def test_broadcast_model_weights_mixed_devices_receiver(mocker): + """Ensure receiver rank with mixed CPU/accelerator tensors correctly receives broadcast data.""" + dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=1) + mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo") + mocker.patch("torch.distributed.barrier") + + def mock_broadcast(tensor, src): + tensor.fill_(7.0) + + mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast) + + model = nn.Module() + model.transformer = nn.Linear(2, 2) # CPU + model.vae = nn.Linear(2, 2).to(dev) # Accelerator + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._broadcast_model_weights(model, target_device=dev, src_rank=0) + + assert torch.allclose(model.transformer.weight, torch.tensor(7.0)) + assert torch.allclose(model.vae.weight, torch.tensor(7.0, device=dev)) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_hsdp_broadcast_weight_load_falls_back_when_online_quantization_enabled(mocker, rank): + """Ensure that when online quantization is enabled, HSDP weight loading falls back to + ordinary per-rank loading on all ranks rather than rank-0 broadcast. + """ + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=4) + mocker.patch("torch.distributed.get_rank", return_value=rank) + + mock_quant_config = SimpleNamespace( + is_checkpoint_quantized=False, + ) + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=4, + ), + enable_broadcast_weight_load=True, + lora_backend=None, + lora_path=None, + quantization_config=mock_quant_config, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._broadcast_model_weights = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, **_kwargs: events.append("broadcast") + ) + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + # Both rank 0 and non-zero ranks MUST execute ordinary load_weights, not broadcast + assert "load" in events + assert "broadcast" not in events + assert events == ["load", "process", "shard"] + + +def _mp_worker_online_quant(rank: int, world_size: int, rendezvous: str, temp_dir: str) -> None: + torch.set_num_threads(1) + import torch.distributed as dist + + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world_size) + try: + + class _MockOnlineQuantModel(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(4, 4) + self.load_called = False + + def load_weights(self, weights): + self.load_called = True + return {"transformer.weight", "transformer.bias"} + + mock_quant_config = SimpleNamespace(is_checkpoint_quantized=False) + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=world_size, + ), + enable_broadcast_weight_load=True, + lora_backend=None, + lora_path=None, + quantization_config=mock_quant_config, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = _MockOnlineQuantModel() + loader._init_from_load_format = lambda *args, **kwargs: model # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *args, **kwargs: None # type: ignore[method-assign] + + with ( + patch.object( + loader_module.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ), + patch.object(loader_module, "apply_hsdp_to_model", return_value=None), + ): + loader._load_model_with_hsdp(torch.device("cpu")) + + if model.load_called: + with open(os.path.join(temp_dir, f"rank_{rank}_success.flag"), "w") as f: + f.write("ok") + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="requires torch.distributed gloo", +) +def test_hsdp_broadcast_weight_load_online_quant_multiprocess(): + """Multi-process regression test verifying that with online quantization and broadcast enabled, + all ranks fall back to independent load_weights in parallel processes without hanging. + """ + import torch.multiprocessing as mp + + with tempfile.TemporaryDirectory() as temp_dir: + rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}" + world_size = 2 + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + mp.spawn( + _mp_worker_online_quant, + args=(world_size, rendezvous, temp_dir), + nprocs=world_size, + join=True, + ) + assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag")) + assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag")) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 250aff884..f20538284 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -82,6 +82,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( "custom_voice_dir", "devices", "disable_autocast", + "enable_broadcast_weight_load", "enable_multithread_weight_load", "env", "has_sampling_extra_args", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 2b016be0d..f0752077b 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -165,6 +165,7 @@ class _ModelEngineOverrides(TypedDict, total=False): max_cudagraph_capture_size: int enable_flashinfer_autotune: bool enable_multithread_weight_load: bool + enable_broadcast_weight_load: bool num_weight_load_threads: int disable_autocast: bool @@ -435,6 +436,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): max_cudagraph_capture_size: int | None = Field(default=None, ge=0) enable_flashinfer_autotune: bool | None = None enable_multithread_weight_load: bool = True + enable_broadcast_weight_load: bool = False num_weight_load_threads: int = Field(default=4, ge=1) disable_autocast: bool = False # Per-stage checkpoint/tokenizer subdirectories under the model root @@ -1028,6 +1030,7 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset( "enable_sleep_mode", "enforce_eager", "enable_multithread_weight_load", + "enable_broadcast_weight_load", "num_weight_load_threads", "disable_autocast", } diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0e962b4cc..063a95b0e 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -472,6 +472,7 @@ class StageDeployConfig: # Runtime optimizations used by diffusion loading/execution. enable_multithread_weight_load: bool | None = None + enable_broadcast_weight_load: bool | None = None num_weight_load_threads: int | None = None diffusion_offload_config: dict[str, Any] | None = None # Compatibility aliases for existing callers and model-specific stage diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d27cc36e1..7a7de3bb1 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -920,6 +920,7 @@ class OmniDiffusionConfig: # Parallel weight loading (for faster diffusion model startup) enable_multithread_weight_load: bool = True + enable_broadcast_weight_load: bool = False num_weight_load_threads: int = 4 # Enable sleep mode diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index a13363c7a..eb3f35188 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -43,6 +43,7 @@ def _prepare_lora_delta( lora_a_suffix: str = "lora_A.weight", lora_b_suffix: str = "lora_B.weight", lora_bias_suffix: str = "bias", + compute_device: torch.device | None = None, ): used_keys: set[str] = set() # stacked_params_mapping param_to_weight_names @@ -70,6 +71,8 @@ def _prepare_lora_delta( if lora_bias_key not in lora_state_dict: return None, used_keys delta = lora_state_dict[lora_bias_key] + if compute_device is not None and compute_device.type != "cpu": + delta = delta.to(compute_device) stacked_deltas.append(delta) used_keys.add(lora_bias_key) else: @@ -80,6 +83,9 @@ def _prepare_lora_delta( return None, used_keys a = lora_state_dict[lora_a_key] b = lora_state_dict[lora_b_key] + if compute_device is not None and compute_device.type != "cpu": + a = a.to(compute_device) + b = b.to(compute_device) delta = torch.matmul(b, a) stacked_deltas.append(delta) used_keys.add(lora_a_key) @@ -94,7 +100,10 @@ def _prepare_lora_delta( if lora_bias_key not in lora_state_dict: return None, used_keys used_keys.add(lora_bias_key) - return lora_state_dict[lora_bias_key], used_keys + delta = lora_state_dict[lora_bias_key] + if compute_device is not None and compute_device.type != "cpu": + delta = delta.to(compute_device) + return delta, used_keys lora_a_key = f"{base_key}.{lora_a_suffix}" lora_b_key = f"{base_key}.{lora_b_suffix}" @@ -102,9 +111,13 @@ def _prepare_lora_delta( return None, used_keys a = lora_state_dict[lora_a_key] b = lora_state_dict[lora_b_key] + if compute_device is not None and compute_device.type != "cpu": + a = a.to(compute_device) + b = b.to(compute_device) + delta = torch.matmul(b, a) used_keys.add(lora_a_key) used_keys.add(lora_b_key) - return torch.matmul(b, a), used_keys + return delta, used_keys def _load_lora_state_dict( @@ -267,6 +280,7 @@ class LoraLoaderMixin: else: continue + compute_dev = params.device if params.device.type != "cpu" else None delta, used_keys = _prepare_lora_delta( state_dict, base_key, @@ -275,6 +289,7 @@ class LoraLoaderMixin: lora_a_suffix, lora_b_suffix, lora_bias_suffix, + compute_device=compute_dev, ) if delta is None: continue @@ -319,6 +334,7 @@ class LoraLoaderMixin: else: continue + compute_dev = param.device if param.device.type != "cpu" else None delta, used_keys = _prepare_lora_delta( state_dict, base_key, @@ -327,6 +343,7 @@ class LoraLoaderMixin: lora_a_suffix, lora_b_suffix, lora_bias_suffix, + compute_device=compute_dev, ) if delta is None: continue diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 27774aa37..78247e7b9 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -509,11 +509,6 @@ class DiffusersPipelineLoader(HWRLoaderMixin): if getattr(model, "lora_is_fused", False): return - # A warm HWR restore already brings back the fused final-layout weights. - if self._hwr_state is not None and self._hwr_state.get("warm_snapshot") is not None: - setattr(model, "lora_is_fused", True) - return - lora_path = getattr(self.od_config, "lora_path", None) if not lora_path: return @@ -531,6 +526,111 @@ class DiffusersPipelineLoader(HWRLoaderMixin): else: logger.warning("Pipeline %s does not support loading distilled LoRA weights.", model.__class__.__name__) + def _broadcast_model_weights( + self, + model: nn.Module, + target_device: torch.device, + src_rank: int = 0, + bucket_cap_bytes: int = 256 * 1024 * 1024, + ) -> None: + """Broadcast model parameters and buffers from src_rank to all other ranks using coalesced bucketing.""" + if not torch.distributed.is_initialized() or torch.distributed.get_world_size() <= 1: + return + + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + backend = torch.distributed.get_backend() + is_cpu_backend = backend == "gloo" + + logger.info( + "Worker %d: %s full model weights via %s across %d ranks", + rank, + "Broadcasting" if rank == src_rank else "Receiving", + backend, + world_size, + ) + t0 = time.perf_counter() + + all_tensors = [p.data for _, p in model.named_parameters() if p.numel() > 0] + [ + b.data for _, b in model.named_buffers() if b.numel() > 0 + ] + if not all_tensors: + return + + count = len(all_tensors) + total_bytes = sum(t.numel() * t.element_size() for t in all_tensors) + + # Group tensors into buckets by dtype and maximum byte size + buckets: list[list[torch.Tensor]] = [] + curr_bucket: list[torch.Tensor] = [] + curr_bytes = 0 + + for tensor in all_tensors: + t_bytes = tensor.numel() * tensor.element_size() + if curr_bucket and (curr_bucket[0].dtype != tensor.dtype or curr_bytes + t_bytes > bucket_cap_bytes): + buckets.append(curr_bucket) + curr_bucket = [] + curr_bytes = 0 + curr_bucket.append(tensor) + curr_bytes += t_bytes + + if curr_bucket: + buckets.append(curr_bucket) + + for bucket in buckets: + tot_numel = sum(t.numel() for t in bucket) + dtype = bucket[0].dtype + + if is_cpu_backend: + if rank == src_rank: + flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket]) + else: + flat_cpu = torch.empty(tot_numel, dtype=dtype) + torch.distributed.broadcast(flat_cpu, src=src_rank) + if rank != src_rank: + offset = 0 + for t in bucket: + n = t.numel() + t.copy_(flat_cpu[offset : offset + n].view_as(t)) + offset += n + del flat_cpu + else: + if rank == src_rank: + flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket]) + dev_flat = flat_cpu.to(target_device, non_blocking=False) + del flat_cpu + torch.distributed.broadcast(dev_flat, src=src_rank) + del dev_flat + else: + dev_flat = torch.empty(tot_numel, dtype=dtype, device=target_device) + torch.distributed.broadcast(dev_flat, src=src_rank) + flat_cpu = dev_flat.to("cpu", non_blocking=False) + del dev_flat + offset = 0 + for t in bucket: + n = t.numel() + t.copy_(flat_cpu[offset : offset + n].view_as(t), non_blocking=False) + offset += n + del flat_cpu + + if not is_cpu_backend and target_device.type != "cpu": + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + torch.distributed.barrier() + + elapsed = time.perf_counter() - t0 + gb = total_bytes / 1e9 + logger.info( + "Worker %d: Shared weight broadcast complete (%d tensors, %d buckets, %.2f GB) in %.2fs (%.2f GB/s)", + rank, + count, + len(buckets), + gb, + elapsed, + gb / max(elapsed, 1e-6), + ) + def load_model( self, load_device: str, @@ -654,6 +754,10 @@ class DiffusersPipelineLoader(HWRLoaderMixin): if hwr_active and hwr_state is not None: self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan")) if dit_distributed_offload and not hwr_active and not self._force_canonical_load: + lora_backend = getattr(self.od_config, "lora_backend", None) + has_distilled_lora = lora_backend in (LoRABackend.DISTILL, "distill") and bool( + getattr(self.od_config, "lora_path", None) + ) plan_result = build_checkpoint_mmap_plan( model, dit_modules=tuple(zip(modules.dit_names, modules.dits)), @@ -662,6 +766,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): tensor_parallel_size=tensor_parallel_size, use_hsdp=use_hsdp, online_quantization=any(self._has_online_quant(dit) for dit in modules.dits), + has_distilled_lora=has_distilled_lora, ) self.host_weight_plan = plan_result.plan @@ -689,7 +794,6 @@ class DiffusersPipelineLoader(HWRLoaderMixin): sources=ordinary_sources, planned_weights=host_weight_plan.bindings, ) - self._maybe_fuse_distilled_lora(model) else: if dit_distributed_offload and plan_result is not None: logger.info( @@ -1097,8 +1201,39 @@ class DiffusersPipelineLoader(HWRLoaderMixin): if load_format == "diffusers": raise ValueError("HSDP is not supported with the diffusers adapter load format") model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) - self.load_weights(model) - self._maybe_fuse_distilled_lora(model) + world_size = 1 + rank = 0 + if torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() + + has_online_quant = self._has_online_quant(model) or ( + self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False) + ) + enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1 + + if enable_broadcast and has_online_quant: + logger.info( + "Worker %d: Online quantization detected; falling back to ordinary per-rank weight loading for HSDP", + rank, + ) + enable_broadcast = False + + if enable_broadcast: + if rank == 0: + self.load_weights(model) + self._maybe_fuse_distilled_lora(model) + self._broadcast_model_weights(model, target_device=target_device, src_rank=0) + if ( + rank != 0 + and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill") + and getattr(self.od_config, "lora_path", None) + and hasattr(model, "load_lora_weights") + ): + setattr(model, "lora_is_fused", True) + else: + self.load_weights(model) + self._maybe_fuse_distilled_lora(model) # Quantization methods must finish while parameters are ordinary local # tensors. Some post-load transforms use operations (for example, diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index e6301b0ec..f0ea8024d 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -322,6 +322,7 @@ def build_checkpoint_mmap_plan( tensor_parallel_size: int, use_hsdp: bool, online_quantization: bool, + has_distilled_lora: bool = False, ) -> HostWeightPlanResult: """Build a complete direct-checkpoint plan or return a fallback reason.""" if tensor_parallel_size != 1: @@ -330,6 +331,8 @@ def build_checkpoint_mmap_plan( return HostWeightPlanResult(None, "HSDP requires the ordinary loader") if online_quantization: return HostWeightPlanResult(None, "online quantization requires the ordinary loader") + if has_distilled_lora: + return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader") remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) if not callable(remap_fn): diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index ee5552751..20319de86 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -572,6 +572,7 @@ class OrchestratorArgs: vae_use_slicing: bool = False vae_use_tiling: bool = False enable_multithread_weight_load: bool = True + enable_broadcast_weight_load: bool = False num_weight_load_threads: int = 4 diffusion_offload_config: dict[str, Any] | None = None # Compatibility aliases for existing callers and model-specific stage diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 598cbe893..f9297aa0c 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -726,6 +726,13 @@ class OmniServeCommand(CLISubcommand): default=True, help="Disable multi-threaded safetensors loading (default: enabled with 4 threads).", ) + omni_config_group.add_argument( + "--enable-broadcast-weight-load", + action="store_true", + dest="enable_broadcast_weight_load", + default=False, + help="Enable Rank-0 shared weight broadcast across workers for HSDP (default: disabled).", + ) omni_config_group.add_argument( "--num-weight-load-threads", type=int, From a7ab91fb85b1bb513d34f0a19c3289c55a38b013 Mon Sep 17 00:00:00 2001 From: SuyanLi <126558907+suyanli220@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:47:25 +0800 Subject: [PATCH 020/209] [Example] Migrate HunyuanImage-3.0 to model_extras + shared task examples (#5559) Signed-off-by: suyanli220 Signed-off-by: suyan.li Co-authored-by: Claude Opus 4.8 Co-authored-by: suyan.li Co-authored-by: Hongsheng Liu --- .../offline_inference/hunyuan_image3.md | 35 -- .../hunyuan_image3/README.md | 154 -------- .../hunyuan_image3/end2end.py | 321 ----------------- .../image_to_image/image_edit.py | 126 +++++++ .../text_to_image/text_to_image.py | 126 +++++++ recipes/Tencent/HunyuanImage-3.0-Instruct.md | 52 ++- .../hunyuan_image3/test_prompt_utils.py | 129 ++++++- .../offline_inference/test_hunyuanimage3.py | 122 ++++++- .../openai_api/test_image_server.py | 69 ++++ ...test_serving_chat_multistage_generation.py | 81 +++++ tests/model_extras/test_model_extras.py | 241 +++++++++++++ .../test_shared_script_ar_integration.py | 340 ++++++++++++++++++ .../models/hunyuan_image3/prompt_utils.py | 136 +++++++ vllm_omni/entrypoints/openai/serving_chat.py | 11 +- vllm_omni/model_extras/__init__.py | 4 + vllm_omni/model_extras/hunyuan_image3.py | 171 ++++++++- vllm_omni/model_extras/registry.py | 53 +++ 17 files changed, 1627 insertions(+), 544 deletions(-) delete mode 100644 docs/user_guide/examples/offline_inference/hunyuan_image3.md delete mode 100644 examples/offline_inference/hunyuan_image3/README.md delete mode 100644 examples/offline_inference/hunyuan_image3/end2end.py create mode 100644 tests/model_extras/test_shared_script_ar_integration.py diff --git a/docs/user_guide/examples/offline_inference/hunyuan_image3.md b/docs/user_guide/examples/offline_inference/hunyuan_image3.md deleted file mode 100644 index 636aa0c11..000000000 --- a/docs/user_guide/examples/offline_inference/hunyuan_image3.md +++ /dev/null @@ -1,35 +0,0 @@ -# HunyuanImage-3.0 Image-to-Text Inference - -Source . - - -This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni. - -## Local CLI Usage - -Download the example image: - -```bash -wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg -``` - -Run example: - -```bash -python image_to_text.py \ - --image cherry_blossom.jpg \ - --prompt "<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." -``` - -Key arguments: - -- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional). -- `--image`: Path to input image (required). -- `--prompt`: Text description used to guide image understanding (required). - -## Example materials - -??? abstract "image_to_text.py" - ``````py - --8<-- "examples/offline_inference/hunyuan_image3/image_to_text.py" - `````` diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md deleted file mode 100644 index b4d5af2c7..000000000 --- a/examples/offline_inference/hunyuan_image3/README.md +++ /dev/null @@ -1,154 +0,0 @@ -# HunyuanImage-3.0-Instruct - -This example runs HunyuanImage-3.0-Instruct offline with the unified deploy -YAMLs under `vllm_omni/deploy/`. - -## Deploy Configs - -| File | Topology | Default use | -| :--- | :--- | :--- | -| `vllm_omni/deploy/hunyuan_image_3_moe.yaml` | AR + DiT | Default for `text2img` and `img2img`. | -| `vllm_omni/deploy/hunyuan_image3_ar.yaml` | AR only | Default for `img2text` and `text2text`. | -| `vllm_omni/deploy/hunyuan_image3_dit.yaml` | DiT only | Standalone diffusion stage. Pass it explicitly with `--deploy-config`. | - -The example chooses a deploy config automatically when `--deploy-config` is -omitted: - -| `--modality` | `mode` passed to Omni | Default deploy | -| :--- | :--- | :--- | -| `text2img` | `text-to-image` | `hunyuan_image_3_moe.yaml` | -| `img2img` | `image-editing` | `hunyuan_image_3_moe.yaml` | -| `img2text` | `image-to-text` | `hunyuan_image3_ar.yaml` | -| `text2text` | `text-to-text` | `hunyuan_image3_ar.yaml` | - -`--modality` is an offline example convenience flag. It maps to the internal -`mode` argument passed to `Omni(...)` by this script. HunyuanImage3 uses -separate deploy YAMLs for AR + DiT, AR-only, and DiT-only topologies, so the -stage topology is selected by the deploy file rather than by YAML mode -overrides. - -Online serving does not expose a `--modality` flag or accept `mode` as an API -request field. Choose the deploy topology when starting the server with -`--deploy-config`, then use the OpenAI-compatible endpoint and request shape for -the scenario. The `modalities` request field is used by the chat completions -path; the image endpoints infer the image task from the endpoint and payload. - -| Online scenario | Server deploy | Request | -| :--- | :--- | :--- | -| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. | -| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/edits`. | -| Image/text to text | `--deploy-config vllm_omni/deploy/hunyuan_image3_ar.yaml` | `POST /v1/chat/completions` for text output, for example with `"modalities": ["text"]`. | -| DiT-only image generation | `--deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml` | `POST /v1/images/generations`. | - -## Run Examples - -Text to image, using the default AR + DiT deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --prompts "A cute cat sitting on a windowsill watching the sunset" -``` - -Image editing, using the default AR + DiT deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality img2img \ - --image-path /path/to/image.png \ - --prompts "Make the petals neon pink" -``` - -Image to text, using the AR-only deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "Describe the content of the picture." -``` - -Text to text, using the AR-only deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2text \ - --prompts "What is the capital of France?" -``` - -Standalone DiT, using the DiT-only deploy explicitly: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml \ - --prompts "A cinematic portrait of an astronaut in a greenhouse" -``` - -Override the default full AR + DiT deploy explicitly: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ - --prompts "A cute cat" -``` - -## Additional Config - -You can pass diffusion worker `additional_config` from the offline example as a JSON object. -This maps to the upstream vLLM `VllmConfig.additional_config` platform extension field: -https://docs.vllm.ai/en/stable/api/vllm/config/#vllm.config.VllmConfig.additional_config - -```bash -python end2end.py --modality text2img \ - --prompts "A cute cat" \ - --additional-config '{"torchair_graph_config":{"enabled":true}}' -``` - - -## Key Arguments - -| Argument | Description | -| :--- | :--- | -| `--deploy-config` | Optional path for unified deploy YAML overrides. | -| `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. | -| `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. | -| `--steps` | Number of diffusion inference steps for image generation. | -| `--guidance-scale` | Classifier-free guidance scale for image generation. | -| `--height`, `--width` | Output image size for `text2img`. | -| `--bot-task` | Override prompt mode. `none`, `think`, `recaption`, `think_recaption`, or `vanilla`. | -| `--sys-type` | Override the system prompt type, for example `en_unified` or `en_vanilla`. | -| `--vae-use-tiling` | Enable VAE tiling for memory reduction. | - -## Notes - -- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy. -- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`. -- Platform-specific overrides are included in the unified deploy YAMLs. - -## Prompt Format - -HunyuanImage-3.0-Instruct uses an instruct chat template: - -```text -<|startoftext|>{system_prompt} - -User: {?}{user_prompt} - -Assistant: {trigger_tag?} -``` - -- ``: Placeholder for each input image (single token; expanded by the multimodal pipeline). -- Trigger tags: `` for CoT and `` for recaptioning, placed after `Assistant: `. -- System prompt: Auto-selected from `task` and `bot_task`. -- `bot_task='vanilla'` with `task='t2i'` uses the bare pretrain template. - -The shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils.build_prompt_tokens()` -helper handles segment-by-segment tokenization and matches HF `apply_chat_template`. diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py deleted file mode 100644 index fbe4289be..000000000 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ /dev/null @@ -1,321 +0,0 @@ -""" -HunyuanImage-3.0-Instruct unified end-to-end inference script. -""" - -import argparse -import json -import os -from pathlib import Path - -from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( - MAX_IMAGES_PER_REQUEST, - build_prompt_tokens, - resolve_stop_token_ids, - resolve_sys_type, -) -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniPromptType - -_REPO_ROOT = Path(__file__).resolve().parents[3] -_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml") -_DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml") - -_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = { - "text2img": ("t2i", "think"), - "img2img": ("it2i", "think"), - "img2text": ("i2t", None), - "text2text": ("t2t", None), -} - -_MODALITY_DEFAULT_DEPLOY_CONFIG = { - "text2img": _DEFAULT_DEPLOY_CONFIG, - "img2img": _DEFAULT_DEPLOY_CONFIG, - "img2text": _DEFAULT_AR_DEPLOY_CONFIG, - "text2text": _DEFAULT_AR_DEPLOY_CONFIG, -} - -_MODALITY_MODE = { - "text2img": "text-to-image", - "img2img": "image-editing", - "img2text": "image-to-text", - "text2text": "text-to-text", -} - - -def parse_args(): - parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.") - parser.add_argument("--model", default="tencent/HunyuanImage-3.0-Instruct", help="Model name or local path.") - parser.add_argument( - "--modality", - default="text2img", - choices=list(_MODALITY_TASK_MAP), - ) - parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") - parser.add_argument( - "--image-path", - type=str, - default=None, - help="Input image path(s) for img2img/img2text. Comma-separated for multi-image (up to 3).", - ) - parser.add_argument("--output", type=str, default=".", help="Output directory to save results.") - parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") - parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.") - parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--height", type=int, default=None, help="Output image height.") - parser.add_argument("--width", type=int, default=None, help="Output image width.") - parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.") - parser.add_argument( - "--bot-task", - type=str, - default=None, - choices=["none", "think", "recaption", "think_recaption", "vanilla"], - help="Override prompt mode. Default: auto from --modality.", - ) - parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.") - parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.") - parser.add_argument( - "--stream", - action="store_true", - default=False, - help="Enable streaming CoT display; print AR text token-by-token in real time.", - ) - parser.add_argument("--log-stats", action="store_true", default=False) - parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") - parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") - parser.add_argument( - "--diffusion-kv-cache-dtype", - type=str, - default=None, - help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.", - ) - parser.add_argument( - "--diffusion-kv-cache-skip-steps", - type=str, - default=None, - help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.", - ) - parser.add_argument( - "--diffusion-kv-cache-skip-layers", - type=str, - default=None, - help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.", - ) - parser.add_argument( - "--additional-config", - type=str, - default=None, - help=( - "JSON object forwarded to Omni/additional_config, for example " - '\'{"torchair_graph_config":{"enabled":true}}\'. Different platforms may support different ' - "configs. Make sure the configs are valid for the platform you are using. " - "Contents must be hashable." - ), - ) - - return parser.parse_args() - - -def parse_additional_config(raw_value: str | None) -> dict | None: - """Parse a JSON string into an additional_config mapping.""" - if raw_value is None: - return None - - try: - additional_config = json.loads(raw_value) - except json.JSONDecodeError as exc: - raise ValueError(f"Invalid --additional-config JSON: {exc}") from exc - - if additional_config is None: - return None - if not isinstance(additional_config, dict): - raise ValueError(f"--additional-config must decode to a JSON object, got {type(additional_config).__name__}") - return additional_config - - -def main(): - args = parse_args() - os.makedirs(args.output, exist_ok=True) - additional_config = parse_additional_config(args.additional_config) - - task, default_bot_task = _MODALITY_TASK_MAP[args.modality] - if args.bot_task is None: - bot_task: str | None = default_bot_task - elif args.bot_task == "none": - bot_task = None - else: - bot_task = args.bot_task - - deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] - - omni_kwargs = { - "model": args.model, - "vae_use_tiling": args.vae_use_tiling, - "log_stats": args.log_stats, - "init_timeout": args.init_timeout, - "enforce_eager": args.enforce_eager, - "mode": _MODALITY_MODE[args.modality], - "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype, - "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps, - "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers, - } - - if additional_config is not None: - omni_kwargs["additional_config"] = additional_config - omni_kwargs["deploy_config"] = deploy_config - - omni = Omni(**omni_kwargs) - - prompts = args.prompts or ["A cute cat"] - input_images: list = [] - if args.modality in ("img2img", "img2text"): - if not args.image_path: - raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}") - from PIL import Image - - image_paths = [p.strip() for p in args.image_path.split(",") if p.strip()] - if len(image_paths) > MAX_IMAGES_PER_REQUEST: - raise ValueError( - f"--image-path accepts at most {MAX_IMAGES_PER_REQUEST} images for " - f"HunyuanImage-3.0 IT2I, got {len(image_paths)}: {args.image_path}" - ) - for image_path in image_paths: - if not os.path.exists(image_path): - raise ValueError(f"Image path does not exist: {image_path}") - input_images.append(Image.open(image_path).convert("RGB")) - if not input_images: - raise ValueError(f"--image-path produced no usable paths: {args.image_path!r}") - - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) - mm_image_payload = (input_images[0] if len(input_images) == 1 else input_images) if input_images else None - - formatted_prompts: list[OmniPromptType] = [] - for prompt in prompts: - build_kwargs: dict = {"task": task, "bot_task": bot_task, "sys_type": args.sys_type} - if input_images: - build_kwargs["num_images"] = len(input_images) - result = build_prompt_tokens(prompt, tokenizer, **build_kwargs) - token_ids = result.token_ids - effective_sys_type = args.sys_type or resolve_sys_type(bot_task) - - prompt_dict: dict = { - "prompt_token_ids": token_ids, - "prompt": prompt, - "use_system_prompt": effective_sys_type, - } - if args.modality == "text2img": - prompt_dict["modalities"] = ["image"] - elif args.modality == "img2img": - prompt_dict["modalities"] = ["image"] - prompt_dict["multi_modal_data"] = {"image": mm_image_payload} - prompt_dict["height"] = input_images[0].height - prompt_dict["width"] = input_images[0].width - elif args.modality == "img2text": - prompt_dict["modalities"] = ["text"] - prompt_dict["multi_modal_data"] = {"image": mm_image_payload} - else: - prompt_dict["modalities"] = ["text"] - formatted_prompts.append(prompt_dict) - - from vllm.sampling_params import RequestOutputKind - - params_list = list(omni.default_sampling_params_list) - - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - if (args.height is None) != (args.width is None): - raise ValueError("--height and --width must both be specified or both omitted.") - user_specified_size = args.height is not None and args.width is not None - if args.modality in ("img2text", "text2text"): - ar_image_size = "auto" - elif user_specified_size: - ar_image_size = f"{args.width}x{args.height}" - else: - ar_image_size = None - ar_stop_token_ids = resolve_stop_token_ids( - task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size - ) - print( - f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}" - ) - for sp in params_list: - if isinstance(sp, OmniDiffusionSamplingParams): - sp.num_inference_steps = args.steps - sp.guidance_scale = args.guidance_scale - sp.guidance_scale_provided = True - if args.seed is not None: - sp.seed = args.seed - if args.modality == "text2img": - sp.height = args.height - sp.width = args.width - elif hasattr(sp, "stop_token_ids"): - sp.stop_token_ids = ar_stop_token_ids - # When --stream is set, request DELTA output from the AR stage - # so we can display CoT text token-by-token in real time. - if args.stream and hasattr(sp, "output_kind"): - sp.output_kind = RequestOutputKind.DELTA - - print(f"\n{'=' * 60}") - print("HunyuanImage-3.0 Generation Configuration:") - print(f" Model: {args.model}") - print(f" Modality: {args.modality}") - print(f" Prompt task: {task}") - print(f" Bot task: {bot_task}") - print(f" Deploy config: {deploy_config}") - print(f" Num stages: {omni.num_stages}") - if args.modality in ("text2img", "img2img"): - print(f" Inference steps: {args.steps}") - print(f" Guidance scale: {args.guidance_scale}") - print(f" Seed: {args.seed}") - print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}") - print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}") - print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}") - if args.modality == "text2img": - print(f" Output size: {args.width}x{args.height}") - if args.image_path: - print(f" Input image: {args.image_path}") - if additional_config is not None: - print(f" Additional config: {additional_config}") - print(f" Prompts: {prompts}") - print(f"{'=' * 60}\n") - - # When --stream is set, print AR CoT text token-by-token in real time. - # Otherwise, collect and print the full AR text once when stage 0 finishes. - omni_outputs = omni.generate( - prompts=formatted_prompts, - sampling_params_list=params_list, - py_generator=True, - use_tqdm=False, - ) - img_idx = 0 - for req_output in omni_outputs: - ro = req_output - stage_id = getattr(req_output, "stage_id", None) - - # AR stage text — each CompletionOutput.text is already a delta when - # output_kind=DELTA, so we can print it directly (matching the pattern - # in serving_chat.py). - if stage_id == 0 and ro and getattr(ro, "outputs", None): - for o in ro.outputs: - text = getattr(o, "text", "") or "" - if text: - print(text, end="", flush=True) - # Non-streaming: one shot with full text — emit a trailing newline. - if not args.stream: - print(flush=True) - - # Collect images from diffusion stage - images = getattr(req_output, "images", None) - if not images and ro and hasattr(ro, "images"): - images = ro.images - if images: - for j, img in enumerate(images): - save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") - img.save(save_path) - print(f"\n[Output] Saved image to {save_path}") - img_idx += 1 - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 136d99857..cec7014de 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -96,6 +96,7 @@ from typing import Any import torch from PIL import Image +from vllm_omni.diffusion.data import logger from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni @@ -103,6 +104,8 @@ from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_extras import ( build_image_to_image_prompt, + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_model_class_name, should_init_extra_args_for_non_diffusion_stages, @@ -110,6 +113,91 @@ from vllm_omni.model_extras import ( from vllm_omni.platforms import current_omni_platform +def _apply_ar_stage_inputs( + ar_input_builder: Any, + *, + model: str, + prompt_text: str, + extra_body: dict[str, Any], + num_images: int, + height: int | None, + width: int | None, + prompt_dict: dict[str, Any], + sampling_params_list: list[Any], + text_output: bool = False, + trust_remote_code: bool = False, + validate_tokenizer: Any | None = None, + allow_tokenizer_fallback: bool = False, +) -> None: + """Apply a model's declared AR-stage inputs to the request in place. + + Mirrors the helper in ``text_to_image.py``: loads the tokenizer for + byte-for-byte HF-parity AR tokenization, asks the model's + ``ar_input_builder`` for the AR prefill + stop tokens, and writes them + onto ``prompt_dict`` and the non-diffusion (AR) stage params. + Model-agnostic -- only the declared-hook contract is assumed. + + ``trust_remote_code`` must be threaded in from the caller's own resolved + ``--trust-remote-code`` flag -- this is a separate tokenizer load outside + the main engine, so it needs the same explicit user opt-in rather than + defaulting to trusting whatever ``--model`` was passed. + + ``validate_tokenizer``, when the model declares one via + ``get_ar_tokenizer_validator``, is called on a successfully-loaded real + tokenizer -- outside the load's try/except, so a validation failure (a + model/tokenizer revision drifting from hardcoded special-token ids) + raises instead of being swallowed into the string-prompt fallback. + + ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a + tokenizer load failure normally raises, since silently degrading to the + string-prompt form can produce a request that completes with subtly + wrong stop tokens instead of surfacing the real problem (missing + ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only + for explicit unit tests or a deliberate offline-compat run where the + caller has already decided a degraded prompt is acceptable. + """ + try: + from transformers import AutoTokenizer + + ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) + except Exception as exc: # noqa: BLE001 - re-raised unless the caller opts into the fallback + if not allow_tokenizer_fallback: + raise + logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).") + ar_tokenizer = None + + if ar_tokenizer is not None and validate_tokenizer is not None: + validate_tokenizer(ar_tokenizer) + + ar_inputs = ar_input_builder( + prompt=prompt_text, + tokenizer=ar_tokenizer, + extra_body=extra_body, + num_images=num_images, + height=height, + width=width, + text_output=text_output, + ) + + if ar_inputs.prompt_token_ids is not None: + prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids + elif ar_inputs.prompt is not None: + prompt_dict["prompt"] = ar_inputs.prompt + if ar_inputs.use_system_prompt: + prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt + prompt_dict["modalities"] = ar_inputs.modalities + + # Apply stop_token_ids to exactly the stage(s) the builder names -- not + # "whichever stage isn't a diffusion stage," which would misfire on any + # future topology with more than one non-diffusion stage. + for stage_index in ar_inputs.stage_indices: + if stage_index >= len(sampling_params_list): + continue + stage_params = sampling_params_list[stage_index] + if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"): + stage_params.stop_token_ids = ar_inputs.stop_token_ids + + def parse_profiler_config(value: str) -> dict[str, Any]: try: config = json.loads(value) @@ -419,6 +507,22 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).", + ) + parser.add_argument( + "--allow-tokenizer-fallback", + action="store_true", + help=( + "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading " + "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead " + "of failing the run. Off by default -- a tokenizer load failure usually means a " + "missing --trust-remote-code or a network/cache issue that's worth surfacing, not " + "silently masking with a possibly-wrong prompt." + ), + ) return parser.parse_args() @@ -492,6 +596,8 @@ def main(): ) if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager + if args.trust_remote_code: + omni_kwargs["trust_remote_code"] = True if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config omni = Omni(**omni_kwargs) @@ -584,6 +690,26 @@ def main(): if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] + # Models with an AR text stage (e.g. HunyuanImage3 image-editing) declare an + # ar_input_builder; build the AR prefill + stop tokens declaratively from the + # plain prompt + reference image count + extra_args. Others are untouched. + ar_input_builder = get_ar_input_builder(model_class_name) + if ar_input_builder is not None: + _apply_ar_stage_inputs( + ar_input_builder, + model=args.model, + prompt_text=args.prompt, + extra_body=extra_args_from_cli, + num_images=len(input_images), + height=args.height, + width=args.width, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=args.trust_remote_code, + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + allow_tokenizer_fallback=args.allow_tokenizer_fallback, + ) + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2ee09e178..adcc10669 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -24,6 +24,8 @@ from vllm_omni.model_extras import ( build_text_to_image_prompt as build_model_text_to_image_prompt, ) from vllm_omni.model_extras import ( + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_model_class_name, should_init_extra_args_for_non_diffusion_stages, @@ -377,10 +379,111 @@ def parse_args() -> argparse.Namespace: default=None, help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).", ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).", + ) + parser.add_argument( + "--allow-tokenizer-fallback", + action="store_true", + help=( + "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading " + "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead " + "of failing the run. Off by default -- a tokenizer load failure usually means a " + "missing --trust-remote-code or a network/cache issue that's worth surfacing, not " + "silently masking with a possibly-wrong prompt." + ), + ) current_omni_platform.pre_register_and_update(parser) return parser.parse_args() +def _apply_ar_stage_inputs( + ar_input_builder: Any, + *, + model: str, + prompt_text: str, + extra_body: dict[str, Any], + num_images: int, + height: int | None, + width: int | None, + prompt_dict: dict[str, Any], + sampling_params_list: list[Any], + text_output: bool = False, + trust_remote_code: bool = False, + validate_tokenizer: Any | None = None, + allow_tokenizer_fallback: bool = False, +) -> None: + """Apply a model's declared AR-stage inputs to the request in place. + + Loads the model tokenizer (for byte-for-byte HF-parity segment + tokenization) and asks the model's ``ar_input_builder`` for the AR + prefill + stop tokens, then writes them onto ``prompt_dict`` and the + non-diffusion (AR) stage sampling params. Kept model-agnostic: the + example only knows the declared-hook contract. + + ``trust_remote_code`` must be threaded in from the caller's own resolved + ``--trust-remote-code`` flag -- this is a separate tokenizer load outside + the main engine, so it needs the same explicit user opt-in rather than + defaulting to trusting whatever ``--model`` was passed. + + ``validate_tokenizer``, when the model declares one via + ``get_ar_tokenizer_validator``, is called on a successfully-loaded real + tokenizer -- outside the load's try/except, so a validation failure (a + model/tokenizer revision drifting from hardcoded special-token ids) + raises instead of being swallowed into the string-prompt fallback. + + ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a + tokenizer load failure normally raises, since silently degrading to the + string-prompt form can produce a request that completes with subtly + wrong stop tokens instead of surfacing the real problem (missing + ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only + for explicit unit tests or a deliberate offline-compat run where the + caller has already decided a degraded prompt is acceptable. + """ + try: + from transformers import AutoTokenizer + + ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) + except Exception as exc: # noqa: BLE001 - re-raised unless the caller opts into the fallback + if not allow_tokenizer_fallback: + raise + logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).") + ar_tokenizer = None + + if ar_tokenizer is not None and validate_tokenizer is not None: + validate_tokenizer(ar_tokenizer) + + ar_inputs = ar_input_builder( + prompt=prompt_text, + tokenizer=ar_tokenizer, + extra_body=extra_body, + num_images=num_images, + height=height, + width=width, + text_output=text_output, + ) + + if ar_inputs.prompt_token_ids is not None: + prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids + elif ar_inputs.prompt is not None: + prompt_dict["prompt"] = ar_inputs.prompt + if ar_inputs.use_system_prompt: + prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt + prompt_dict["modalities"] = ar_inputs.modalities + + # Apply stop_token_ids to exactly the stage(s) the builder names -- not + # "whichever stage isn't a diffusion stage," which would misfire on any + # future topology with more than one non-diffusion stage. + for stage_index in ar_inputs.stage_indices: + if stage_index >= len(sampling_params_list): + continue + stage_params = sampling_params_list[stage_index] + if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"): + stage_params.stop_token_ids = ar_inputs.stop_token_ids + + def main(): args = parse_args() generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) @@ -469,6 +572,8 @@ def main(): omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager + if args.trust_remote_code: + omni_kwargs["trust_remote_code"] = True if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config if use_nextstep: @@ -613,6 +718,27 @@ def main(): if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] + # Models with an AR text stage (e.g. HunyuanImage3) declare an + # ar_input_builder. When present, build the AR prefill token-ids and AR + # stop-token-ids declaratively from the plain prompt + extra_body, so this + # example stays model-agnostic. Models without one are untouched. + ar_input_builder = get_ar_input_builder(model_class_name) + if ar_input_builder is not None: + _apply_ar_stage_inputs( + ar_input_builder, + model=args.model, + prompt_text=args.prompt, + extra_body=user_extra, + num_images=0, + height=args.height, + width=args.width, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=args.trust_remote_code, + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + allow_tokenizer_fallback=args.allow_tokenizer_fallback, + ) + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index 4614df382..eb5ccedf6 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -8,7 +8,7 @@ - Vendor: Tencent Hunyuan - Model: `tencent/HunyuanImage-3.0-Instruct` -- Task: Text-to-image generation; text-to-text and image-to-text understanding +- Task: Text-to-image generation, image editing; text-to-text and image-to-text understanding - Mode: Offline understanding, online serving, and performance benchmarking - Maintainer: Community @@ -31,8 +31,10 @@ FP8/NVFP4 configuration: ## References - Model: -- Existing model-specific offline example: - [`examples/offline_inference/hunyuan_image3`](../../examples/offline_inference/hunyuan_image3) +- Shared T2I example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Shared IT2I example: + [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py) - Shared T2T/I2T example: [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) - Related PRs: @@ -66,8 +68,48 @@ python examples/offline_inference/x_to_text/x_to_text.py \ ``` The AR-only default uses four GPUs. Pass `--deploy-config` to override the -layout. Image-output tasks remain documented by the existing model-specific -example and the DiT sections below. +layout. + +### Shared T2I/IT2I offline examples + +Text-to-image and image-editing route through the shared task examples, with +all HunyuanImage-3.0-specific knobs (`bot_task`, `use_system_prompt`, +`system_prompt`, `negative_prompt`) declared in +`vllm_omni/model_extras/hunyuan_image3.py` and passed via `--extra-body` / +`--extra-args`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ + --trust-remote-code \ + --prompt "A cute cat sitting on a windowsill watching the sunset" \ + --height 1024 --width 1024 \ + --guidance-scale 5.0 --num-inference-steps 50 --seed 42 \ + --extra-body '{"bot_task": "think", "use_system_prompt": "en_recaption"}' \ + --output hunyuan_t2i.png +``` + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ + --trust-remote-code \ + --image /path/to/image.png \ + --prompt "Make the petals neon pink" \ + --guidance-scale 5.0 --num-inference-steps 50 --seed 42 \ + --extra-args '{"bot_task": "think"}' \ + --output hunyuan_edit.png +``` + +`image_edit.py` accepts up to 3 reference images (repeat `--image`) for +HunyuanImage-3.0's multi-image fusion. All four shared-script paths (t2i, +it2i, i2t, t2t) build the AR prefill and stop tokens through the same +declarative seam (`build_ar_stage_inputs` / `build_x_to_text_prompt`, both +wrapping `prompt_utils.build_ar_prompt_inputs`), so prompt formatting is +identical across them. The OpenAI server's `serving_chat.py` does not go +through this seam yet -- see that function's docstring for the resulting +`bot_task`-resolution divergence risk when omitted from `extra_body`. ### 4x H100/H800 80GB diff --git a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py index 7cb7ea83a..7a467e359 100644 --- a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py +++ b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py @@ -13,9 +13,11 @@ from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, available_bot_tasks, available_tasks, + build_ar_prompt_inputs, build_prompt, build_prompt_tokens, resolve_stop_token_ids, + validate_special_token_ids, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -296,19 +298,42 @@ def _repo_root() -> pathlib.Path: return pathlib.Path(__file__).resolve().parents[4] -def test_end2end_routes_through_shared_prompt_utils(): - end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py" - tree = ast.parse(end2end_path.read_text(encoding="utf-8")) +def test_example_routes_through_shared_seam(): + """Neither the shared task examples nor the ``model_extras`` seam itself + should duplicate AR prompt tokenization. - local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} - assert not (local_func_names & {"build_prompt", "build_prompt_tokens"}) + After the model_extras migration, t2i / image-editing run through the + shared ``text_to_image.py`` / ``image_edit.py`` examples (both call + ``build_ar_stage_inputs``), and t2t / i2t run through the shared + ``x_to_text.py`` example (which calls ``build_x_to_text_prompt``). This + test pins that both ``model_extras.hunyuan_image3`` seam functions + delegate to ``prompt_utils.build_ar_prompt_inputs`` / ``build_ar_stage_inputs`` + rather than re-implementing tokenization inline. + """ + seam_path = _repo_root() / "vllm_omni" / "model_extras" / "hunyuan_image3.py" + seam_tree = ast.parse(seam_path.read_text(encoding="utf-8")) - imported_from_prompt_utils: set[str] = set() - for node in ast.walk(tree): + local_func_names = {n.name for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef)} + assert not (local_func_names & {"build_prompt", "build_prompt_tokens", "resolve_stop_token_ids"}) + + # build_x_to_text_prompt must route through build_ar_stage_inputs, not + # re-derive the AR prefill / stop-token logic itself. + build_x_to_text_prompt_node = next( + n for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef) and n.name == "build_x_to_text_prompt" + ) + called_names = { + node.func.id + for node in ast.walk(build_x_to_text_prompt_node) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "build_ar_stage_inputs" in called_names + + # The model_extras seam itself must delegate to the shared prompt_utils. + routed_from_prompt_utils: set[str] = set() + for node in ast.walk(seam_tree): if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"): - imported_from_prompt_utils.update(alias.name for alias in node.names) - expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"} - assert expected_imports <= imported_from_prompt_utils + routed_from_prompt_utils.update(alias.name for alias in node.names) + assert "build_ar_prompt_inputs" in routed_from_prompt_utils _HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct" @@ -331,3 +356,87 @@ def test_segment_tokenize_diverges_from_full_string_encode(): full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False) assert seg_ids != full_ids assert len(seg_ids) >= len(full_ids) + + +def test_build_ar_prompt_inputs_tokenizer_path_matches_build_prompt_tokens(): + """With a tokenizer, the helper delegates to build_prompt_tokens (parity).""" + tok = FakeTokenizer() + ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="think", tokenizer=tok) + expected = build_prompt_tokens("a red panda", FakeTokenizer(), task="t2i", bot_task="think") + assert ar.prompt is None + assert ar.prompt_token_ids == expected.token_ids + assert ar.system_prompt_type == expected.system_prompt_type + # think + no explicit size -> AR predicts ratio, so stop range is non-trivial. + assert len(ar.stop_token_ids) > 1 + + +def test_build_ar_prompt_inputs_string_fallback_without_tokenizer(): + """Without a tokenizer, the helper returns the build_prompt string form.""" + ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="recaption") + assert ar.prompt_token_ids is None + assert ar.prompt == build_prompt("a red panda", task="t2i", bot_task="recaption") + assert ar.system_prompt_type + assert ar.stop_token_ids + + +def test_build_ar_prompt_inputs_explicit_size_stops_at_terminator(): + """Explicit image_size makes 'recaption' stop at only.""" + ar = build_ar_prompt_inputs("x", task="t2i", bot_task="recaption", image_size="1024x1024") + assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +def test_build_ar_prompt_inputs_comprehension_stops_on_answer(): + """Text-output tasks (t2t / i2t) stop on .""" + ar = build_ar_prompt_inputs("hi", task="t2t") + assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +class _MatchingTokenizer: + """Reports exactly HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS -- a "real" tokenizer + that hasn't drifted from the hardcoded table.""" + + def convert_tokens_to_ids(self, tok: str) -> int: + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +class _DriftedTokenizer: + """Reports a different id for one token -- simulates a model/tokenizer + revision bump shifting special-token ids out from under the hardcoded + table.""" + + def convert_tokens_to_ids(self, tok: str) -> int: + if tok == "": + return 999999 + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +class _BaseCheckpointTokenizer: + """Simulates tencent/HunyuanImage-3.0 (base): reports the same + model_class_name as -Instruct, but its vocab doesn't include the + Instruct-only - tokens -- convert_tokens_to_ids + falls back to unk_token_id for those, like a real HF tokenizer would.""" + + unk_token_id = 0 + _INSTRUCT_ONLY = {"", ""} + + def convert_tokens_to_ids(self, tok: str) -> int: + if tok in self._INSTRUCT_ONLY: + return self.unk_token_id + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +def test_validate_special_token_ids_passes_for_matching_tokenizer(): + validate_special_token_ids(_MatchingTokenizer()) # must not raise + + +def test_validate_special_token_ids_raises_on_drift(): + with pytest.raises(ValueError, match=".*expected 128023.*got 999999"): + validate_special_token_ids(_DriftedTokenizer()) + + +def test_validate_special_token_ids_tolerates_base_checkpoint_missing_tokens(): + """The base checkpoint's tokenizer not recognizing Instruct-only ratio + tokens (unk_token_id, not a real drifted id) must not be flagged as a + mismatch -- both checkpoints share model_class_name, so this hook fires + for either one.""" + validate_special_token_ids(_BaseCheckpointTokenizer()) # must not raise diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index 97179d963..8655bd63f 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -1,5 +1,8 @@ # ruff: noqa: E501 +import importlib.util +import sys from collections.abc import Generator +from pathlib import Path import pytest import torch @@ -10,13 +13,15 @@ from transformers import CLIPModel, CLIPProcessor from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path from vllm_omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name from vllm_omni.platforms import current_omni_platform PROMPT = "A brown and white dog is running on the grass" MODEL_NAME = "tencent/HunyuanImage-3.0" LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" -DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3.yaml") +DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] @@ -328,6 +333,15 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: return _extract_generated_image(outputs) +# A floor for "this is recognizably the same generation" against each +# case's precomputed reference embedding, not a quality bar -- catches a +# gross migration regression (wrong/blank/garbage image, or the wrong +# system prompt applied) without being flaky on run-to-run sampling noise. +# Matches the threshold test_shared_script_ar_path_reaches_generation uses +# for its own (seed=1234, en_recaption) comparison below. +MIN_SEMANTIC_SIMILARITY = 0.5 + + @pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES) def test_system_prompt_scores( @@ -342,3 +356,109 @@ def test_system_prompt_scores( score = compare_semantic(expected_embedding, generated_image, clip_model, clip_processor) print(f"{system_prompt_name}: CLIP cosine similarity = {score:.6f}") + assert score >= MIN_SEMANTIC_SIMILARITY, ( + f"{system_prompt_name}: CLIP similarity {score:.4f} is below " + f"{MIN_SEMANTIC_SIMILARITY} against the reference (seed=1234) embedding -- " + "the migrated path may have diverged from the reference HunyuanImage-3.0 " + f"output for use_system_prompt={use_system_prompt!r}." + ) + + +def _load_text_to_image_module(): + """Dynamically load the shared example script by file path (``examples/`` + is not an installed package) so this test can call its real + ``_apply_ar_stage_inputs`` instead of re-implementing its logic.""" + repo_root = Path(__file__).resolve().parents[3] + path = repo_root / "examples/offline_inference/text_to_image/text_to_image.py" + spec = importlib.util.spec_from_file_location("_e2e_shared_text_to_image", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +def test_shared_script_ar_path_reaches_generation( + omni: Omni, + clip_bundle: tuple[CLIPModel, CLIPProcessor], +) -> None: + """Exercise the path this PR actually migrates HunyuanImage-3.0 onto: + + get_model_class_name -> get_ar_input_builder -> _apply_ar_stage_inputs + -> Omni.generate() + + using the *real* model_extras registry lookup (keyed on the engine's + live ``model_class_name``) and the shared ``text_to_image.py``'s own + helper -- not the hand-constructed ``sampling_params.extra_args`` that + ``_generate_image``/``test_system_prompt_scores`` above use, which never + touch the registry or the AR-tokenizer-loading path at all. This is the + one gap those tests leave: a regression in the registry key (the bug + this PR fixes), in ``_apply_ar_stage_inputs``, or in the + ``ar_tokenizer_validator`` wiring would not fail any of them. + + Reuses the ``en_recaption`` reference embedding from + ``test_system_prompt_scores`` so this also asserts the migrated path + produces the *same* output as the hand-constructed one, not just "didn't + crash." + """ + text_to_image = _load_text_to_image_module() + clip_model, clip_processor = clip_bundle + + model_class_name = get_model_class_name(omni) + ar_input_builder = get_ar_input_builder(model_class_name) + assert ar_input_builder is not None, ( + f"get_ar_input_builder({model_class_name!r}) returned None for the running " + "engine -- this is exactly the registry-key regression this PR fixes " + "(spec registered under the wrong model_class_name), and it would " + "silently skip AR-stage input building entirely." + ) + + generator_device = current_omni_platform.device_type or "cuda" + diffusion_params = OmniDiffusionSamplingParams( + seed=1234, + generator=torch.Generator(device=generator_device).manual_seed(1234), + num_outputs_per_prompt=1, + ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + sampling_params_list[idx] = diffusion_params + + prompt_dict: dict[str, object] = {"prompt": PROMPT, "modalities": ["image"]} + text_to_image._apply_ar_stage_inputs( + ar_input_builder, + model=MODEL_NAME, + prompt_text=PROMPT, + extra_body={"use_system_prompt": "en_recaption"}, + num_images=0, + height=None, + width=None, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=True, + # validate_special_token_ids now tolerates tokens the loaded tokenizer + # doesn't recognize at all (unk_token_id) instead of requiring every + # entry in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS to be present, so this is + # safe against MODEL_NAME's base checkpoint missing the Instruct-only + # / tokens. Wiring it here exercises the + # real production path (get_ar_tokenizer_validator via the registry), + # not just the isolated fake-tokenizer coverage in test_prompt_utils.py + # / test_shared_script_ar_integration.py. + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + ) + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) + generated_image = _extract_generated_image(outputs) + score = compare_semantic(SYSTEM_EN_RECAPTION, generated_image, clip_model, clip_processor) + print(f"shared-script AR path (en_recaption): CLIP cosine similarity = {score:.6f}") + # test_system_prompt_scores now asserts this same (seed=1234, en_recaption) + # case through the hand-constructed extra_args path against the same + # module-level MIN_SEMANTIC_SIMILARITY floor, so this reuses it too rather + # than picking an independent threshold for the same comparison. + assert score >= MIN_SEMANTIC_SIMILARITY, ( + f"shared-script AR path (en_recaption) CLIP similarity {score:.4f} is below " + f"{MIN_SEMANTIC_SIMILARITY} -- the migrated get_ar_input_builder path may have " + "diverged from the reference HunyuanImage-3.0 output for this prompt/system-prompt." + ) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index c45563617..60850248e 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -2314,3 +2314,72 @@ def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_on assert captured_prompt["prompt"].count("") == 2, ( f"N=2 reference images must emit 2 placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}" ) + + +@pytest.mark.xfail( + strict=True, + reason=( + "Pre-existing gap, not introduced by this change: the /v1/images/edits route in " + "api_server.py never calls resolve_stop_token_ids at all (grep: api_server.py has " + "no stop_token_ids reference), so the AR stage keeps SamplingParams' default empty " + "list. The omitted-bot_task fix landed in serving_chat.py, which is the only " + "production caller; that path is covered by " + "test_serving_chat_multistage_generation.py::" + "test_build_multistage_generation_inputs_omitted_bot_task_matches_prompt_default. " + "Wiring the images/edits route through the same seam needs a tokenizer in that " + "scope and a decision on whether HunyuanImage3 it2i should resolve AR stop tokens " + "unconditionally (today the whole AR block is gated on an explicit bot_task / " + "use_system_prompt / system_prompt), so it is tracked separately." + ), +) +def test_image_edits_omitted_bot_task_stop_tokens_match_prompt_default( + async_omni_stage_configs_only_client, +): + """Regression: an omitted bot_task must resolve identically for the AR + prompt and its stop_token_ids. + + build_prompt/build_prompt_tokens default an omitted bot_task per-task + (e.g. "think" for the base "it2i" task, since it isn't itself a key in + _TASK_PRESETS). resolve_stop_token_ids must land on that same default + to compute the matching stop set. Passing the raw (still-None) outer + bot_task variable to resolve_stop_token_ids -- instead of mirroring + build_kwargs's own omitted-or-not "bot_task" entry -- made it normalize + bot_task=None instead of "think", so it fell through to the full + stop range instead of the think/recaption-only pair. + + An explicit (non-"auto") size is required to reach this: with + need_ratio=True (size="auto"), resolve_stop_token_ids returns the full + ratio range regardless of bot_task, so the two code paths only visibly + disagree once a concrete size selects the narrower think/recaption stop + set. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + + img = make_test_image_bytes((64, 64)) + response = async_omni_stage_configs_only_client.post( + "/v1/images/edits", + files=[("image", img)], + data={ + "prompt": "make it neon", + "size": "512x512", + }, + ) + assert response.status_code == 200, response.text + + engine = async_omni_stage_configs_only_client.app.state.engine_client + captured = engine.captured_sampling_params_list + assert captured is not None + + ar_params = captured[0] + expected_stop_token_ids = [ + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + ] + assert ar_params.stop_token_ids == expected_stop_token_ids, ( + f"omitted bot_task with an explicit size must resolve stop_token_ids " + f"for the default 'think' bot_task ({expected_stop_token_ids}); got " + f"{ar_params.stop_token_ids} -- this is the full ratio range, meaning " + "resolve_stop_token_ids disagreed with build_prompt_tokens's default." + ) diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py index 0dea97179..639ddf6c6 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -538,3 +538,84 @@ def test_build_multistage_generation_inputs_no_stop_token_ids_without_size(servi assert sampling_params_list[0].stop_token_ids == [], ( "Without bot_task, AR stage stop_token_ids must be the default empty list" ) + + +def test_build_multistage_generation_inputs_omitted_bot_task_matches_prompt_default(serving_chat): + """Regression: an omitted bot_task must resolve identically for the AR + prompt and for its stop_token_ids. + + ``build_prompt_tokens``/``build_prompt`` default an omitted bot_task + per-task via the ``_DEFAULT_BOT_TASK`` sentinel (``"think"`` for the base + ``it2i`` task, which isn't itself a key in ``_TASK_PRESETS``). + ``resolve_stop_token_ids`` must normalize from that same omitted-or-not + starting point. Passing the raw -- still ``None`` -- outer ``bot_task`` + variable instead of mirroring ``build_kwargs``'s own omitted-or-not + ``"bot_task"`` entry made it normalize ``bot_task=None`` (plain mode), so + it fell through to the full ```` range instead of the + think/recaption terminator pair. + + An explicit (non-"auto") size is required to observe this: with + ``need_ratio=True`` the ratio range is returned regardless of bot_task, so + the two paths only visibly disagree once a concrete size selects the + narrower stop set. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + class FakeTokenizer: + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 4, + } + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return list(range(100, 100 + len(text))) + + def _build(extra_body): + engine = SimpleNamespace( + stage_configs=[ + SimpleNamespace(stage_type="llm", is_comprehension=True), + SimpleNamespace(stage_type="diffusion", is_comprehension=False), + ], + default_sampling_params_list=[ + SamplingParams(temperature=0.0), + OmniDiffusionSamplingParams(), + ], + ) + _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs( + serving_chat, + engine=engine, + prompt="draw a cat", + extra_body=extra_body, + reference_images=[Image.new("RGB", (32, 32), color="red")], + gen_params=OmniDiffusionSamplingParams(height=768, width=1024), + tokenizer=FakeTokenizer(), + ) + return sampling_params_list[0].stop_token_ids + + # bot_task omitted entirely; use_system_prompt is what carries the request + # into the AR-prompt path. + omitted = _build({"use_system_prompt": "en_recaption"}) + # Same request with the default spelled out explicitly. + explicit = _build({"bot_task": "think", "use_system_prompt": "en_recaption"}) + + expected = [ + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + ] + assert omitted == expected, ( + f"omitted bot_task with an explicit size must resolve stop_token_ids for the " + f"default 'think' bot_task ({expected}); got {omitted}. A long list here is the " + f"full range, meaning resolve_stop_token_ids normalized bot_task=None " + f"(plain mode) instead of agreeing with build_prompt_tokens's default." + ) + assert omitted == explicit, ( + "omitting bot_task must resolve the same stop_token_ids as passing the default " + f"explicitly; omitted={omitted} explicit={explicit}" + ) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index e55691e10..dcdf13678 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sys from typing import Any import pytest @@ -592,6 +593,193 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None: } +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("HunyuanImage3Pipeline") == frozenset( + { + "bot_task", + "use_system_prompt", + "system_prompt", + "negative_prompt", + } + ) + assert get_extra_output_params("HunyuanImage3Pipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("HunyuanImage3Pipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_extra_registry_matches_runtime_model_class_name() -> None: + # od_config.model_class_name resolves to "HunyuanImage3ForCausalMM" at + # runtime; "HunyuanImage3Pipeline" is kept only as a compatibility alias. + assert get_extra_body_params("HunyuanImage3ForCausalMM") == get_extra_body_params("HunyuanImage3Pipeline") + assert get_extra_output_params("HunyuanImage3ForCausalMM") == get_extra_output_params("HunyuanImage3Pipeline") + assert should_init_extra_args_for_non_diffusion_stages( + "HunyuanImage3ForCausalMM" + ) == should_init_extra_args_for_non_diffusion_stages("HunyuanImage3Pipeline") + + +class _FakeHunyuanTokenizer: + """Minimal tokenizer with distinct ids so AR token assertions are unambiguous.""" + + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 3, + "": 4, + } + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + # Offset well above the special-id range to avoid collisions. + return [500 + len(text)] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_ar_input_builder_registered() -> None: + from vllm_omni.model_extras import get_ar_input_builder + + builder = get_ar_input_builder("HunyuanImage3Pipeline") + assert builder is not None + assert get_ar_input_builder("HunyuanImage3ForCausalMM") is builder + # Models without the hook return None. + assert get_ar_input_builder("WanVACEPipeline") is None + assert get_ar_input_builder(None) is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_ar_tokenizer_validator_registered() -> None: + from vllm_omni.model_extras import get_ar_tokenizer_validator + + validator = get_ar_tokenizer_validator("HunyuanImage3Pipeline") + assert validator is not None + assert get_ar_tokenizer_validator("HunyuanImage3ForCausalMM") is validator + # Models without the hook return None. + assert get_ar_tokenizer_validator("WanVACEPipeline") is None + assert get_ar_tokenizer_validator(None) is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_t2i_string_fallback() -> None: + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="a red panda", + tokenizer=None, + extra_body={"bot_task": "recaption"}, + num_images=0, + height=1024, + width=1024, + ) + # No tokenizer -> string prompt form; image modality; recaption + explicit + # size stops at only. + assert ar.prompt_token_ids is None + assert ar.prompt and "a red panda" in ar.prompt + assert ar.modalities == ["image"] + assert ar.use_system_prompt + assert len(ar.stop_token_ids) == 1 + # stop_token_ids belong to stage 0 (HunyuanImage3's single AR stage) -- + # declared explicitly, not inferred by the caller from stage type. + assert ar.stage_indices == [0] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_t2i_tokenizer_parity() -> None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="a red panda", + tokenizer=_FakeHunyuanTokenizer(), + extra_body={"bot_task": "think"}, + num_images=0, + ) + expected = build_prompt_tokens("a red panda", _FakeHunyuanTokenizer(), task="t2i", bot_task="think") + assert ar.prompt is None + assert ar.prompt_token_ids == expected.token_ids + assert ar.modalities == ["image"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_it2i_inserts_image_placeholders() -> None: + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="edit it", + tokenizer=_FakeHunyuanTokenizer(), + extra_body={}, + num_images=2, + ) + # Two reference images -> two placeholder ids (2) in the AR prefill. + assert ar.prompt_token_ids is not None + assert ar.prompt_token_ids.count(_FakeHunyuanTokenizer.SPECIAL[""]) == 2 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_comprehension_text_output() -> None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + # t2t (no image, text output) -> text modality, stop on . + ar_t2t = build_ar_stage_inputs("hi", tokenizer=None, extra_body={}, num_images=0, text_output=True) + assert ar_t2t.modalities == ["text"] + assert ar_t2t.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + # i2t (image, text output) -> single placeholder in the prefill. + ar_i2t = build_ar_stage_inputs( + "describe", tokenizer=_FakeHunyuanTokenizer(), extra_body={}, num_images=1, text_output=True + ) + assert ar_i2t.modalities == ["text"] + assert ar_i2t.prompt_token_ids.count(_FakeHunyuanTokenizer.SPECIAL[""]) == 1 + + +class _RealIdHunyuanTokenizer: + """Mirrors HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS exactly, so validate_ar_tokenizer + accepts it -- unlike _FakeHunyuanTokenizer above, whose ids are deliberately + wrong to exercise the drift-detection failure path.""" + + unk_token_id = -1 + + def convert_tokens_to_ids(self, tok: str) -> int: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.get(tok, self.unk_token_id) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [900 + len(text)] + + +def _patch_hunyuan_autotokenizer(monkeypatch: pytest.MonkeyPatch, tokenizer: object) -> list[str]: + """Patch the ``AutoTokenizer`` that ``hunyuan_image3.build_x_to_text_prompt`` + imports locally, recording the models it was asked to load.""" + calls: list[str] = [] + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> object: + calls.append(model) + assert trust_remote_code is True + return tokenizer + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FakeAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + return calls + + @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: @@ -619,6 +807,59 @@ def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> N assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_t2t(monkeypatch: pytest.MonkeyPatch) -> None: + """t2t (no image) routes through build_ar_stage_inputs and stops on .""" + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + calls = _patch_hunyuan_autotokenizer(monkeypatch, _RealIdHunyuanTokenizer()) + prompt_dict, stop_token_ids = hunyuan_build_x_to_text_prompt( + "tencent/HunyuanImage-3.0-Instruct", "What is the capital of France?", has_image=False + ) + + assert calls == ["tencent/HunyuanImage-3.0-Instruct"] + assert prompt_dict["modalities"] == ["text"] + assert prompt_dict["prompt_token_ids"] is not None + assert prompt_dict["prompt_token_ids"].count(HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]) == 0 + assert stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_i2t(monkeypatch: pytest.MonkeyPatch) -> None: + """i2t (image) inserts exactly one placeholder and still stops on .""" + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + _patch_hunyuan_autotokenizer(monkeypatch, _RealIdHunyuanTokenizer()) + prompt_dict, stop_token_ids = hunyuan_build_x_to_text_prompt( + "tencent/HunyuanImage-3.0-Instruct", "Describe this image.", has_image=True + ) + + assert prompt_dict["modalities"] == ["text"] + assert prompt_dict["prompt_token_ids"].count(HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]) == 1 + assert stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_validates_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None: + """A tokenizer whose special-token ids have drifted must be rejected here too -- + x_to_text.py (unlike text_to_image.py/image_edit.py) has no registry-driven + ar_tokenizer_validator wiring of its own, so this seam must self-validate.""" + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + _patch_hunyuan_autotokenizer(monkeypatch, _FakeHunyuanTokenizer()) + with pytest.raises(ValueError, match="no longer match"): + hunyuan_build_x_to_text_prompt("tencent/HunyuanImage-3.0-Instruct", "Hello.", has_image=False) + + @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_text_to_image_prompt_builder() -> None: diff --git a/tests/model_extras/test_shared_script_ar_integration.py b/tests/model_extras/test_shared_script_ar_integration.py new file mode 100644 index 000000000..f781ca5fb --- /dev/null +++ b/tests/model_extras/test_shared_script_ar_integration.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration coverage for the shared task examples' ``_apply_ar_stage_inputs``. + +The unit tests in ``test_model_extras.py`` cover ``get_ar_input_builder`` and +``build_ar_stage_inputs`` in isolation, but nothing previously exercised the +*glue* in ``text_to_image.py`` / ``image_edit.py`` that wires them together: +resolving the model's declared ``ar_input_builder`` via the registry, loading +the AR tokenizer, and writing the result onto ``prompt_dict`` / +``sampling_params_list``. A regression in that glue -- e.g. the registry key +mismatch fixed in this PR (``HunyuanImage3ForCausalMM`` vs +``HunyuanImage3Pipeline``), or the ``trust_remote_code`` opt-in being +silently bypassed -- would not have been caught by the isolated unit tests +above, since they call ``build_ar_stage_inputs`` directly rather than going +through the shared script's own code path. + +These tests dynamically load the example scripts by file path (``examples/`` +is not an installed package) and call ``_apply_ar_stage_inputs`` the same way +``main()`` does, with a fake tokenizer standing in for a real HF download. +""" + +from __future__ import annotations + +import importlib.util +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _load_module(name: str, relpath: str): + path = _REPO_ROOT / relpath + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +class _FakeHunyuanTokenizer: + """Mirrors test_model_extras.py's fixture: distinct ids, no network I/O.""" + + SPECIAL = {"<|startoftext|>": 1, "": 2, "": 3, "": 4} + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [500 + len(text)] + + +@dataclass +class _FakeARSamplingParams: + """Stands in for an LLM (non-diffusion) stage's sampling params.""" + + stop_token_ids: list[int] | None = None + + +def _patch_tokenizer_loader(monkeypatch: pytest.MonkeyPatch, module: Any) -> list[dict[str, Any]]: + """Patch the ``AutoTokenizer`` the module imports locally, recording calls.""" + calls: list[dict[str, Any]] = [] + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> _FakeHunyuanTokenizer: + calls.append({"model": model, "trust_remote_code": trust_remote_code}) + return _FakeHunyuanTokenizer() + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FakeAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + return calls + + +def _patch_failing_tokenizer_loader(monkeypatch: pytest.MonkeyPatch) -> None: + """Patch AutoTokenizer.from_pretrained to always raise, simulating a + missing --trust-remote-code / network / cache failure.""" + + class _FailingAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> Any: + raise OSError("simulated tokenizer load failure") + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FailingAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_routes_through_real_registry( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """The shared script's glue must resolve HunyuanImage3's *real* registry + entry (keyed on the runtime ``model_class_name``) and write its output + onto the request -- not a stub, not a hardcoded id.""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module) + + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + assert ar_input_builder is not None, ( + "get_ar_input_builder('HunyuanImage3ForCausalMM') returned None -- " + "if this regresses, the shared script silently skips AR-stage " + "input building entirely (this is exactly how the registry-key bug " + "this PR fixes went undetected)." + ) + + prompt_dict: dict[str, Any] = {} + ar_params = _FakeARSamplingParams() + diffusion_params = OmniDiffusionSamplingParams() + sampling_params_list: list[Any] = [ar_params, diffusion_params] + + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={"bot_task": "think", "use_system_prompt": "en_recaption"}, + num_images=0, + height=1024, + width=1024, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=False, + ) + + # The tokenizer was loaded exactly once, for the model the user asked for. + assert len(tokenizer_calls) == 1 + assert tokenizer_calls[0]["model"] == "tencent/HunyuanImage-3.0-Instruct" + + # Byte-for-byte tokenizer path: prompt_token_ids populated, not the string form. + assert prompt_dict.get("prompt_token_ids"), "AR prefill token-ids were not written onto the request" + assert "prompt" not in prompt_dict + assert prompt_dict.get("use_system_prompt") == "en_recaption" + assert prompt_dict.get("modalities") == ["image"] + + # stop_token_ids reached the AR-stage params, and only those. The diffusion + # stage must be left untouched: OmniDiffusionSamplingParams has no + # stop_token_ids field, so the helper skipping it means the attribute is + # never created (reading it would raise AttributeError, not return None). + assert ar_params.stop_token_ids, "AR stage did not receive stop_token_ids" + assert not hasattr(diffusion_params, "stop_token_ids"), "diffusion stage params must not receive stop_token_ids" + + +@pytest.mark.parametrize( + "module_name,relpath,cli_flag", + [ + ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", False), + ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", True), + ], +) +def test_apply_ar_stage_inputs_respects_trust_remote_code_option( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str, cli_flag: bool +) -> None: + """The tokenizer load must use the caller's resolved --trust-remote-code + value, not silently default to True regardless of user intent.""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={}, + num_images=0, + height=1024, + width=1024, + prompt_dict={}, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=cli_flag, + ) + + assert tokenizer_calls[0]["trust_remote_code"] is cli_flag + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image_validate", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit_validate", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_runs_declared_tokenizer_validator( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """get_ar_tokenizer_validator's hook must run on a real tokenizer load, + and a validation failure must propagate -- not get swallowed by the + tokenizer-load try/except (which would silently fall back to the string + prompt instead of surfacing the drift).""" + from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator + + module = _load_module(module_name, relpath) + _patch_tokenizer_loader(monkeypatch, module) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + validator = get_ar_tokenizer_validator("HunyuanImage3ForCausalMM") + assert validator is not None + + # The fake tokenizer's ids intentionally don't match + # HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, so the real validator must reject it. + with pytest.raises(ValueError, match="no longer match"): + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={}, + num_images=0, + height=1024, + width=1024, + prompt_dict={}, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=False, + validate_tokenizer=validator, + ) + + +@pytest.mark.parametrize( + "module_name,relpath,required_argv", + [ + ("_shared_text_to_image_cli", "examples/offline_inference/text_to_image/text_to_image.py", []), + # image_edit.py's parser declares --image/--prompt as required (no defaults); + # omitting them makes parse_args() raise SystemExit(2) before the assertion + # below is ever reached. + ( + "_shared_image_edit_cli", + "examples/offline_inference/image_to_image/image_edit.py", + ["--image", "dummy.png", "--prompt", "x"], + ), + ], +) +@pytest.mark.parametrize("argv,expected", [([], False), (["--trust-remote-code"], True)]) +def test_trust_remote_code_flag_parses_from_real_cli( + monkeypatch: pytest.MonkeyPatch, + module_name: str, + relpath: str, + required_argv: list[str], + argv: list[str], + expected: bool, +) -> None: + """The other tests in this file call _apply_ar_stage_inputs directly with + an already-resolved trust_remote_code bool, which never exercises + argparse itself. Closing that specific gap doesn't need a GPU or a model + load -- parse_args() only touches sys.argv -- so cover it here instead of + the (expensive, one-model-load-per-invocation) full CLI subprocess a + stricter GPU e2e test would need.""" + module = _load_module(module_name, relpath) + monkeypatch.setattr(sys, "argv", ["prog", *required_argv, *argv]) + args = module.parse_args() + assert args.trust_remote_code is expected + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image_failfast", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit_failfast", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_tokenizer_load_failure_raises_by_default( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """A tokenizer load failure must surface as an error by default, not + silently degrade to the string-prompt form -- that degraded form can + complete "successfully" with subtly wrong stop tokens, masking the real + problem (e.g. a missing --trust-remote-code).""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + _patch_failing_tokenizer_loader(monkeypatch) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + + with pytest.raises(OSError, match="simulated tokenizer load failure"): + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={}, + num_images=0, + height=1024, + width=1024, + prompt_dict={}, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=False, + ) + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image_fallback", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit_fallback", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_allow_tokenizer_fallback_option( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """With allow_tokenizer_fallback=True (the CLI's --allow-tokenizer-fallback), + a tokenizer load failure degrades to the string-prompt form instead of + raising -- for a deliberate offline-compat run, not the default.""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + _patch_failing_tokenizer_loader(monkeypatch) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + + prompt_dict: dict[str, Any] = {} + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={"bot_task": "recaption"}, + num_images=0, + height=1024, + width=1024, + prompt_dict=prompt_dict, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=False, + allow_tokenizer_fallback=True, + ) + + # No tokenizer -> string-prompt form, not the byte-for-byte token-ids form. + assert "prompt_token_ids" not in prompt_dict + assert prompt_dict.get("prompt") and "a red panda" in prompt_dict["prompt"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py index b64fe2e97..b4ecd7dca 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py @@ -60,6 +60,48 @@ HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS: dict[str, int] = { "": 130106, } + +def validate_special_token_ids(tokenizer: Any) -> None: + """Fail fast if a real tokenizer's ids drift from the hardcoded table. + + ``HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS`` exists so stop-token / trigger-tag + resolution doesn't need a tokenizer round-trip, but that only holds as + long as the ids match ``tokenizer.json`` for whatever model repo/revision + is actually loaded. Call this once a real tokenizer is available (the + example scripts do, right after loading it) to catch drift immediately + instead of silently producing a request that runs to completion with the + wrong stop tokens. Not wired into ``build_prompt_tokens`` itself: callers + that intentionally pass a stand-in tokenizer with different ids (e.g. + this module's own unit tests) must not be forced through this check. + + Both ``tencent/HunyuanImage-3.0`` (base) and ``-Instruct`` report the + same ``model_class_name`` (``HunyuanImage3ForCausalMM``), so this hook + fires for either checkpoint -- but the base tokenizer's vocab doesn't + include the Instruct-only ````-```` tokens + (added for it2i's second aspect-ratio range). A token the tokenizer + doesn't recognize at all (``convert_tokens_to_ids`` returning + ``unk_token_id``, or ``None`` for tokenizers without an unk concept) is + "not part of this checkpoint's vocab," not "drifted"; only an id that + resolves to something *else* indicates the hardcoded table is stale. + """ + unk_id = getattr(tokenizer, "unk_token_id", None) + mismatches: dict[str, tuple[int, int]] = {} + for tok, expected_id in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.items(): + actual_id = tokenizer.convert_tokens_to_ids(tok) + if actual_id is None or (unk_id is not None and actual_id == unk_id): + continue # not in this checkpoint's vocab -- nothing to compare + if actual_id != expected_id: + mismatches[tok] = (expected_id, actual_id) + if mismatches: + details = ", ".join(f"{tok!r}: expected {exp}, got {act}" for tok, (exp, act) in mismatches.items()) + raise ValueError( + "HunyuanImage-3.0 tokenizer special-token ids no longer match " + f"HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS ({details}). The loaded " + "tokenizer/model revision has likely changed; update the hardcoded " + "table in prompt_utils.py to match." + ) + + # bot_task -> (sys_type, trigger_tag). # ``vanilla`` is special-cased downstream: it bypasses the chat template # (no ``User:`` / ``Assistant:`` framing) and is only valid with @@ -317,12 +359,106 @@ def build_prompt_tokens( ) +@dataclass +class ARPromptInputs: + """Fully-resolved AR-stage inputs for one HunyuanImage3 request. + + Consolidates the three things every caller (offline example, OpenAI + server) must compute before submitting an AR request: the AR prefill + (as segment-tokenized ``prompt_token_ids`` when a tokenizer is + available, else a formatted ``prompt`` string), the resolved DiT + system-prompt type, and the AR ``stop_token_ids``. + """ + + prompt: str | None + prompt_token_ids: list[int] | None + system_prompt_type: str + stop_token_ids: list[int] + + +def build_ar_prompt_inputs( + user_prompt: str, + task: str = "t2i", + bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK, + sys_type: str | None = None, + custom_system_prompt: str | None = None, + num_images: int = 1, + tokenizer: Any | None = None, + image_size: str | None = None, +) -> ARPromptInputs: + """Build the complete AR-stage inputs for a HunyuanImage3 request. + + Wraps :func:`build_prompt_tokens` (byte-for-byte HF parity when + ``tokenizer`` is given), the :func:`build_prompt` string fallback (when + no tokenizer is available), and :func:`resolve_stop_token_ids` behind + one call, so a caller doesn't have to re-derive the AR prefill/stop-token + logic -- and, critically, doesn't have to independently resolve + ``bot_task`` for the prompt-building call vs. the stop-token call, which + can silently diverge if a caller passes the same nominal ``bot_task`` + value to both but resolves "omitted" vs. "explicit None" differently for + each (this seam always forwards the identical resolved value to both). + + Used by all four offline shared task examples: ``text_to_image.py`` / + ``image_edit.py`` (via ``model_extras.hunyuan_image3.build_ar_stage_inputs``) + and ``x_to_text.py`` (via ``build_x_to_text_prompt``, which also calls + ``build_ar_stage_inputs``). The OpenAI server's ``serving_chat.py`` does + not currently call this seam -- it still independently calls + ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids``, and + can hit exactly the divergence described above when ``bot_task`` is + omitted from ``extra_body`` (pre-existing, unrelated to this module). + ``image_size`` follows the same convention as :func:`resolve_stop_token_ids` + (``None``/``"auto"`` lets the AR predict the aspect ratio; an explicit + ``"{w}x{h}"`` makes it stop at the terminator). + """ + if tokenizer is not None: + result = build_prompt_tokens( + user_prompt, + tokenizer, + task=task, + bot_task=bot_task, + sys_type=sys_type, + custom_system_prompt=custom_system_prompt, + num_images=num_images, + ) + prompt_str: str | None = None + prompt_token_ids: list[int] | None = result.token_ids + system_prompt_type = result.system_prompt_type + else: + prompt_str = build_prompt( + user_prompt, + task=task, + bot_task=bot_task, + sys_type=sys_type, + custom_system_prompt=custom_system_prompt, + num_images=num_images, + ) + prompt_token_ids = None + _, resolved_bot_task = _normalize_task_and_bot_task(task, bot_task) + system_prompt_type = sys_type or resolve_sys_type(resolved_bot_task) + + stop_token_ids = resolve_stop_token_ids( + task=task, + bot_task=bot_task, + tokenizer=tokenizer, + image_size=image_size, + ) + return ARPromptInputs( + prompt=prompt_str, + prompt_token_ids=prompt_token_ids, + system_prompt_type=system_prompt_type, + stop_token_ids=stop_token_ids, + ) + + __all__ = [ "HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS", + "validate_special_token_ids", "MAX_IMAGES_PER_REQUEST", "_TASK_PRESETS", + "ARPromptInputs", "available_bot_tasks", "available_tasks", + "build_ar_prompt_inputs", "build_prompt", "build_prompt_tokens", "resolve_stop_token_ids", diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f5cb24efd..20bddaba6 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3080,11 +3080,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ar_image_size: str | None = None if height is not None and width is not None: ar_image_size = f"{width}x{height}" + # Reuse build_kwargs's own (possibly-omitted) "bot_task" entry rather + # than the raw outer `bot_task` variable: build_prompt_tokens/build_prompt + # above default an omitted bot_task per-task (e.g. "think" for t2i), and + # resolve_stop_token_ids must normalize from that same omitted-or-not + # starting point to agree on which stop tokens apply -- passing the raw + # `bot_task` (still None when omitted) made the two calls disagree. + ar_stop_kwargs: dict[str, Any] = {} + if "bot_task" in build_kwargs: + ar_stop_kwargs["bot_task"] = build_kwargs["bot_task"] ar_stop_token_ids = resolve_stop_token_ids( task=ar_task, - bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size, + **ar_stop_kwargs, ) engine_prompt: OmniTextPrompt = {"prompt": prompt} diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index 143a9b224..2e1146a82 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -6,6 +6,8 @@ from vllm_omni.model_extras.registry import ( build_image_to_video_prompt, build_text_to_image_prompt, build_x_to_text_prompt, + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_extra_output_params, get_model_class_name, @@ -22,6 +24,8 @@ __all__ = [ "build_image_to_video_prompt", "build_text_to_image_prompt", "build_x_to_text_prompt", + "get_ar_input_builder", + "get_ar_tokenizer_validator", "get_extra_body_params", "get_extra_output_params", "get_model_class_name", diff --git a/vllm_omni/model_extras/hunyuan_image3.py b/vllm_omni/model_extras/hunyuan_image3.py index 8a53dfd54..46d48b513 100644 --- a/vllm_omni/model_extras/hunyuan_image3.py +++ b/vllm_omni/model_extras/hunyuan_image3.py @@ -3,6 +3,8 @@ from __future__ import annotations +from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Any @@ -11,27 +13,162 @@ def build_x_to_text_prompt( prompt: str, has_image: bool, ) -> tuple[dict[str, Any], list[int] | None]: - """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T.""" + """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T. + + Routes through :func:`build_ar_stage_inputs`, the same declarative seam + the shared text_to_image / image_edit examples use, instead of + re-deriving the AR prefill and stop-token logic here. Unlike those + examples -- generic scripts that reach ``validate_ar_tokenizer`` only + indirectly via ``get_ar_tokenizer_validator`` -- this function is already + HunyuanImage3-specific, so it calls the validator directly rather than + requiring the (model-agnostic) ``x_to_text.py`` caller to know about it. + """ from transformers import AutoTokenizer - from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( - build_prompt_tokens, - resolve_stop_token_ids, - resolve_sys_type, - ) - - task = "i2t" if has_image else "t2t" tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) - build_kwargs: dict[str, Any] = {"task": task, "bot_task": None} - if has_image: - build_kwargs["num_images"] = 1 - built = build_prompt_tokens(prompt, tokenizer, **build_kwargs) + validate_ar_tokenizer(tokenizer) + ar_inputs = build_ar_stage_inputs( + prompt=prompt, + tokenizer=tokenizer, + extra_body=None, + num_images=1 if has_image else 0, + text_output=True, + ) return ( { - "prompt": prompt, - "prompt_token_ids": built.token_ids, - "modalities": ["text"], - "use_system_prompt": resolve_sys_type(None), + "prompt": ar_inputs.prompt, + "prompt_token_ids": ar_inputs.prompt_token_ids, + "modalities": ar_inputs.modalities, + "use_system_prompt": ar_inputs.use_system_prompt, }, - resolve_stop_token_ids(task=task, bot_task=None, tokenizer=tokenizer, image_size="auto"), + ar_inputs.stop_token_ids, ) + + +@dataclass +class ARStageInputs: + """AR-stage inputs the shared example applies to a HunyuanImage3 request. + + ``prompt_token_ids`` (byte-for-byte HF parity) is preferred; ``prompt`` + is the string fallback used when no tokenizer is available. ``modalities`` + picks the AR output type (``["image"]`` for generation, ``["text"]`` for + comprehension). ``use_system_prompt`` feeds the DiT system-prompt prefix + and ``stop_token_ids`` terminate the AR decode. + + ``stage_indices`` names which stage(s) in the request's + ``sampling_params_list`` these ``stop_token_ids`` belong to -- the shared + scripts apply them by explicit index, not by scanning for "whichever + stage isn't a diffusion stage." HunyuanImage-3.0's topology always puts + its single AR stage at index 0 (see + ``vllm_omni.model_executor.models.hunyuan_image3.pipeline``); a future + model with multiple AR/understanding stages would declare its own + indices here instead of relying on type-based inference. + """ + + prompt: str | None + prompt_token_ids: list[int] | None + modalities: list[str] + use_system_prompt: str + stop_token_ids: list[int] + stage_indices: list[int] = field(default_factory=lambda: [0]) + + +def build_ar_stage_inputs( + prompt: str, + tokenizer: Any | None, + extra_body: Mapping[str, Any] | None, + *, + num_images: int = 0, + height: int | None = None, + width: int | None = None, + text_output: bool = False, +) -> ARStageInputs: + """Resolve HunyuanImage-3.0 AR-stage inputs declaratively. + + Invoked generically by the shared task examples (``text_to_image`` / + ``image_to_image`` / ``x_to_text``) whenever the model declares an + ``ar_input_builder``. All model-specific knobs (``bot_task`` / + ``use_system_prompt`` / ``system_prompt``) arrive via ``extra_body``, so + the example itself stays model-agnostic. The heavy lifting is delegated to + :func:`build_ar_prompt_inputs`. The OpenAI server's ``serving_chat.py`` + does not go through this seam -- it independently calls the same + underlying ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids`` + primitives (see that function's docstring for the resulting divergence risk). + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + build_ar_prompt_inputs, + ) + + extra_body = extra_body or {} + has_image = num_images > 0 + if text_output: + task = "i2t" if has_image else "t2t" + else: + task = "it2i" if has_image else "t2i" + + # ar_image_size: None -> AR predicts aspect ratio; explicit "{w}x{h}" -> + # AR stops at the terminator. Comprehension (text) tasks always auto-size. + ar_image_size: str | None = None + if not text_output and height is not None and width is not None: + ar_image_size = f"{width}x{height}" + + kwargs: dict[str, Any] = { + "task": task, + "sys_type": extra_body.get("use_system_prompt"), + "custom_system_prompt": extra_body.get("system_prompt"), + "num_images": max(num_images, 1) if has_image else 1, + "tokenizer": tokenizer, + "image_size": ar_image_size, + } + # Only forward bot_task when the caller set it: present (incl. explicit + # ``None`` -> plain mode) overrides; absent lets each task pick its default + # trigger via build_ar_prompt_inputs' own sentinel default. + if "bot_task" in extra_body: + kwargs["bot_task"] = extra_body["bot_task"] + + resolved = build_ar_prompt_inputs(prompt, **kwargs) + return ARStageInputs( + prompt=resolved.prompt, + prompt_token_ids=resolved.prompt_token_ids, + modalities=["text"] if text_output else ["image"], + use_system_prompt=resolved.system_prompt_type, + stop_token_ids=resolved.stop_token_ids, + # HunyuanImage-3.0's AR stage is always stage 0 (see pipeline.py's + # HUNYUAN_IMAGE3_PIPELINE / _AR_PIPELINE topologies). + stage_indices=[0], + ) + + +def validate_ar_tokenizer(tokenizer: Any) -> None: + """Registry-declared ``ar_tokenizer_validator`` hook for HunyuanImage3. + + Called by the shared task examples right after they load a *real* + tokenizer for the AR stage, so a model/tokenizer revision bump that + silently shifts special-token ids fails loudly instead of producing a + request that completes with the wrong stop tokens. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + validate_special_token_ids, + ) + + validate_special_token_ids(tokenizer) + + +# Per-request, model-specific knobs (non-standard sampling-param fields). +# Standard knobs (guidance_scale / num_inference_steps / seed / height / width) +# stay on OmniDiffusionSamplingParams; engine-level knobs +# (diffusion_kv_cache_* / vae_use_tiling) stay as Omni() init args. +HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS = frozenset( + { + "bot_task", + "use_system_prompt", + "system_prompt", + "negative_prompt", + } +) +# Text outputs for i2t/t2t are surfaced through the standard AR text path; +# no diffusion custom-output params are declared. +HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS = frozenset() +# HunyuanImage3 runs an AR stage before DiT; its extra_args must be initialised +# so bot_task / use_system_prompt reach the AR input path. +HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 4d2ad010a..f21691eab 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -28,7 +28,18 @@ from vllm_omni.model_extras.helios import ( HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_OUTPUT_PARAMS, ) +from vllm_omni.model_extras.hunyuan_image3 import ( + HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS, + HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS, + HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.hunyuan_image3 import ( + build_ar_stage_inputs as build_hunyuan_image3_ar_stage_inputs, +) from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt +from vllm_omni.model_extras.hunyuan_image3 import ( + validate_ar_tokenizer as validate_hunyuan_image3_ar_tokenizer, +) from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.ltx2 import ( LTX_EXTRA_BODY_PARAMS, @@ -233,6 +244,18 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "LTX2DistilledTwoStagePipeline", ) }, + # HunyuanImage3 keys on the model architecture name that ``get_model_class_name`` + # reports at runtime for its AR+DiT / AR-only / DiT-only deploys + # (``od_config.model_class_name`` == the diffusion registry arch key). The + # "HunyuanImage3Pipeline" alias is kept for any path that surfaces the + # pipeline class name instead. + "HunyuanImage3ForCausalMM": { + "extra_body_params": HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS, + "extra_output_params": HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "ar_input_builder": build_hunyuan_image3_ar_stage_inputs, + "ar_tokenizer_validator": validate_hunyuan_image3_ar_tokenizer, + }, "SanaVideoPipeline": { "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, }, @@ -263,6 +286,10 @@ for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"): _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder +# Alias: some code paths surface HunyuanImage3's pipeline class name rather than +# the architecture name; point both at the same spec. +_EXTRA_SPECS["HunyuanImage3Pipeline"] = _EXTRA_SPECS["HunyuanImage3ForCausalMM"] + # Multi-stage discovery reports the top-level wrapper rather than its DiT # submodule, so both names must resolve to the same request builders. _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] @@ -358,6 +385,32 @@ def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) +def get_ar_input_builder(model_class_name: str | None) -> Callable[..., Any] | None: + """Return a model's AR-stage input builder, or ``None`` if undeclared. + + Models with a text/AR stage that needs template-formatted prompt tokens + and AR stop tokens (e.g. HunyuanImage3) declare an ``ar_input_builder``. + The shared task examples call it generically when present, so the example + scripts stay model-agnostic; models without one are unaffected. + """ + spec = _get_spec(model_class_name) + return spec.get("ar_input_builder") if spec is not None else None + + +def get_ar_tokenizer_validator(model_class_name: str | None) -> Callable[[Any], None] | None: + """Return a model's AR-tokenizer validator, or ``None`` if undeclared. + + Models whose AR prompt/stop-token logic depends on hardcoded special + token ids (e.g. HunyuanImage3) declare an ``ar_tokenizer_validator`` to + check those ids against whatever tokenizer actually loads at runtime. + The shared task examples call it generically, right after loading a real + tokenizer, so model/tokenizer revision drift fails loudly instead of + silently producing a request with the wrong stop tokens. + """ + spec = _get_spec(model_class_name) + return spec.get("ar_tokenizer_validator") if spec is not None else None + + def build_text_to_image_prompt( model_class_name: str | None, prompt: dict[str, Any], From 3f7c2174574f14a9767c022b01834fa2c52420df Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:09:26 +0800 Subject: [PATCH 021/209] [Model] Avoid scalar synchronizations in GLM-Image preparation (#7172) Signed-off-by: hyw --- .../glm_image/test_glm_image_quantization.py | 5 ++- .../models/glm_image/glm_image_transformer.py | 32 ++++++------------- 2 files changed, 11 insertions(+), 26 deletions(-) diff --git a/tests/diffusion/models/glm_image/test_glm_image_quantization.py b/tests/diffusion/models/glm_image/test_glm_image_quantization.py index 1a6170ca2..dfcd15b1a 100644 --- a/tests/diffusion/models/glm_image/test_glm_image_quantization.py +++ b/tests/diffusion/models/glm_image/test_glm_image_quantization.py @@ -399,15 +399,14 @@ class TestGlmImagePrepareModule: prepare = GlmImagePrepare( image_projector=projector, rope=rope, - patch_size=2, ) # Test forward pass hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W] result = prepare(hidden_states) - assert len(result) == 5 - hidden_out, rope_cos, rope_sin, height, width = result + assert len(result) == 3 + hidden_out, rope_cos, rope_sin = result assert hidden_out.shape[0] == 1 # batch size assert hidden_out.shape[1] == 1024 # seq_len (32 * 32) assert hidden_out.shape[2] == 2560 # hidden_dim diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 75ba60203..1b0868400 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math from collections.abc import Iterable @@ -124,8 +124,8 @@ class GlmImageImageProjector(nn.Module): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, channel, height, width = hidden_states.shape - post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) - post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p] hidden_states = hidden_states.reshape( @@ -189,18 +189,16 @@ class GlmImagePrepare(nn.Module): self, image_projector: nn.Module, rope: GlmImageRotaryPosEmbed, - patch_size: int, ): super().__init__() self.image_projector = image_projector self.rope = rope - self.patch_size = patch_size def forward( self, hidden_states: torch.Tensor, prior_hidden_states: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Process hidden_states and compute RoPE embeddings. Args: @@ -211,14 +209,7 @@ class GlmImagePrepare(nn.Module): hidden_states: Patched hidden states [B, seq_len, D] rope_cos: RoPE cos embeddings [seq_len, dim] rope_sin: RoPE sin embeddings [seq_len, dim] - post_patch_height: Scalar tensor for height after patching - post_patch_width: Scalar tensor for width after patching """ - batch_size, num_channels, height, width = hidden_states.shape - - post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) - post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) - # Compute RoPE (uses original 4D hidden_states shape) image_rotary_emb = self.rope(hidden_states) rope_cos = image_rotary_emb[0].to(hidden_states.device) @@ -231,7 +222,7 @@ class GlmImagePrepare(nn.Module): if prior_hidden_states is not None: hidden_states = hidden_states + prior_hidden_states - return hidden_states, rope_cos, rope_sin, post_patch_height, post_patch_width + return hidden_states, rope_cos, rope_sin class GlmImageAdaLayerNormZero(nn.Module): @@ -961,7 +952,6 @@ class GlmImageTransformer2DModel(CachedTransformer): 1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), # RoPE sin: [seq_len, dim] - shard along sequence dimension 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), - # post_patch_height and post_patch_width are scalars, not sharded }, # Gather output at proj_out "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), @@ -1044,7 +1034,7 @@ class GlmImageTransformer2DModel(CachedTransformer): ) # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan) - self.prepare = GlmImagePrepare(self.image_projector, self.rope, patch_size) + self.prepare = GlmImagePrepare(self.image_projector, self.rope) self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( embedding_dim=time_embed_dim, @@ -1137,12 +1127,10 @@ class GlmImageTransformer2DModel(CachedTransformer): # 1. Prepare hidden_states and RoPE via GlmImagePrepare module # _sp_plan will shard hidden_states and RoPE together via split_output=True - hidden_states, rope_cos, rope_sin, post_patch_height_t, post_patch_width_t = self.prepare( - hidden_states, prior_hidden_states - ) + hidden_states, rope_cos, rope_sin = self.prepare(hidden_states, prior_hidden_states) image_rotary_emb = (rope_cos, rope_sin) - post_patch_height = int(post_patch_height_t.item()) - post_patch_width = int(post_patch_width_t.item()) + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size # Timestep conditioning temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) @@ -1163,8 +1151,6 @@ class GlmImageTransformer2DModel(CachedTransformer): device=hidden_states.device, ) hidden_states_mask[:, ctx.sp_original_seq_len :] = False - if hidden_states_mask.all(): - hidden_states_mask = None # 2. Transformer blocks for layer_idx, block in enumerate(self.transformer_blocks): From f62d4ad4a794bd0706340d8dcc9d92e67672af1c Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:15:32 +0800 Subject: [PATCH 022/209] [Model][ERNIE-Image] Delay AdaLN modulation broadcast (#7171) Signed-off-by: hyw --- .../models/ernie_image/ernie_image_transformer.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py index 244aeeb95..bd96630e8 100644 --- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py +++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ernie_image.py from collections.abc import Iterable @@ -599,17 +599,12 @@ class ErnieImageTransformer2DModel(nn.Module): hidden_states, freqs_cos, freqs_sin, attention_mask = self.unified_prepare(hidden_states, text_bth, text_lens) attention_mask = self._slice_attention_mask_for_ring(attention_mask) rotary_pos_emb = (freqs_cos, freqs_sin) - S = hidden_states.shape[0] - sample = self.time_proj(timestep) sample = sample.to(dtype=dtype) c = self.time_embedding(sample) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [ - t.unsqueeze(0).expand(S, -1, -1).contiguous() for t in self.adaLN_modulation(c).chunk(6, dim=-1) - ] + temb = self.adaLN_modulation(c).chunk(6, dim=-1) for layer in self.layers: - temb = [shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( layer, From 43b8de9b0ecdeb78ec3585222db02164a087103a Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:18:12 +0800 Subject: [PATCH 023/209] [Kernel][MiniMax-H3] Run Q/K RMSNorm-RoPE in one launch (#7167) Signed-off-by: hyw --- .../models/minimax_h3/ops/vae/qk_norm_rope.py | 80 +++++++++++-------- 1 file changed, 47 insertions(+), 33 deletions(-) diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py index 40b7069f6..f964c1577 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py @@ -11,18 +11,26 @@ if HAS_TRITON: @triton.jit def _qk_rms_norm_rope_exact_kernel( - x_ptr, + q_ptr, + k_ptr, cos_ptr, sin_ptr, - output_ptr, - x_stride_token, - x_stride_head, - x_stride_dim, + q_output_ptr, + k_output_ptr, + q_stride_token, + q_stride_head, + q_stride_dim, + k_stride_token, + k_stride_head, + k_stride_dim, rope_stride_token, rope_stride_dim, - output_stride_token, - output_stride_head, - output_stride_dim, + q_output_stride_token, + q_output_stride_head, + q_output_stride_dim, + k_output_stride_token, + k_output_stride_head, + k_output_stride_dim, num_heads: tl.constexpr, head_dim: tl.constexpr, rotary_dim: tl.constexpr, @@ -31,12 +39,15 @@ if HAS_TRITON: ): token = tl.program_id(0) head_group = tl.program_id(1) + qk_domain = tl.program_id(2) heads = head_group * heads_per_program + tl.arange(0, heads_per_program) dims = tl.arange(0, head_dim) mask = heads[:, None] < num_heads - offsets = token * x_stride_token + heads[:, None] * x_stride_head + dims[None, :] * x_stride_dim - x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + q_offsets = token * q_stride_token + heads[:, None] * q_stride_head + dims[None, :] * q_stride_dim + k_offsets = token * k_stride_token + heads[:, None] * k_stride_head + dims[None, :] * k_stride_dim + x_ptrs = tl.where(qk_domain == 0, q_ptr + q_offsets, k_ptr + k_offsets) + x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) normalized = (x * inv_rms[:, None]).to(tl.float16) @@ -46,8 +57,10 @@ if HAS_TRITON: dims + rotary_half, tl.where(dims < rotary_dim, dims - rotary_half, dims), ) - pair_offsets = token * x_stride_token + heads[:, None] * x_stride_head + pair_dims[None, :] * x_stride_dim - pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + q_pair_offsets = token * q_stride_token + heads[:, None] * q_stride_head + pair_dims[None, :] * q_stride_dim + k_pair_offsets = token * k_stride_token + heads[:, None] * k_stride_head + pair_dims[None, :] * k_stride_dim + pair_x_ptrs = tl.where(qk_domain == 0, q_ptr + q_pair_offsets, k_ptr + k_pair_offsets) + pair_x = tl.load(pair_x_ptrs, mask=mask, other=0.0).to(tl.float32) pair_normalized = (pair_x * inv_rms[:, None]).to(tl.float16) rope_mask = dims < rotary_dim @@ -79,10 +92,18 @@ if HAS_TRITON: ) output = tl.where(rope_mask[None, :], rotated, normalized) - output_offsets = ( - token * output_stride_token + heads[:, None] * output_stride_head + dims[None, :] * output_stride_dim + q_output_offsets = ( + token * q_output_stride_token + heads[:, None] * q_output_stride_head + dims[None, :] * q_output_stride_dim ) - tl.store(output_ptr + output_offsets, output, mask=mask) + k_output_offsets = ( + token * k_output_stride_token + heads[:, None] * k_output_stride_head + dims[None, :] * k_output_stride_dim + ) + output_ptrs = tl.where( + qk_domain == 0, + q_output_ptr + q_output_offsets, + k_output_ptr + k_output_offsets, + ) + tl.store(output_ptrs, output, mask=mask) def _supported_inputs( @@ -143,11 +164,11 @@ def try_qk_norm_rope_exact( sin = sin.reshape(tokens, sin.shape[-1]) q_output = torch.empty_like(q) k_output = torch.empty_like(k) - # This launch layout is part of the exactness contract on the validated - # SM90 and SM103 targets: changing the warp-to-row mapping changes the - # FP32 RMS reduction order. + # The first two launch dimensions are part of the exactness contract on + # the validated SM90 and SM103 targets: changing the warp-to-row mapping + # changes the FP32 RMS reduction order. heads_per_program = 8 - grid = (tokens, triton.cdiv(heads, heads_per_program)) + grid = (tokens, triton.cdiv(heads, heads_per_program), 2) launch_args = { "num_heads": heads, "head_dim": head_dim, @@ -158,29 +179,22 @@ def try_qk_norm_rope_exact( } _qk_rms_norm_rope_exact_kernel[grid]( q, - cos, - sin, - q_output, - q.stride(0), - q.stride(1), - q.stride(2), - cos.stride(0), - cos.stride(1), - q_output.stride(0), - q_output.stride(1), - q_output.stride(2), - **launch_args, - ) - _qk_rms_norm_rope_exact_kernel[grid]( k, cos, sin, + q_output, k_output, + q.stride(0), + q.stride(1), + q.stride(2), k.stride(0), k.stride(1), k.stride(2), cos.stride(0), cos.stride(1), + q_output.stride(0), + q_output.stride(1), + q_output.stride(2), k_output.stride(0), k_output.stride(1), k_output.stride(2), From dff8f8f4cfe9ba6e6eaeeac455f8aee66f95c4ae Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:41:45 -0700 Subject: [PATCH 024/209] [CI][ROCm] Align AMD image with vLLM 0.29 (#7395) Signed-off-by: andyluo7 Co-authored-by: Hongsheng Liu --- docker/Dockerfile.rocm | 24 +++--- .../installation/gpu/rocm.inc.md | 44 ++++++++++- tests/buildkite/test_rocm_dockerfile.py | 77 +++++++++++++++++++ 3 files changed, 132 insertions(+), 13 deletions(-) create mode 100644 tests/buildkite/test_rocm_dockerfile.py diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 597f5d134..adc767691 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,16 +1,18 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0 +# Keep the ROCm runtime aligned with the vLLM release targeted by +# docker/Dockerfile.ci. A stale base image can build successfully while failing +# at runtime when vLLM-Omni imports APIs added by the newer vLLM release. +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.29.0 FROM ${BASE_IMAGE} AS base -# Declare a variable to know if we want to use the nightly build or the stable build. # NOTE: REMINDER to vLLM-Omni rebase maintainer -# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image -# 1. If vLLM-omni maintainer is forced to use custom commits -# during rebasing, they can change the two variables -# 2. Whenever vLLM upstream has released a stable version, -# we should swap over to use stable release ASAP. -# We should avoid relying on custom commits. +# BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0. +# VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when +# USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's +# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An +# intentional custom source pin must update that regression in the same change, +# and should return to the stable release as soon as one is available. ARG USE_NIGHTLY_BUILD=0 -ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0 +ARG VLLM_VERSION_OR_COMMIT_HASH=v0.29.0 ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} @@ -38,6 +40,10 @@ RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \ rm -rf vllm/.git; \ fi +# Fail during image construction instead of letting every AMD test fail at +# runtime when the base image and vLLM-Omni target incompatible vLLM APIs. +RUN python3 -c "import vllm; from vllm.v1.kv_cache_interface import compute_layout_strides; from vllm.v1.kv_cache_layout import KVCacheLayout; print(f'Validated vLLM {vllm.__version__}')" + # Build TorchCodec after any optional nightly vLLM reinstall so it targets the # final torch/ROCm ABI. Keep this layer before COPY . so source changes do not # invalidate the expensive native build. FFmpeg development headers are diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index d15bba2eb..83efe0c62 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -1,11 +1,23 @@ + + # --8<-- [start:requirements] - GPU: Validated on gfx942 (It should be supported on the AMD GPUs that are supported by vLLM.) # --8<-- [end:requirements] + # --8<-- [start:set-up-using-python] -vLLM-Omni current recommends the steps in under setup through Docker Images. +For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images. + +vLLM-Omni depends on the matching major/minor release of vLLM. The 0.29 +development line uses vLLM 0.29.x. Published 0.28.0 wheels and images use vLLM +0.28.x. + +The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The +`vllm-omni` package does not install vLLM as a dependency, so non-Docker source +installs must install the matching ROCm vLLM release explicitly before +installing vLLM-Omni, as shown below. # --8<-- [start:pre-built-wheels] @@ -14,6 +26,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below. vLLM-Omni is built based on vLLM. Please install it with command below. + ```bash uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 ``` @@ -36,6 +49,7 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM + If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library ```bash @@ -45,7 +59,9 @@ uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/roc The pre-built 0.29.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni + Since vllm-omni is rapidly evolving, it's recommended to install it from source + ```bash git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -71,15 +87,31 @@ python3 setup.py develop #### Build docker image +The source build defaults to the published upstream base image +`vllm/vllm-openai-rocm:v0.29.0`, aligned with the vLLM release used by CI. +Older or custom bases must provide the vLLM APIs checked by the Dockerfile's +image-build canary. This upstream base is distinct from the prebuilt +`vllm/vllm-omni-rocm` images discussed below; their published-tag availability +is tracked in [#7405](https://github.com/vllm-project/vllm-omni/issues/7405). + ```bash DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm . ``` +To select the upstream ROCm base explicitly, pass `BASE_IMAGE`: + +```bash +DOCKER_BUILDKIT=1 docker build \ + -f docker/Dockerfile.rocm \ + --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.29.0 \ + -t vllm-omni-rocm . +``` + #### Launch the docker image ##### Launch with OpenAI API Server -``` +```bash docker run --rm \ --group-add=video \ --ipc=host \ @@ -97,7 +129,7 @@ vllm-omni-rocm \ ##### Launch with interactive session for development -``` +```bash docker run --rm -it \ --network=host \ --group-add=video \ @@ -119,7 +151,9 @@ vllm-omni-rocm vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni-rocm](https://hub.docker.com/r/vllm/vllm-omni-rocm/tags). The version of vLLM-Omni indicates which release of vLLM it is based on. #### Launch vLLM-Omni Server + Here's an example deployment command that has been verified on 2 x MI300's: + ```bash docker run --rm \ --group-add=video \ @@ -136,8 +170,10 @@ docker run --rm \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` -#### Launch an interactive terminal with prebuilt docker image. +#### Launch an interactive terminal with prebuilt docker image + If you want to run in dev environment you can launch the docker image as follows: + ```bash docker run --rm -it \ --network=host \ diff --git a/tests/buildkite/test_rocm_dockerfile.py b/tests/buildkite/test_rocm_dockerfile.py new file mode 100644 index 000000000..9171b3d53 --- /dev/null +++ b/tests/buildkite/test_rocm_dockerfile.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REPO_ROOT = Path(__file__).resolve().parents[2] +AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2" +CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci" +ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm" + + +def _docker_arg(path: Path, name: str) -> str: + prefix = f"ARG {name}=" + matches = [ + line.removeprefix(prefix) for line in path.read_text(encoding="utf-8").splitlines() if line.startswith(prefix) + ] + assert len(matches) == 1, f"expected one {name} declaration in {path}, found {len(matches)}" + return matches[0] + + +def _line_index(lines: list[str], prefix: str) -> int: + matches = [index for index, line in enumerate(lines) if line.startswith(prefix)] + assert len(matches) == 1, f"expected one line starting with {prefix!r}, found {len(matches)}" + return matches[0] + + +def test_rocm_base_tracks_ci_vllm_release() -> None: + ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG") + rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE") + + image_ref, separator, image_tag = rocm_base.rpartition(":") + assert separator, f"expected a tagged ROCm base image, got {rocm_base}" + assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref + assert image_tag == ci_release + + +def test_rocm_defaults_to_prebuilt_base_image() -> None: + assert _docker_arg(ROCM_DOCKERFILE, "USE_NIGHTLY_BUILD") == "0" + + +def test_amd_build_uses_rocm_dockerfile_defaults() -> None: + template = AMD_TEMPLATE.read_text(encoding="utf-8") + build_commands = [line.strip() for line in template.splitlines() if '"docker build ' in line] + + assert len(build_commands) == 1, f"expected one AMD image build command, found {len(build_commands)}" + build_command = build_commands[0] + assert "-f docker/Dockerfile.rocm" in build_command + for arg_name in ("BASE_IMAGE", "USE_NIGHTLY_BUILD"): + assert f"--build-arg {arg_name}" not in build_command + assert f"--build-arg={arg_name}" not in build_command + + +def test_rocm_source_ref_tracks_ci_vllm_release() -> None: + ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG") + rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH") + + assert rocm_source_ref == ci_release + + +def test_rocm_dockerfile_contains_vllm_api_canary() -> None: + dockerfile = ROCM_DOCKERFILE.read_text(encoding="utf-8") + lines = dockerfile.splitlines() + nightly_start = _line_index(lines, 'RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then') + canary = ( + 'RUN python3 -c "import vllm; ' + "from vllm.v1.kv_cache_interface import compute_layout_strides; " + "from vllm.v1.kv_cache_layout import KVCacheLayout;" + ) + canary_index = _line_index(lines, canary) + + assert any(line.strip() == "fi" for line in lines[nightly_start + 1 : canary_index]), ( + "vLLM API canary must follow the optional nightly source reinstall block" + ) From de9a133d48cbcb35c7165c772524be20e0fd1bdd Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:05:12 +0800 Subject: [PATCH 025/209] [Bugfix] Add embed_multimodal to MiniCPM-o 4.5 omni LLM class (#7384) Signed-off-by: Guangjian --- .../test_minicpmo_4_5_omni_llm_embed.py | 33 +++++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 4 +++ 2 files changed, 37 insertions(+) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py new file mode 100644 index 000000000..9f53fb7ec --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_embed_multimodal_delegates_to_get_multimodal_embeddings(monkeypatch): + """vLLM V1 encoder profiling calls embed_multimodal; without an override the + SupportsMultiModal protocol stub returns None and startup fails downstream.""" + sentinel = (object(),) + seen: dict[str, object] = {} + + def fake_get_multimodal_embeddings(self, **kwargs): + seen.update(kwargs) + return sentinel + + monkeypatch.setattr( + MiniCPMO45OmniLLMForConditionalGeneration, + "get_multimodal_embeddings", + fake_get_multimodal_embeddings, + ) + model = object.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + + out = model.embed_multimodal(pixel_values="pv", audio_features="af") + + assert out is sentinel + assert seen == {"pixel_values": "pv", "audio_features": "af"} diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index d70bc30d0..6d1c5de32 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -4580,6 +4580,10 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + """vLLM V1 encoder profiling calls this; the inherited Protocol stub returns None.""" + return self.get_multimodal_embeddings(**kwargs) + def get_input_embeddings( self, input_ids: torch.Tensor, From 02aaa34059478280eb3e6a5ee05a9484f0fa23c4 Mon Sep 17 00:00:00 2001 From: wtz2333 <2955110911@qq.com> Date: Fri, 11 Sep 2026 18:37:12 +0800 Subject: [PATCH 026/209] [Model] Add LingBot World Ulysses sequence parallelism (#6841) Signed-off-by: wtz2333 <2955110911@qq.com> Co-authored-by: Zhou Taichang --- docs/design/feature/realtime_ar_diffusion.md | 18 + .../diffusion/lingbot_world_v2.py | 6 +- .../test_lingbot_world_attention.py | 84 ++++- .../test_lingbot_world_transformer.py | 312 +++++++++++++++++- .../test_pipeline_lingbot_world.py | 51 ++- .../test_lingbot_world_v2.py | 24 ++ .../models/lingbot_world/pipeline.py | 42 ++- .../models/lingbot_world/transformer.py | 167 +++++++++- 8 files changed, 659 insertions(+), 45 deletions(-) create mode 100644 tests/examples/offline_inference/test_lingbot_world_v2.py diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md index 5963befe7..426ee5afd 100644 --- a/docs/design/feature/realtime_ar_diffusion.md +++ b/docs/design/feature/realtime_ar_diffusion.md @@ -178,6 +178,24 @@ For LingBot World v2: - `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and produces one block plus the standard metadata envelope. +## Ulysses sequence parallelism + +LingBot supports pure Ulysses sequence parallelism for both direct and +AR-Diffusion execution. With Ulysses degree greater than one, hidden tokens, +camera features, token-expanded timestep modulation, and RoPE tables are sharded +together. Without SP, timestep modulation retains the frame-broadcast path. +Self-attention performs the sequence-to-head all-to-all before reading or writing +paged KV, while static text K/V uses the same local head shard. Text K/V shards +own compact storage; cross-attention exchanges query/output layouts to use these +shards. This retains the shared cache geometry and reduces text K/V storage at +the cost of two all-to-all calls per layer. + +The output head projects local tokens before gathering the flow values. For +the 14B model this reduces the gathered width from 5120 to 64; frame modulation +uses each shard's global token offset, including shards that split a frame. Only +`ulysses_mode="strict"` is supported; `advanced_uaa`, Ring, and AllGather-KV modes +remain unsupported for this model. + ## Non-goals This contract does not currently provide: diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py index 590442c4d..4bbea2f4f 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Generate a LingBot-World v2 video from an image and camera trajectory. The official checkpoint is licensed separately under CC BY-NC-SA and is @@ -99,6 +99,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism inside the DiT.", ) + parser.add_argument("--ulysses-degree", type=int, default=1, help="Pure Ulysses sequence parallel degree.") parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.") parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.") parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.") @@ -158,6 +159,8 @@ def build_omni_kwargs( if args.tensor_parallel_size <= 0: raise ValueError("--tensor-parallel-size must be a positive integer.") + if args.ulysses_degree <= 0: + raise ValueError("--ulysses-degree must be a positive integer.") flow_shift = _positive_finite(args.flow_shift, "--flow-shift") model_path = Path(args.model).expanduser() model = str(model_path.resolve()) if model_path.exists() else args.model @@ -165,6 +168,7 @@ def build_omni_kwargs( "model": model, "flow_shift": flow_shift, "tensor_parallel_size": args.tensor_parallel_size, + "ulysses_degree": args.ulysses_degree, "enforce_eager": args.enforce_eager, "model_config": {"lingbot_action_root": str(paths.action_root)}, } diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py index f5567a706..3d883414e 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -7,7 +7,7 @@ import importlib.util import math import sys from pathlib import Path -from types import ModuleType +from types import ModuleType, SimpleNamespace import pytest import torch @@ -33,6 +33,10 @@ _STUBBED_MODULE_NAMES = ( "vllm_omni.diffusion", "vllm_omni.diffusion.attention", "vllm_omni.diffusion.attention.layer", + "vllm_omni.diffusion.distributed", + "vllm_omni.diffusion.distributed.comm", + "vllm_omni.diffusion.distributed.parallel_state", + "vllm_omni.diffusion.distributed.sp_plan", "vllm_omni.diffusion.layers", "vllm_omni.diffusion.layers.norm", "vllm_omni.diffusion.layers.rope", @@ -58,6 +62,47 @@ def _install_vllm_stubs() -> None: distributed.get_tensor_model_parallel_world_size = lambda: 1 distributed.tensor_model_parallel_all_reduce = lambda value: value + class _SeqAllToAll4D: + @staticmethod + def apply(group, value, scatter_idx, gather_idx, use_sync=False): + del group, scatter_idx, gather_idx, use_sync + return value + + setattr( + sys.modules["vllm_omni.diffusion.distributed.comm"], + "SeqAllToAll4D", + _SeqAllToAll4D, + ) + + def get_sp_group(): + return SimpleNamespace( + ulysses_world_size=1, + ulysses_rank=0, + ulysses_group=None, + ) + + setattr( + sys.modules["vllm_omni.diffusion.distributed.parallel_state"], + "get_sp_group", + get_sp_group, + ) + + class _SequenceParallelInput: + def __init__(self, split_dim, expected_dims=None, split_output=False, auto_pad=False): + self.split_dim = split_dim + self.expected_dims = expected_dims + self.split_output = split_output + self.auto_pad = auto_pad + + class _SequenceParallelOutput: + def __init__(self, gather_dim, expected_dims=None): + self.gather_dim = gather_dim + self.expected_dims = expected_dims + + sp_plan = sys.modules["vllm_omni.diffusion.distributed.sp_plan"] + setattr(sp_plan, "SequenceParallelInput", _SequenceParallelInput) + setattr(sp_plan, "SequenceParallelOutput", _SequenceParallelOutput) + def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None: for name, value in attrs.items(): setattr(weight, name, value) @@ -541,3 +586,38 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0])) torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0])) + + +def test_single_token_text_kv_shard_owns_compact_storage(monkeypatch): + module = _load_module() + monkeypatch.setattr( + module, + "get_sp_group", + lambda: SimpleNamespace( + ulysses_world_size=2, + ulysses_rank=1, + ulysses_group=None, + ), + ) + attention = module.LingBotCrossAttention(dim=8, num_heads=4) + full = torch.arange(8, dtype=torch.float32).reshape(1, 1, 4, 2) + shard = attention.shard_kv_heads(full) + torch.testing.assert_close(shard, full[:, :, 2:], rtol=0, atol=0) + assert shard.is_contiguous() + assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size() + + +@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention", "LingBotCrossAttention"]) +def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class): + module = _load_module() + monkeypatch.setattr( + module, + "get_sp_group", + lambda: SimpleNamespace( + ulysses_world_size=3, + ulysses_rank=0, + ulysses_group=None, + ), + ) + with pytest.raises(ValueError, match="heads must be divisible"): + getattr(module, attention_class)(dim=8, num_heads=4) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py index e2e2a08b6..b4bf7805f 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -7,14 +7,17 @@ import hashlib import inspect import json import math +from datetime import timedelta from pathlib import Path +from types import SimpleNamespace import pytest import torch from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests +from tests.helpers.mark import hardware_test -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] _FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture" _SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture" @@ -117,6 +120,7 @@ def _assert_self_cache_unchanged(cache, snapshot) -> None: attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot) +@pytest.mark.cpu def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None: torch.manual_seed(7) module = attention_tests._load_module() @@ -191,6 +195,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera assert not torch.equal(camera_output, alternate_output) +@pytest.mark.cpu def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=2).eval() @@ -229,6 +234,7 @@ def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() - @pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks")) +@pytest.mark.cpu def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None: module = attention_tests._load_module() model = _tiny_model( @@ -257,6 +263,7 @@ def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: assert cache.cross_attention == [None] +@pytest.mark.cpu def test_forward_accepts_exactly_one_configured_frame_block() -> None: module = attention_tests._load_module() model = _tiny_model( @@ -279,6 +286,7 @@ def test_forward_accepts_exactly_one_configured_frame_block() -> None: assert output.shape == (1, 2, 3, 4, 4) +@pytest.mark.cpu def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None: module = attention_tests._load_module() model = _tiny_model( @@ -303,6 +311,7 @@ def test_transformer_allocates_request_cache_from_its_configured_geometry() -> N assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention) +@pytest.mark.cpu def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1, num_frames_per_block=2) @@ -323,6 +332,7 @@ def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None: torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0])) +@pytest.mark.cpu def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1, num_frames_per_block=2) @@ -363,6 +373,7 @@ def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None: torch.testing.assert_close(output, expected) +@pytest.mark.cpu def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1, num_frames_per_block=2) @@ -381,6 +392,7 @@ def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None: torch.testing.assert_close(output, expected) +@pytest.mark.cpu def test_constructor_defaults_match_official_checkpoint_config() -> None: module = attention_tests._load_module() parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters @@ -406,6 +418,7 @@ def test_constructor_defaults_match_official_checkpoint_config() -> None: assert {name: parameters[name].default for name in expected} == expected +@pytest.mark.cpu def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None: module = attention_tests._load_module() model = _tiny_model(module) @@ -413,12 +426,14 @@ def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None: assert model.dtype == next(model.parameters()).dtype +@pytest.mark.cpu def test_transformer_declares_regional_compile_block() -> None: module = attention_tests._load_module() assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"] +@pytest.mark.cpu def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None: module = attention_tests._load_module() reduced_values: list[torch.Tensor] = [] @@ -440,6 +455,7 @@ def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None: torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]])) +@pytest.mark.cpu def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None: module = attention_tests._load_module() monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3) @@ -448,6 +464,7 @@ def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None: module.LingBotSelfAttention(dim=4, num_heads=2) +@pytest.mark.cpu def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None: module = attention_tests._load_module() @@ -456,6 +473,7 @@ def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None: @pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"]) +@pytest.mark.cpu def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None: module = attention_tests._load_module() @@ -463,6 +481,7 @@ def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None module.CausalLingBotWorldTransformer3DModel(**{field: 4}) +@pytest.mark.cpu def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None: module = attention_tests._load_module() @@ -470,6 +489,7 @@ def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> No module.CausalLingBotWorldTransformer3DModel(quant_config=object()) +@pytest.mark.cpu def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None: module = attention_tests._load_module() @@ -522,6 +542,7 @@ def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> N ("sink_size", 0), ], ) +@pytest.mark.cpu def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None: module = attention_tests._load_module() config = { @@ -554,6 +575,7 @@ def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object module.CausalLingBotWorldTransformer3DModel.from_config(config) +@pytest.mark.cpu def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None: module = attention_tests._load_module() config = { @@ -569,6 +591,7 @@ def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None: assert model.config.num_layers == 40 +@pytest.mark.cpu def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1) @@ -606,6 +629,7 @@ def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> model.load_weights([("unexpected_model.weight", torch.ones(1))]) +@pytest.mark.cpu def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1) @@ -619,6 +643,7 @@ def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None: assert model.load_weights(weights()) == {first_name} +@pytest.mark.cpu def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None: utils = pytest.importorskip("vllm.model_executor.models.utils") module = attention_tests._load_module() @@ -638,6 +663,7 @@ def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> No assert set(dict(pipeline.named_parameters())) <= loaded +@pytest.mark.cpu def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None: module = attention_tests._load_module() fixture = json.loads(_FIXTURE_PATH.read_text()) @@ -667,6 +693,7 @@ def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None: assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values()) +@pytest.mark.cpu def test_official_default_shapes_match_public_safetensors_header_fixture() -> None: module = attention_tests._load_module() fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text()) @@ -683,3 +710,284 @@ def test_official_default_shapes_match_public_safetensors_header_fixture() -> No shape, dtype = parameter_specs[name] assert tuple(shape) == tuple(expected_shape), name assert dtype == expected_dtype + + +@pytest.mark.cpu +def test_timestep_expansion_only_runs_with_ulysses(monkeypatch): + dtype = torch.bfloat16 + module = attention_tests._load_module() + hidden = torch.randn(2, 12, 4).to(dtype) + camera = torch.randn_like(hidden) + table = torch.randn(2, 3, 6, 4).to(dtype) + rotary = (torch.randn(12, 2), torch.randn(12, 2)) + prepare = module._LingBotSPPrepare() + assert prepare(hidden, camera, table, rotary)[2] is table + monkeypatch.setattr(module, "get_sp_group", lambda: SimpleNamespace(ulysses_world_size=2)) + expanded = prepare(hidden, camera, table, rotary)[2] + assert expanded.shape == (2, 12, 6, 4) + torch.testing.assert_close(expanded, table.repeat_interleave(4, dim=1), rtol=0, atol=0) + + +@pytest.mark.parametrize("unsharded", [None, 3], ids=["missing-hook", "unsharded-rope"]) +@pytest.mark.cpu +def test_missing_or_partial_sp_split_fails_before_attention(monkeypatch, unsharded): + module = attention_tests._load_module() + model = _tiny_model(module, num_frames_per_block=3, sliding_window_num_frames=6) + cache = _cache(module, model) + for block in model.blocks: + block.self_attn.ulysses_world_size = 2 + monkeypatch.setattr(module, "get_sp_group", lambda: SimpleNamespace(ulysses_world_size=2)) + original = model.sp_prepare.forward + + def partial_split(*args): + values = original(*args) + if unsharded is None: + return values + return tuple( + value if i == unsharded else value.chunk(2, dim=1 if i < 3 else 0)[0] for i, value in enumerate(values) + ) + + monkeypatch.setattr(model.sp_prepare, "forward", partial_split) + + def forbidden(*args, **kwargs): + raise AssertionError("attention/collectives must not run before shard validation") + + monkeypatch.setattr(model.blocks[0], "forward", forbidden) + with pytest.raises(RuntimeError, match="SP input hooks"): + model( + torch.randn(1, 36, 3, 4, 4), + torch.tensor([1.0]), + torch.randn(1, 3, 6), + torch.randn(1, 384, 3, 4, 4), + cache=cache, + start_frame=0, + update_cache=False, + ) + + +# Real multi-rank regression: synthetic weights, native SP/TP and attention. +_HEADS, _HEAD_DIM, _LAYERS = 8, 32, 2 +_FRAMES, _SIDE, _TOKENS_PER_FRAME = 3, 32, 256 + + +def _model(dtype): + from vllm_omni.diffusion.models.lingbot_world.transformer import CausalLingBotWorldTransformer3DModel + + model = CausalLingBotWorldTransformer3DModel( + num_attention_heads=_HEADS, + attention_head_dim=_HEAD_DIM, + num_layers=_LAYERS, + in_channels=4, + out_channels=4, + text_dim=8, + freq_dim=16, + ffn_dim=512, + patch_size=(1, 2, 2), + sink_size=1, + num_frames_per_block=_FRAMES, + sliding_window_num_frames=6, + rope_max_seq_len=32, + ).eval() + return model.to(device=torch.device("cuda", torch.accelerator.current_device_index()), dtype=dtype) + + +def _checkpoint(model): + weights = {} + generator = torch.Generator().manual_seed(17) + for name, param in sorted(model.named_parameters()): + data = torch.randn(param.shape, generator=generator) * 0.02 + if "norm" in name and name.endswith("weight"): + data.fill_(1) + if name.endswith("modulation") and data.shape[-2] == 6: + data[..., 2, :] = 1 + data[..., 5, :] = 1 + data = data.to(param.dtype) + param.copy_(data) + if ".self_attn.qkv." in name: + for component, chunk in zip(("q", "k", "v"), data.chunk(3), strict=True): + weights[name.replace(".self_attn.qkv.", f".self_attn.{component}.")] = chunk + else: + weights[name] = data + return weights + + +def _rollout(model, mode, dtype, batch): + from vllm_omni.diffusion.models.lingbot_world.transformer import LingBotTransformerCache + from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec + from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig + from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + + device = next(model.parameters()).device + state = None + if mode == "paged": + kv = ARDiffusionKVCache( + ARDiffusionKVConfig(enable=True, chunk_size=_TOKENS_PER_FRAME, window_chunks=5, sink_chunks=1), + num_layers=_LAYERS, + num_kv_heads=model.blocks[0].self_attn.num_sp_heads, + head_size=_HEAD_DIM, + dtype=dtype, + block_size=_TOKENS_PER_FRAME, + max_model_len=4096, + available_bytes=1 << 27, + kv_branches=(ARDiffusionKVBranchSpec("main", 0),), + session_capacity=1, + frames_per_block=_FRAMES, + max_scratch_tokens_per_branch=_FRAMES * _TOKENS_PER_FRAME, + cross_attention_lengths={"text": 5}, + device=device, + ) + state = ARDiffusionKVState(kv, "numeric", {"main": kv.begin_request("numeric")}, num_layers=_LAYERS) + cache = LingBotTransformerCache(self_attention=[], cross_attention=[None] * _LAYERS) + else: + cache = model.allocate_cache( + batch_size=batch, latent_height=_SIDE, latent_width=_SIDE, device=device, dtype=dtype + ) + generator = torch.Generator().manual_seed(123) + text = torch.randn(batch, 5, 8, generator=generator).to(device, dtype) + outputs = [] + try: + for start in range(0, 12, _FRAMES): + latent = torch.randn(batch, 4, _FRAMES, _SIDE, _SIDE, generator=generator).to(device, dtype) + camera = torch.randn(batch, 384, _FRAMES, _SIDE, _SIDE, generator=generator).to(device, dtype) + for commit in (False, False, True): + if state is not None: + cache.self_attention = state.get_kv_caches( + "main", seq_len=_FRAMES * _TOKENS_PER_FRAME, commit_current=commit + ) + output = model( + latent, + torch.tensor([100.0, 400.0, 900.0], device=device).expand(batch, -1), + text, + camera, + cache=cache, + start_frame=start, + update_cache=commit, + ) + outputs.append(output.float().cpu()) + if state is not None: + if not state.is_cross_attention_populated("main", "text"): + state.populate_cross_attention( + "main", "text", [(layer.key, layer.value) for layer in cache.cross_attention] + ) + for layer, pool in zip( + cache.cross_attention, state.get_cross_attention_kv("main", "text"), strict=True + ): + layer.key, layer.value = pool["k"], pool["v"] + state.commit_paged_context("main") + cross = [ + (layer.key.detach().float().cpu(), layer.value.detach().float().cpu()) for layer in cache.cross_attention + ] + return torch.stack(outputs), cross + finally: + if state is not None: + state.close() + + +def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous): + from vllm.config import VllmConfig + from vllm.config.vllm import set_current_vllm_config + from vllm.distributed import get_tensor_model_parallel_rank + + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + destroy_model_parallel, + get_sp_group, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm_omni.diffusion.forward_context import set_forward_context + from vllm_omni.diffusion.registry import _apply_sequence_parallel_if_enabled + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.set_device(torch.device("cuda", rank)) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.distributed.init_process_group( + "nccl", init_method=f"file://{rendezvous}", world_size=world_size, rank=rank, timeout=timedelta(seconds=90) + ) + init_distributed_environment(world_size=world_size, rank=rank, local_rank=rank, backend="nccl") + try: + with torch.inference_mode(), set_current_vllm_config(VllmConfig()): + for baseline in (True, False): + sp, tp = (1, 1) if baseline else (sp_size, tp_size) + initialize_model_parallel( + data_parallel_size=world_size // (sp * tp), + sequence_parallel_size=sp, + ulysses_degree=sp, + tensor_parallel_size=tp, + ) + config = OmniDiffusionConfig( + model=str(Path(rendezvous).parent), + dtype=dtype, + enforce_eager=True, + parallel_config=DiffusionParallelConfig( + data_parallel_size=world_size // (sp * tp), + sequence_parallel_size=sp, + ulysses_degree=sp, + tensor_parallel_size=tp, + ), + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")), + ) + with set_current_diffusion_config(config), set_forward_context(omni_diffusion_config=config): + model = _model(dtype) + if baseline: + weights = _checkpoint(model) + else: + model.load_weights(iter(weights.items())) + _apply_sequence_parallel_if_enabled(SimpleNamespace(transformer=model), config) + result, cross = _rollout(model, mode, dtype, batch) + if baseline: + expected, expected_cross = result, cross + else: + assert expected.abs().max() > 1e-3 + bound = 1e-5 if dtype == torch.float32 else 1e-2 + error = (result - expected).double().norm() / expected.double().norm() + assert error <= bound, f"rank={rank}, TP{tp} SP{sp}: relative L2 {error.item():.3g} > {bound}" + local_heads = _HEADS // tp // sp + first = ( + get_tensor_model_parallel_rank() * (_HEADS // tp) + + get_sp_group().ulysses_rank * local_heads + ) + for (k, v), (ek, ev) in zip(cross, expected_cross, strict=True): + tolerance = 1e-5 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close( + k, ek[:, :, first : first + local_heads], rtol=tolerance, atol=tolerance + ) + torch.testing.assert_close( + v, ev[:, :, first : first + local_heads], rtol=tolerance, atol=tolerance + ) + if rank == 0: + print(f"{mode} {dtype} TP{tp} SP{sp}: relative L2={error.item():.3g}", flush=True) + del model + torch.distributed.barrier() + destroy_model_parallel() + finally: + destroy_distributed_env() + + +def _run(tmp_path: Path, sp, tp, mode, dtype, batch): + if torch.accelerator.device_count() < sp * tp: + pytest.skip(f"requires {sp * tp} GPUs") + torch.multiprocessing.spawn( + _worker, args=(sp * tp, sp, tp, mode, dtype, batch, str(tmp_path / "rendezvous")), nprocs=sp * tp + ) + + +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=2) +def test_sp2_direct_matches_sp1_fp32(tmp_path): + _run(tmp_path, 2, 1, "direct", torch.float32, 2) + + +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +def test_sp4_direct_matches_sp1_bf16(tmp_path): + _run(tmp_path, 4, 1, "direct", torch.bfloat16, 1) + + +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +def test_tp2_sp2_paged_matches_sp1_bf16(tmp_path): + _run(tmp_path, 2, 2, "paged", torch.bfloat16, 1) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index 52113208a..a58eee935 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -126,8 +126,6 @@ class _RecordingTransformer(nn.Module): sink_size=3, ) self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()]) - for block in self.blocks: - block.self_attn = SimpleNamespace(num_local_heads=2, head_dim=4) self.calls: list[dict] = [] self.cache_allocations: list[dict] = [] self.raise_on_call = raise_on_call @@ -387,6 +385,9 @@ def _od_config(**overrides): "parallel_config": SimpleNamespace( pipeline_parallel_size=1, sequence_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + allgather_degree=1, cfg_parallel_size=1, vae_patch_parallel_size=1, use_hsdp=False, @@ -438,15 +439,15 @@ def test_pipeline_respects_loader_managed_component_placement(offload_field: str assert getattr(pipeline.vae, "to_calls", []) == [] -def test_ar_diffusion_capability_uses_fixed_tp_local_lingbot_geometry() -> None: +def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None: module = _load_pipeline_module() - module.get_tensor_model_parallel_world_size = lambda: 1 pipeline = _pipeline(module) - + # Use the constructed head count even when the config describes a different geometry. + pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1) spec = pipeline.ar_diffusion_kv_cache_spec() assert spec.num_layers == 2 - assert spec.num_kv_heads == 2 + assert spec.num_kv_heads == 1 assert spec.head_size == 4 assert spec.tokens_per_frame == 1 assert spec.frames_per_block == 3 @@ -604,7 +605,6 @@ def test_component_discovery_uses_official_checkpoint_contract() -> None: ("field", "value", "feature"), [ ("pipeline_parallel_size", 2, "pipeline parallelism"), - ("sequence_parallel_size", 2, "sequence parallelism"), ("cfg_parallel_size", 2, "CFG parallelism"), ("vae_patch_parallel_size", 2, "VAE parallelism"), ("use_hsdp", True, "HSDP"), @@ -622,6 +622,43 @@ def test_unsupported_parallel_modes_fail_before_component_loading(field: str, va assert module._loader_state.prefetch_calls == [] +def test_pure_ulysses_parallel_config_is_supported() -> None: + module = _load_pipeline_module() + parallel_config = _od_config().parallel_config + parallel_config.sequence_parallel_size = 2 + parallel_config.ulysses_degree = 2 + + pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=parallel_config)) + + assert pipeline.transformer is not None + + +@pytest.mark.parametrize( + "overrides", + [ + {"sequence_parallel_size": 4, "ring_degree": 2}, # Normalized hybrid. + {"ulysses_degree": 1, "allgather_degree": 2}, # Normalized AllGather-KV. + {"ring_degree": 2}, # Isolate each clause from the SP-size mismatch. + {"allgather_degree": 2}, + {"ulysses_mode": "advanced_uaa"}, + {"ulysses_a2a_permute": True}, + {"ulysses_degree": None}, # Missing degree must not imply pure Ulysses. + ], +) +def test_unsupported_sp_config_fails_before_component_loading(overrides): + module = _load_pipeline_module() + config = _od_config().parallel_config + config.sequence_parallel_size = config.ulysses_degree = 2 + for name, value in overrides.items(): + if value is None: + delattr(config, name) + else: + setattr(config, name, value) + with pytest.raises(NotImplementedError, match="pure Ulysses"): + module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=config)) + assert module._loader_state.prefetch_calls == [] + + def test_unsupported_quantization_fails_before_component_loading() -> None: module = _load_pipeline_module() diff --git a/tests/examples/offline_inference/test_lingbot_world_v2.py b/tests/examples/offline_inference/test_lingbot_world_v2.py new file mode 100644 index 000000000..41d2e73b7 --- /dev/null +++ b/tests/examples/offline_inference/test_lingbot_world_v2.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest + +from examples.offline_inference.diffusion import lingbot_world_v2 as example + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_offline_ulysses_argument(tmp_path): + argv = ["--image", "frame.jpg", "--action-dir", "forward", "--prompt", "A lake"] + paths = example.LingBotPaths( + tmp_path / "frame.jpg", tmp_path / "forward", tmp_path, Path("forward"), 9, tmp_path / "out.mp4" + ) + assert example.parse_args(argv).ulysses_degree == 1 + args = example.parse_args([*argv, "--ulysses-degree", "4"]) + kwargs = example.build_omni_kwargs(args, paths) + assert kwargs["ulysses_degree"] == 4 and kwargs["tensor_parallel_size"] == 1 + args.ulysses_degree = 0 + with pytest.raises(ValueError, match="--ulysses-degree"): + example.build_omni_kwargs(args, paths) diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index b41abd861..d2dfde3d4 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -18,7 +18,6 @@ import torch.nn.functional as F from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel -from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -203,7 +202,6 @@ def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: return unsupported_sizes = { "pipeline_parallel_size": "pipeline parallelism", - "sequence_parallel_size": "sequence parallelism", "cfg_parallel_size": "CFG parallelism", "vae_patch_parallel_size": "VAE parallelism", } @@ -211,6 +209,25 @@ def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: size = getattr(parallel_config, field, 1) or 1 if size > 1: raise NotImplementedError(f"LingBot World v1 does not support {feature} ({field}={size}).") + sequence_parallel_size = getattr(parallel_config, "sequence_parallel_size", 1) or 1 + ulysses_degree = getattr(parallel_config, "ulysses_degree", 1) or 1 + ring_degree = getattr(parallel_config, "ring_degree", 1) or 1 + allgather_degree = getattr(parallel_config, "allgather_degree", 1) or 1 + ulysses_mode = getattr(parallel_config, "ulysses_mode", "strict") + ulysses_a2a_permute = bool(getattr(parallel_config, "ulysses_a2a_permute", False)) + if ( + sequence_parallel_size != ulysses_degree + or ring_degree != 1 + or allgather_degree != 1 + or ulysses_mode != "strict" + or ulysses_a2a_permute + ): + raise NotImplementedError( + "LingBot World sequence parallelism requires pure Ulysses with ulysses_a2a_permute disabled: " + f"sequence_parallel_size={sequence_parallel_size}, ulysses_degree={ulysses_degree}, " + f"ring_degree={ring_degree}, allgather_degree={allgather_degree}, ulysses_mode={ulysses_mode!r}, " + f"ulysses_a2a_permute={ulysses_a2a_permute}." + ) if getattr(parallel_config, "use_hsdp", False): raise NotImplementedError("LingBot World v1 does not support HSDP.") if getattr(parallel_config, "enable_expert_parallel", False): @@ -557,8 +574,6 @@ class LingBotWorldCausalDMDPipeline( latent_height = self._ar_height // spatial latent_width = self._ar_width // spatial tokens_per_frame = (latent_height // patch_height) * (latent_width // patch_width) - tp_size = get_tensor_model_parallel_world_size() - num_local_heads = int(self.transformer.config.num_attention_heads) // tp_size total_window_frames = ( int(self.transformer.config.local_attn_size) if int(self.transformer.config.local_attn_size) != -1 @@ -579,7 +594,7 @@ class LingBotWorldCausalDMDPipeline( ) return ARDiffusionKVCacheSpec( num_layers=int(self.transformer.config.num_layers), - num_kv_heads=num_local_heads, + num_kv_heads=int(self.transformer.blocks[0].self_attn.num_sp_heads), head_size=int(self.transformer.config.attention_head_dim), tokens_per_frame=tokens_per_frame, frames_per_block=int(self.transformer.config.num_frames_per_block), @@ -960,20 +975,11 @@ class LingBotWorldCausalDMDPipeline( def layer_kv() -> Iterator[tuple[torch.Tensor, torch.Tensor]]: for block in self.transformer.blocks: cross_attention = block.cross_attn - key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten( - 2, - ( - cross_attention.num_local_heads, - cross_attention.head_dim, - ), - ) - value = cross_attention.v(projected_text).unflatten( - 2, - ( - cross_attention.num_local_heads, - cross_attention.head_dim, - ), + shape = (cross_attention.num_local_heads, cross_attention.head_dim) + key = cross_attention.shard_kv_heads( + cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(2, shape) ) + value = cross_attention.shard_kv_heads(cross_attention.v(projected_text).unflatten(2, shape)) yield key, value state.populate_cross_attention( diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py index 35ec18801..32931db0c 100644 --- a/vllm_omni/diffusion/models/lingbot_world/transformer.py +++ b/vllm_omni/diffusion/models/lingbot_world/transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Checkpoint-compatible causal DiT for the LingBot World v2 model package.""" from __future__ import annotations @@ -24,6 +24,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.utils import set_weight_attrs from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D +from vllm_omni.diffusion.distributed.parallel_state import get_sp_group +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.layers.norm import LayerNorm from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( @@ -127,6 +130,37 @@ def _projection_prefix(prefix: str, name: str) -> str: return f"{prefix}.{name}" if prefix else name +def _ulysses_state() -> tuple[int, int, torch.distributed.ProcessGroup | None]: + coordinator = get_sp_group() + return ( + int(coordinator.ulysses_world_size), + int(coordinator.ulysses_rank), + coordinator.ulysses_group, + ) + + +class _LingBotSPPrepare(nn.Module): + """Expand frame conditioning only when Ulysses shards the token sequence.""" + + def forward( + self, + hidden_states: torch.Tensor, + camera_hidden_states: torch.Tensor, + timestep_projection: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + num_frames = timestep_projection.shape[1] + if hidden_states.shape[1] % num_frames: + raise ValueError("LingBot token count must be divisible by the latent frame count.") + tokens_per_frame = hidden_states.shape[1] // num_frames + if get_sp_group().ulysses_world_size > 1: + timestep_projection = ( + timestep_projection.unsqueeze(2).expand(-1, -1, tokens_per_frame, -1, -1).flatten(1, 2) + ) + cosine, sine = rotary_emb + return hidden_states, camera_hidden_states, timestep_projection, cosine, sine + + class LingBotSelfAttention(nn.Module): """Block-causal self-attention over retained history and one full chunk.""" @@ -157,6 +191,13 @@ class LingBotSelfAttention(nn.Module): prefix=_projection_prefix(prefix, "qkv"), ) self.num_local_heads = self.qkv.num_heads + self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state() + if self.num_local_heads % self.ulysses_world_size: + raise ValueError( + "LingBot local attention heads must be divisible by the Ulysses degree: " + f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}." + ) + self.num_sp_heads = self.num_local_heads // self.ulysses_world_size self.tp_inner_dim = self.num_local_heads * self.head_dim self.o = RowParallelLinear( dim, @@ -170,9 +211,9 @@ class LingBotSelfAttention(nn.Module): self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps) self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) self.attn = Attention( - num_heads=self.num_local_heads, + num_heads=self.num_sp_heads, head_size=self.head_dim, - num_kv_heads=self.num_local_heads, + num_kv_heads=self.num_sp_heads, softmax_scale=self.head_dim**-0.5, causal=False, role="self", @@ -287,6 +328,11 @@ class LingBotSelfAttention(nn.Module): query = self.rotary_embedding(query, cos, sin) key = self.rotary_embedding(key, cos, sin) + if self.ulysses_world_size > 1: + query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False) + key = SeqAllToAll4D.apply(self.ulysses_group, key, 2, 1, False) + value = SeqAllToAll4D.apply(self.ulysses_group, value, 2, 1, False) + if isinstance(cache, ARDiffusionPagedLayerInputs): if query.shape[0] != 1: raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.") @@ -347,6 +393,8 @@ class LingBotSelfAttention(nn.Module): ) else: output = self.attn(query, visible_key, visible_value) + if self.ulysses_world_size > 1: + output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False) return self.o(output.flatten(2, 3)) @@ -373,6 +421,13 @@ class LingBotCrossAttention(nn.Module): self.num_heads = num_heads self.head_dim = dim // num_heads self.num_local_heads = num_heads // tp_size + self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state() + if self.num_local_heads % self.ulysses_world_size: + raise ValueError( + "LingBot local cross-attention heads must be divisible by the Ulysses degree: " + f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}." + ) + self.num_sp_heads = self.num_local_heads // self.ulysses_world_size self.tp_inner_dim = self.num_local_heads * self.head_dim self.q = ColumnParallelLinear( @@ -410,9 +465,9 @@ class LingBotCrossAttention(nn.Module): self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps) self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps) self.attn = Attention( - num_heads=self.num_local_heads, + num_heads=self.num_sp_heads, head_size=self.head_dim, - num_kv_heads=self.num_local_heads, + num_kv_heads=self.num_sp_heads, softmax_scale=self.head_dim**-0.5, causal=False, role="cross", @@ -422,6 +477,14 @@ class LingBotCrossAttention(nn.Module): disable_kv_quant=True, ) + def shard_kv_heads(self, value: torch.Tensor) -> torch.Tensor: + """Select this Ulysses rank's projected K/V heads.""" + + if self.ulysses_world_size == 1: + return value + start = self.ulysses_rank * self.num_sp_heads + return value[:, :, start : start + self.num_sp_heads].clone(memory_format=torch.contiguous_format) + def forward( self, hidden_states: torch.Tensor, @@ -431,6 +494,8 @@ class LingBotCrossAttention(nn.Module): ) -> tuple[torch.Tensor, LingBotAttentionCache]: query = self.norm_q(self.q(hidden_states)) query = query.unflatten(2, (self.num_local_heads, self.head_dim)) + if self.ulysses_world_size > 1: + query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False) # Text K/V is constant within a request and is projected once per layer. if cache is None: @@ -438,8 +503,8 @@ class LingBotCrossAttention(nn.Module): raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.") key = self.norm_k(self.k(encoder_hidden_states)) value = self.v(encoder_hidden_states) - key = key.unflatten(2, (self.num_local_heads, self.head_dim)) - value = value.unflatten(2, (self.num_local_heads, self.head_dim)) + key = self.shard_kv_heads(key.unflatten(2, (self.num_local_heads, self.head_dim))) + value = self.shard_kv_heads(value.unflatten(2, (self.num_local_heads, self.head_dim))) cache = LingBotAttentionCache( key=key, value=value, @@ -452,6 +517,8 @@ class LingBotCrossAttention(nn.Module): value = cache.value[:, : cache.end] output = self.attn(query, key, value) + if self.ulysses_world_size > 1: + output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False) return self.o(output.flatten(2, 3)), cache @@ -540,7 +607,18 @@ class LingBotAttentionBlock(nn.Module): rotary_emb: tuple[torch.Tensor, torch.Tensor], ) -> tuple[torch.Tensor, LingBotAttentionCache]: batch_size, token_count, dim = hidden_states.shape + if ( + timestep_projection.ndim != 4 + or timestep_projection.shape[0] != batch_size + or timestep_projection.shape[2:] != (6, dim) + ): + raise ValueError("LingBot timestep projection must have shape (batch, frames or tokens, 6, dim).") num_frames = timestep_projection.shape[1] + if num_frames == 0 or token_count % num_frames: + raise ValueError("LingBot timestep projection must align with hidden tokens.") + if self.self_attn.ulysses_world_size > 1 and num_frames != token_count: + raise ValueError("LingBot Ulysses requires one timestep projection per local token.") + # Without SP, broadcast per frame. With SP, each entry represents one token. tokens_per_frame = token_count // num_frames # Timestep, camera, and text remain separate conditioning paths. modulation = self.modulation.unsqueeze(1) + timestep_projection.float() @@ -635,13 +713,26 @@ class _LingBotHead(nn.Module): self, hidden_states: torch.Tensor, timestep_embedding: torch.Tensor, + *, + tokens_per_frame: int | None = None, + token_offset: int = 0, ) -> torch.Tensor: num_frames = timestep_embedding.shape[1] - tokens_per_frame = hidden_states.shape[1] // num_frames modulation = self.modulation.unsqueeze(1) + timestep_embedding.unsqueeze(2).float() shift, scale = modulation.chunk(2, dim=2) - normalized = self.norm(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame)) - normalized = (normalized * (1 + scale) + shift).flatten(1, 2).to(hidden_states.dtype) + normalized = self.norm(hidden_states.float()) + if tokens_per_frame is None: + normalized = normalized.unflatten(1, (num_frames, -1)) + normalized = (normalized * (1 + scale) + shift).flatten(1, 2) + else: + # A sequence shard can begin/end inside a frame. + frames = ( + torch.arange(hidden_states.shape[1], device=hidden_states.device) + token_offset + ) // tokens_per_frame + scale = scale.squeeze(2).index_select(1, frames) + shift = shift.squeeze(2).index_select(1, frames) + normalized = normalized * (1 + scale) + shift + normalized = normalized.to(hidden_states.dtype) return self.head(normalized) @@ -678,6 +769,16 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): _repeated_blocks = ["LingBotAttentionBlock"] packed_modules_mapping = {"qkv": ["q", "k", "v"]} _layerwise_offload_blocks_attrs = ["blocks"] + _sp_plan = { + "sp_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + 4: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + }, + "sp_output_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } def __init__( self, @@ -768,6 +869,8 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): stride=patch_size, ) self.patch_embedding_wancamctrl = _LingBotCameraPatchEmbedding(6 * 8 * 8, dim, patch_size) + self.sp_prepare = _LingBotSPPrepare() + self.sp_output_gather = nn.Identity() self.c2ws_hidden_states_layer1 = ColumnParallelLinear( dim, dim, @@ -904,8 +1007,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames ) max_tokens = int(window_frames * post_patch_height * post_patch_width) - tp_size = get_tensor_model_parallel_world_size() - num_local_heads = self.config.num_attention_heads // tp_size + num_local_heads = self.blocks[0].self_attn.num_sp_heads return allocate_lingbot_cache( batch_size=batch_size, num_layers=self.config.num_layers, @@ -1138,6 +1240,33 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): query_len=hidden_states.shape[1], ) cache.self_attention = [layer_context.to_layer_inputs() for layer_context in cache.self_attention] + sp_size = self.blocks[0].self_attn.ulysses_world_size + global_tokens = patched_frames * tokens_per_frame + hidden_dim = hidden_states.shape[-1] + rope_dim = rotary_emb[0].shape[-1] + if global_tokens % sp_size: + raise ValueError("LingBot Ulysses requires the token count to be divisible by its degree.") + hidden_states, camera_hidden_states, timestep_projection, cosine, sine = self.sp_prepare( + hidden_states, + camera_hidden_states, + timestep_projection, + rotary_emb, + ) + if sp_size > 1: + local_tokens = global_tokens // sp_size + expected_hidden = (batch_size, local_tokens, hidden_dim) + if ( + hidden_states.shape != expected_hidden + or camera_hidden_states.shape != expected_hidden + or timestep_projection.shape != (batch_size, local_tokens, 6, hidden_dim) + or cosine.shape != (local_tokens, rope_dim) + or sine.shape != (local_tokens, rope_dim) + ): + raise RuntimeError( + "LingBot SP input hooks did not shard all conditioning tensors consistently; " + f"expected {local_tokens} tokens per rank. Check SP hook registration and input dimensions." + ) + rotary_emb = (cosine, sine) # Phase 3: each layer receives its own cache entry. Text K/V is passed # only when absent; the returned cache is stored for subsequent DMD # steps and causal blocks in this request. @@ -1156,9 +1285,17 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): ) cache.cross_attention[index] = cross_cache - # Phase 4: map tokens to per-patch 16-channel flow values and restore - # [B, C, F, H, W] for the Pipeline's sampler update. - hidden_states = self.head(hidden_states, timestep_embedding) + # Phase 4: project local tokens before gathering the much narrower flow values. + hidden_states = self.head( + hidden_states, + timestep_embedding, + tokens_per_frame=tokens_per_frame if sp_size > 1 else None, + token_offset=get_sp_group().ulysses_rank * (global_tokens // sp_size) if sp_size > 1 else 0, + ) + hidden_states = self.sp_output_gather(hidden_states) + projected_dim = self.config.out_channels * math.prod(self.config.patch_size) + if sp_size > 1 and hidden_states.shape != (batch_size, global_tokens, projected_dim): + raise RuntimeError("LingBot SP output hook did not gather the full projected token sequence.") return self._unpatchify( hidden_states, batch_size=batch_size, From 33fff2155e0bf9bc8688f7ffb92ef508efdf6954 Mon Sep 17 00:00:00 2001 From: GXIN <37653830+gxxx-hum@users.noreply.github.com> Date: Fri, 11 Sep 2026 21:22:15 +0800 Subject: [PATCH 027/209] [Feature][TTS] Add Speech API streaming metrics (#6853) Signed-off-by: XIN GAO <1037396230@qq.com> Co-authored-by: Hongsheng Liu --- docs/design/metrics.md | 45 +- docs/usage/metrics.md | 25 +- .../openai_api/test_serving_speech.py | 3 +- .../openai_api/test_serving_speech_metrics.py | 465 ++++++++++++++++++ .../openai_api/test_serving_speech_stream.py | 73 ++- tests/metrics/test_modality.py | 46 +- .../entrypoints/openai/serving_speech.py | 281 ++++++++--- .../openai/serving_speech_stream.py | 21 +- vllm_omni/metrics/definitions.py | 2 + vllm_omni/metrics/modality.py | 26 + 10 files changed, 895 insertions(+), 92 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_serving_speech_metrics.py diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 85eb7ac38..874cbd310 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -60,7 +60,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by There are five independent paths for metric collection. -**Path 1: Pipeline-level metrics (`vllm_omni:*`)** +#### Path 1: Pipeline-level metrics (`vllm_omni:*`) `OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress: @@ -69,7 +69,7 @@ There are five independent paths for metric collection. - `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter. - `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`. -**Path 2: Image and diffusion metrics** +#### Path 2: Image and diffusion metrics Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice. @@ -81,15 +81,15 @@ Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetri Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active. -**Path 3: Audio modality metrics (`vllm_omni:audio_*`)** +#### Path 3: Audio modality metrics (`vllm_omni:audio_*`) `OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites: - `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path). -- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry. -- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`. +- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` and the streaming TTS output paths in `serving_speech.py` on the first non-empty audio payload for a request. Streaming Speech observes the first PCM payload rather than a header-only WAV chunk. HTTP carries the middleware request timestamp through response generation; WebSocket uses the start of each sentence synthesis request. Non-streaming Speech does not emit TTFP and uses `e2e_request_latency_s` for overall latency instead. +- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` and `serving_speech.py` after a streaming chunk loop exhausts normally. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same output paths that update TTFP. Speech records raw PCM payloads before raw/SSE/WebSocket encoding, excludes WAV headers and empty chunks, and does not emit continuity samples for failed, cancelled, or non-streaming requests. -**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)** +#### Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`) `OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution. @@ -100,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected). -**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)** +#### Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap) The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics. @@ -174,7 +174,7 @@ second. ### Pipeline (4) | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | | `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | | `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | @@ -183,7 +183,7 @@ second. ### Image and diffusion service-level metrics | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing | | `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded | | `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage | @@ -200,7 +200,7 @@ second. Labels: `{model_name, stage, replica}`. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request | | `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count | | `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time | @@ -218,7 +218,7 @@ The optional breakdown and memory families are sparse by design: when an engine Labels: `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | -|--------|------|-------------|-------------| +| -------- | ------ | ------------- | ------------- | | `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | | `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | | `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | @@ -227,12 +227,29 @@ Labels: `{model_name, stage, replica}` plus the listed extra label. | `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | | `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | +### Speech streaming (2) + +Labels: `{model_name}` plus the listed extra label. + +| Metric | Type | Extra label | Description | +| -------- | ------ | ------------- | ------------- | +| `vllm_omni:speech_stream_aborted_total` | Counter | `reason` | Interrupted audio generators, including before first PCM; reasons: `cancelled`, `closed`, `engine_dead`, `error` | +| `vllm_omni:speech_stream_completed_total` | Counter | — | Normally completed audio generators; does not confirm client receipt | + +Scope: started Speech audio generators (raw/SSE/WebSocket), counted once per +generation; WebSocket counts each sentence. Chat, non-streaming Speech, and +failures before generator execution are excluded. Interrupted streams do not +emit continuity or underrun samples. + +Interruption ratio: `aborted / (aborted + completed)`, using increments over the +same time window and summing all reasons per model. + ### Cross-stage transfer (4) Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | | `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | | `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | @@ -245,8 +262,8 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, ## Naming Convention - All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used: - - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values. - - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. + - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values. + - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. - Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Sizes use the `_bytes` suffix. - All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`. diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index 184a76eef..1e0c1d068 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -12,7 +12,7 @@ curl http://localhost:8000/metrics ## Metric Namespaces | Prefix | Source | Present when | -|--------|--------|--------------| +| -------- | -------- | -------------- | | `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent | | `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage | | `http_` / `process_` | Uvicorn / Python runtime | Always | @@ -24,7 +24,7 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the ### Request counts | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) | | `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue | | `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | @@ -32,15 +32,15 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the ### Latency | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds | ## Audio Modality Metrics (`vllm_omni:`) -Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label. +Emitted at request finalize, except for `audio_ttfp_s` (the first audio packet/frame from streaming Chat or Speech APIs) and `audio_underrun_s` / `audio_continuity_ok_total` (Chat or Speech streaming finalize, after the chunk stream is exhausted). Non-streaming Speech latency is represented by `e2e_request_latency_s`, not TTFP. All audio metrics carry `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | -|--------|------|-------------|-------------| +| -------- | ------ | ------------- | ------------- | | `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | | `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | | `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` | @@ -50,13 +50,22 @@ Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the fi | `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition. +For Speech, continuity applies to raw audio, SSE, and WebSocket streaming. WAV headers and empty chunks are excluded; non-streaming Speech requests do not emit continuity samples. + +Continuity rate for the default 100 ms threshold can be derived from the successful-continuity counter and the underrun Histogram's request count: + +```promql +sum by (model_name, stage, replica) (rate(vllm_omni:audio_continuity_ok_total{threshold_ms="100"}[5m])) +/ +sum by (model_name, stage, replica) (rate(vllm_omni:audio_underrun_s_count[5m])) +``` ## Diffusion Metrics (`vllm_omni:`) Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds | | `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) | | `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds | @@ -80,7 +89,7 @@ rate(vllm_omni:diffusion_preprocess_s_sum[5m]) Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | | `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | | `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | @@ -106,7 +115,7 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl ## Metric Availability by Pipeline Type | Metric group | Multi-stage LLM (Qwen3-Omni) | -|---|---| +| --- | --- | | `vllm_omni:` request tracking + latency | With `--log-stats` | | `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage | | `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages | diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 6040bede3..732f86f2f 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -3287,7 +3287,7 @@ class TestStreamingResponse: finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]} captured: dict = {} - async def prepare(_request, request_id=None): + async def prepare(_request, request_id=None, arrival_time=None): return request_id, object(), finalized_tts_params async def generate_chunks( @@ -3296,6 +3296,7 @@ class TestStreamingResponse: _response_format="pcm", raw_request=None, request_start_s=None, + request_arrival_ts=None, include_sample_rate=False, usage_acc=None, tts_params=None, diff --git a/tests/entrypoints/openai_api/test_serving_speech_metrics.py b/tests/entrypoints/openai_api/test_serving_speech_metrics.py new file mode 100644 index 000000000..a5e284a01 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_speech_metrics.py @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Prometheus metric coverage for the Speech API audio path.""" + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.entrypoints.openai import serving_speech as speech_module +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _MetricsStub: + def __init__(self) -> None: + self.ttfp_calls: list[tuple[str, str, float]] = [] + self.underrun_calls: list[tuple[str, str, float]] = [] + self.continuity_calls: list[tuple[str, str, int]] = [] + self.abort_calls: list[str] = [] + self.completed = 0 + + def inc_speech_stream_completed(self) -> None: + self.completed += 1 + + def inc_speech_stream_aborted(self, reason: str) -> None: + self.abort_calls.append(reason) + + def observe_audio_ttfp(self, stage: str, replica: str, seconds: float) -> None: + self.ttfp_calls.append((stage, replica, seconds)) + + def observe_audio_underrun(self, stage: str, replica: str, seconds: float) -> None: + self.underrun_calls.append((stage, replica, seconds)) + + def inc_audio_continuity_ok(self, stage: str, replica: str, threshold_ms: int) -> None: + self.continuity_calls.append((stage, replica, threshold_ms)) + + +def _serving(metrics: _MetricsStub, *, adapter=None) -> OmniOpenAIServingSpeech: + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + serving._tts_model_type = "qwen3_tts" + serving.engine_client = SimpleNamespace(mod_metrics=metrics, request_states={}) + serving._get_tts_adapter = lambda: adapter + serving.create_audio = lambda audio_obj: SimpleNamespace( + audio_data=b"\0\0" * int(audio_obj.audio_tensor.size), + media_type="audio/pcm", + ) + serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None + serving._discard_ref_audio_artifact_warmup = lambda request_id: None + return serving + + +def _result(samples: int = 320, *, stage_id: int = 1, replica_id: int | None = 2) -> SimpleNamespace: + return SimpleNamespace( + multimodal_output={"audio": torch.zeros(samples), "sr": 16000}, + stage_id=stage_id, + replica_id=replica_id, + ) + + +async def _generate(*results): + for result in results: + yield result + + +@pytest.mark.asyncio +async def test_streaming_speech_observes_ttfp_once_on_first_pcm_payload(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + chunks = serving._generate_audio_chunks( + _generate(_result(), _result()), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 2 + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] + assert metrics.continuity_calls == [("1", "2", 100)] + assert metrics.abort_calls == [] + assert metrics.completed == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("before_pcm", [True, False]) +@pytest.mark.parametrize( + ("error_type", "reason"), + [(asyncio.CancelledError, "cancelled"), (speech_module.EngineDeadError, "engine_dead"), (RuntimeError, "error")], +) +async def test_speech_stream_abort_counted_once_without_continuity(before_pcm, error_type, reason): + metrics = _MetricsStub() + serving = _serving(metrics) + + async def results(): + if not before_pcm: + yield _result() + raise error_type() + + chunks = serving._generate_audio_chunks(results(), request_id="speech-test", request_arrival_ts=100.0) + with pytest.raises(error_type): + _ = [chunk async for chunk in chunks] + + assert metrics.abort_calls == [reason] + assert metrics.completed == 0 + assert len(metrics.ttfp_calls) == (0 if before_pcm else 1) + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_speech_stream_close_counted_once_without_continuity(): + metrics = _MetricsStub() + serving = _serving(metrics) + chunks = serving._generate_audio_chunks( + _generate(_result(), _result()), request_id="speech-test", request_arrival_ts=100.0 + ) + assert await anext(chunks) + await chunks.aclose() + await chunks.aclose() + assert metrics.abort_calls == ["closed"] + assert metrics.completed == 0 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_retries_ttfp_labels_without_moving_first_packet_time(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + clock = {"now": 100.25, "perf": 0.25} + monkeypatch.setattr(speech_module.time, "time", lambda: clock["now"]) + monkeypatch.setattr(speech_module.time, "perf_counter", lambda: clock["perf"]) + + async def results(): + yield _result(replica_id=None) + clock["now"] = 101.0 + clock["perf"] = 0.26 + yield _result(replica_id=2) + clock["now"] = 101.25 + clock["perf"] = 0.27 + yield _result(replica_id=2) + + chunks = serving._generate_audio_chunks( + results(), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 3 + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] + assert metrics.continuity_calls == [("1", "2", 100)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["raw", "sse", "pcm"]) +async def test_speech_disconnect_during_send_closes_entire_generator_chain(mode): + metrics = _MetricsStub() + serving = _serving(metrics) + closed = [] + + async def source(): + try: + yield _result() + yield _result() + finally: + closed.append(True) + + # Retain the source to prove cleanup does not depend on garbage collection. + generator = source() + method = { + "raw": serving._generate_audio_chunks, + "sse": serving._generate_audio_sse_events, + "pcm": serving._generate_pcm_chunks, + }[mode] + stream = method(generator, request_id="speech-test", request_arrival_ts=100.0) + + async def send(message): + if message["type"] == "http.response.body": + raise OSError("client disconnected") + + if mode == "pcm": + # WebSocket handlers already explicitly close their PCM iterator. + assert await anext(stream) + await stream.aclose() + else: + response = speech_module._SpeechStreamingResponse(stream) + with pytest.raises(OSError, match="client disconnected"): + await response.stream_response(send) + assert closed == [True] + assert metrics.abort_calls == ["closed"] + assert metrics.completed == 0 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.parametrize("arrival_ts", [0.0, -1.0]) +def test_speech_ttfp_invalid_arrival_does_not_consume_guard(arrival_ts): + metrics = _MetricsStub() + serving = _serving(metrics) + state = SimpleNamespace( + external_request_id="speech-test", + request_arrival_ts=arrival_ts, + first_audio_ts=None, + audio_emit_stage_id=None, + audio_emit_replica_id=None, + ) + serving.engine_client.request_states = {"internal-test": state} + + assert serving._observe_speech_audio_ttfp(request_id="speech-test", result=_result(), first_packet_ts=100.25) == ( + 1, + 2, + False, + ) + assert metrics.ttfp_calls == [] + assert state.first_audio_ts is None + assert state.audio_emit_stage_id is None + assert state.audio_emit_replica_id is None + + state.request_arrival_ts = 100.0 + assert serving._observe_speech_audio_ttfp(request_id="speech-test", result=_result(), first_packet_ts=100.25) == ( + 1, + 2, + True, + ) + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert state.first_audio_ts == 100.25 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("during_send", [False, True]) +async def test_speech_response_task_cancellation_closes_source(during_send): + metrics = _MetricsStub() + serving = _serving(metrics) + waiting = asyncio.Event() + blocker = asyncio.Event() + closed = [] + + async def source(): + try: + if not during_send: + waiting.set() + await blocker.wait() + yield _result() + finally: + closed.append(True) + + async def send(message): + if during_send and message["type"] == "http.response.body": + waiting.set() + await blocker.wait() + + generator = source() + stream = serving._generate_audio_chunks(generator, request_id="speech-test", request_arrival_ts=100.0) + response = speech_module._SpeechStreamingResponse(stream) + task = asyncio.create_task(response.stream_response(send)) + try: + await asyncio.wait_for(waiting.wait(), timeout=5) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert closed == [True] + assert metrics.abort_calls == ["closed" if during_send else "cancelled"] + assert metrics.completed == 0 + assert metrics.underrun_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_does_not_count_empty_payload_as_first_packet(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + chunks = serving._generate_audio_chunks( + _generate(_result(samples=0), _result()), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 2 + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] + assert metrics.continuity_calls == [("1", "2", 100)] + + +@pytest.mark.asyncio +async def test_streaming_speech_does_not_finalize_continuity_on_error(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + async def failing_stream(): + yield _result() + raise RuntimeError("stream failed") + + chunks = serving._generate_audio_chunks( + failing_stream(), + request_id="speech-test", + request_arrival_ts=100.0, + ) + with pytest.raises(RuntimeError, match="stream failed"): + _ = [chunk async for chunk in chunks] + + assert len(metrics.ttfp_calls) == 1 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_does_not_finalize_continuity_on_validation_error(monkeypatch): + metrics = _MetricsStub() + + def reject_generation(_tts_params, **_kwargs): + raise RuntimeError("generation validation failed") + + adapter = SimpleNamespace(validates_generation=True, validate_generation=reject_generation) + serving = _serving(metrics, adapter=adapter) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + chunks = serving._generate_audio_chunks( + _generate(_result()), + request_id="speech-test", + request_arrival_ts=100.0, + tts_params={"task_type": ["Base"]}, + ) + with pytest.raises(RuntimeError, match="generation validation failed"): + _ = [chunk async for chunk in chunks] + + assert len(metrics.ttfp_calls) == 1 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_reports_late_chunk_as_underrun(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + perf_times = iter((0.0, 0.1, 0.1, 2.0, 2.0)) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.perf_counter", lambda: next(perf_times)) + + chunks = serving._generate_audio_chunks( + _generate(_result(), _result()), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 2 + assert metrics.underrun_calls[0][:2] == ("1", "2") + assert metrics.underrun_calls[0][2] > 0.1 + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_non_streaming_speech_does_not_observe_ttfp(): + metrics = _MetricsStub() + serving = _serving(metrics) + serving._audio_encode_speed = lambda _request: 1.0 + + async def prepare(_request, **_kwargs): + result = _result() + result.metrics = {} + return "speech-test", _generate(result), {} + + serving._prepare_speech_generation = prepare + request = OpenAICreateSpeechRequest(input="hello", response_format="pcm") + + audio_data, media_type = await serving._generate_audio_bytes(request, request_arrival_ts=100.0) + + assert audio_data + assert media_type == "audio/pcm" + assert metrics.ttfp_calls == [] + assert metrics.completed == 0 + assert metrics.abort_calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flush_only", [False, True]) +@pytest.mark.parametrize("response_format", ["pcm", "wav"]) +async def test_resampled_speech_records_flush_with_audio_producer(monkeypatch, flush_only, response_format): + metrics = _MetricsStub() + serving = _serving(metrics) + clock = {"now": 0.0} + finalized = [] + original_finalize = speech_module.observe_audio_streaming_finalize + + def capture_finalize(*args, **kwargs): + finalized.append(kwargs) + original_finalize(*args, **kwargs) + + class BufferedResampler: + def __init__(self, source_rate, target_rate): + assert (source_rate, target_rate) == (16000, 24000) + + def process(self, chunk, *, final=False): + clock["now"] = 0.5 if final else 0.25 + if not final and flush_only: + return np.empty(0, dtype=np.float32) + return np.zeros(2400, dtype=np.float32) + + monkeypatch.setattr(speech_module, "StreamingAudioResampler", BufferedResampler) + monkeypatch.setattr(speech_module, "observe_audio_streaming_finalize", capture_finalize) + monkeypatch.setattr(speech_module.time, "time", lambda: 100.0 + clock["now"]) + monkeypatch.setattr(speech_module.time, "perf_counter", lambda: clock["now"]) + # The last result is not audio and must not supply the flush metric labels. + non_audio = SimpleNamespace(multimodal_output={"timestamps": []}, stage_id=9, replica_id=8) + chunks = [ + chunk + async for chunk in serving._generate_audio_chunks( + _generate(_result(), non_audio), + request_id="speech-test", + response_format=response_format, + request_start_s=0.0, + request_arrival_ts=100.0, + target_sample_rate=24000, + ) + ] + + if response_format == "wav": + assert chunks.pop(0).startswith(b"RIFF") + expected_arrivals = [0.5] if flush_only else [0.25, 0.5] + assert [len(chunk) for chunk in chunks] == [4800] * len(expected_arrivals) + assert metrics.ttfp_calls == [("1", "2", pytest.approx(expected_arrivals[0]))] + assert len(finalized) == 1 + assert finalized[0]["sample_rate"] == 24000 + assert finalized[0]["channels"] == 1 + assert finalized[0]["chunk_bytes"] == [4800] * len(expected_arrivals) + assert finalized[0]["chunk_arrival_times_s"] == expected_arrivals + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0 if flush_only else 0.15))] + assert metrics.continuity_calls == ([("1", "2", 100)] if flush_only else []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_pcm", [False, True]) +async def test_resampled_speech_without_pcm_does_not_emit_metrics(monkeypatch, empty_pcm): + metrics = _MetricsStub() + serving = _serving(metrics) + + class EmptyResampler: + def __init__(self, source_rate, target_rate): + pass + + def process(self, chunk, *, final=False): + # Also cover a nonempty flush waveform whose encoder emits no PCM. + return np.zeros(10 if final and empty_pcm else 0, dtype=np.float32) + + monkeypatch.setattr(speech_module, "StreamingAudioResampler", EmptyResampler) + serving.create_audio = lambda audio_obj: SimpleNamespace(audio_data=b"", media_type="audio/pcm") + chunks = [ + chunk + async for chunk in serving._generate_audio_chunks( + _generate(_result()), + request_id="speech-test", + request_arrival_ts=100.0, + target_sample_rate=24000, + ) + ] + assert not any(chunks) + assert metrics.ttfp_calls == [] + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 6414ab4ff..39c9e66d9 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -48,7 +48,14 @@ def _build_test_app( speech_service.forced_aligner_enabled = False async def mock_generate_pcm_chunks( - _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + collect=None, ): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -249,6 +256,7 @@ class TestStreamingSpeechWebSocket: def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] + captured_timing = {} captured_tts_params = [] speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) @@ -257,13 +265,27 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_prepare_speech_generation(request): + async def mock_prepare_speech_generation(request, *, arrival_time=None): captured_requests.append(request) + assert arrival_time is not None + captured_timing["prepare_arrival"] = arrival_time return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]} speech_service._prepare_speech_generation = mock_prepare_speech_generation - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + async def mock_generate_pcm_chunks( + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + ): + assert request_start_s is not None + assert request_arrival_ts is not None + captured_timing["chunk_start"] = request_start_s + captured_timing["chunk_arrival"] = request_arrival_ts captured_tts_params.append(tts_params) for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): yield (chunk, 24000) if include_sample_rate else chunk @@ -309,6 +331,8 @@ class TestStreamingSpeechWebSocket: assert captured_requests[0].stream is True assert captured_requests[0].response_format == "pcm" assert captured_requests[0].initial_codec_chunk_frames == 12 + assert captured_timing["prepare_arrival"] == captured_timing["chunk_arrival"] + assert captured_timing["chunk_start"] > 0 assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}] assert speech_service._generate_audio_bytes.await_count == 0 @@ -341,8 +365,9 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = True - async def mock_prepare_speech_generation(request): + async def mock_prepare_speech_generation(request, *, arrival_time=None): captured_requests.append(request) + assert arrival_time is not None return "req-stream", object(), {} speech_service._prepare_speech_generation = mock_prepare_speech_generation @@ -353,8 +378,17 @@ class TestStreamingSpeechWebSocket: # The forced-aligner stage rides the same generator: its pooling output # is surfaced via the ``collect`` channel once the audio has streamed. async def mock_generate_pcm_chunks( - _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + collect=None, ): + assert request_start_s is not None + assert request_arrival_ts is not None for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -438,7 +472,14 @@ class TestStreamingSpeechWebSocket: speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) async def mock_generate_pcm_chunks( - _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + collect=None, ): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk @@ -620,7 +661,15 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + async def mock_generate_pcm_chunks( + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + ): yield b"\x01\x02" raise RuntimeError("stream boom") @@ -716,7 +765,15 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + async def mock_generate_pcm_chunks( + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + ): yield b"\x01\x02" speech_service._generate_pcm_chunks = mock_generate_pcm_chunks diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py index 14bb4f309..a918dce53 100644 --- a/tests/metrics/test_modality.py +++ b/tests/metrics/test_modality.py @@ -48,6 +48,8 @@ _EXPECTED_FAMILIES = [ defs.AUDIO_UNDERRUN_S, defs.AUDIO_CONTINUITY_OK_METRIC, defs.AUDIO_SKIPPED_REQUESTS_METRIC, + defs.SPEECH_STREAM_ABORTED_METRIC, + defs.SPEECH_STREAM_COMPLETED_METRIC, defs.DIFFUSION_EXEC_S, defs.DIFFUSION_EXEC_PER_STEP_S, defs.DIFFUSION_PREPROCESS_S, @@ -70,6 +72,8 @@ class TestRegistration: mod.observe_audio_underrun("s", "r", 0.01) mod.inc_audio_continuity_ok("s", "r", 100) mod.inc_audio_skipped("s", "r", "malformed_codec") + mod.inc_speech_stream_aborted("error") + mod.inc_speech_stream_completed() mod.observe_diffusion_exec("s", "r", 0.5) mod.observe_diffusion_exec_per_step("s", "r", 0.01) mod.observe_diffusion_preprocess("s", "r", 0.01) @@ -91,6 +95,23 @@ class TestRegistration: class TestAudio: + def test_speech_stream_abort_counter_and_disabled_logging(self): + model = "test-speech-stream-abort-counter" + metrics = OmniModalityMetrics(model_name=model) + disabled = OmniModalityMetrics(model_name=model, log_stats=False) + metrics.inc_speech_stream_completed() + disabled.inc_speech_stream_completed() + assert REGISTRY.get_sample_value(defs.SPEECH_STREAM_COMPLETED_METRIC + "_total", {"model_name": model}) == 1.0 + for reason in ("cancelled", "closed", "engine_dead", "error"): + metrics.inc_speech_stream_aborted(reason) + disabled.inc_speech_stream_aborted(reason) + assert ( + REGISTRY.get_sample_value( + defs.SPEECH_STREAM_ABORTED_METRIC + "_total", {"model_name": model, "reason": reason} + ) + == 1.0 + ) + def test_audio_ttfp_observed(self, mod: OmniModalityMetrics) -> None: stage, replica = "talker_ttfp", "0" mod.observe_audio_ttfp(stage, replica, 0.42) @@ -242,6 +263,9 @@ class _StubModMetrics: def observe_audio_rtf(self, s, r, rtf): self.calls.append(("observe_audio_rtf", s, r, rtf)) + def observe_audio_ttfp(self, s, r, t): + self.calls.append(("observe_audio_ttfp", s, r, t)) + def observe_audio_underrun(self, s, r, u): self.calls.append(("observe_audio_underrun", s, r, u)) @@ -599,7 +623,6 @@ class TestObserveDiffusionFinalize: class TestObserveAudioFirstPacket: def test_observes_with_valid_inputs(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet( stub, @@ -612,19 +635,16 @@ class TestObserveAudioFirstPacket: def test_replica_none_skipped(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet(stub, stage_id=1, replica_id=None, arrival_ts=100.0, now_ts=100.5) assert stub.calls == [] def test_arrival_ts_zero_skipped(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=0.0, now_ts=100.5) assert stub.calls == [] def test_clock_skew_clamped_to_zero(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=100.5, now_ts=100.0) assert stub.calls == [("observe_audio_ttfp", "1", "0", 0.0)] @@ -668,6 +688,24 @@ class TestObserveAudioStreamingFinalize: assert underrun_calls[0][-1] > 0.1 assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls) + def test_stereo_channel_count_is_used_for_playback_rate(self): + stub = _StubModMetrics() + # At 100 Hz, s16le stereo consumes 400 bytes/s. The first 200-byte + # chunk buffers 0.5s, so a second chunk arriving at 1.0s underruns by 0.5s. + observe_audio_streaming_finalize( + stub, + stage_id=1, + replica_id=0, + chunk_arrival_times_s=[0.0, 1.0], + chunk_bytes=[200, 200], + sample_rate=100, + channels=2, + threshold_s=0.1, + ) + underrun_calls = [c for c in stub.calls if c[0] == "observe_audio_underrun"] + assert underrun_calls == [("observe_audio_underrun", "1", "0", pytest.approx(0.5))] + assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls) + def test_empty_arrivals_skipped(self): stub = _StubModMetrics() observe_audio_streaming_finalize( diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 2f4e472f8..d4aa72818 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -13,12 +13,14 @@ import struct import time from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor +from contextlib import aclosing from http import HTTPStatus from pathlib import Path -from typing import Any +from typing import Any, cast from urllib.parse import urlparse from urllib.request import url2pathname +import anyio import numpy as np import soundfile as sf import torch @@ -59,6 +61,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( tts_entry_stage_archs, ) from vllm_omni.entrypoints.utils import coerce_param_message_types +from vllm_omni.metrics.modality import observe_audio_first_packet, observe_audio_streaming_finalize from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -194,6 +197,17 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool: return False +class _SpeechStreamingResponse(StreamingResponse): + async def stream_response(self, send) -> None: + try: + await super().stream_response(send) + finally: + # A disconnect during send leaves the iterator suspended at yield. + # Close it explicitly, even inside Starlette's cancelled scope. + with anyio.CancelScope(shield=True): + await cast(Any, self.body_iterator).aclose() + + class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _diffusion_mode: bool = False _media_connector: MediaConnector | None = None @@ -451,9 +465,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """Find and return the TTS stage config, or None if not found.""" tts_stage_keys = all_tts_stage_keys() entry_stage_archs = tts_entry_stage_archs() - all_stages = frozenset( - getattr(stage.engine_args, "model_stage", None) for stage in self.engine_client.stage_configs - ) + all_stages = frozenset(stage.engine_args.model_stage for stage in self.engine_client.stage_configs) for stage in self.engine_client.stage_configs: engine_args = stage.engine_args model_stage = engine_args.model_stage @@ -1083,7 +1095,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "sample_rate is not supported by the current TTS model" return None - def _validate_ref_audio_format(self, ref_audio: str) -> str | None: + def _validate_ref_audio_format(self, ref_audio: str | list[str] | None) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" if not isinstance(ref_audio, str): return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" @@ -1362,6 +1374,57 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): resolved.append((wav_list, sr)) return resolved + def _observe_speech_audio_ttfp( + self, + *, + request_id: str, + result: Any, + request_arrival_ts: float | None = None, + first_packet_ts: float | None = None, + ) -> tuple[int | None, int | None, bool]: + """Emit the Speech API TTFP sample once the first PCM payload exists.""" + engine_client = getattr(self, "engine_client", None) + mod_metrics = getattr(engine_client, "mod_metrics", None) + if mod_metrics is None: + return None, None, False + + req_state = next( + ( + state + for state in getattr(engine_client, "request_states", {}).values() + if state.external_request_id == request_id + ), + None, + ) + if req_state is not None and req_state.first_audio_ts is not None: + return req_state.audio_emit_stage_id, req_state.audio_emit_replica_id, True + + arrival_ts = request_arrival_ts + if arrival_ts is None and req_state is not None: + arrival_ts = req_state.request_arrival_ts + stage_id = getattr(result, "stage_id", None) + replica_id = getattr(result, "replica_id", None) + if replica_id is None and req_state is not None and stage_id is not None: + stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None) + if stage_pools is not None and 0 <= stage_id < len(stage_pools): + replica_id = stage_pools[stage_id].get_bound_replica_id(req_state.request_id) + if arrival_ts is None or arrival_ts <= 0 or stage_id is None or replica_id is None: + return stage_id, replica_id, False + + now_ts = first_packet_ts if first_packet_ts is not None else time.time() + observe_audio_first_packet( + mod_metrics, + stage_id=stage_id, + replica_id=replica_id, + arrival_ts=arrival_ts, + now_ts=now_ts, + ) + if req_state is not None: + req_state.first_audio_ts = now_ts + req_state.audio_emit_stage_id = stage_id + req_state.audio_emit_replica_id = replica_id + return stage_id, replica_id, True + async def _generate_audio_chunks( self, generator, @@ -1369,6 +1432,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format: str = "pcm", raw_request: Request | None = None, request_start_s: float | None = None, + request_arrival_ts: float | None = None, include_sample_rate: bool = False, usage_acc: SpeechOutputTokenCounter | None = None, tts_params: dict[str, Any] | None = None, @@ -1395,10 +1459,42 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): sample_rate_val = 24000 first_chunk = True first_audio_chunk_s: float | None = None + first_audio_packet_ts: float | None = None + ttfp_observed = False stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() artifact_ready = False + audio_chunk_arrivals_s: list[float] = [] + audio_chunk_bytes: list[int] = [] + audio_stage_id: int | None = None + audio_replica_id: int | None = None + audio_channels = 1 source_sample_rate: int | None = None resampler: StreamingAudioResampler | None = None + last_audio_result: Any = None + + def record_stream_abort(reason: str) -> None: + mod_metrics = getattr(getattr(self, "engine_client", None), "mod_metrics", None) + if mod_metrics is not None: + mod_metrics.inc_speech_stream_aborted(reason) + + def record_audio_chunk(audio_bytes: bytes, chunk_np: np.ndarray, result: Any) -> None: + nonlocal first_audio_chunk_s, first_audio_packet_ts, ttfp_observed + nonlocal audio_stage_id, audio_replica_id, audio_channels + if not audio_bytes: + return + if first_audio_chunk_s is None: + first_audio_chunk_s = time.perf_counter() + first_audio_packet_ts = time.time() + audio_channels = _infer_audio_num_channels(np.asarray(chunk_np)) + if not ttfp_observed: + audio_stage_id, audio_replica_id, ttfp_observed = self._observe_speech_audio_ttfp( + request_id=request_id, + result=result, + request_arrival_ts=request_arrival_ts, + first_packet_ts=first_audio_packet_ts, + ) + audio_chunk_arrivals_s.append(max(time.perf_counter() - stream_start_s, 0.0)) + audio_chunk_bytes.append(len(audio_bytes)) # SSE supplies an accumulator for usage output. Raw-audio and WebSocket # streams retain terminal metrics only when their model adapter needs @@ -1419,6 +1515,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if collect is not None and self._is_timestamps_output(res): collect["aligner_res"] = res continue + audio_output = cast(dict, audio_output) sr_raw = audio_output.get("sr") if sr_raw is not None: @@ -1452,6 +1549,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): output_sample_rate = target_sample_rate or sample_rate_val for chunk_tensor in new_chunks: + # Flush may run after a non-audio (e.g. aligner) result. + # Retain the producer of the audio buffered by the resampler. + last_audio_result = res chunk_np = ( chunk_tensor.float().detach().cpu().numpy() if hasattr(chunk_tensor, "float") else chunk_tensor ) @@ -1466,6 +1566,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # as first audio; the post-loop guard below needs to # see an audio-less stream to fail the request. continue + wav_header: bytes | None = None # For WAV format, emit header before first audio chunk if response_format == "wav" and first_chunk: num_channels = _infer_audio_num_channels(np.asarray(chunk_np)) @@ -1474,7 +1575,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): num_channels=num_channels, bits_per_sample=16, ) - yield wav_header first_chunk = False # Convert audio to PCM bytes @@ -1485,9 +1585,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): speed=1.0, base64_encode=False, ) - if first_audio_chunk_s is None: - first_audio_chunk_s = time.perf_counter() - audio_bytes = self.create_audio(audio_obj).audio_data + audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data) + record_audio_chunk(audio_bytes, chunk_np, res) + if wav_header is not None: + yield wav_header if include_sample_rate: yield audio_bytes, output_sample_rate else: @@ -1497,8 +1598,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): final_chunk = resampler.process(np.empty((0,), dtype=np.float32), final=True) if final_chunk.size: output_sample_rate = target_sample_rate or sample_rate_val + wav_header = None if response_format == "wav" and first_chunk: - yield _create_wav_header( + wav_header = _create_wav_header( sample_rate=output_sample_rate, num_channels=1, bits_per_sample=16, @@ -1511,9 +1613,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): speed=1.0, base64_encode=False, ) - if first_audio_chunk_s is None: - first_audio_chunk_s = time.perf_counter() - audio_bytes = self.create_audio(audio_obj).audio_data + audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data) + record_audio_chunk(audio_bytes, final_chunk, last_audio_result) + if wav_header is not None: + yield wav_header if include_sample_rate: yield audio_bytes, output_sample_rate else: @@ -1527,8 +1630,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # bytes, but they must terminate as an error rather than cleanly. if tts_params is not None and usage_acc is not None: self._validate_tts_generation(tts_params, usage_acc) + mod_metrics = getattr(getattr(self, "engine_client", None), "mod_metrics", None) + if mod_metrics is not None and audio_stage_id is not None and audio_replica_id is not None: + observe_audio_streaming_finalize( + mod_metrics, + stage_id=audio_stage_id, + replica_id=audio_replica_id, + chunk_arrival_times_s=audio_chunk_arrivals_s, + chunk_bytes=audio_chunk_bytes, + sample_rate=target_sample_rate or sample_rate_val, + channels=audio_channels, + ) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True + if mod_metrics is not None: + mod_metrics.inc_speech_stream_completed() total_ms = (time.perf_counter() - stream_start_s) * 1000.0 if first_audio_chunk_s is not None: first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0 @@ -1544,7 +1660,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id, total_ms, ) + except GeneratorExit: + record_stream_abort("closed") + raise except asyncio.CancelledError: + record_stream_abort("cancelled") total_ms = (time.perf_counter() - stream_start_s) * 1000.0 logger.info( "[SpeechE2E] request_id=%s stream=true status=cancelled total_ms=%.2f", @@ -1554,6 +1674,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): logger.info("Streaming request %s cancelled by client", request_id) raise except EngineDeadError as e: + record_stream_abort("engine_dead") total_ms = (time.perf_counter() - stream_start_s) * 1000.0 logger.error( "[SpeechE2E] request_id=%s stream=true status=engine_dead total_ms=%.2f", @@ -1573,6 +1694,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) raise except Exception as e: + record_stream_abort("error") total_ms = (time.perf_counter() - stream_start_s) * 1000.0 logger.exception( "[SpeechE2E] request_id=%s stream=true status=error total_ms=%.2f error=%s", @@ -1585,6 +1707,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): finally: if not artifact_ready: self._discard_ref_audio_artifact_warmup(request_id) + close = getattr(generator, "aclose", None) + if close is not None: + with anyio.CancelScope(shield=True): + await close() async def _generate_audio_sse_events( self, @@ -1593,6 +1719,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format: str = "pcm", raw_request: Request | None = None, request_start_s: float | None = None, + request_arrival_ts: float | None = None, request: OpenAICreateSpeechRequest | None = None, tts_params: dict[str, Any] | None = None, ): @@ -1613,24 +1740,28 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): usage_acc = SpeechOutputTokenCounter() emitted_audio = False try: - async for chunk in self._generate_audio_chunks( - generator, - request_id, - response_format, - raw_request=raw_request, - request_start_s=request_start_s, - usage_acc=usage_acc, - tts_params=tts_params, - target_sample_rate=request.sample_rate if request is not None else None, - ): - payload = { - "type": "speech.audio.delta", - "audio": base64.b64encode(chunk).decode("ascii"), - "response_format": response_format, - } - data = json.dumps(payload, separators=(",", ":")) - emitted_audio = True - yield f"event: speech.audio.delta\ndata: {data}\n\n" + async with aclosing( + self._generate_audio_chunks( + generator, + request_id, + response_format, + raw_request=raw_request, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, + usage_acc=usage_acc, + tts_params=tts_params, + target_sample_rate=request.sample_rate if request is not None else None, + ) + ) as chunks: + async for chunk in chunks: + payload = { + "type": "speech.audio.delta", + "audio": base64.b64encode(chunk).decode("ascii"), + "response_format": response_format, + } + data = json.dumps(payload, separators=(",", ":")) + emitted_audio = True + yield f"event: speech.audio.delta\ndata: {data}\n\n" done_payload: dict[str, Any] = {"type": "speech.audio.done"} if request is not None: # Streaming path: output_tokens = sum of stage-0 deltas. @@ -1695,6 +1826,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request: OpenAICreateSpeechRequest, request_id: str | None = None, has_inline_ref_audio: bool | None = None, + arrival_time: float | None = None, ) -> tuple[str, Any, dict[str, Any]]: if self.engine_client.errored: raise self.engine_client.dead_error @@ -1823,6 +1955,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id=request_id, sampling_params_list=sampling_params_list, output_modalities=output_modalities, + arrival_time=arrival_time, ) self._track_ref_audio_artifact_warmup( request_id, @@ -1836,6 +1969,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): generator, request_id: str, *, + request_start_s: float | None = None, + request_arrival_ts: float | None = None, include_sample_rate: bool = False, tts_params: dict[str, Any] | None = None, collect: dict | None = None, @@ -1848,28 +1983,36 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ``collect`` (when given) receives the forced-aligner stage's pooling output under ``"aligner_res"`` for downstream word-timestamp extraction. """ - async for chunk in self._generate_audio_chunks( - generator, - request_id, - response_format="pcm", - include_sample_rate=include_sample_rate, - tts_params=tts_params, - collect=collect, - target_sample_rate=target_sample_rate, - ): - yield chunk + async with aclosing( + self._generate_audio_chunks( + generator, + request_id, + response_format="pcm", + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, + include_sample_rate=include_sample_rate, + tts_params=tts_params, + collect=collect, + target_sample_rate=target_sample_rate, + ) + ) as chunks: + async for chunk in chunks: + yield chunk async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest): """Yield raw PCM bytes for a speech request as soon as chunks are decoded.""" request_id, generator, tts_params = await self._prepare_speech_generation(request) try: - async for chunk in self._generate_pcm_chunks( - generator, - request_id, - tts_params=tts_params, - target_sample_rate=request.sample_rate, - ): - yield chunk + async with aclosing( + self._generate_pcm_chunks( + generator, + request_id, + tts_params=tts_params, + target_sample_rate=request.sample_rate, + ) + ) as chunks: + async for chunk in chunks: + yield chunk finally: self._discard_ref_audio_artifact_warmup(request_id) @@ -1881,14 +2024,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): usage_out: list[SpeechTokenUsage] | None = None, has_inline_ref_audio: bool | None = None, collect: dict | None = None, + request_arrival_ts: float | None = None, ) -> tuple[bytes | str, str]: # ``usage_out`` is an opt-in output channel: when a list is passed, the # computed SpeechTokenUsage is appended to it. The return stays a # 2-tuple so existing callers (and their test mocks) are unaffected; # batch and non-streaming response-header paths opt in when surfacing # usage outside the raw audio body. + request_arrival_ts = request_arrival_ts if request_arrival_ts is not None else time.time() request_id, generator, bytes_tts_params = await self._prepare_speech_generation( - request, request_id=request_id, has_inline_ref_audio=has_inline_ref_audio + request, + request_id=request_id, + has_inline_ref_audio=has_inline_ref_audio, + arrival_time=request_arrival_ts, ) artifact_ready = False @@ -1927,6 +2075,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): continue if step_key is None: continue + step_audio = cast(dict, step_audio) chunk = step_audio[step_key] candidates = chunk if isinstance(chunk, list) else [chunk] for cand in candidates: @@ -1948,6 +2097,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): audio_output, audio_key = self._extract_audio_output(audio_source) if audio_key is None: raise ValueError("TTS model did not produce audio output.") + audio_output = cast(dict, audio_output) # Surface forced-aligner word timestamps to the caller (set as a # response header) when requested and an aligner stage produced them. @@ -2074,7 +2224,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id = f"speech-{random_uuid()}" prompt: dict[str, Any] = {"input": request.input} if request.ref_audio: - wav, sr, _ = await self._resolve_ref_audio(request.ref_audio) + wav, sr, _ = await self._resolve_ref_audio(cast(str, request.ref_audio)) prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr) if request.ref_text: prompt["ref_text"] = request.ref_text @@ -2151,6 +2301,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): audio_output, audio_key = self._extract_audio_output(final_output) if audio_key is None: raise ValueError("TTS model did not produce audio output.") + audio_output = cast(dict, audio_output) audio_tensor = audio_output[audio_key] sr_raw = audio_output.get("sr", 24000) @@ -2265,6 +2416,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id = f"speech-{random_uuid()}" request_start_s = time.perf_counter() + request_arrival_ts = ( + float(getattr(raw_request.state, "request_timestamp", time.time())) + if raw_request is not None + else time.time() + ) if raw_request: raw_request.state.request_metadata = RequestResponseMetadata( request_id=request_id, @@ -2292,14 +2448,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return error media_type = "audio/wav" if response_format == "wav" else "audio/pcm" - _, generator, raw_tts_params = await self._prepare_speech_generation(request, request_id=request_id) - return StreamingResponse( + _, generator, raw_tts_params = await self._prepare_speech_generation( + request, + request_id=request_id, + arrival_time=request_arrival_ts, + ) + return _SpeechStreamingResponse( self._generate_audio_chunks( generator, request_id, response_format, raw_request=raw_request, request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, tts_params=raw_tts_params, target_sample_rate=request.sample_rate, ), @@ -2314,14 +2475,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if error is not None: return error - _, generator, sse_tts_params = await self._prepare_speech_generation(request, request_id=request_id) - return StreamingResponse( + _, generator, sse_tts_params = await self._prepare_speech_generation( + request, + request_id=request_id, + arrival_time=request_arrival_ts, + ) + return _SpeechStreamingResponse( self._generate_audio_sse_events( generator, request_id, response_format, raw_request=raw_request, request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, request=request, tts_params=sse_tts_params, ), @@ -2332,7 +2498,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): usage_box: list[SpeechTokenUsage] = [] try: audio_bytes, media_type = await self._generate_audio_bytes( - request, request_id=request_id, usage_out=usage_box, collect=collect + request, + request_id=request_id, + usage_out=usage_box, + collect=collect, + request_arrival_ts=request_arrival_ts, ) except TTSGenerationError as error: # An adapter can reject otherwise completed audio. Retry only @@ -2358,6 +2528,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id=retry_request_id, usage_out=usage_box, collect=collect, + request_arrival_ts=request_arrival_ts, ) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 32e9fa40d..ceeb208c9 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -55,6 +55,7 @@ Protocol: import asyncio import base64 import json +import time from contextlib import aclosing from fastapi import WebSocket, WebSocketDisconnect @@ -284,6 +285,8 @@ class OmniStreamingSpeechHandler: ``utterance_index`` identifies the flush this sentence belongs to and ``sentence_index`` its position inside that flush. """ + request_arrival_ts = time.time() + request_start_s = time.perf_counter() response_format = config.response_format or "wav" # Reject unmet word-timestamps preconditions early with a clear reason. @@ -346,7 +349,10 @@ class OmniStreamingSpeechHandler: request_id = None try: if config.stream_audio: - request_id, generator, tts_params = await self._speech_service._prepare_speech_generation(request) + request_id, generator, tts_params = await self._speech_service._prepare_speech_generation( + request, + arrival_time=request_arrival_ts, + ) if config.word_timestamps: total_bytes = await self._stream_audio_with_alignments( websocket=websocket, @@ -356,6 +362,8 @@ class OmniStreamingSpeechHandler: utterance_index=utterance_index, sentence_index=sentence_index, language=config.language, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, tts_params=tts_params, ) else: @@ -363,6 +371,8 @@ class OmniStreamingSpeechHandler: self._speech_service._generate_pcm_chunks( generator, request_id, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, tts_params=tts_params, ) ) as stream: @@ -370,7 +380,10 @@ class OmniStreamingSpeechHandler: total_bytes += len(chunk) await websocket.send_bytes(chunk) else: - audio_bytes, _ = await self._speech_service._generate_audio_bytes(request) + audio_bytes, _ = await self._speech_service._generate_audio_bytes( + request, + request_arrival_ts=request_arrival_ts, + ) total_bytes = len(audio_bytes) await websocket.send_bytes(audio_bytes) except WebSocketDisconnect: @@ -416,6 +429,8 @@ class OmniStreamingSpeechHandler: sentence_text: str, utterance_index: int, sentence_index: int, + request_start_s: float, + request_arrival_ts: float, language: str | None = None, tts_params: dict | None = None, ) -> int: @@ -463,6 +478,8 @@ class OmniStreamingSpeechHandler: self._speech_service._generate_pcm_chunks( generator, request_id, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, include_sample_rate=True, tts_params=tts_params, collect=collect, diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 6cc2825c1..a09fbcc2f 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -160,6 +160,8 @@ AUDIO_FRAMES_METRIC = METRIC_PREFIX + AUDIO_FRAMES AUDIO_UNDERRUN_S = METRIC_PREFIX + AUDIO_UNDERRUN + "_s" AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS +SPEECH_STREAM_ABORTED_METRIC = METRIC_PREFIX + "speech_stream_aborted" +SPEECH_STREAM_COMPLETED_METRIC = METRIC_PREFIX + "speech_stream_completed" # Realtime Server VAD serving metrics. REALTIME_VAD_ACTIVE_SESSIONS = METRIC_PREFIX + "realtime_vad_active_sessions" diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py index e0d51e2b0..bb7842e36 100644 --- a/vllm_omni/metrics/modality.py +++ b/vllm_omni/metrics/modality.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """OmniModalityMetrics — per-modality Prometheus families (audio path only). 7 audio business-semantic metric families. Text-path metrics (TTFT / ITL / @@ -76,6 +79,16 @@ _audio_skipped_family = Counter( "Silent-loss counter — code2wav rejected malformed codec input and returned 200 OK with empty audio.", labelnames=list(defs.AUDIO_SKIPPED_LABELS), ) +_speech_stream_aborted_family = Counter( + defs.SPEECH_STREAM_ABORTED_METRIC, + "Speech audio generators terminated before normal completion, including before the first PCM payload.", + labelnames=["model_name", "reason"], +) +_speech_stream_completed_family = Counter( + defs.SPEECH_STREAM_COMPLETED_METRIC, + "Speech audio generators that completed normally; does not confirm client receipt.", + labelnames=["model_name"], +) # ---------------------------------------------------------------------------- @@ -153,6 +166,17 @@ class OmniModalityMetrics: # ---- Audio ------------------------------------------------------------ + def inc_speech_stream_aborted(self, reason: str) -> None: + if not self._log_stats: + return + if reason not in {"cancelled", "closed", "engine_dead", "error"}: + reason = "error" + _speech_stream_aborted_family.labels(model_name=self._model_name, reason=reason).inc() + + def inc_speech_stream_completed(self) -> None: + if self._log_stats: + _speech_stream_completed_family.labels(model_name=self._model_name).inc() + def observe_audio_ttfp(self, stage: str, replica: str, ttfp_seconds: float) -> None: if not self._log_stats: return @@ -335,6 +359,7 @@ def observe_audio_streaming_finalize( chunk_arrival_times_s: list[float], chunk_bytes: list[int], sample_rate: int, + channels: int = defs.DEFAULT_AUDIO_CHANNELS, threshold_s: float = defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S, ) -> None: """Emit audio_underrun_s + audio_continuity_ok_total at request end. @@ -353,6 +378,7 @@ def observe_audio_streaming_finalize( chunk_arrival_times_s=chunk_arrival_times_s, chunk_bytes=chunk_bytes, sample_rate=sample_rate, + channels=channels, threshold_s=threshold_s, ) stage_label = str(stage_id) From 78e1a0ae5ce6bec260043a085df2d35d42ce1d03 Mon Sep 17 00:00:00 2001 From: Qi Jia Date: Fri, 11 Sep 2026 22:18:07 +0800 Subject: [PATCH 028/209] [Bugfix][Model] Fix FLUX.2 Klein multi-image edit metadata (#7430) Signed-off-by: QI JIA Co-authored-by: QI JIA Co-authored-by: Cursor --- .../diffusion/test_diffusion_config_propagation.py | 13 +++++++++++++ vllm_omni/diffusion/model_metadata.py | 6 ++++++ 2 files changed, 19 insertions(+) diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 394f7c04f..23ff796e2 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -17,6 +17,7 @@ from vllm_omni.diffusion.data import ( ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.model_metadata import ( + FLUX2_KLEIN_MAX_INPUT_IMAGES, HUNYUAN_IMAGE3_MAX_INPUT_IMAGES, QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, ) @@ -196,6 +197,18 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES +def test_flux2_klein_sets_generic_multimodal_limit(): + od_config = OmniDiffusionConfig( + model="black-forest-labs/FLUX.2-klein-9B", + model_class_name="Flux2KleinPipeline", + ) + + od_config.update_multimodal_support() + + assert od_config.supports_multimodal_inputs is True + assert od_config.max_multimodal_image_inputs == FLUX2_KLEIN_MAX_INPUT_IMAGES + + def test_task_type_roundtrip(): od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task") assert od.task_type == "model-defined-task" diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index c91d2ead5..2ee1b1f4b 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -19,6 +19,8 @@ class DiffusionModelMetadata: final_output_type: str | None = None +# FLUX.2 Klein supports up to four reference images. +FLUX2_KLEIN_MAX_INPUT_IMAGES = 4 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3. HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3 @@ -27,6 +29,10 @@ BOOGU_IMAGE_MAX_INPUT_IMAGES = 1 _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { + "Flux2KleinPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=FLUX2_KLEIN_MAX_INPUT_IMAGES, + ), "QwenImageEditPlusPipeline": DiffusionModelMetadata( supports_multimodal_inputs=True, max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, From d3493384cfc6c4b9ba1fc4e534e91fe588ecb4f1 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Fri, 11 Sep 2026 10:28:55 -0400 Subject: [PATCH 029/209] [BugFix] Fix leftovers of the legacy OpenAI realtime API event names (#7426) Signed-off-by: Nick Cao Co-authored-by: Codex --- docs/serving/realtime_api.md | 6 +-- docs/serving/video_stream_api.md | 4 +- examples/online_serving/qwen3_omni/README.md | 36 ++++++++------ .../qwen3_omni/openai_realtime_client.py | 14 +++--- .../openai_api/test_duplex_handler.py | 4 +- .../test_qwen3_omni_realtime_websocket.py | 16 +++--- .../openai_api/test_serving_video_stream.py | 49 +++++++++++++++++++ .../entrypoints/openai/realtime_connection.py | 8 +-- .../openai/serving_video_stream.py | 4 +- .../entrypoints/openai/video_stream_base.py | 10 ++-- 10 files changed, 103 insertions(+), 48 deletions(-) diff --git a/docs/serving/realtime_api.md b/docs/serving/realtime_api.md index 153564140..e44ff7ed0 100644 --- a/docs/serving/realtime_api.md +++ b/docs/serving/realtime_api.md @@ -43,8 +43,8 @@ event is base64-encoded raw PCM16. | Client to server | `input_audio_buffer.commit` with `final: true` | Marks the end of input | | Server to client | `transcription.delta` | Carries incremental response text | | Server to client | `transcription.done` | Carries final text and token usage | -| Server to client | `response.audio.delta` | Carries incremental PCM16 response audio | -| Server to client | `response.audio.done` | Marks the end of response audio | +| Server to client | `response.output_audio.delta` | Carries incremental PCM16 response audio | +| Server to client | `response.output_audio.done` | Marks the end of response audio | | Server to client | `error` | Reports an invalid event, model, or audio payload | A minimal client sends events in this order: @@ -63,7 +63,7 @@ engine is consuming the stream. ## Audio Handling - Input is mono PCM16 at 16 kHz for the Qwen3-Omni example. -- `response.audio.delta.audio` contains base64-encoded PCM16 bytes. +- `response.output_audio.delta.audio` contains base64-encoded PCM16 bytes. - Read `sample_rate_hz` from each audio event instead of assuming an output rate. Qwen3-Omni output is typically 24 kHz. - Concatenate audio deltas in receive order to construct the output waveform. diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index e1f583991..78575e0fc 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -45,8 +45,8 @@ WebSocket /v1/video/chat/stream | Server -> Client | `response.start` | none | Query generation started. | | Server -> Client | `response.text.delta` | `delta` | Incremental text output. | | Server -> Client | `response.text.done` | `text` | Final text output for the query. | -| Server -> Client | `response.audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. | -| Server -> Client | `response.audio.done` | none | Audio output finished. | +| Server -> Client | `response.output_audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. | +| Server -> Client | `response.output_audio.done` | none | Audio output finished. | | Server -> Client | `session.done` | none | Session closed. | | Server -> Client | `error` | `message` | Recoverable protocol or generation error. | diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index ef2367412..107e1f098 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -1,5 +1,7 @@ # Qwen3-Omni: Online serving + + ## 🛠️ Installation Please refer to [README.md](../../../README.md) @@ -19,6 +21,7 @@ Additionally, NPU, ROCm, and XPU per-platform configuration deltas are determini `platforms`: section of the corresponding YAML. To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly: + ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/your_deploy_config.yaml @@ -53,7 +56,7 @@ Use the stage-based CLI when you want to run one stage per process. The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via `CUDA_VISIBLE_DEVICES`. -**1. Stage 0 (Thinker + API server)** +#### 1. Stage 0 (Thinker + API server) ```bash CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -63,7 +66,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --omni-master-port 26000 ``` -**2. Stage 1 (Talker)** +#### 2. Stage 1 (Talker) ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -73,7 +76,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --omni-master-port 26000 ``` -**3. Stage 2 (Code2Wav)** +#### 3. Stage 2 (Code2Wav) ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -252,11 +255,12 @@ pattern works for Qwen2.5-Omni — replace `base_config:` with the path to ### Send Multi-modal Request Get into the example folder + ```bash cd examples/online_serving/qwen3_omni ``` -#### Send request via python +#### Send request via python ```bash python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost" @@ -264,7 +268,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene #### Realtime WebSocket client (`openai_realtime_client.py`) -[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. +[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.output_audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation. @@ -289,7 +293,7 @@ python openai_realtime_client.py \ **Arguments:** | Flag | Default | Description | -|------|---------|-------------| +| ------ | --------- | ------------- | | `--url` | `ws://localhost:8091/v1/realtime` | Full WebSocket URL including path | | `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (sent in `session.update`) | | `--input-wav` | *(required)* | Input WAV: mono, 16-bit PCM, **16 kHz** | @@ -297,7 +301,7 @@ python openai_realtime_client.py \ | `--output-text` | *(optional)* | If set, write final transcription text to this path | | `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) | | `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) | -| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files | +| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.output_audio.delta` WAV files | | `--server-vad` | disabled | Enable server-side Silero VAD instead of sending explicit commits | | `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) | | `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 | @@ -334,7 +338,6 @@ The Python client supports the following command-line arguments: - `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"` - `--speaker`: TTS speaker/voice for audio output when requesting audio (e.g. `ethan`, `chelsie`, `aiden`). Omit to use the model default. Example: `--speaker "chelsie"` - For example, to use a local video file with custom prompt: ```bash @@ -345,22 +348,22 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene --prompt "What are the main activities shown in this video?" ``` -#### Send request via curl +#### Send request via curl ```bash bash run_curl_multimodal_generation.sh use_image ``` - ### FAQ ## Modality control + You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance. ### Supported modalities | Modalities | Output | -|------------|--------| +| ------------ | -------- | | `["text"]` | Text only | | `["audio"]` | Audio only | | `["text", "audio"]` | Text + Audio | @@ -496,7 +499,9 @@ print(response.choices[0].message.audio) Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default. ## Streaming Output + If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally. + ```bash python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ --query-type use_image \ @@ -521,12 +526,14 @@ The convenience script launches both the vLLM server and Gradio demo together: ``` This script will: + 1. Start the vLLM server in the background 2. Wait for the server to be ready 3. Launch the Gradio demo 4. Handle cleanup when you press Ctrl+C The script supports the following arguments: + - `--model`: Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct) - `--server-port`: Port for vLLM server (default: 8091) - `--gradio-port`: Port for Gradio demo (default: 7861) @@ -537,18 +544,19 @@ The script supports the following arguments: #### Option 2: Manual Launch (Two-Step Process) -**Step 1: Launch the vLLM API server** +##### Step 1: Launch the vLLM API server ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` If you have a custom deploy config file: + ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file ``` -**Step 2: Run the Gradio demo** +##### Step 2: Run the Gradio demo In a separate terminal: @@ -561,7 +569,7 @@ Then open `http://localhost:7861/` on your local browser to interact with the we The gradio script supports the following arguments: - `--model`: Model name/path (should match the server model) -- `--api-base`: Base URL for the vLLM API server (default: http://localhost:8091/v1) +- `--api-base`: Base URL for the vLLM API server (default: ) - `--ip`: Host/IP for Gradio server (default: 127.0.0.1) - `--port`: Port for Gradio server (default: 7861) - `--share`: Share the Gradio demo publicly (creates a public link) diff --git a/examples/online_serving/qwen3_omni/openai_realtime_client.py b/examples/online_serving/qwen3_omni/openai_realtime_client.py index 8a6871580..0a98d69a0 100644 --- a/examples/online_serving/qwen3_omni/openai_realtime_client.py +++ b/examples/online_serving/qwen3_omni/openai_realtime_client.py @@ -6,14 +6,14 @@ This client: 1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz), 2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events, -3) Receives response.audio.*, response.output_text.*, and transcript events, +3) Receives response.output_audio.*, response.output_text.*, and transcript events, 4) Saves synthesized audio to an output WAV file and optional text file. -By default each ``response.audio.delta`` is treated as an **incremental PCM** +By default each ``response.output_audio.delta`` is treated as an **incremental PCM** chunk and all chunks are concatenated into the final ``--output-wav``. Optional debugging: pass ``--delta-dump-dir DIR`` to write every -``response.audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, … +``response.output_audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, … Usage: python openai_realtime_client.py \ @@ -187,7 +187,7 @@ async def run_client( if event_type == "session.created": continue - if event_type == "response.audio.delta": + if event_type == "response.output_audio.delta": sr = event.get("sample_rate_hz") if isinstance(sr, int) and sr > 0: output_sample_rate = sr @@ -223,11 +223,11 @@ async def run_client( print(f"{log_prefix}text usage: {usage}") continue - if event_type == "response.audio_transcript.done": + if event_type == "response.output_audio_transcript.done": final_audio_transcript = event.get("transcript", "") continue - if event_type == "response.audio.done": + if event_type == "response.output_audio.done": if not server_vad: break continue @@ -339,7 +339,7 @@ def main() -> None: "--delta-dump-dir", type=Path, default=None, - help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory", + help="If set, each response.output_audio.delta is saved as delta_NNNNNN.wav under this directory", ) parser.add_argument( "--server-vad", diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 0f9698457..206a38917 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -2197,9 +2197,7 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ) audio_events = [ - payload - for payload in payloads - if payload["type"] in {"response.output_audio.delta", "response.output_audio.delta"} + payload for payload in payloads if payload["type"] in {"response.output_audio.delta", "response.audio.delta"} ] assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} assert {payload["format"] for payload in audio_events} == {"pcm16"} diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 406668a4c..be83ea9db 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -175,7 +175,7 @@ async def _run_realtime_audio_roundtrip( if event_type == "session.created": continue - if event_type == "response.audio.delta": + if event_type == "response.output_audio.delta": delta_events += 1 sr = event.get("sample_rate_hz") if isinstance(sr, int) and sr > 0: @@ -195,7 +195,7 @@ async def _run_realtime_audio_roundtrip( final_text = event.get("text", "") or "".join(text_chunks) continue - if event_type == "response.audio.done": + if event_type == "response.output_audio.done": break if event_type == "error": @@ -315,7 +315,7 @@ def _server_vad_pcm16_input() -> bytes: def _assert_realtime_smoke(result: dict) -> None: out_pcm = result["output_pcm"] assert result["delta_events"] >= 1 - assert out_pcm, "No output PCM from response.audio.delta" + assert out_pcm, "No output PCM from response.output_audio.delta" assert len(out_pcm) % 2 == 0 assert len(out_pcm) >= 4096, "Output audio unexpectedly small" assert result["output_sample_rate"] > 0 @@ -388,7 +388,7 @@ class TestQwen3OmniRealtimeWebSocket: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True) def test_long_audio_response_completes_async_chunk(self, omni_server) -> None: - """Regression for #6474: long async-chunk audio must emit response.audio.done.""" + """Regression for #6474: long async-chunk audio must emit response.output_audio.done.""" issue_6474_pcm16 = _synthetic_pcm16_input( phrase_text=ISSUE_6474_SYNTH_PHRASE_TEXT, ) @@ -450,8 +450,8 @@ class TestQwen3OmniRealtimeWebSocket: "input_audio_buffer.speech_stopped", "input_audio_buffer.committed", "response.created", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", "response.done", ] input_item_ids: list[str] = [] @@ -459,7 +459,7 @@ class TestQwen3OmniRealtimeWebSocket: for events in turns: event_types = [event["type"] for event in events] for event_type in required_sequence: - if event_type != "response.audio.delta": + if event_type != "response.output_audio.delta": assert event_types.count(event_type) == 1, event_types positions = [event_types.index(event_type) for event_type in required_sequence] assert positions == sorted(positions) @@ -486,7 +486,7 @@ class TestQwen3OmniRealtimeWebSocket: output_pcm = b"".join( base64.b64decode(event["delta"]) for event in events - if event["type"] == "response.audio.delta" and event.get("delta") + if event["type"] == "response.output_audio.delta" and event.get("delta") ) assert output_pcm assert created["id"] == done["id"] diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index b534d6c67..58ea2b57f 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -14,6 +14,7 @@ import weakref from typing import Any import pytest +import torch from PIL import Image from vllm_omni.entrypoints.openai import video_stream_base, video_stream_envs @@ -457,6 +458,54 @@ async def test_async_chunk_mode_is_read_by_engine_path_at_runtime(monkeypatch): assert not any(m.get("type") == "response.text.delta" for m in ws_off.sent) +@pytest.mark.asyncio +@pytest.mark.parametrize("async_chunk_mode", ["on", "off"]) +async def test_audio_events_use_output_audio_names(monkeypatch, async_chunk_mode): + class AudioEngine: + def generate(self, **_kwargs): + async def _gen(): + yield _audio_result([torch.zeros(1)]) + + return _gen() + + class AudioHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt": "x"} + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", async_chunk_mode) + monkeypatch.setattr( + OmniStreamingVideoHandler, + "_encode_audio_wav_b64", + staticmethod(lambda _audio: "audio-b64"), + ) + + ws = MockWebSocket() + handler = AudioHandler(chat_service=object(), engine_client=AudioEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"]) + + await handler._process_query_engine( + ws, + config, + [_b64(_make_jpeg())], + bytearray(), + [], + "describe", + "req-audio-events", + asyncio.Event(), + {}, + ) + + audio_events = [message for message in ws.sent if "audio" in message.get("type", "")] + assert audio_events == [ + { + "type": "response.output_audio.delta", + "data": "audio-b64", + "format": "wav", + }, + {"type": "response.output_audio.done"}, + ] + + @pytest.mark.asyncio async def test_query_without_engine_client_sends_error(): ws = MockWebSocket() diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 86a8ba031..3d37fbe37 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -176,7 +176,7 @@ class RealtimeConnection(VllmRealtimeConnection): sent_audio = True await self.send_json( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "audio": self._pcm16_b64(chunk), "format": "pcm16", "sample_rate_hz": sample_rate, @@ -194,7 +194,7 @@ class RealtimeConnection(VllmRealtimeConnection): await self.send(TranscriptionDone(text=full_text, usage=usage)) if sent_audio: - await self.send_json({"type": "response.audio.done", "has_audio": True}) + await self.send_json({"type": "response.output_audio.done", "has_audio": True}) audio_done_sent = True except Exception as e: logger.exception("Error in generation: %s", e) @@ -213,9 +213,9 @@ class RealtimeConnection(VllmRealtimeConnection): # Always send terminal event so clients don't hang forever. if self._is_connected and not audio_done_sent: try: - await self.send_json({"type": "response.audio.done", "has_audio": sent_audio}) + await self.send_json({"type": "response.output_audio.done", "has_audio": sent_audio}) except Exception: - logger.exception("Failed to send response.audio.done") + logger.exception("Failed to send response.output_audio.done") while not self.audio_queue.empty(): self.audio_queue.get_nowait() diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index e05a40eef..5b30ff0fa 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -18,8 +18,8 @@ Protocol: {"type": "response.start"} {"type": "response.text.delta", "delta": "..."} {"type": "response.text.done", "text": "..."} - {"type": "response.audio.delta", "data": "...", "format": "wav"} - {"type": "response.audio.done"} + {"type": "response.output_audio.delta", "data": "...", "format": "wav"} + {"type": "response.output_audio.done"} {"type": "session.done"} {"type": "error", "message": "..."} """ diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index ab5d51c33..977a0b065 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -21,8 +21,8 @@ Protocol: {"type": "response.start"} {"type": "response.text.delta", "delta": "..."} {"type": "response.text.done", "text": "..."} - {"type": "response.audio.delta", "data": "...", "format": "wav"} - {"type": "response.audio.done"} + {"type": "response.output_audio.delta", "data": "...", "format": "wav"} + {"type": "response.output_audio.done"} {"type": "session.done"} {"type": "error", "message": "..."} """ @@ -763,7 +763,7 @@ class OmniStreamingVideoHandler: if b64: await websocket.send_json( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "data": b64, "format": "wav", } @@ -807,7 +807,7 @@ class OmniStreamingVideoHandler: if b64: await websocket.send_json( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "data": b64, "format": "wav", } @@ -816,7 +816,7 @@ class OmniStreamingVideoHandler: logger.exception("Failed to coalesce off-path audio") if audio_chunk_count > 0: - await websocket.send_json({"type": "response.audio.done"}) + await websocket.send_json({"type": "response.output_audio.done"}) response_text = "".join(text_parts) self.on_turn_complete(message_history, user_message, response_text) From c111f3a2d6048848cd239b9cbce37cd55001d57c Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:09:58 -0700 Subject: [PATCH 030/209] [Model] Add Tencent AuK speech generation and editing (encoder + diffusion pipeline) (#7385) Signed-off-by: Yueqian Lin Co-authored-by: Sy03 <1370724210@qq.com> --- docs/models/supported_models.md | 1 + .../text_to_speech/README.md | 1 + recipes/README.md | 1 + recipes/Tencent/AuK-H100.md | 167 +++++ .../models/auk/test_auk_transformer_parity.py | 439 ++++++++++++ .../models/auk/test_auk_vae_parity.py | 254 +++++++ .../diffusion/models/auk/test_pipeline_auk.py | 514 +++++++++++++ tests/e2e/offline_inference/test_auk.py | 121 ++++ .../stage_input_processors/test_auk.py | 188 +++++ tests/model_extras/test_auk_requests.py | 157 ++++ tests/model_tests/diffusion/test_alignment.py | 4 + tools/auk_parity_reference.py | 227 ++++++ tools/prepare_auk_checkpoint.py | 369 ++++++++++ vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/auk.yaml | 51 ++ vllm_omni/diffusion/models/auk/__init__.py | 16 + .../diffusion/models/auk/auk_transformer.py | 677 ++++++++++++++++++ vllm_omni/diffusion/models/auk/auk_vae.py | 505 +++++++++++++ .../diffusion/models/auk/pipeline_auk.py | 486 +++++++++++++ vllm_omni/diffusion/registry.py | 6 + .../model_executor/models/auk/__init__.py | 5 + vllm_omni/model_executor/models/auk/auk.py | 304 ++++++++ .../model_executor/models/auk/pipeline.py | 60 ++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_input_processors/auk.py | 148 ++++ vllm_omni/model_extras/auk.py | 145 ++++ .../transformers_utils/configs/__init__.py | 3 + vllm_omni/transformers_utils/configs/auk.py | 130 ++++ 28 files changed, 4986 insertions(+) create mode 100644 recipes/Tencent/AuK-H100.md create mode 100644 tests/diffusion/models/auk/test_auk_transformer_parity.py create mode 100644 tests/diffusion/models/auk/test_auk_vae_parity.py create mode 100644 tests/diffusion/models/auk/test_pipeline_auk.py create mode 100644 tests/e2e/offline_inference/test_auk.py create mode 100644 tests/model_executor/stage_input_processors/test_auk.py create mode 100644 tests/model_extras/test_auk_requests.py create mode 100644 tools/auk_parity_reference.py create mode 100755 tools/prepare_auk_checkpoint.py create mode 100644 vllm_omni/deploy/auk.yaml create mode 100644 vllm_omni/diffusion/models/auk/__init__.py create mode 100644 vllm_omni/diffusion/models/auk/auk_transformer.py create mode 100644 vllm_omni/diffusion/models/auk/auk_vae.py create mode 100644 vllm_omni/diffusion/models/auk/pipeline_auk.py create mode 100644 vllm_omni/model_executor/models/auk/__init__.py create mode 100644 vllm_omni/model_executor/models/auk/auk.py create mode 100644 vllm_omni/model_executor/models/auk/pipeline.py create mode 100644 vllm_omni/model_executor/stage_input_processors/auk.py create mode 100644 vllm_omni/model_extras/auk.py create mode 100644 vllm_omni/transformers_utils/configs/auk.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6b0952eac..1f4093828 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -68,6 +68,7 @@ th { | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — | +| `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | | `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/dots.tts-soar` | ✅︎ | | | | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index a8a273026..89ff2dbce 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -14,6 +14,7 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| +| AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | diff --git a/recipes/README.md b/recipes/README.md index ea4576201..d9f4b6822 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -84,6 +84,7 @@ recipes/ | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X | | [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB | | [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X | +| [`Tencent/AuK-H100.md`](./Tencent/AuK-H100.md) | Offline instruction-driven TTS, speech editing, enhancement and separation (AuK base and Flash) | 1x H100 80GB | | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB | | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) | diff --git a/recipes/Tencent/AuK-H100.md b/recipes/Tencent/AuK-H100.md new file mode 100644 index 000000000..5acd01cb0 --- /dev/null +++ b/recipes/Tencent/AuK-H100.md @@ -0,0 +1,167 @@ +# AuK speech generation and editing — H100 + +> Instruction-driven zero-shot TTS, instruct TTS, speech editing, enhancement +> and separation with Tencent's AuK, served as a two-stage vLLM-Omni pipeline +> (Qwen2.5-Omni thinker encoder, then a rectified-flow DiT with a BigVGAN-flow +> VAE) on one H100. + +## Summary + +- Vendor: Tencent Hunyuan +- Model: `tencent/AuK` (32-step base) and `tencent/AuK-Flash` (distilled 4-step), + both with `Qwen/Qwen2.5-Omni-3B` as the frozen encoder +- Task: zero-shot TTS, instruct TTS, content and acoustic editing, + paralinguistic editing, speech enhancement, speaker and music separation +- Mode: offline inference (online chat serving is not yet qualified) +- Hardware: 1x H100 80GB +- Maintainer: Community + +## When to use this recipe + +Use this recipe to run AuK through vLLM-Omni for any of its instruction-driven +tasks with a single shared request shape: an instruction, an optional source or +reference clip, and a target duration. The task is chosen by the instruction +text, exactly as in the upstream repository's cookbook. + +## Supported model contract + +| Task family | Instruction | Audio input | Duration | +| --- | --- | --- | --- | +| Zero-shot TTS | "Say the following with the same voice: '...'" | reference clip | `gen_seconds` required | +| Instruct TTS | `Generate speech based on the following description: "{voice description}". The content to speak is: "{text}".` (upstream template; other orderings can make the model speak the description) | none | `gen_seconds` required | +| Content / lyric editing | "Replace '...' with '...'" and similar | source clip | defaults to the source length | +| Acoustic / paralinguistic editing | pitch, speed, volume, emotion, timbre, de-accent, nonverbal, whisper | source clip | defaults to the source length (speed edits scale it) | +| Enhancement / separation | "Keep pure speech voice, remove noise and reverberation." and similar | source clip | defaults to the source length | + +| Property | Value | +| --- | --- | +| Input audio | mono or stereo (averaged), any sample rate (resampled to 24 kHz for the VAE and to 16 kHz for the encoder), one clip per request | +| Output | mono float32 waveform at 24 kHz, `ceil(gen_seconds * 50)` latent frames | +| Sequence limit | reference plus target latents up to 65536 frames (about 21 minutes) | +| Sampling knobs | `num_inference_steps` (default 32), `guidance_scale` (default 2.0), `seed`; Flash pins 4 steps, CFG 0 and its own time grid | +| Request knobs | `gen_seconds`, `sway` (default -1.0), `t_grid`, `vae_sample` via `additional_information["auk"]` | +| Concurrency | one request per DiT forward; the encoder stage runs with `max_num_seqs: 1`, prefix caching and chunked prefill off | +| Streaming | none: the ODE runs over the whole target | + +## References + +- Upstream repository and cookbook: +- Technical report: +- Weights: , , + +- Pipeline topology: `vllm_omni/model_executor/models/auk/pipeline.py`; + deploy defaults: `vllm_omni/deploy/auk.yaml` +- Request helpers: `vllm_omni/model_extras/auk.py` +- Test: `tests/e2e/offline_inference/test_auk.py` + +## Hardware + +- Accelerator model and per-device memory: NVIDIA H100 80GB +- Number of devices: 1 (both stages share the device) +- Device interconnect: not applicable +- Host memory: 64 GB is ample; the assembled directory is symlinks only +- Qualification scope: offline `Omni` inference for the four cookbook tasks + below, base and Flash checkpoints, seed-reproducible, four concurrent + requests bit-identical to sequential runs + +## Software environment + +- OS: Linux +- Python: 3.12 +- Driver / runtime: CUDA 13 driver with the vLLM 0.29.0 wheel (torch 2.13) +- vLLM version: 0.29.0 +- vLLM-Omni version or commit: branch `auk-support` on top of `aff7d6494` + +## Command + +Assemble the checkpoint directory once (symlinks by default): + +```bash +hf download tencent/AuK --local-dir ckpts/AuK +hf download Qwen/Qwen2.5-Omni-3B --local-dir ckpts/Qwen2.5-Omni-3B +python tools/prepare_auk_checkpoint.py --auk-dir ckpts/AuK --qwen-dir ckpts/Qwen2.5-Omni-3B --out ckpts/auk-omni +# For AuK-Flash: --auk-dir ckpts/AuK-Flash --out ckpts/auk-omni-flash +``` + +Offline inference (any task; the instruction selects it): + +```python +import soundfile as sf +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params + +omni = Omni(model="ckpts/auk-omni") # loads vllm_omni/deploy/auk.yaml +wav, sr = load_audio("ref.wav", sr=None) +prompt = auk_prompt( + "Say the following with the same voice: 'Ladies and gentlemen, welcome.'", + (wav, sr), + gen_seconds=4.0, +) +out = omni.generate(prompt, auk_sampling_params(nfe=32, cfg=2.0, seed=0))[0] +sf.write("out.wav", out.multimodal_output["audio"], out.multimodal_output["audio_sample_rate"]) +``` + +Text-only instruct TTS passes `None` for the audio and must set `gen_seconds`. +Editing and enhancement requests pass the source clip and may omit +`gen_seconds` to keep the source length. + +## Verification + +```bash +VLLM_OMNI_AUK_MODEL_DIR=ckpts/auk-omni python -m pytest tests/e2e/offline_inference/test_auk.py -q +``` + +The parity suites compare the port against the upstream implementation and +need reference artifacts produced by the upstream package (install it with +`pip install -e /path/to/AuK` in its own environment; it pins its own torch): + +```bash +python tools/auk_parity_reference.py --auk-repo /path/to/AuK --ckpt-dir ckpts/AuK \ + --qwen-dir ckpts/Qwen2.5-Omni-3B --out auk-parity +AUK_CKPT_DIR=ckpts/AuK AUK_REF_WAV=/path/to/AuK/assets/demo-input-audio/zero-shot-tts/ref.wav \ + python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py tests/diffusion/models/auk/test_auk_vae_parity.py -q +AUK_OMNI_CKPT_DIR=ckpts/auk-omni AUK_PARITY_REF=auk-parity/parity_ref/base \ + python -m pytest tests/diffusion/models/auk/test_pipeline_auk.py -q +``` + +Measured on one H100 against the upstream implementation with the same seeds +and mean VAE latents (four cookbook cases, base checkpoint): identical +transcripts, speaker similarity to the upstream output 0.993 to 0.999, log-mel +L1 0.11 to 0.47 (the upstream VAE-resampling noise floor is 0.39). Flash: +0.989 to 0.999 and 0.24 to 0.51. Wall per request on the second call: base +1.0 to 2.0 s for 3.5 to 10.9 s of audio, Flash 0.17 to 0.30 s; engine start +about 47 s with a warm page cache. + +## Notes + +- Memory usage: about 25 GB peak on the device for a 12 s generation with + both stages resident (deploy defaults: encoder 0.45, diffusion stage 0.35 + of device memory). +- Key flags: `enforce_eager` on both stages (the encoder walks the decoder + layers itself for the layer fusion). `enable_prefix_caching` must stay off + for the encoder: a cache hit skips prompt positions that the fused + condition needs. `enable_chunked_prefill` is off by default: forcing it + (128-token chunks, so two to three chunks per prompt) reproduces the + unchunked condition to bf16 rounding (per-token cosine 0.99999) and leaves + request walls unchanged, because the encoder is about 3% of a base request, + while concurrent runs stop being bit-identical to sequential ones. +- Known limitations: the encoder's audio tower runs without `flash_attn` in a + plain vLLM install, which shifts the encoder output on audio token positions + (per-token cosine 0.96 vs the upstream fp32 fusion; text positions 0.9999) + with waveform-level effects at the VAE noise floor; installing `flash-attn` + removes the gap. No streaming. Digits are read the way upstream reads them + (spell numbers out in the text). Instruct TTS without a reference can leak + the description into the speech, as upstream does. + +## Supported features + +| Feature | Status | Guide | +| --- | --- | --- | +| Offline `Omni.generate` | supported (base and Flash) | `docs/getting_started/quickstart.md` | +| Online `/v1/chat/completions` with audio | not yet qualified | `docs/serving/` | +| `/v1/audio/speech` | not supported (needs a TTS adapter) | `docs/contributing/model/adding_tts_model.md` | +| Streaming / async chunk | not supported | `docs/design/feature/async_chunk.md` | +| Batching across requests | one request per DiT forward | `docs/user_guide/diffusion/` | +| Tensor / sequence parallelism | not supported | `docs/configuration/composable_parallel.md` | diff --git a/tests/diffusion/models/auk/test_auk_transformer_parity.py b/tests/diffusion/models/auk/test_auk_transformer_parity.py new file mode 100644 index 000000000..937369d05 --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_transformer_parity.py @@ -0,0 +1,439 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK transformer parity: the correctness oracle for the in-tree port. + +Checks that :class:`AuKTransformer` reproduces the upstream ``auk`` reference backbone +on the released ``auk_base`` weights. Both models are built from the same +state dict and fed the same inputs, under the reference's own inference +configuration: fp32 weights with ``torch.autocast(bfloat16)``. That +configuration matters. The reference keeps its rotary frequencies in a buffer, +so casting its weights to bf16 would round both the frequencies and the +integer positions and move the numbers for reasons unrelated to the port. + +Cases: + +* ``single`` - one conditional forward +* ``cfg`` - the doubled ``cfg_infer`` forward, both branches +* ``masked`` - batch of 2 with padded targets, references and text +* ``rope`` - the local rotary implementation against ``x_transformers``, and + its survival of a cast to bf16 +* ``sample`` - an 8-step Euler trajectory against the reference backbone + driven by a transcription of upstream's ODE loop +* ``generator`` - the sampler's two noise paths agreeing +* ``cache`` - the text projection cache and ``clear_cache()`` +* ``mask flag`` - ``attn_mask_enabled=False`` reproducing the reference with + its own flag off + +This directory deliberately has no ``__init__.py``, unlike its siblings. +pytest's prepend import mode puts the nearest non-package directory on +``sys.path``, and an ``auk`` package here would shadow the upstream ``auk`` +distribution this test compares against. + +Needs the released checkpoint and a GPU, so it runs locally rather than in +CI. Point ``AUK_CKPT_DIR`` at a directory holding ``auk_base.safetensors`` and +``config.yaml``; the test skips when that is unset or when the upstream +``auk`` package is not importable:: + + AUK_CKPT_DIR=/path/to/ckpts/AuK CUDA_VISIBLE_DEVICES=6 \ + python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py -v -s +""" + +from __future__ import annotations + +import importlib.util +import math +import os +from pathlib import Path + +import pytest +import torch +import yaml + +from vllm_omni.diffusion.models.auk.auk_transformer import ( + AuKTransformer, + Rotary, + _apply_rope, + dit_state_dict, + sample_latents, +) + +CKPT_DIR = os.environ.get("AUK_CKPT_DIR") + +# Tolerances. The port evaluates rotary embeddings in fp32 and skips the +# reference's redundant uncond text projection, so the two graphs are not +# bit-identical under autocast; they agree well inside bf16 noise. +FORWARD_ATOL = 3e-2 +FORWARD_COS = 0.999 +SAMPLE_ATOL = 5e-2 +SAMPLE_COS = 0.999 +ROPE_ATOL = 1e-5 + +TARGET_FRAMES = 300 +REF_FRAMES = 150 +TEXT_TOKENS = 40 +SHORT_LENS = (220, 90, 25) # target, reference, text for the padded batch row +TIME = 0.37 +SAMPLE_STEPS = 8 +SWAY_COEF = -1.0 +SEED = 1234 +DEVICE = "cuda:0" +AUTOCAST_DTYPE = torch.bfloat16 + + +def _skip_reason() -> str | None: + if CKPT_DIR is None: + return "AUK_CKPT_DIR is not set" + if not Path(CKPT_DIR, "auk_base.safetensors").is_file(): + return f"no auk_base.safetensors under {CKPT_DIR}" + try: + # Resolving the submodule, not just "auk", also catches the shadow an + # __init__.py in this directory would create. + upstream = importlib.util.find_spec("auk.model") + except ModuleNotFoundError: + upstream = None + if upstream is None: + return "the upstream 'auk' package is not importable or is shadowed by this directory" + if not torch.cuda.is_available(): + return "no CUDA device" + return None + + +SKIP_REASON = _skip_reason() + +pytestmark = [ + pytest.mark.local_model, + pytest.mark.diffusion, + pytest.mark.gpu, + pytest.mark.cuda, + pytest.mark.skipif(SKIP_REASON is not None, reason=str(SKIP_REASON)), +] + + +def _arch() -> dict: + """Read the backbone geometry out of the checkpoint's config.""" + config = yaml.safe_load(Path(CKPT_DIR, "config.yaml").read_text())["model"] + return { + "dim": config["arch"]["dim"], + "heads": config["arch"]["heads"], + "dim_head": config["arch"].get("dim_head", 64), + "ff_mult": config["arch"]["ff_mult"], + "text_hidden_dim": config["arch"]["text_hidden_dim"], + "num_layers": config["arch"]["num_layers"], + "num_single_layers": config["arch"]["num_single_layers"], + "latent_dim": config["vae"]["latent_dim"], + } + + +def _stats(port: torch.Tensor, ref: torch.Tensor) -> tuple[float, float, float]: + """Return max abs diff, cosine similarity and the reference's RMS.""" + a = port.flatten().float() + b = ref.flatten().float() + return ( + (a - b).abs().max().item(), + torch.nn.functional.cosine_similarity(a, b, dim=0).item(), + b.pow(2).mean().sqrt().item(), + ) + + +def _assert_close(name: str, port: torch.Tensor, ref: torch.Tensor, atol: float, min_cos: float) -> None: + max_abs, cos, rms = _stats(port, ref) + print(f"[{name}] max|Δ|={max_abs:.4e} cos={cos:.8f} ref_rms={rms:.4f} rel={max_abs / rms:.3e}") + assert max_abs < atol, f"{name}: max|Δ|={max_abs:.4e} >= {atol:.1e} (reference rms {rms:.4f})" + assert cos > min_cos, f"{name}: cos={cos:.8f} <= {min_cos}" + + +def _autocast(): + return torch.autocast("cuda", dtype=AUTOCAST_DTYPE) + + +@pytest.fixture(scope="module") +def models() -> tuple[torch.nn.Module, AuKTransformer]: + """The reference and the port, same weights, fp32 on the accelerator.""" + from auk.model.flux2_edit import Flux2Edit + from safetensors.torch import load_file + + arch = _arch() + device = torch.device(DEVICE) + state = dit_state_dict(load_file(str(Path(CKPT_DIR, "auk_base.safetensors")), device="cpu")) + + reference = Flux2Edit(dropout=0.0, attn_backend="torch", attn_mask_enabled=True, **arch) + reference.load_state_dict(state, strict=True) + reference = reference.to(device).eval().requires_grad_(False) + + port = AuKTransformer(**arch) + # strict=True is the point of the port: every released tensor lands. + port.load_state_dict(state, strict=True) + port = port.to(device).eval().requires_grad_(False) + + del state + return reference, port + + +@pytest.fixture(scope="module") +def inputs() -> dict[str, torch.Tensor]: + """Fixed random inputs, as a batch of 2 whose second row is padded.""" + device = torch.device(DEVICE) + generator = torch.Generator(device="cpu").manual_seed(SEED) + arch = _arch() + + def randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device) + + def pad_mask(width: int, short: int) -> torch.Tensor: + mask = torch.zeros(2, width, dtype=torch.bool, device=device) + mask[0, :] = True + mask[1, :short] = True + return mask + + mask2 = pad_mask(TARGET_FRAMES, SHORT_LENS[0]) + ref_mask2 = pad_mask(REF_FRAMES, SHORT_LENS[1]) + c_mask2 = pad_mask(TEXT_TOKENS, SHORT_LENS[2]) + + # Padding rows are all-zero in the text encoder's output, so keep the + # inputs consistent with the masks. + x2 = randn(2, TARGET_FRAMES, arch["latent_dim"]) * mask2.unsqueeze(-1) + ref2 = randn(2, REF_FRAMES, arch["latent_dim"]) * ref_mask2.unsqueeze(-1) + text2 = randn(2, TEXT_TOKENS, arch["text_hidden_dim"]) * c_mask2.unsqueeze(-1) + + return { + "x": x2[:1].contiguous(), + "ref": ref2[:1].contiguous(), + "text": text2[:1].contiguous(), + "c_mask": c_mask2[:1].contiguous(), + "x2": x2, + "ref2": ref2, + "text2": text2, + "mask2": mask2, + "ref_mask2": ref_mask2, + "c_mask2": c_mask2, + "time": torch.tensor(TIME, dtype=torch.float32, device=device), + } + + +def test_single_forward(models, inputs) -> None: + """One conditional forward: reference against port.""" + reference, port = models + kwargs = dict(text=inputs["text"], time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"]) + with torch.no_grad(), _autocast(): + want = reference(x=inputs["x"], **kwargs) + got = port(inputs["x"], **kwargs) + + assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim) + _assert_close("single", got, want, FORWARD_ATOL, FORWARD_COS) + + +def test_cfg_forward(models, inputs) -> None: + """The doubled cfg_infer forward, cond branch first.""" + reference, port = models + kwargs = dict( + text=inputs["text"], + time=inputs["time"], + c_mask=inputs["c_mask"], + ref=inputs["ref"], + cfg_infer=True, + ) + with torch.no_grad(), _autocast(): + want = reference(x=inputs["x"], **kwargs) + got = port(inputs["x"], **kwargs) + + assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim) + _assert_close("cfg-cond", got[:1], want[:1], FORWARD_ATOL, FORWARD_COS) + _assert_close("cfg-uncond", got[1:], want[1:], FORWARD_ATOL, FORWARD_COS) + + +def test_masked_batch_forward(models, inputs) -> None: + """Batch of 2 with padded targets, references and text. + + Only valid target frames are compared. Padded frames are unconstrained: + the reference leaves them to whatever masked attention emits there. + """ + reference, port = models + kwargs = dict( + text=inputs["text2"], + time=inputs["time"], + mask=inputs["mask2"], + c_mask=inputs["c_mask2"], + ref=inputs["ref2"], + ref_mask=inputs["ref_mask2"], + ) + with torch.no_grad(), _autocast(): + want = reference(x=inputs["x2"], **kwargs) + got = port(inputs["x2"], **kwargs) + + assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim) + valid = inputs["mask2"].unsqueeze(-1).expand_as(got) + _assert_close("masked", got[valid], want[valid], FORWARD_ATOL, FORWARD_COS) + + +def _set_mask_flag(reference, port, enabled: bool) -> None: + """Flip attn_mask_enabled on both models in place, so no second copy is loaded.""" + for module in port.modules(): + if hasattr(module, "attn_mask_enabled"): + module.attn_mask_enabled = enabled + for blocks in (reference.transformer_blocks, reference.single_transformer_blocks): + for block in blocks: + block.attn.processor.attn_mask_enabled = enabled + + +def test_mask_flag_disabled(models, inputs) -> None: + """``attn_mask_enabled=False`` must track the reference with its own flag off. + + With the flag off both models run attention unmasked and still zero the + padded outputs, so the padded batch is compared in full. + """ + reference, port = models + kwargs = dict( + text=inputs["text2"], + time=inputs["time"], + mask=inputs["mask2"], + c_mask=inputs["c_mask2"], + ref=inputs["ref2"], + ref_mask=inputs["ref_mask2"], + ) + with torch.no_grad(), _autocast(): + masked = port(inputs["x2"], **kwargs) + _set_mask_flag(reference, port, False) + try: + want = reference(x=inputs["x2"], **kwargs) + got = port(inputs["x2"], **kwargs) + finally: + _set_mask_flag(reference, port, True) + + _assert_close("mask-flag-off", got, want, FORWARD_ATOL, FORWARD_COS) + # Guard against a vacuous pass: the flag has to actually change something. + assert (got.float() - masked.float()).abs().max().item() > 1e-3 + + +def test_sampler_generator_matches_seed(models, inputs) -> None: + """A generator seeded with 0 must draw exactly what ``seed=0`` draws.""" + _, port = models + common = dict( + gen_frames=TARGET_FRAMES, + nfe=2, + cfg_strength=0.0, + sway_sampling_coef=SWAY_COEF, + text=inputs["text"], + c_mask=inputs["c_mask"], + ref=inputs["ref"], + ref_mask=None, + ) + with torch.no_grad(), _autocast(): + by_seed = sample_latents(port, seed=0, **common) + by_generator = sample_latents(port, generator=torch.Generator(device=DEVICE).manual_seed(0), **common) + + print(f"[generator] max|Δ|={(by_generator.float() - by_seed.float()).abs().max().item():.4e}") + torch.testing.assert_close(by_generator, by_seed, rtol=0, atol=0) + + +def test_cache_needs_clearing(models, inputs) -> None: + """The cached text projection is reused until ``clear_cache()``. + + ``sample_latents`` clears on the way out, so a pipeline that reuses the + module across requests is safe without doing anything. A caller driving + ``forward(cache=True)`` itself has to clear when the text changes. + """ + _, port = models + other_text = inputs["text"] * 0.5 + kwargs = dict(time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"], cfg_infer=True) + + port.clear_cache() + with torch.no_grad(), _autocast(): + want = port(inputs["x"], inputs["text"], **kwargs) + port(inputs["x"], other_text, cache=True, **kwargs) + stale = port(inputs["x"], inputs["text"], cache=True, **kwargs) + port.clear_cache() + fresh = port(inputs["x"], inputs["text"], cache=True, **kwargs) + port.clear_cache() + + # The cache is real: the second call reused the first call's projection. + assert (stale.float() - want.float()).abs().max().item() > 1e-3 + torch.testing.assert_close(fresh, want, rtol=0, atol=0) + assert port.text_cond is None and port.text_uncond is None + + +def test_rope_matches_x_transformers() -> None: + """The local rotary implementation against x_transformers, in fp32.""" + from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + + device = torch.device(DEVICE) + dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES + q = torch.randn(2, 24, seq_len, dim_head, dtype=torch.float32, device=device) + + freqs = RotaryEmbedding(dim_head).to(device).forward_from_seq_len(seq_len)[0] + want = apply_rotary_pos_emb(q, freqs) + got = _apply_rope(q, Rotary(dim_head).to(device)(seq_len)) + + max_abs, cos, _ = _stats(got, want) + print(f"[rope] max|Δ|={max_abs:.4e} cos={cos:.10f}") + assert max_abs < ROPE_ATOL, f"rope: max|Δ|={max_abs:.4e} >= {ROPE_ATOL:.1e}" + + +def test_rope_survives_half_precision_cast() -> None: + """Casting the module to bf16 must not round the rotary frequencies. + + The reference stores them in a buffer, so ``.to(bfloat16)`` rounds both the + frequencies and the integer positions and silently changes every rotation. + The port recomputes them instead. + """ + dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES + exact = Rotary(dim_head)(seq_len) + got = Rotary(dim_head).to(torch.bfloat16)(seq_len) + + assert got.dtype == torch.float32 + torch.testing.assert_close(got, exact, rtol=0, atol=0) + + +def _reference_euler(reference, *, x, text, c_mask, ref, steps, cfg_strength, sway) -> torch.Tensor: + """Upstream's ODE loop, transcribed. + + ``CFMEdit.sample`` hands the backbone to ``torchdiffeq.odeint(method= + "euler")``, which on a fixed grid is ``y += (t1 - t0) * f(t0, y)``. + Constructing a real ``CFMEdit`` would drag in the Qwen-Omni text encoder, + so only the backbone here comes from upstream. + """ + t = torch.linspace(0, 1, steps + 1, device=x.device, dtype=torch.float32) + if sway is not None: + t = t + sway * (torch.cos(math.pi / 2 * t) - 1 + t) + + shared = dict(text=text, c_mask=c_mask, ref=ref, mask=None, ref_mask=None, cache=True) + for i in range(steps): + if cfg_strength < 1e-5: + v = reference(x=x, time=t[i], **shared) + else: + pred = reference(x=x, time=t[i], cfg_infer=True, **shared) + v_cond, v_uncond = torch.chunk(pred, 2, dim=0) + v = v_cond + (v_cond - v_uncond) * cfg_strength + x = x + (t[i + 1] - t[i]) * v + reference.clear_cache() + return x + + +@pytest.mark.parametrize("cfg_strength", [0.0, 2.0]) +def test_euler_sample(models, inputs, cfg_strength: float) -> None: + """An 8-step Euler trajectory, with and without guidance.""" + reference, port = models + shared = dict(text=inputs["text"], c_mask=inputs["c_mask"], ref=inputs["ref"]) + + with torch.no_grad(), _autocast(): + got = sample_latents( + port, + gen_frames=TARGET_FRAMES, + ref_mask=None, + nfe=SAMPLE_STEPS, + cfg_strength=cfg_strength, + sway_sampling_coef=SWAY_COEF, + seed=SEED, + **shared, + ) + # Re-seeding reproduces the noise sample_latents drew for itself. + torch.manual_seed(SEED) + noise = torch.randn( + TARGET_FRAMES, port.latent_dim, device=inputs["ref"].device, dtype=inputs["ref"].dtype + ).unsqueeze(0) + want = _reference_euler( + reference, x=noise, steps=SAMPLE_STEPS, cfg_strength=cfg_strength, sway=SWAY_COEF, **shared + ) + + assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim) + _assert_close(f"sample-cfg{cfg_strength:g}", got, want, SAMPLE_ATOL, SAMPLE_COS) diff --git a/tests/diffusion/models/auk/test_auk_vae_parity.py b/tests/diffusion/models/auk/test_auk_vae_parity.py new file mode 100644 index 000000000..cafc99614 --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_vae_parity.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK VAE parity: the numeric oracle for the vendored BigVGAN-flow codec. + +Checks ``AuKVAE`` against the AuK reference implementation on real speech, with +both models loaded from the same released ``vae.safetensors``: + + * encoder statistics (mean / log-std) agree, + * ``encode(sample=False)`` reproduces the reference mean path, + * ``encode(sample=True)`` reproduces ``encoding_and_normalization`` for a + matching RNG seed, + * ``decode`` reproduces ``denormalize`` + ``inference_from_latents``, in fp32 + and (reported, not gated tightly) under bf16 autocast, + * the waveform survives a mean-path round trip. + +Runs on GPU in fp32 and needs both the released checkpoint and the reference +``auk`` package importable, so it is a local_model test rather than CI. Point +``AUK_CKPT_DIR`` at the directory holding ``vae.safetensors`` and ``AUK_REF_WAV`` +at a speech clip (default ``$AUK_CKPT_DIR/ref.wav``):: + + AUK_CKPT_DIR=... AUK_REF_WAV=... python -m pytest \ + tests/diffusion/models/auk/test_auk_vae_parity.py -v -s + +Unlike its sibling directories this one carries no ``__init__.py`` on purpose: it would put +``tests/diffusion/models`` on ``sys.path`` ahead of site-packages, and this directory's own name +would then shadow the reference ``auk`` package the parity check imports. +""" + +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path + +import pytest +import soundfile as sf +import torch +import torchaudio.functional as AF + +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE + +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.gpu, pytest.mark.cuda] + +# The released model.vae.model_init_kwargs, shared by both models under test. +VAE_CONFIG = { + "upsample_rates": [5, 4, 3, 2, 2, 2], + "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4], + "upsample_initial_channel": 1536, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "downsample_rates": [2, 2, 2, 3, 4, 5], + "downsample_channels": [12, 24, 48, 96, 192, 384, 768], + "snake_logscale": True, + "latent_dim": 64, + "use_vae": True, + "causal": True, + "flow_hidden_channels": 256, + "act_causal": True, +} + +SAMPLE_RATE = 24000 +HOP_SIZE = 480 +LATENT_DIM = 64 +SEED = 1234 +# Keep the clip short: the parity run shares its GPU with other jobs. +MAX_SECONDS = 4.0 +TOLERANCE = 1e-5 +DEVICE = "cuda" + +_CKPT_DIR = os.environ.get("AUK_CKPT_DIR") +_VAE_PATH = Path(_CKPT_DIR) / "vae.safetensors" if _CKPT_DIR else None +_REF_WAV = Path(os.environ.get("AUK_REF_WAV") or (Path(_CKPT_DIR) / "ref.wav" if _CKPT_DIR else "unset")) + + +def _reference_importable() -> bool: + """Resolve ``auk.model``, not ``auk``. + + An ``__init__.py`` in this directory would make its own name resolve as ``auk`` and satisfy a + top-level probe, leaving the fixtures to fail at setup instead of skipping. The submodule + probe fails in that case, which is the honest answer. + """ + try: + return importlib.util.find_spec("auk.model") is not None + except (ImportError, ValueError): + return False + + +def _skip_reason() -> str | None: + if not _CKPT_DIR: + return "set AUK_CKPT_DIR to the directory holding vae.safetensors" + if _VAE_PATH is None or not _VAE_PATH.is_file(): + return f"no vae.safetensors under {_CKPT_DIR}" + if not _REF_WAV.is_file(): + return f"no reference clip at {_REF_WAV} (set AUK_REF_WAV)" + if not _reference_importable(): + return "cannot import the reference auk.model package (an __init__.py here would shadow it)" + if not torch.cuda.is_available(): + return "parity runs in fp32 on GPU" + return None + + +pytestmark.append(pytest.mark.skipif(_skip_reason() is not None, reason=_skip_reason() or "")) + + +def _max_abs(lhs: torch.Tensor, rhs: torch.Tensor) -> float: + return (lhs.float() - rhs.float()).abs().max().item() + + +def _snr_db(signal: torch.Tensor, estimate: torch.Tensor) -> float: + noise = (signal.float() - estimate.float()).pow(2).sum() + return (10.0 * torch.log10(signal.float().pow(2).sum() / noise)).item() + + +@pytest.fixture(scope="module") +def clip() -> torch.Tensor: + """Mono 24 kHz speech, trimmed to a whole number of latent frames.""" + audio, sample_rate = sf.read(str(_REF_WAV), dtype="float32", always_2d=True) # (T, C) + wav = torch.from_numpy(audio).mean(dim=1, keepdim=True).transpose(0, 1) + if sample_rate != SAMPLE_RATE: + wav = AF.resample(wav, sample_rate, SAMPLE_RATE) + samples = min(int(MAX_SECONDS * SAMPLE_RATE), wav.size(-1)) // HOP_SIZE * HOP_SIZE + wav = wav[:, :samples].contiguous().to(DEVICE) + print(f"\n[clip] {_REF_WAV.name}: {samples} samples, {samples / SAMPLE_RATE:.2f} s, {samples // HOP_SIZE} frames") + return wav + + +@pytest.fixture(scope="module") +def reference(clip: torch.Tensor) -> dict[str, torch.Tensor]: + """Reference tensors from the upstream model, which is then released so one model fits at a time.""" + from auk.model.vae import load_vae_model + from auk.model.vae.bigvgan_flow_vae import BigVGANFlowVAEConfig + + model = load_vae_model( + vae_name="BigVGANFlowVAE", + vae_cfg=BigVGANFlowVAEConfig.from_dict(VAE_CONFIG), + vae_ckpt=str(_VAE_PATH), + map_location="cpu", + ) + model = model.to(DEVICE).eval() + model.requires_grad_(False) + + with torch.no_grad(): + mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1) + scale = torch.sqrt(model.global_log_std.float()) + mean_latents = (mean.transpose(1, 2).float() - model.global_mean.float()) / scale + torch.manual_seed(SEED) + sampled_latents, _ = model.encoding_and_normalization(clip.unsqueeze(1)) + decoded = model.inference_from_latents(model.denormalize(mean_latents).permute(0, 2, 1)) + + tensors = { + "mean": mean.clone(), + "log_std": log_std.clone(), + "mean_latents": mean_latents.clone(), + "sampled_latents": sampled_latents.clone(), + "decoded": decoded.squeeze(1).clone(), + } + del model + torch.accelerator.empty_cache() + return tensors + + +@pytest.fixture(scope="module") +def port(reference: dict[str, torch.Tensor]) -> tuple[AuKVAE, list[str], list[str]]: + """The vendored module, loaded from the same checkpoint. Depends on ``reference`` for load ordering. + + Placed on the accelerator before loading so weight norm is folded there; see ``load_weights``. + """ + model = AuKVAE.from_config(VAE_CONFIG).to(DEVICE) + missing, unexpected = model.load_weights(_VAE_PATH) + model = model.eval() + model.requires_grad_(False) + return model, missing, unexpected + + +def test_checkpoint_covers_every_inference_tensor(port: tuple[AuKVAE, list[str], list[str]]) -> None: + model, missing, unexpected = port + print(f"[keys] missing={len(missing)} unexpected={len(unexpected)}") + print(f"[keys] unexpected prefixes={sorted({key.split('.')[0] for key in unexpected})}") + assert missing == [] + assert unexpected and all(key.startswith("flow.") for key in unexpected) + assert (model.hop_size, model.sample_rate, model.latent_dim) == (HOP_SIZE, SAMPLE_RATE, LATENT_DIM) + + +def test_encoder_statistics_match( + clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]] +) -> None: + model = port[0] + with torch.no_grad(): + mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1) + mean_diff = _max_abs(mean, reference["mean"]) + log_std_diff = _max_abs(log_std, reference["log_std"]) + print(f"[encoder] max|dmean|={mean_diff:.3e} max|dlog_std|={log_std_diff:.3e}") + assert mean_diff < TOLERANCE + assert log_std_diff < TOLERANCE + # Folded and hook-computed weights are bit-identical on one device, so the whole encoder is. + assert (mean_diff, log_std_diff) == (0.0, 0.0) + + +def test_encode_mean_path_matches( + clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]] +) -> None: + model = port[0] + with torch.no_grad(): + latents = model.encode(clip) + assert latents.shape == (1, clip.size(-1) // HOP_SIZE, LATENT_DIM) + diff = _max_abs(latents, reference["mean_latents"]) + print(f"[encode mean] shape={tuple(latents.shape)} max|d|={diff:.3e}") + assert diff == 0.0 + + +def test_encode_sampled_matches_seeded_reference( + clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]] +) -> None: + model = port[0] + generator = torch.Generator(device=DEVICE) + generator.manual_seed(SEED) + with torch.no_grad(): + latents = model.encode(clip, sample=True, generator=generator) + diff = _max_abs(latents, reference["sampled_latents"]) + spread = _max_abs(latents, reference["mean_latents"]) + print(f"[encode sample] max|d vs seeded reference|={diff:.3e} max|d vs mean path|={spread:.3e}") + assert diff < TOLERANCE + assert diff == 0.0 # a fresh device Generator replays the default one after torch.manual_seed + assert spread > TOLERANCE # the draw actually moved the latents + + +def test_decode_matches(reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]) -> None: + model = port[0] + latents = reference["mean_latents"] + with torch.no_grad(): + wav = model.decode(latents) + assert wav.shape == (1, latents.size(1) * HOP_SIZE) + diff = _max_abs(wav, reference["decoded"]) + print(f"[decode fp32] shape={tuple(wav.shape)} max|d|={diff:.3e}") + assert diff < TOLERANCE + assert diff == 0.0 + + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + wav_bf16 = model.decode(latents) + bf16_diff = _max_abs(wav_bf16, reference["decoded"]) + print(f"[decode bf16 autocast] max|d| vs fp32 reference={bf16_diff:.3e} snr={_snr_db(wav, wav_bf16):.2f} dB") + assert torch.isfinite(wav_bf16).all() + assert bf16_diff < 0.5 # sanity only: bf16 is reported, not held to fp32 parity + + +def test_round_trip_snr(clip: torch.Tensor, port: tuple[AuKVAE, list[str], list[str]]) -> None: + model = port[0] + with torch.no_grad(): + wav = model.decode(model.encode(clip)) + assert wav.shape == clip.shape + snr = _snr_db(clip, wav) + print(f"[round trip] snr={snr:.2f} dB peak={wav.abs().max().item():.4f}") + assert snr > 10.0 diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py new file mode 100644 index 000000000..261c5240d --- /dev/null +++ b/tests/diffusion/models/auk/test_pipeline_auk.py @@ -0,0 +1,514 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK stage-1 pipeline: request parsing on CPU, generation parity on a GPU. + +The CPU tests stub the transformer, the codec and the ODE so they exercise only +what the pipeline itself owns: how a request turns into a target length, a +sampling schedule and a noise generator. + +The parity test replays the saved upstream reference for the zero-shot TTS case +and reports latent, per-frame and waveform distances rather than asserting bit +equality, because two differences live in the conditioning rather than in the +port: the saved text condition is an fp32 thinker forward while the reference +generated under bf16 autocast, and the reference's stochastic VAE draw came from +the process-global RNG seeded at 1234. + +It replays with ``vae_sample=False``. That is the deliberate choice: the +reference re-seeds the global RNG immediately before drawing the ODE noise, so a +per-request generator reproduces that noise exactly (verified bit-identical) as +long as nothing else has drawn from it. Sampling the VAE posterior consumes the +generator first, which moves the noise and produces a different, equally valid +realization; comparing latents then measures the noise, not the port. The +posterior mean costs a reference latent that is 0.045 MSE (0.969 cosine) away +from the draw the reference used, which is the smaller of the two errors. +Generate the reference artifacts once with the upstream package installed +(``tools/auk_parity_reference.py --auk-repo ... --ckpt-dir ... --qwen-dir ... +--out /path/to/auk-parity``; it writes ``parity_ref/`` and +``fusion_ref``), then run:: + + AUK_OMNI_CKPT_DIR=/path/to/auk-omni-base \ + AUK_PARITY_REF=/path/to/auk-parity/parity_ref/base \ + python -m pytest -s tests/diffusion/models/auk/test_pipeline_auk.py + +``AUK_OMNI_CKPT_DIR`` is the assembled directory this pipeline loads, which is +not the same thing as the released AuK snapshot that the transformer and codec +parity tests in this directory read from ``AUK_CKPT_DIR``. ``AUK_CKPT_DIR`` is +accepted as a fallback, and the test skips rather than fails when the directory +it resolves to is not an assembled one. +""" + +from __future__ import annotations + +import json +import math +import os +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import soundfile as sf +import torch +import torch.nn.functional as F +import torchaudio +from torch import nn + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.auk import pipeline_auk +from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +LATENT_DIM = 64 +HOP = 480 +SAMPLE_RATE = 24000 +TEXT_HIDDEN_DIM = 2048 + +# The distilled student's grid, exact values from the released checkpoint. +FLASH_T_GRID = [0.0, 0.07612049579620361, 0.2928932309150696, 0.6173166036605835, 1.0] + +CKPT_DIR = os.environ.get("AUK_OMNI_CKPT_DIR") or os.environ.get("AUK_CKPT_DIR") +PARITY_REF = os.environ.get("AUK_PARITY_REF") +# An assembled directory is the only thing this pipeline can load. +IS_ASSEMBLED = bool(CKPT_DIR) and (Path(CKPT_DIR).expanduser() / "config.json").is_file() + +# Measured on the released base checkpoint: mean frame cosine 0.961 to 0.968, +# relative latent MSE 0.051 to 0.063, log-mel distance 0.34. That spread is +# bf16 against fp32 DiT weights plus where the codec folds its weight norm, +# neither of which is worth gating on. A port that has lost the conditioning +# sits at 0.110 cosine, 1.82 relative MSE and 2.88 log-mel, so these gates +# still separate the two by an order of magnitude. +PARITY_MIN_FRAME_COSINE = 0.93 +PARITY_MAX_RELATIVE_MSE = 0.15 +PARITY_MAX_MEL_DISTANCE = 0.6 + + +class _StubTransformer(nn.Module): + """Records its construction arguments; owns one weight so strict loading runs.""" + + def __init__( + self, + *, + dim: int, + heads: int, + dim_head: int, + ff_mult: float, + latent_dim: int, + text_hidden_dim: int, + num_layers: int, + num_single_layers: int, + attn_mask_enabled: bool = True, + ) -> None: + super().__init__() + self.init_kwargs = { + "dim": dim, + "heads": heads, + "dim_head": dim_head, + "ff_mult": ff_mult, + "latent_dim": latent_dim, + "text_hidden_dim": text_hidden_dim, + "num_layers": num_layers, + "num_single_layers": num_single_layers, + "attn_mask_enabled": attn_mask_enabled, + } + self.latent_dim = latent_dim + self.proj = nn.Linear(latent_dim, latent_dim, bias=False) + + +class _StubVAE(nn.Module): + """Deterministic stand-in for the codec, one latent frame per hop.""" + + hop_size = HOP + sample_rate = SAMPLE_RATE + latent_dim = LATENT_DIM + + def __init__(self) -> None: + super().__init__() + self.encode_calls: list[dict[str, Any]] = [] + self.decode_calls = 0 + self.weights_path: str | None = None + self.config: dict[str, Any] = {} + + @classmethod + def from_config(cls, cfg: dict[str, Any]) -> _StubVAE: + vae = cls() + vae.config = dict(cfg) + return vae + + def load_weights(self, path: str, **_: Any) -> tuple[list[str], list[str]]: + self.weights_path = str(path) + return [], [] + + def encode( + self, + wav: torch.Tensor, + *, + sample: bool = False, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + frames = wav.shape[-1] // HOP + self.encode_calls.append({"samples": int(wav.shape[-1]), "sample": sample, "generator": generator}) + if sample: + return torch.randn(1, frames, LATENT_DIM, generator=generator, device=wav.device) + return torch.zeros(1, frames, LATENT_DIM, device=wav.device) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + self.decode_calls += 1 + return torch.zeros(1, latents.shape[1] * HOP, device=latents.device) + + +def _stub_sampler(calls: list[dict[str, Any]]): + """Record the ODE arguments and return noise drawn from the request generator.""" + + def sample_latents(dit: nn.Module, **kwargs: Any) -> torch.Tensor: + calls.append(kwargs) + shape = (1, kwargs["gen_frames"], kwargs["latent_dim"]) + generator = kwargs.get("generator") + if generator is None: + return torch.zeros(*shape, device=kwargs["device"], dtype=kwargs["dtype"]) + return torch.randn(*shape, generator=generator, device=kwargs["device"], dtype=kwargs["dtype"]) + + return sample_latents + + +def _write_checkpoint(root: Path, variant: str) -> Path: + """Write a tiny assembled checkpoint the stubs can load.""" + + from safetensors.torch import save_file + + model_dir = root / f"auk-{variant}" + model_dir.mkdir(parents=True, exist_ok=True) + config = { + "model_type": "auk", + "architectures": ["AuKForConditionalGeneration"], + "dit": { + "dim": 16, + "heads": 2, + "dim_head": 8, + "ff_mult": 2, + "text_hidden_dim": TEXT_HIDDEN_DIM, + "num_layers": 1, + "num_single_layers": 1, + "attn_mask_enabled": True, + }, + "vae": { + "latent_dim": LATENT_DIM, + "downsample_rate": HOP, + "target_sample_rate": SAMPLE_RATE, + "model_init_kwargs": {"latent_dim": LATENT_DIM}, + }, + "variant": variant, + "flash_t_grid": FLASH_T_GRID, + "defaults": {"nfe": 32, "cfg": 2.0, "sway": -1.0}, + } + (model_dir / "config.json").write_text(json.dumps(config)) + save_file( + { + "transformer.proj.weight": torch.zeros(LATENT_DIM, LATENT_DIM), + "layer_weights": torch.zeros(36), + "layer_scale": torch.ones(1), + }, + str(model_dir / "auk.safetensors"), + ) + save_file({"global_mean": torch.zeros(LATENT_DIM)}, str(model_dir / "vae.safetensors")) + return model_dir + + +@pytest.fixture +def build_pipeline(tmp_path, monkeypatch): + """Build an AuKPipeline whose transformer, codec and ODE are stubs.""" + + monkeypatch.setattr(pipeline_auk, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(pipeline_auk, "AuKTransformer", _StubTransformer) + monkeypatch.setattr(pipeline_auk, "AuKVAE", _StubVAE) + + def _build(variant: str = "base") -> tuple[AuKPipeline, list[dict[str, Any]]]: + calls: list[dict[str, Any]] = [] + monkeypatch.setattr(pipeline_auk, "sample_latents", _stub_sampler(calls)) + od_config = OmniDiffusionConfig( + model=str(_write_checkpoint(tmp_path, variant)), + dtype=torch.float32, + model_class_name="AuKPipeline", + ) + return AuKPipeline(od_config=od_config), calls + + return _build + + +def _prompt( + *, + tokens: int = 8, + audio: Any = None, + knobs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build the diffusion prompt the stage input processor would emit.""" + + prompt: dict[str, Any] = { + "prompt": "", + "prompt_embeds": torch.zeros(tokens, TEXT_HIDDEN_DIM), + "additional_information": {"auk": dict(knobs or {})}, + } + if audio is not None: + prompt["multi_modal_data"] = {"audio": audio} + return prompt + + +def _batch(prompt: dict[str, Any], **sampling: Any) -> DiffusionRequestBatch: + request = OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(**sampling), + request_id="auk-test-0", + ) + return DiffusionRequestBatch(requests=[request]) + + +def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> tuple[np.ndarray, int]: + return np.zeros(int(seconds * sample_rate), dtype=np.float32), sample_rate + + +@pytest.mark.core_model +@pytest.mark.cpu +class TestRequestParsing: + def test_pipeline_declares_audio_output(self, build_pipeline): + pipeline, _ = build_pipeline() + + assert pipeline.support_audio_output is True + assert pipeline.audio_sample_rate == SAMPLE_RATE + assert pipeline.supports_request_batch is False + # The warmup request cannot carry an encoder-stage text condition. + assert pipeline.dummy_run_num_frames == 0 + + def test_config_reaches_the_transformer(self, build_pipeline): + pipeline, _ = build_pipeline() + + assert pipeline.dit.init_kwargs["latent_dim"] == LATENT_DIM + assert pipeline.dit.init_kwargs["text_hidden_dim"] == TEXT_HIDDEN_DIM + assert pipeline.dit.init_kwargs["attn_mask_enabled"] is True + + def test_gen_seconds_sets_the_target_length(self, build_pipeline): + pipeline, calls = build_pipeline() + + outputs = pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": 6.0}), seed=1)) + + assert calls[0]["gen_frames"] == math.ceil(6.0 * SAMPLE_RATE / HOP) + assert calls[0]["ref"].shape == (1, 100, LATENT_DIM) + assert calls[0]["ref_mask"].shape == (1, 100) + assert calls[0]["c_mask"].shape == (1, 8) + assert outputs[0].output.shape == (300 * HOP,) + assert outputs[0].output.dtype is torch.float32 + + def test_target_length_falls_back_to_the_source_clip(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": None}), seed=1)) + + assert calls[0]["gen_frames"] == 100 + + def test_text_only_request_without_a_duration_is_rejected(self, build_pipeline): + pipeline, calls = build_pipeline() + + with pytest.raises(ValueError, match="gen_seconds"): + pipeline.forward(_batch(_prompt(knobs={"gen_seconds": None}), seed=1)) + assert calls == [] + + def test_text_only_request_gets_an_empty_reference(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 3.5}), seed=1)) + + assert calls[0]["ref"].shape == (1, 0, LATENT_DIM) + assert calls[0]["gen_frames"] == 175 + + def test_source_clip_is_resampled_to_the_codec_rate(self, build_pipeline): + pipeline, _ = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(1.0, 16000), knobs={"gen_seconds": 1.0}), seed=1)) + + assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE + + def test_stereo_source_is_mixed_to_mono(self, build_pipeline): + pipeline, _ = build_pipeline() + stereo = np.zeros((2, SAMPLE_RATE), dtype=np.float32) + + pipeline.forward(_batch(_prompt(audio=(stereo, SAMPLE_RATE), knobs={"gen_seconds": 1.0}), seed=1)) + + assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE + + def test_missing_text_condition_is_rejected(self, build_pipeline): + pipeline, _ = build_pipeline() + prompt = _prompt(knobs={"gen_seconds": 1.0}) + del prompt["prompt_embeds"] + + with pytest.raises(ValueError, match="prompt_embeds"): + pipeline.forward(_batch(prompt, seed=1)) + + def test_base_variant_uses_the_requested_schedule(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward( + _batch( + _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}), + seed=1, + num_inference_steps=16, + guidance_scale=3.0, + ) + ) + + assert calls[0]["nfe"] == 16 + assert calls[0]["cfg_strength"] == 3.0 + assert calls[0]["sway_sampling_coef"] == -1.0 + assert calls[0]["t_grid"] is None + + def test_schedule_defaults_come_from_the_checkpoint(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": None}), seed=1)) + + assert calls[0]["nfe"] == 32 + assert calls[0]["cfg_strength"] == 2.0 + assert calls[0]["sway_sampling_coef"] == -1.0 + + def test_flash_variant_locks_the_distilled_recipe(self, build_pipeline): + pipeline, calls = build_pipeline("flash") + + pipeline.forward( + _batch( + _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}), + seed=1, + num_inference_steps=32, + guidance_scale=2.0, + ) + ) + + assert calls[0]["nfe"] == 4 + assert calls[0]["cfg_strength"] == 0.0 + assert calls[0]["sway_sampling_coef"] is None + assert calls[0]["t_grid"] == FLASH_T_GRID + + def test_noise_comes_from_a_seeded_request_generator(self, build_pipeline): + pipeline, calls = build_pipeline() + + outputs = pipeline.forward( + _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=0, output_type="latent"), + ) + + expected = torch.randn(1, 50, LATENT_DIM, generator=torch.Generator().manual_seed(0)) + assert torch.equal(outputs[0].output, expected) + # The global RNG is never seeded on the pipeline's behalf. + assert calls[0]["seed"] is None + assert isinstance(calls[0]["generator"], torch.Generator) + + def test_latent_output_type_skips_the_decoder(self, build_pipeline): + pipeline, _ = build_pipeline() + + outputs = pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1, output_type="latent")) + + assert outputs[0].output.shape == (1, 50, LATENT_DIM) + assert pipeline.vae.decode_calls == 0 + + def test_vae_sample_draws_from_the_request_generator(self, build_pipeline): + pipeline, _ = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)) + + call = pipeline.vae.encode_calls[0] + assert call["sample"] is True + assert isinstance(call["generator"], torch.Generator) + + def test_one_request_per_forward(self, build_pipeline): + pipeline, _ = build_pipeline() + batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1) + batch.requests.append(batch.requests[0]) + + with pytest.raises(AssertionError, match="one request per forward"): + pipeline.forward(batch) + + +def _reference_audio_path(messages: list[dict[str, Any]]) -> str: + for message in messages: + for item in message.get("content") or (): + if isinstance(item, dict) and item.get("type") == "audio": + path = item.get("audio") or item.get("audio_url") + if path: + return str(path) + raise AssertionError("the saved reference case carries no source clip") + + +def _log_mel_distance(left: torch.Tensor, right: torch.Tensor, sample_rate: int) -> float: + """Mean absolute log-mel difference between two mono waveforms.""" + + mel = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate, n_fft=1024, hop_length=256, n_mels=80) + length = min(left.shape[-1], right.shape[-1]) + spectra = [torch.log(mel(wav[..., :length].float()) + 1e-5) for wav in (left, right)] + return float((spectra[0] - spectra[1]).abs().mean()) + + +@pytest.mark.local_model +@pytest.mark.diffusion +@pytest.mark.skipif( + not (IS_ASSEMBLED and PARITY_REF), + reason="needs AUK_PARITY_REF and AUK_OMNI_CKPT_DIR pointing at an assembled checkpoint directory", +) +def test_zero_shot_tts_tracks_the_upstream_reference(): + """Replay the saved zero-shot TTS case and report the distances.""" + + parity_dir = Path(PARITY_REF) + fusion_dir = Path(os.environ.get("AUK_FUSION_REF") or parity_dir.parent.parent / "fusion_ref") + reference = torch.load(parity_dir / "zs_tts_en.pt", map_location="cpu", weights_only=False) + fusion = torch.load(fusion_dir / "zs_tts_audio.pt", map_location="cpu", weights_only=False) + + # soundfile rather than torchaudio.load: the repo's test convention, and + # torchaudio's decoder needs an ffmpeg the port environment does not have. + samples, source_rate = sf.read(_reference_audio_path(reference["messages"]), dtype="float32", always_2d=True) + source = torch.from_numpy(samples).transpose(0, 1) + knobs = {"gen_seconds": reference["gen_seconds"], "sway": -1.0, "t_grid": None, "vae_sample": False} + prompt = { + "prompt": "", + # The encoder stage emits bf16; the saved fusion is an fp32 forward. + "prompt_embeds": fusion["fused"].to(torch.bfloat16), + "multi_modal_data": {"audio": (source, int(source_rate))}, + "additional_information": {"auk": knobs}, + } + od_config = OmniDiffusionConfig(model=CKPT_DIR, dtype=torch.bfloat16, model_class_name="AuKPipeline") + pipeline = AuKPipeline(od_config=od_config) + + def run(output_type: str) -> torch.Tensor: + batch = _batch( + dict(prompt), + seed=reference["seed"], + num_inference_steps=32, + guidance_scale=2.0, + output_type=output_type, + ) + return pipeline.forward(batch)[0].output + + # Two deterministic runs: the ODE is seeded, so the latents belong to the + # waveform even though the decode happens in the second call. + latents = run("latent") + waveform = run("pt") + + ref_frames = int(reference["ref_latent_lens"][0]) + expected = reference["generated"][:, ref_frames:, :].float() + assert latents.shape == expected.shape + assert torch.isfinite(latents).all() + assert torch.isfinite(waveform).all() + assert waveform.shape[-1] == expected.shape[1] * HOP + + cosine = F.cosine_similarity(latents[0], expected[0], dim=-1) + metrics = { + "gen_frames": int(expected.shape[1]), + "ref_frames": ref_frames, + "latent_mse": float(torch.mean((latents - expected) ** 2)), + "latent_relative_mse": float(torch.mean((latents - expected) ** 2) / expected.var()), + "frame_cosine_mean": float(cosine.mean()), + "frame_cosine_min": float(cosine.min()), + "log_mel_distance": _log_mel_distance(waveform, reference["audio"].reshape(-1), int(reference["sr"])), + "rms_port": float(waveform.pow(2).mean().sqrt()), + "rms_reference": float(reference["audio"].reshape(-1).pow(2).mean().sqrt()), + } + print("\nAuK zero-shot TTS parity:", json.dumps(metrics, indent=1)) + + assert metrics["frame_cosine_mean"] > PARITY_MIN_FRAME_COSINE + assert metrics["latent_relative_mse"] < PARITY_MAX_RELATIVE_MSE + assert metrics["log_mel_distance"] < PARITY_MAX_MEL_DISTANCE diff --git a/tests/e2e/offline_inference/test_auk.py b/tests/e2e/offline_inference/test_auk.py new file mode 100644 index 000000000..746af7bd2 --- /dev/null +++ b/tests/e2e/offline_inference/test_auk.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Offline E2E smoke test for AuK (encoder stage -> diffusion stage). + +AuK ships as ``config.yaml`` plus loose safetensors and needs the Qwen2.5-Omni-3B +snapshot next to it, so the test expects an assembled directory produced by +``tools/prepare_auk_checkpoint.py`` and skips when it is not available: + + VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk.py + +Two requests are exercised: zero-shot TTS from a reference clip and text-only +instruct TTS. The checks are structural (duration, sample rate, level, +finiteness); fidelity against the upstream implementation is covered by the +parity suites under ``tests/diffusion/models/auk``. +""" + +from __future__ import annotations + +import functools +import math +import os +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.model_extras.auk import HOP, SAMPLE_RATE, auk_prompt, auk_sampling_params + +MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR" + +# Any clean speech clip works as the cloning reference; reuse the vendored +# CosyVoice3 prompt so no new binary asset is needed. +REFERENCE_WAV_PATH = get_asset_path("cosyvoice3/zero_shot_prompt.wav") +ZERO_SHOT_TEXT = "The quick brown fox jumps over the lazy dog." +ZERO_SHOT_SECONDS = 3.0 +# The upstream cookbook template: description first, then the content. +INSTRUCT_TEXT = ( + 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". ' + 'The content to speak is: "Welcome back, how was your day?".' +) +INSTRUCT_SECONDS = 2.5 +# 8 Euler steps keep the test short; the base recipe is 32. +TEST_NFE = 8 + +_model_dir = os.environ.get(MODEL_DIR_ENV) +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"), + pytest.mark.parametrize( + "omni_runner", + [pytest.param((str(Path(_model_dir or ".").resolve()), get_deploy_config_path("auk.yaml"), {}), id="auk")], + indirect=True, + ), +] + + +@functools.lru_cache(maxsize=1) +def _load_reference_wav() -> tuple[np.ndarray, int]: + audio, sr = sf.read(str(REFERENCE_WAV_PATH), dtype="float32", always_2d=False) + if isinstance(audio, np.ndarray) and audio.ndim > 1: + audio = np.mean(audio, axis=-1) + return np.asarray(audio, dtype=np.float32), int(sr) + + +def _audio_from_output(output) -> tuple[np.ndarray, int]: + mm = getattr(output, "multimodal_output", None) or {} + assert "audio" in mm, f"no audio in the final output; keys={sorted(mm)}" + audio = mm["audio"] + if isinstance(audio, list): + audio = np.concatenate([np.asarray(a, dtype=np.float32).reshape(-1) for a in audio]) + if hasattr(audio, "detach"): + audio = audio.detach().cpu().float().numpy() + audio = np.asarray(audio, dtype=np.float32).reshape(-1) + sr = mm.get("audio_sample_rate", SAMPLE_RATE) + if isinstance(sr, list) and sr: + sr = sr[-1] + if hasattr(sr, "item"): + sr = sr.item() + return audio, int(sr) + + +def _check_waveform(audio: np.ndarray, sr: int, *, gen_seconds: float) -> None: + assert sr == SAMPLE_RATE + expected = math.ceil(gen_seconds * SAMPLE_RATE / HOP) * HOP + assert audio.size == expected, f"expected {expected} samples for {gen_seconds}s, got {audio.size}" + assert np.isfinite(audio).all(), "non-finite samples" + assert np.abs(audio).max() <= 1.0 + 1e-6, "audio exceeds full scale" + assert float(np.sqrt(np.mean(audio**2))) > 1e-3, "audio is silent" + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_auk_zero_shot_tts(omni_runner: OmniRunner) -> None: + """A reference clip plus an instruction yields speech of the requested length.""" + wav, sr = _load_reference_wav() + prompt = auk_prompt( + f"Say the following with the same voice: '{ZERO_SHOT_TEXT}'", + (wav, sr), + gen_seconds=ZERO_SHOT_SECONDS, + ) + outputs = omni_runner.omni.generate(prompt, auk_sampling_params(nfe=TEST_NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + audio, out_sr = _audio_from_output(outputs[0]) + _check_waveform(audio, out_sr, gen_seconds=ZERO_SHOT_SECONDS) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_auk_instruct_tts_text_only(omni_runner: OmniRunner) -> None: + """A text-only request (no reference audio) takes the no-prompt-audio path.""" + prompt = auk_prompt(INSTRUCT_TEXT, None, gen_seconds=INSTRUCT_SECONDS) + assert prompt["prompt"].rstrip().endswith("<|im_start|>assistant\n".rstrip()) + assert "||" in prompt["prompt"] + outputs = omni_runner.omni.generate(prompt, auk_sampling_params(nfe=TEST_NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + audio, out_sr = _audio_from_output(outputs[0]) + _check_waveform(audio, out_sr, gen_seconds=INSTRUCT_SECONDS) diff --git a/tests/model_executor/stage_input_processors/test_auk.py b/tests/model_executor/stage_input_processors/test_auk.py new file mode 100644 index 000000000..b2904e1f7 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_auk.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the AuK encoder -> diffusion stage input processor.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.auk import encoder2dit + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +HIDDEN = 2048 + + +def _request_output( + mm_output: Any, + completion_mm_output: Any = None, + request_id: str | None = None, +) -> SimpleNamespace: + """Build a minimal encoder-stage output mock.""" + return SimpleNamespace( + outputs=[SimpleNamespace(token_ids=[0], multimodal_output=completion_mm_output)], + multimodal_output=mm_output, + request_id=request_id, + ) + + +def _payload(text_cond: Any) -> dict[str, Any]: + """The slot the encoder actually emits: hidden_states.output.""" + return {"hidden_states": {"output": text_cond}} + + +def _prompt( + *, + audio: bool = False, + gen_seconds: float | None = None, + has_audio: bool | None = None, +) -> dict[str, Any]: + knobs: dict[str, Any] = { + "gen_seconds": gen_seconds, + "sway": -1.0, + "t_grid": None, + "vae_sample": False, + } + if has_audio is not None: + knobs["has_audio"] = has_audio + prompt: dict[str, Any] = { + "prompt": "<|im_start|>system\n", + "additional_information": {"auk": knobs}, + } + if audio: + prompt["multi_modal_data"] = {"audio": (np.zeros(24000, dtype=np.float32), 24000)} + return prompt + + +class TestTextCondExtraction: + def test_tensor_output(self): + cond = torch.zeros(7, HIDDEN, dtype=torch.bfloat16) + out = encoder2dit([_request_output(_payload(cond))], _prompt()) + assert out is not None + assert out["prompt"] == "" + assert out["prompt_embeds"].shape == (7, HIDDEN) + assert out["prompt_embeds"].dtype == torch.bfloat16 + assert not out["prompt_embeds"].is_cuda + + def test_list_of_tensors_is_concatenated(self): + parts = [torch.zeros(3, HIDDEN), torch.ones(4, HIDDEN)] + out = encoder2dit([_request_output(_payload(parts))], _prompt()) + assert out["prompt_embeds"].shape == (7, HIDDEN) + assert torch.equal(out["prompt_embeds"][3:], torch.ones(4, HIDDEN)) + + def test_single_element_list(self): + out = encoder2dit([_request_output(_payload([torch.zeros(5, HIDDEN)]))], _prompt()) + assert out["prompt_embeds"].shape == (5, HIDDEN) + + def test_numpy_output(self): + cond = np.zeros((4, HIDDEN), dtype=np.float32) + out = encoder2dit([_request_output(_payload(cond))], _prompt()) + assert isinstance(out["prompt_embeds"], torch.Tensor) + assert out["prompt_embeds"].shape == (4, HIDDEN) + + def test_completion_output_fallback(self): + cond = torch.zeros(6, HIDDEN) + out = encoder2dit( + [_request_output(None, completion_mm_output=_payload(cond))], + _prompt(), + ) + assert out["prompt_embeds"].shape == (6, HIDDEN) + + def test_hidden_states_as_a_struct(self): + """A deserialized payload arrives as a struct, not a mapping.""" + cond = torch.zeros(9, HIDDEN) + mm_output = {"hidden_states": SimpleNamespace(output=cond, layers=None)} + out = encoder2dit([_request_output(mm_output)], _prompt()) + assert out["prompt_embeds"].shape == (9, HIDDEN) + + def test_flattened_payload_key(self): + cond = torch.zeros(8, HIDDEN) + out = encoder2dit([_request_output({"hidden_states.output": cond})], _prompt()) + assert out["prompt_embeds"].shape == (8, HIDDEN) + + def test_direct_text_cond_key(self): + cond = torch.zeros(2, HIDDEN) + out = encoder2dit([_request_output({"text_cond": cond})], _prompt()) + assert out["prompt_embeds"].shape == (2, HIDDEN) + + def test_no_source_outputs_returns_none(self): + assert encoder2dit([], _prompt()) is None + + def test_missing_payload_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output({})], _prompt()) + + def test_empty_hidden_states_slot_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output({"hidden_states": {}})], _prompt()) + + def test_empty_multimodal_output_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output(None)], _prompt()) + + def test_empty_list_raises(self): + with pytest.raises(ValueError, match="empty text condition"): + encoder2dit([_request_output(_payload([]))], _prompt()) + + def test_wrong_rank_raises(self): + with pytest.raises(ValueError, match=r"\[tokens, hidden\]"): + encoder2dit([_request_output(_payload(torch.zeros(2, 7, HIDDEN)))], _prompt()) + + +class TestKnobPropagation: + def test_knobs_are_forwarded(self): + prompt = _prompt(audio=True, gen_seconds=6.0) + prompt["additional_information"]["auk"]["t_grid"] = [0.0, 0.5, 1.0] + prompt["additional_information"]["auk"]["vae_sample"] = True + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + knobs = out["additional_information"]["auk"] + assert knobs["gen_seconds"] == 6.0 + assert knobs["sway"] == -1.0 + assert knobs["t_grid"] == [0.0, 0.5, 1.0] + assert knobs["vae_sample"] is True + + def test_sampling_knobs_are_not_copied(self): + """nfe/cfg/seed ride the stage-1 sampling params, not the prompt.""" + prompt = _prompt() + prompt["additional_information"]["auk"].update({"nfe": 8, "cfg": 1.0, "seed": 3}) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + knobs = out["additional_information"]["auk"] + assert set(knobs) == {"gen_seconds", "sway", "t_grid", "vae_sample", "has_audio"} + + def test_multi_modal_data_is_forwarded(self): + prompt = _prompt(audio=True) + out = encoder2dit( + [_request_output(_payload(torch.zeros(3, HIDDEN)))], + prompt, + requires_multimodal_data=True, + ) + assert out["multi_modal_data"] is prompt["multi_modal_data"] + + def test_text_only_request_has_no_multi_modal_data(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], _prompt()) + assert out["multi_modal_data"] is None + assert out["additional_information"]["auk"]["has_audio"] is False + + def test_has_audio_declared_by_the_prompt_wins(self): + prompt = _prompt(audio=True, has_audio=True) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + assert out["additional_information"]["auk"]["has_audio"] is True + + def test_has_audio_inferred_from_multi_modal_data(self): + prompt = _prompt(audio=True) + prompt["additional_information"]["auk"].pop("has_audio", None) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + assert out["additional_information"]["auk"]["has_audio"] is True + + def test_prompt_as_list(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], [_prompt(gen_seconds=2.0)]) + assert out["additional_information"]["auk"]["gen_seconds"] == 2.0 + + def test_missing_prompt_yields_empty_knobs(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], None) + knobs = out["additional_information"]["auk"] + assert knobs["gen_seconds"] is None + assert knobs["has_audio"] is False diff --git a/tests/model_extras/test_auk_requests.py b/tests/model_extras/test_auk_requests.py new file mode 100644 index 000000000..1f319d092 --- /dev/null +++ b/tests/model_extras/test_auk_requests.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the AuK request helpers.""" + +from __future__ import annotations + +import numpy as np +import pytest +from vllm import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras.auk import ( + NO_PROMPT_AUDIO, + auk_prompt, + auk_sampling_params, + resolve_gen_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +INSTRUCTION = "Say the following with the same voice: 'Ladies and gentlemen'" + +EXPECTED_WITH_AUDIO = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + f"{INSTRUCTION}<|audio_bos|><|AUDIO|><|audio_eos|><|im_end|>\n" + "<|im_start|>assistant\n" +) + +EXPECTED_TEXT_ONLY = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + f"{INSTRUCTION}{NO_PROMPT_AUDIO}<|im_end|>\n" + "<|im_start|>assistant\n" +) + + +def _audio() -> tuple[np.ndarray, int]: + return np.zeros(24000, dtype=np.float32), 24000 + + +class TestPromptRendering: + def test_with_audio(self): + audio = _audio() + prompt = auk_prompt(INSTRUCTION, audio, gen_seconds=6.0) + assert prompt["prompt"] == EXPECTED_WITH_AUDIO + assert prompt["multi_modal_data"] == {"audio": audio} + assert prompt["additional_information"]["auk"]["has_audio"] is True + + def test_text_only_appends_the_marker(self): + prompt = auk_prompt(INSTRUCTION, gen_seconds=6.0) + assert prompt["prompt"] == EXPECTED_TEXT_ONLY + assert "multi_modal_data" not in prompt + assert prompt["additional_information"]["auk"]["has_audio"] is False + + def test_marker_is_not_appended_twice(self): + once = auk_prompt(INSTRUCTION, gen_seconds=6.0) + twice = auk_prompt(INSTRUCTION + NO_PROMPT_AUDIO, gen_seconds=6.0) + assert once["prompt"] == twice["prompt"] + assert twice["prompt"].count(NO_PROMPT_AUDIO) == 1 + + def test_marker_is_not_added_when_audio_is_present(self): + prompt = auk_prompt(INSTRUCTION, _audio()) + assert NO_PROMPT_AUDIO not in prompt["prompt"] + + def test_knob_defaults(self): + knobs = auk_prompt(INSTRUCTION, _audio())["additional_information"]["auk"] + assert knobs == { + "gen_seconds": None, + "sway": -1.0, + "t_grid": None, + "vae_sample": False, + "has_audio": True, + } + + def test_knobs_are_carried(self): + grid = [0.0, 0.25, 1.0] + knobs = auk_prompt( + INSTRUCTION, + _audio(), + gen_seconds=3.5, + sway=0.0, + t_grid=grid, + vae_sample=True, + )["additional_information"]["auk"] + assert knobs["gen_seconds"] == 3.5 + assert knobs["sway"] == 0.0 + assert knobs["vae_sample"] is True + assert knobs["t_grid"] == grid + # A copy, so a caller mutating its list cannot reach the request. + assert knobs["t_grid"] is not grid + + def test_sampling_knobs_are_absent_from_the_prompt(self): + knobs = auk_prompt(INSTRUCTION, _audio())["additional_information"]["auk"] + assert not {"nfe", "cfg", "seed"} & set(knobs) + + +class TestSamplingParams: + def test_two_stages_in_order(self): + params = auk_sampling_params() + assert len(params) == 2 + assert isinstance(params[0], SamplingParams) + assert isinstance(params[1], OmniDiffusionSamplingParams) + + def test_encoder_stage_is_prefill_only(self): + encoder = auk_sampling_params()[0] + assert encoder.max_tokens == 1 + assert encoder.temperature == 0.0 + assert encoder.detokenize is False + + def test_defaults_are_the_base_recipe(self): + diffusion = auk_sampling_params()[1] + assert diffusion.num_inference_steps == 32 + assert diffusion.guidance_scale == 2.0 + assert diffusion.seed == 0 + + def test_overrides(self): + encoder, diffusion = auk_sampling_params(nfe=4, cfg=0.0, seed=None) + assert diffusion.num_inference_steps == 4 + assert diffusion.guidance_scale == 0.0 + assert diffusion.seed is None + assert encoder.seed is None + + +class TestDurationPolicy: + @pytest.mark.parametrize( + ("gen_seconds", "expected"), + [ + (6.0, 300), + (0.02, 1), + (0.001, 1), + (5.59, 280), + (1.0, 50), + ], + ) + def test_explicit_duration_rounds_up(self, gen_seconds: float, expected: int): + assert resolve_gen_frames(gen_seconds, 0) == expected + + def test_explicit_duration_wins_over_the_source(self): + assert resolve_gen_frames(2.0, 500) == 100 + + def test_source_length_is_the_default(self): + assert resolve_gen_frames(None, 137) == 137 + + def test_text_only_without_a_duration_raises(self): + with pytest.raises(ValueError, match="gen_seconds is required"): + resolve_gen_frames(None, 0) + + @pytest.mark.parametrize("gen_seconds", [0.0, -1.0]) + def test_non_positive_duration_raises(self, gen_seconds: float): + with pytest.raises(ValueError, match="must be positive"): + resolve_gen_frames(gen_seconds, 100) + + def test_custom_rate(self): + assert resolve_gen_frames(1.0, 0, sample_rate=16000, hop=320) == 50 diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 183ef8659..f0a1285bc 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -83,6 +83,10 @@ EXCLUDED_MODELS = [ "Gr00tN1d7Pipeline", "Pi0Pipeline", "SanaWmPipeline", + # Audio pipeline conditioned by a separate Qwen2.5-Omni encoder stage; + # covered by tests/diffusion/models/auk (CPU pipeline tests, parity) and + # the real-checkpoint E2E test instead of a tiny model. + "AuKPipeline", ] diff --git a/tools/auk_parity_reference.py b/tools/auk_parity_reference.py new file mode 100644 index 000000000..33b664232 --- /dev/null +++ b/tools/auk_parity_reference.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generate the upstream AuK reference artifacts consumed by the parity tests. + +This runs the *upstream* implementation, so it must be executed in an +environment where the upstream package is installed +(``pip install -e /path/to/AuK``; it pins its own torch/transformers line and +does not need vLLM-Omni). It writes: + + /parity_ref//.pt, .wav, meta.json + Upstream generations for four cookbook cases (zero-shot TTS, instruct + TTS, content edit, denoise) with the reference latents, the generated + latents and the waveform. Point ``AUK_PARITY_REF`` at + ``/parity_ref/`` for tests/diffusion/models/auk/test_pipeline_auk.py. + /fusion_ref/zs_tts_audio.pt, instruct_text_only.pt + The fp32 layer-fused text condition from the HF thinker for the + zero-shot and text-only prompts (``AUK_FUSION_REF``). + +By default the VAE encode uses the posterior mean, matching the port's default +and making the artifacts deterministic; ``--vae-sample`` reproduces the +upstream stochastic draw (seeded through ``--global-seed``). The ODE noise is +drawn from ``--seed`` on both sides. + +Usage: + python tools/auk_parity_reference.py \\ + --auk-repo /path/to/AuK --ckpt-dir /path/to/AuK \\ + --qwen-dir /path/to/Qwen2.5-Omni-3B --out /path/to/auk-parity \\ + [--variant base|flash] [--vae-sample] [--fusion-device cpu|cuda] [--skip-fusion] +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F + +WEIGHTS = {"base": "auk_base.safetensors", "flash": "auk_flash.safetensors"} + +# Cookbook cases; audio paths are relative to /assets/demo-input-audio. +CASES: dict[str, dict[str, Any]] = { + "zs_tts_en": { + "text": ( + "Say the following with the same voice: 'Ladies and gentlemen, it's an honor to have the " + "opportunity to address such a distinguished audience'" + ), + "audio": "zero-shot-tts/ref.wav", + "gen_seconds": 6.0, + }, + "instruct_en": { + "text": ( + 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". ' + 'The content to speak is: "Welcome back, how was your day at work?".' + ), + "audio": None, + "gen_seconds": 3.5, + }, + "content_edit": { + "text": "Replace 'but accepting what we cannot have' with 'and living well with dreams unmet'.", + "audio": "content-edit/content.wav", + "gen_seconds": 7.0, + }, + "denoise": { + "text": "Keep pure speech voice, remove noise and reverberation.", + "audio": "vocal-extraction/vocal-1-input.wav", + "gen_seconds": None, + }, +} +# The two prompts the fusion reference covers (they must match the cases above). +FUSION_CASES = {"zs_tts_audio": "zs_tts_en", "instruct_text_only": "instruct_en"} +NO_PROMPT_AUDIO = "||" + + +def _messages(case: dict[str, Any], assets: Path) -> list[dict[str, Any]]: + content: list[dict[str, Any]] = [{"type": "text", "text": case["text"]}] + if case["audio"]: + content.append({"type": "audio", "audio": str(assets / case["audio"])}) + return [{"role": "user", "content": content}] + + +def _mean_encode(vae): + """Posterior-mean version of the upstream ``encoding_and_normalization``.""" + + def encode(sample, sample_lengths=None): + latent_stats = vae.audio_encoder(sample) + if sample_lengths is None: + sample_lengths = torch.LongTensor([sample.size(-1)] * sample.size(0)).to(sample.device) + latent_lens = sample_lengths // vae.hop_size + mean, _log_std = latent_stats.chunk(2, 1) + latents = mean.transpose(1, 2).float() + latents = (latents - vae.global_mean.float()) / torch.sqrt(vae.global_log_std.float()) + return latents, torch.clamp(latent_lens, max=latents.size(1)) + + return encode + + +def generate_references(args: argparse.Namespace) -> None: + from auk.infer.infer_auk import AukInfer, save_audio + + ckpt = Path(args.ckpt_dir) + out = Path(args.out) / "parity_ref" / args.variant + out.mkdir(parents=True, exist_ok=True) + assets = Path(args.auk_repo) / "assets" / "demo-input-audio" + + engine = AukInfer(str(ckpt / "config.yaml"), str(ckpt / WEIGHTS[args.variant]), qwen_path=args.qwen_dir) + encode = engine.vae_model.encoding_and_normalization if args.vae_sample else _mean_encode(engine.vae_model) + recorded: dict[str, torch.Tensor] = {} + + def encode_hook(sample, sample_lengths=None): + latents, lens = encode(sample, sample_lengths) + recorded["ref_latents"] = latents.detach().float().cpu() + recorded["ref_latent_lens"] = lens.detach().cpu() + return latents, lens + + engine.vae_model.encoding_and_normalization = encode_hook + upstream_sample = engine.model.sample + + def sample_hook(*a, **k): + latents, trajectory = upstream_sample(*a, **k) + recorded["generated"] = latents.detach().float().cpu() + return latents, trajectory + + engine.model.sample = sample_hook + + meta: dict[str, Any] = {} + for name, case in CASES.items(): + messages = _messages(case, assets) + recorded.clear() + torch.manual_seed(args.global_seed) + started = time.perf_counter() + audio, sr = engine.generate(messages, gen_seconds=case["gen_seconds"], seed=args.seed) + wall = time.perf_counter() - started + save_audio(audio, sr, str(out / f"{name}.wav")) + torch.save( + { + "messages": messages, + "gen_seconds": case["gen_seconds"], + "seed": args.seed, + "global_seed": args.global_seed, + "vae_sample": bool(args.vae_sample), + "audio": audio.float().cpu(), + "sr": sr, + **recorded, + }, + out / f"{name}.pt", + ) + ref_frames = int(recorded["ref_latent_lens"][0]) if "ref_latent_lens" in recorded else 0 + meta[name] = { + "wall_s": wall, + "out_s": audio.shape[-1] / sr, + "sr": sr, + "ref_frames": ref_frames, + "gen_frames": int(recorded["generated"].shape[1]) - ref_frames, + } + print(name, meta[name], flush=True) + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + + +def generate_fusion(args: argparse.Namespace) -> None: + from auk.model.cfm_edit import CFMEdit + from safetensors import safe_open + from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration + + out = Path(args.out) / "fusion_ref" + out.mkdir(parents=True, exist_ok=True) + assets = Path(args.auk_repo) / "assets" / "demo-input-audio" + with safe_open(str(Path(args.ckpt_dir) / WEIGHTS[args.variant]), "pt") as f: + layer_weights = f.get_tensor("layer_weights").float() + layer_scale = f.get_tensor("layer_scale").float() + + processor = Qwen2_5OmniProcessor.from_pretrained(args.qwen_dir) + thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(args.qwen_dir, torch_dtype=torch.float32) + thinker.visual = None + thinker = thinker.to(args.fusion_device).eval() + + for out_name, case_name in FUSION_CASES.items(): + case = dict(CASES[case_name]) + if not case["audio"] and not case["text"].endswith(NO_PROMPT_AUDIO): + case["text"] = case["text"] + NO_PROMPT_AUDIO + cond = CFMEdit.build_cond_inputs([_messages(case, assets)], processor).to(args.fusion_device) + with torch.no_grad(): + hidden_states = thinker(**cond, output_hidden_states=True).hidden_states + width = hidden_states[0].shape[-1] + stacked = torch.stack([F.layer_norm(h.float(), [width]) for h in hidden_states[1:]], 0) + weights = F.softmax(layer_weights.to(stacked.device), 0) + fused = (stacked * weights[:, None, None, None]).sum(0) * layer_scale.to(stacked.device) + torch.save( + { + "input_ids": cond["input_ids"].cpu(), + "attention_mask": cond["attention_mask"].cpu(), + "fused": fused[0].float().cpu(), + "final_normed": hidden_states[-1][0].float().cpu(), + "text": case["text"], + }, + out / f"{out_name}.pt", + ) + print(out_name, {"tokens": int(cond["input_ids"].shape[1]), "fused_norm": float(fused.norm())}, flush=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--auk-repo", required=True, help="checkout of the upstream AuK repository (for demo clips)") + parser.add_argument("--ckpt-dir", required=True, help="tencent/AuK or tencent/AuK-Flash snapshot") + parser.add_argument("--qwen-dir", required=True, help="Qwen/Qwen2.5-Omni-3B snapshot") + parser.add_argument("--out", required=True, help="output root; writes parity_ref/ and fusion_ref") + parser.add_argument("--variant", choices=sorted(WEIGHTS), default="base") + parser.add_argument("--seed", type=int, default=0, help="ODE noise seed passed to the upstream sampler") + parser.add_argument("--global-seed", type=int, default=1234, help="global RNG seed set before each generate()") + parser.add_argument("--vae-sample", action="store_true", help="draw the VAE posterior sample as upstream does") + parser.add_argument("--fusion-device", default="cpu", help="device for the fp32 HF thinker forward") + parser.add_argument("--skip-fusion", action="store_true", help="only write the generation references") + args = parser.parse_args() + + generate_references(args) + if not args.skip_fusion: + generate_fusion(args) + print(f"reference artifacts written under {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/prepare_auk_checkpoint.py b/tools/prepare_auk_checkpoint.py new file mode 100755 index 000000000..7f55dab8a --- /dev/null +++ b/tools/prepare_auk_checkpoint.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Assemble an AuK checkpoint directory that ``--model `` can load. + +The released AuK repos ship ``config.yaml`` plus loose safetensors and depend +on a separate Qwen2.5-Omni-3B snapshot for the frozen text encoder. This tool +writes the ``config.json`` the vLLM-Omni pipeline expects and links every +artifact both stages need into one flat directory: + + / + config.json the Qwen2.5-Omni-3B config with model_type "auk", + AuK's architecture, and the dit / vae / variant / + flash_t_grid / defaults sections added at the top + level (see transformers_utils/configs/auk.py) + auk.safetensors -> auk_base.safetensors or auk_flash.safetensors + vae.safetensors -> the AuK VAE weights + thinker/ -> the Qwen2.5-Omni-3B snapshot, for reference only + model-0000*.safetensors, model.safetensors.index.json, + tokenizer*, vocab.json, merges.txt, added_tokens.json, + special_tokens_map.json, chat_template.json, + preprocessor_config.json, generation_config.json, spk_dict.pt + -> the same snapshot's entries + +The Qwen entries are linked at the top level rather than left under +``thinker/`` because both stages load from this one directory: the encoder +reads the thinker shards, the tokenizer and the audio processor from here, +and takes the layer-fusion tensors out of ``auk.safetensors``. The weight +index keeps ``auk.safetensors`` and ``vae.safetensors`` out of the thinker's +own weight load, and the DiT and VAE weights are read by model code directly. + +Usage: + python tools/prepare_auk_checkpoint.py \ + --auk-dir /path/to/AuK \ + --qwen-dir /path/to/Qwen2.5-Omni-3B \ + --out /path/to/auk-omni-base + +Re-running against an existing directory is safe: config.json is rewritten +and the links are repointed. Pass ``--copy`` on filesystems without symlinks. +""" + +from __future__ import annotations + +import argparse +import ast +import json +import re +import shutil +import sys +from pathlib import Path +from typing import Any + +# The distilled student's 4-step grid. Duplicated from +# vllm_omni/transformers_utils/configs/auk.py so this tool stays runnable +# without importing the package (and its torch dependency). +FLASH_T_GRID = [ + 0.0, + 0.07612049579620361, + 0.2928932309150696, + 0.6173166036605835, + 1.0, +] + +# Names the assembled directory owns; everything else in it is a Qwen entry. +WEIGHTS_LINK = "auk.safetensors" +VAE_LINK = "vae.safetensors" +THINKER_LINK = "thinker" +CONFIG_NAME = "config.json" + +# Snapshot entries that are not reused: the config is rewritten, and the repo +# paperwork would misdescribe the assembled directory. +_SKIP_QWEN_ENTRIES = frozenset({CONFIG_NAME, "README.md", "LICENSE"}) + +DIT_KEYS = ("dim", "heads", "ff_mult", "text_hidden_dim", "num_layers", "num_single_layers") +VAE_KEYS = ("latent_dim", "downsample_rate", "target_sample_rate", "model_init_kwargs") + +_INTERPOLATION = re.compile(r"^\$\{([A-Za-z0-9_.]+)\}$") + + +def _strip_comment(line: str) -> str: + """Drop a trailing ``#`` comment that is not inside a quoted scalar.""" + out: list[str] = [] + quote: str | None = None + for ch in line: + if quote is not None: + if ch == quote: + quote = None + out.append(ch) + continue + if ch in "\"'": + quote = ch + out.append(ch) + continue + if ch == "#": + break + out.append(ch) + return "".join(out).rstrip() + + +def _parse_scalar(text: str) -> Any: + """Parse one YAML scalar or inline sequence.""" + text = text.strip() + if not text: + return None + lowered = text.lower() + if lowered in ("true", "yes"): + return True + if lowered in ("false", "no"): + return False + if lowered in ("null", "~"): + return None + if text[0] in "[{": + normalized = re.sub(r"\b(True|true|yes)\b", "True", text) + normalized = re.sub(r"\b(False|false|no)\b", "False", normalized) + normalized = re.sub(r"\b(null|~)\b", "None", normalized) + try: + return ast.literal_eval(normalized) + except (SyntaxError, ValueError): + return text + try: + return ast.literal_eval(text) + except (SyntaxError, ValueError): + return text.strip("\"'") + + +def _parse_yaml_minimal(text: str) -> dict[str, Any]: + """Parse the small subset of YAML the AuK configs use. + + Nested mappings by indentation, scalars, and inline sequences. Used only + when PyYAML is unavailable. + """ + root: dict[str, Any] = {} + stack: list[tuple[int, dict[str, Any]]] = [(-1, root)] + for raw in text.splitlines(): + line = _strip_comment(raw) + if not line.strip(): + continue + indent = len(line) - len(line.lstrip(" ")) + body = line.strip() + if ":" not in body: + raise ValueError(f"unsupported YAML line (no key): {raw!r}") + key, _, rest = body.partition(":") + key = key.strip() + while stack and indent <= stack[-1][0]: + stack.pop() + if not stack: + raise ValueError(f"unsupported YAML indentation at: {raw!r}") + parent = stack[-1][1] + if rest.strip(): + parent[key] = _parse_scalar(rest) + else: + child: dict[str, Any] = {} + parent[key] = child + stack.append((indent, child)) + return root + + +def _resolve_interpolations(node: Any, root: dict[str, Any]) -> Any: + """Replace ``${a.b.c}`` scalars with the value at that dotted path.""" + if isinstance(node, dict): + return {k: _resolve_interpolations(v, root) for k, v in node.items()} + if isinstance(node, list): + return [_resolve_interpolations(v, root) for v in node] + if isinstance(node, str): + match = _INTERPOLATION.match(node.strip()) + if match is None: + return node + target: Any = root + for part in match.group(1).split("."): + if not isinstance(target, dict) or part not in target: + raise ValueError(f"cannot resolve interpolation ${{{match.group(1)}}}") + target = target[part] + return target + return node + + +def load_auk_yaml(path: Path) -> dict[str, Any]: + """Load the AuK ``config.yaml`` with interpolations resolved.""" + text = path.read_text() + try: + import yaml + + parsed = yaml.safe_load(text) + except ImportError: + parsed = _parse_yaml_minimal(text) + if not isinstance(parsed, dict) or "model" not in parsed: + raise ValueError(f"{path} has no top-level 'model' section") + return _resolve_interpolations(parsed, parsed) + + +def detect_variant(auk_dir: Path, model_section: dict[str, Any]) -> str: + """Infer base vs flash from the weight filename, then the model name.""" + if (auk_dir / "auk_flash.safetensors").exists(): + return "flash" + if (auk_dir / "auk_base.safetensors").exists(): + return "base" + name = str(model_section.get("name", "")).lower() + return "flash" if "flash" in name else "base" + + +def weights_name(variant: str) -> str: + return "auk_flash.safetensors" if variant == "flash" else "auk_base.safetensors" + + +def build_config( + auk_yaml: dict[str, Any], + qwen_config: dict[str, Any], + variant: str, +) -> dict[str, Any]: + """Assemble the ``config.json`` body for the output directory. + + The Qwen config is kept whole, including ``talker_config`` and + ``token2wav_config``: ``AuKConfig`` subclasses ``Qwen2_5OmniConfig`` and + the encoder stage reuses the in-tree thinker processor, which asserts that + type. AuK's own sections are added at the top level. + """ + model = auk_yaml["model"] + arch = model.get("arch", {}) + vae = model.get("vae", {}) + + missing = [key for key in DIT_KEYS if key not in arch] + if missing: + raise ValueError(f"AuK config.yaml arch section is missing {missing}") + missing = [key for key in VAE_KEYS if key not in vae] + if missing: + raise ValueError(f"AuK config.yaml vae section is missing {missing}") + if not isinstance(qwen_config.get("thinker_config"), dict): + raise ValueError("Qwen snapshot config.json has no 'thinker_config' mapping") + + dit = {key: arch[key] for key in DIT_KEYS} + # dim_head is implicit upstream (dim // heads) but the port takes it + # explicitly, so it is recorded rather than recomputed at load time. + dit["dim_head"] = int(arch["dim"]) // int(arch["heads"]) + # Carried from the released config: the DiT builds an attention mask over + # the padded [audio | text] sequence and parity depends on it. + dit["attn_mask_enabled"] = bool(arch.get("attn_mask_enabled", True)) + + config = dict(qwen_config) + config["model_type"] = "auk" + config["architectures"] = ["AuKForConditionalGeneration"] + config["dit"] = dit + config["vae"] = {key: vae[key] for key in VAE_KEYS} + config["variant"] = variant + config["flash_t_grid"] = FLASH_T_GRID + config["defaults"] = {"nfe": 32, "cfg": 2.0, "sway": -1.0} + return config + + +def _replace_path(target: Path) -> None: + """Remove an existing file, directory, or link at ``target``.""" + if target.is_symlink() or target.is_file(): + target.unlink() + elif target.is_dir(): + shutil.rmtree(target) + + +def link_or_copy(source: Path, target: Path, *, copy: bool) -> None: + """Point ``target`` at ``source``, replacing whatever is already there.""" + if not source.exists(): + raise FileNotFoundError(f"missing source artifact: {source}") + if target.is_symlink() and not copy and target.readlink() == source.resolve(): + return + _replace_path(target) + if copy: + if source.is_dir(): + shutil.copytree(source, target) + else: + shutil.copy2(source, target) + return + target.symlink_to(source.resolve(), target_is_directory=source.is_dir()) + + +MANIFEST_NAME = ".auk-assembled.json" + + +def _reject_overlap(out: Path, *inputs: Path) -> None: + """Refuse an output directory that is, contains, or lives inside an input.""" + out_r = out.resolve() + for src in inputs: + src_r = src.resolve() + if out_r == src_r or out_r in src_r.parents or src_r in out_r.parents: + raise ValueError(f"--out {out} overlaps input directory {src}; choose a separate directory") + + +def assemble(auk_dir: Path, qwen_dir: Path, out: Path, variant: str, *, copy: bool) -> None: + """Link the AuK weights and the Qwen snapshot entries into ``out``.""" + _reject_overlap(out, auk_dir, qwen_dir) + out.mkdir(parents=True, exist_ok=True) + manifest_path = out / MANIFEST_NAME + previous: set[str] = set() + if manifest_path.is_file(): + previous = set(json.loads(manifest_path.read_text()).get("entries", [])) + weights = auk_dir / weights_name(variant) + + entries: dict[str, Path] = { + WEIGHTS_LINK: weights, + VAE_LINK: auk_dir / VAE_LINK, + THINKER_LINK: qwen_dir, + } + # The assembled config replaces the Qwen one; every other entry is reused. + for entry in qwen_dir.iterdir(): + if entry.name.startswith(".") or entry.name in _SKIP_QWEN_ENTRIES: + continue + entries[entry.name] = entry + + for name, source in entries.items(): + link_or_copy(source, out / name, copy=copy) + + # Drop only entries this tool created in an earlier run and no longer needs; + # anything else in the directory belongs to the user. + for name in previous - set(entries) - {CONFIG_NAME}: + stale = out / name + if stale.exists() or stale.is_symlink(): + _replace_path(stale) + manifest_path.write_text(json.dumps({"entries": sorted(entries)}, indent=1) + "\n") + + +def print_tree(out: Path) -> None: + """Print the assembled directory, one entry per line.""" + print(f"\n{out}/") + for entry in sorted(out.iterdir()): + if entry.is_symlink(): + print(f" {entry.name} -> {entry.readlink()}") + elif entry.is_dir(): + print(f" {entry.name}/") + else: + print(f" {entry.name} ({entry.stat().st_size / 1e9:.2f} GB)") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--auk-dir", required=True, help="HF snapshot of the AuK or AuK-Flash repo") + parser.add_argument("--qwen-dir", required=True, help="HF snapshot of Qwen2.5-Omni-3B") + parser.add_argument("--out", required=True, help="Directory to assemble (created if absent)") + parser.add_argument( + "--variant", + choices=("base", "flash"), + help="Override the variant; detected from the AuK snapshot by default", + ) + parser.add_argument( + "--copy", + action="store_true", + help="Copy the weights and the snapshot entries instead of symlinking", + ) + args = parser.parse_args() + + auk_dir = Path(args.auk_dir).expanduser().resolve() + qwen_dir = Path(args.qwen_dir).expanduser().resolve() + out = Path(args.out).expanduser().resolve() + _reject_overlap(out, auk_dir, qwen_dir) + + auk_yaml = load_auk_yaml(auk_dir / "config.yaml") + variant = args.variant or detect_variant(auk_dir, auk_yaml["model"]) + qwen_config = json.loads((qwen_dir / CONFIG_NAME).read_text()) + + config = build_config(auk_yaml, qwen_config, variant) + assemble(auk_dir, qwen_dir, out, variant, copy=args.copy) + (out / CONFIG_NAME).write_text(json.dumps(config, indent=2) + "\n") + + print(f"variant: {variant} (weights {weights_name(variant)})") + print_tree(out) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index a69fb08c6..2378ce5b9 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -40,6 +40,7 @@ from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_TTS_PIPELINE, ) from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE +from vllm_omni.model_executor.models.auk.pipeline import AUK_PIPELINE from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_PIPELINE, @@ -164,6 +165,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, "dots_tts": DOTS_TTS_PIPELINE, + "auk": AUK_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, "audex_tta": AUDEX_TTA_PIPELINE, diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml new file mode 100644 index 000000000..bbe544919 --- /dev/null +++ b/vllm_omni/deploy/auk.yaml @@ -0,0 +1,51 @@ +# AuK deploy: thinker encoder (stage 0) + rectified-flow DiT (stage 1). +# Topology declared in vllm_omni/model_executor/models/auk/pipeline.py. +# +# Fields omitted from a stage fall back to StageDeployConfig dataclass +# defaults (see vllm_omni/config/stage_config.py). trust_remote_code, +# enable_prefix_caching and enable_chunked_prefill are pipeline-wide +# DeployConfig fields and apply to both stages. +# +# The ODE runs over the whole target segment, so there is nothing to stream +# between the stages: async_chunk stays off. +async_chunk: false +trust_remote_code: true +# Prefix caching must stay off: a cache hit skips the prefix positions, and +# the fused [tokens, 2048] condition needs every position of the prompt. +# Chunked prefill works (the runner concatenates the per-chunk hidden states) +# but measured no wall-time gain, since the encoder is a few percent of a +# request, and it costs bit-exact reproducibility across scheduling. +enable_prefix_caching: false +enable_chunked_prefill: false + +stages: + # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder. + # Prefill only — one sampled token, discarded downstream. eager because the + # fused-hidden forward walks the decoder layers itself and so bypasses the + # compiled graph. + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.45 + enforce_eager: true + max_num_batched_tokens: 32768 + # AuK takes text plus at most one clip; no image or video tower is built. + limit_mm_per_prompt: {audio: 1} + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + detokenize: false + seed: 0 + + # Stage 1: AuKPipeline (DiT + BigVGAN-flow VAE encoder and decoder). + # Shares GPU 0 with the encoder; 32 Euler steps at CFG 2.0 is the Base + # recipe, and AuK-Flash overrides both with its distilled 4-step grid. + - stage_id: 1 + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.35 + enforce_eager: true + default_sampling_params: + seed: 0 + num_inference_steps: 32 + guidance_scale: 2.0 diff --git a/vllm_omni/diffusion/models/auk/__init__.py b/vllm_omni/diffusion/models/auk/__init__.py new file mode 100644 index 000000000..fd987d42b --- /dev/null +++ b/vllm_omni/diffusion/models/auk/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK diffusion-stage components: the flow transformer, the codec, the pipeline.""" + +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE +from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline, get_auk_post_process_func + +__all__ = [ + "AuKPipeline", + "AuKTransformer", + "AuKVAE", + "dit_state_dict", + "get_auk_post_process_func", + "sample_latents", +] diff --git a/vllm_omni/diffusion/models/auk/auk_transformer.py b/vllm_omni/diffusion/models/auk/auk_transformer.py new file mode 100644 index 000000000..5aca8f39b --- /dev/null +++ b/vllm_omni/diffusion/models/auk/auk_transformer.py @@ -0,0 +1,677 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The AuK audio DiT as a self-contained inference module. + +Two phases. ``num_layers`` double-stream blocks attend jointly over text and +audio while keeping each stream's residual separate, then +``num_single_layers`` single-stream blocks run over the concatenated +``[text | ref audio | target audio]`` sequence. Only the target-audio slice is +projected back to latent space, so the module predicts a flow-matching +velocity for the target frames and :func:`sample_latents` integrates it. + +Parameter names and shapes match the AuK reference backbone, so a released +checkpoint loads with ``load_state_dict(strict=True)`` once the +``transformer.`` prefix is stripped (see :func:`dit_state_dict`). + +Deliberate differences from the reference, all inference-only: + +* ``x_transformers`` and ``torchdiffeq`` are gone. Rotary embeddings are + computed here (interleaved GPT-J pairs, no xpos scaling) and the ODE is an + explicit Euler loop. +* Dropout, activation checkpointing, the ``flash_attn`` backend switch and the + zero-init of the modulation layers are training concerns and are dropped. +* The timestep sinusoid is cast to the time MLP's weight dtype rather than to + the timestep's own dtype, so an fp32 timestep works against half-precision + weights outside autocast. +* ``attn_mask_enabled`` defaults to ``True``, the value the released config + sets, rather than the reference's ``False``. +""" + +import math +from collections.abc import Iterable, Mapping + +import torch +import torch.nn.functional as F +from torch import nn + +__all__ = ["AuKTransformer", "dit_state_dict", "sample_latents"] + + +def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: + """The single attention call site, so the backend can be swapped in one place. + + ``q``/``k``/``v`` are ``[B, H, N, D]``. ``mask`` is a boolean key-padding + mask ``[B, K]`` in which ``True`` marks a valid key; it is broadcast to + ``[B, H, Q, K]`` only when given. + """ + attn_mask = None + if mask is not None: + attn_mask = mask[:, None, None, :].expand(q.shape[0], q.shape[1], q.shape[-2], k.shape[-2]) + return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate interleaved GPT-J pairs: ``(x1, x2) -> (-x2, x1)``.""" + pairs = x.reshape(*x.shape[:-1], -1, 2) + x1, x2 = pairs.unbind(dim=-1) + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary frequencies ``[1, 1, N, D]`` to ``x`` ``[B, H, N, D]`` in fp32.""" + out_dtype = x.dtype + x32 = x.float() + out = x32 * freqs.cos() + _rotate_half(x32) * freqs.sin() + return out.to(out_dtype) + + +class Rotary(nn.Module): + """Rotary position frequencies for one stream, evaluated in fp32. + + The frequencies live in a buffer so that a released checkpoint's own copy + of them loads with ``strict=True``. They are recomputed whenever the + module has been cast below fp32: rounding either the frequencies or the + integer positions to bf16 wrecks the phase at long positions. + """ + + def __init__(self, dim: int, base: float = 10000.0) -> None: + super().__init__() + if dim % 2 != 0: + raise ValueError(f"rotary dim must be even, got {dim}") + self.dim = dim + self.base = base + self.register_buffer("inv_freq", self._frequencies(torch.device("cpu"))) + + def _frequencies(self, device: torch.device) -> torch.Tensor: + exponents = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim + return 1.0 / (self.base**exponents) + + def forward(self, seq_len: int) -> torch.Tensor: + """Return frequencies ``[1, 1, seq_len, dim]`` for positions ``0..seq_len-1``.""" + inv_freq = self.inv_freq + if inv_freq.dtype != torch.float32: + inv_freq = self._frequencies(inv_freq.device) + pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32) + freqs = torch.outer(pos, inv_freq) + return torch.stack((freqs, freqs), dim=-1).flatten(-2)[None, None] + + +class TimeEmbedding(nn.Module): + """Sinusoidal timestep features followed by a two-layer MLP.""" + + def __init__(self, dim: int, freq_embed_dim: int = 256) -> None: + super().__init__() + self.freq_embed_dim = freq_embed_dim + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor, scale: float = 1000.0) -> torch.Tensor: + half = self.freq_embed_dim // 2 + decay = math.log(10000) / (half - 1) + freqs = torch.exp(torch.arange(half, device=timestep.device, dtype=torch.float32) * -decay) + angles = scale * timestep.unsqueeze(1) * freqs.unsqueeze(0) + hidden = torch.cat((angles.sin(), angles.cos()), dim=-1) + return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype)) + + +class ConvPosEmbedding(nn.Module): + """Depthwise-grouped conv stack that adds local position information.""" + + def __init__(self, dim: int, kernel_size: int = 31, groups: int = 16) -> None: + super().__init__() + if kernel_size % 2 == 0: + raise ValueError(f"kernel_size must be odd, got {kernel_size}") + padding = kernel_size // 2 + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=padding), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=padding), + nn.Mish(), + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + keep = None if mask is None else mask.unsqueeze(1) + x = x.transpose(1, 2) + if keep is not None: + x = x.masked_fill(~keep, 0.0) + for layer in self.conv1d: + x = layer(x) + # Re-zero padding after each convolution so it cannot bleed into + # valid frames through the kernel window. + if keep is not None and isinstance(layer, nn.Conv1d): + x = x.masked_fill(~keep, 0.0) + return x.transpose(1, 2) + + +class AudioEmbedding(nn.Module): + """Project audio latents to model width and add conv position information.""" + + def __init__(self, latent_dim: int, dim: int) -> None: + super().__init__() + self.linear = nn.Linear(latent_dim, dim) + self.conv_pos_embed = ConvPosEmbedding(dim) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + x = self.linear(x) + return self.conv_pos_embed(x, mask=mask) + x + + +class AdaLayerNorm(nn.Module): + """Timestep-conditioned modulation for a block: six chunks from one projection.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward( + self, x: torch.Tensor, emb: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.linear(self.silu(emb)).chunk(6, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaLayerNormFinal(nn.Module): + """Timestep-conditioned modulation before the output projection.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + scale, shift = self.linear(self.silu(emb)).chunk(2, dim=1) + return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + + +class FeedForward(nn.Module): + """SwiGLU feed-forward with the gate projection fused into ``linear_in``.""" + + def __init__(self, dim: int, mult: float = 4.0) -> None: + super().__init__() + inner_dim = int(dim * mult) + self.linear_in = nn.Linear(dim, inner_dim * 2, bias=False) + self.linear_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.linear_in(x).chunk(2, dim=-1) + return self.linear_out(F.silu(gate) * up) + + +class Attention(nn.Module): + """Self-attention with per-head QK RMSNorm and rotary positions. + + ``to_out`` stays a ``ModuleList`` because the released checkpoints name the + output projection ``to_out.0``; the reference's trailing dropout carries no + parameters and is dropped. + """ + + def __init__(self, dim: int, heads: int, dim_head: int, attn_mask_enabled: bool = True) -> None: + super().__init__() + self.heads = heads + self.dim_head = dim_head + self.attn_mask_enabled = attn_mask_enabled + inner_dim = heads * dim_head + self.to_qkv = nn.Linear(dim, 3 * inner_dim) + self.q_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.k_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim)]) + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + return x.transpose(1, 2).reshape(x.shape[0], -1, self.heads * self.dim_head) + + def _qkv( + self, + packed: torch.Tensor, + q_norm: nn.Module, + k_norm: nn.Module, + rope: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split a packed QKV projection into heads, RMSNorm Q/K, then rotate.""" + q, k, v = (t.view(t.shape[0], -1, self.heads, self.dim_head).transpose(1, 2) for t in packed.chunk(3, dim=-1)) + q = q_norm(q) + k = k_norm(k) + if rope is not None: + q = _apply_rope(q, rope) + k = _apply_rope(k, rope) + return q, k, v + + def _project(self, x: torch.Tensor, rope: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self._qkv(self.to_qkv(x), self.q_norm, self.k_norm, rope) + + def forward( + self, x: torch.Tensor, mask: torch.Tensor | None = None, rope: torch.Tensor | None = None + ) -> torch.Tensor: + q, k, v = self._project(x, rope) + attn_mask = mask if self.attn_mask_enabled else None + out = self._merge_heads(_sdpa(q, k, v, attn_mask).to(q.dtype)) + out = self.to_out[0](out) + if mask is not None: + out = out.masked_fill(~mask.unsqueeze(-1), 0.0) + return out + + +class JointAttention(Attention): + """Attention over the audio and text streams jointly, with separate projections.""" + + def __init__(self, dim: int, heads: int, dim_head: int, context_dim: int, attn_mask_enabled: bool = True) -> None: + super().__init__(dim, heads, dim_head, attn_mask_enabled) + inner_dim = heads * dim_head + self.to_qkv_c = nn.Linear(context_dim, 3 * inner_dim) + self.c_q_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.c_k_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.to_out_c = nn.Linear(inner_dim, context_dim) + + def _project_context(self, c: torch.Tensor, c_rope: torch.Tensor | None) -> tuple[torch.Tensor, ...]: + return self._qkv(self.to_qkv_c(c), self.c_q_norm, self.c_k_norm, c_rope) + + def forward( # type: ignore[override] # joint attention takes the context stream too + self, + x: torch.Tensor, + c: torch.Tensor, + mask: torch.Tensor | None = None, + rope: torch.Tensor | None = None, + c_rope: torch.Tensor | None = None, + c_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + audio_len = x.shape[1] + q, k, v = self._project(x, rope) + c_q, c_k, c_v = self._project_context(c, c_rope) + + q = torch.cat([q, c_q], dim=2) + k = torch.cat([k, c_k], dim=2) + v = torch.cat([v, c_v], dim=2) + + joint_mask = None + if self.attn_mask_enabled and mask is not None: + if c_mask is not None: + joint_mask = torch.cat([mask, c_mask], dim=1) + else: + joint_mask = F.pad(mask, (0, c.shape[1]), value=True) + + out = self._merge_heads(_sdpa(q, k, v, joint_mask).to(q.dtype)) + x_out = self.to_out[0](out[:, :audio_len]) + c_out = self.to_out_c(out[:, audio_len:]) + + if mask is not None: + x_out = x_out.masked_fill(~mask.unsqueeze(-1), 0.0) + if c_mask is not None: + c_out = c_out.masked_fill(~c_mask.unsqueeze(-1), 0.0) + return x_out, c_out + + +class DoubleBlock(nn.Module): + """MM-DiT block: joint attention, separate modulation and feed-forward per stream.""" + + def __init__(self, dim: int, heads: int, dim_head: int, ff_mult: float, attn_mask_enabled: bool) -> None: + super().__init__() + self.attn_norm_c = AdaLayerNorm(dim) + self.attn_norm_x = AdaLayerNorm(dim) + self.attn = JointAttention( + dim=dim, heads=heads, dim_head=dim_head, context_dim=dim, attn_mask_enabled=attn_mask_enabled + ) + self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_c = FeedForward(dim=dim, mult=ff_mult) + self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_x = FeedForward(dim=dim, mult=ff_mult) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + rope: torch.Tensor, + c_rope: torch.Tensor, + c_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t) + norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t) + + x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask) + + c = c + c_gate_msa[:, None] * c_attn + norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + c = c + c_gate_mlp[:, None] * self.ff_c(norm_c) + + x = x + x_gate_msa[:, None] * x_attn + norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] + x = x + x_gate_mlp[:, None] * self.ff_x(norm_x) + return c, x + + +class SingleBlock(nn.Module): + """DiT block over the concatenated text and audio sequence.""" + + def __init__(self, dim: int, heads: int, dim_head: int, ff_mult: float, attn_mask_enabled: bool) -> None: + super().__init__() + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention(dim=dim, heads=heads, dim_head=dim_head, attn_mask_enabled=attn_mask_enabled) + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult) + + def forward(self, x: torch.Tensor, t: torch.Tensor, mask: torch.Tensor | None, rope: torch.Tensor) -> torch.Tensor: + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t) + x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope) + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + return x + gate_mlp[:, None] * self.ff(norm) + + +class AuKTransformer(nn.Module): + """Flow-matching velocity predictor for AuK audio latents. + + Args: + dim: Model width. + heads: Attention heads. + dim_head: Per-head width. + ff_mult: Feed-forward expansion factor. + latent_dim: Audio VAE latent channels, the input and output width. + text_hidden_dim: Width of the pre-encoded text hidden states. + num_layers: Double-stream (joint attention) block count. + num_single_layers: Single-stream block count. + attn_mask_enabled: Apply padding masks inside attention. When ``False`` + attention runs unmasked, though padded outputs are still zeroed, + which is what the reference does with the flag off. + """ + + def __init__( + self, + *, + dim: int, + heads: int, + dim_head: int, + ff_mult: float, + latent_dim: int, + text_hidden_dim: int, + num_layers: int, + num_single_layers: int, + attn_mask_enabled: bool = True, + ) -> None: + super().__init__() + self.dim = dim + self.latent_dim = latent_dim + self.text_cond: torch.Tensor | None = None + self.text_uncond: torch.Tensor | None = None + + self.time_embed = TimeEmbedding(dim) + self.txt_norm = nn.RMSNorm(dim, elementwise_affine=True) + self.txt_proj = nn.Linear(text_hidden_dim, dim) + self.audio_embed = AudioEmbedding(latent_dim, dim) + self.rotary_embed = Rotary(dim_head) + + self.transformer_blocks = nn.ModuleList( + [ + DoubleBlock( + dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, attn_mask_enabled=attn_mask_enabled + ) + for _ in range(num_layers) + ] + ) + self.single_transformer_blocks = nn.ModuleList( + [ + SingleBlock( + dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, attn_mask_enabled=attn_mask_enabled + ) + for _ in range(num_single_layers) + ] + ) + + self.norm_out = AdaLayerNormFinal(dim) + self.proj_out = nn.Linear(dim, latent_dim) + + def project_text(self, text: torch.Tensor) -> torch.Tensor: + """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width.""" + return self.txt_norm(self.txt_proj(text)) + + def clear_cache(self) -> None: + """Drop the cached text projection. Call this when the conditioning changes.""" + self.text_cond = None + self.text_uncond = None + + def _embed_audio( + self, + x: torch.Tensor, + ref: torch.Tensor | None, + drop_audio_cond: bool, + mask: torch.Tensor | None, + ref_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, int]: + """Embed the target latents, prepending the reference prompt when present. + + Returns the ``[ref | target]`` sequence, its padding mask and the + reference length, which the caller uses to slice the target back out. + """ + if ref is not None and ref.shape[1] == 0: + ref = None + + target = self.audio_embed(x, mask=mask) + if ref is None: + return target, mask, 0 + + if drop_audio_cond: + ref = torch.zeros_like(ref) + prompt = self.audio_embed(ref, mask=ref_mask) + audio = torch.cat([prompt, target], dim=1) + prompt_len = prompt.shape[1] + + audio_mask = None + if mask is not None or ref_mask is not None: + if mask is None: + mask = x.new_ones(x.shape[:2], dtype=torch.bool) + if ref_mask is None: + ref_mask = ref.new_ones(ref.shape[:2], dtype=torch.bool) + audio_mask = torch.cat([ref_mask, mask], dim=1) + return audio, audio_mask, prompt_len + + def forward( + self, + x: torch.Tensor, + text: torch.Tensor, + time: torch.Tensor, + *, + mask: torch.Tensor | None = None, + c_mask: torch.Tensor | None = None, + ref: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + drop_audio_cond: bool = False, + drop_text: bool = False, + cfg_infer: bool = False, + cache: bool = False, + ) -> torch.Tensor: + """Predict the flow-matching velocity for the target frames of ``x``. + + Args: + x: Noised target latents ``[B, n, latent_dim]``. + text: Pre-encoded text hidden states ``[B, nt, text_hidden_dim]``. + time: Flow-matching timestep, a scalar or ``[B]``. + mask: Target padding mask ``[B, n]``, ``True`` where valid. + c_mask: Text padding mask ``[B, nt]``. Inferred from all-zero text + rows when omitted. + ref: Reference prompt latents ``[B, np, latent_dim]``, prepended to + the audio sequence. A zero-length ``ref`` counts as absent. + ref_mask: Reference padding mask ``[B, np]``. + drop_audio_cond: Zero the reference prompt, for the CFG uncond branch. + drop_text: Zero the projected text, for the CFG uncond branch. + cfg_infer: Run the cond and uncond branches as one batch of ``2B``, + cond first. ``drop_audio_cond`` and ``drop_text`` are then + driven per branch and ignored. + cache: Reuse the projected text across calls, for an ODE loop over + fixed conditioning. Only the ``cfg_infer`` path caches, as in + the reference. Call :meth:`clear_cache` when the text changes. + + Returns: + Velocity ``[B, n, latent_dim]``, or ``[2B, n, latent_dim]`` under + ``cfg_infer``. + """ + batch = x.shape[0] + if time.ndim == 0: + time = time.repeat(batch) + t = self.time_embed(time) + + if c_mask is None: + # Padding rows are all-zero in the text encoder's output. + c_mask = text.abs().sum(-1) > 0 + + if cfg_infer: + if cache and self.text_cond is not None: + c_cond, c_uncond = self.text_cond, self.text_uncond + else: + c_cond = self.project_text(text) + # The uncond branch drops the text, which is exactly a zeroed projection. + c_uncond = torch.zeros_like(c_cond) + if cache: + self.text_cond, self.text_uncond = c_cond, c_uncond + x_cond, mask_cond, prompt_len = self._embed_audio(x, ref, False, mask, ref_mask) + x_uncond, mask_uncond, _ = self._embed_audio(x, ref, True, mask, ref_mask) + + x = torch.cat((x_cond, x_uncond), dim=0) + c = torch.cat((c_cond, c_uncond), dim=0) + t = torch.cat((t, t), dim=0) + c_mask = torch.cat((c_mask, c_mask), dim=0) + audio_mask = None + if mask_cond is not None and mask_uncond is not None: + audio_mask = torch.cat((mask_cond, mask_uncond), dim=0) + else: + c = self.project_text(text) + if drop_text: + c = torch.zeros_like(c) + x, audio_mask, prompt_len = self._embed_audio(x, ref, drop_audio_cond, mask, ref_mask) + + seq_len = x.shape[1] + text_len = c.shape[1] + rope_audio = self.rotary_embed(seq_len) + rope_text = self.rotary_embed(text_len) + + for block in self.transformer_blocks: + c, x = block(x, c, t, mask=audio_mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask) + + x = torch.cat([c, x], dim=1) + rope = self.rotary_embed(text_len + seq_len) + single_mask = None if audio_mask is None else torch.cat([c_mask, audio_mask], dim=1) + + for block in self.single_transformer_blocks: + x = block(x, t, mask=single_mask, rope=rope) + + x = x[:, text_len + prompt_len :] + return self.proj_out(self.norm_out(x, t)) + + +def dit_state_dict( + checkpoint: Mapping[str, torch.Tensor] | Iterable[tuple[str, torch.Tensor]], + prefix: str = "transformer.", +) -> dict[str, torch.Tensor]: + """Select the DiT tensors out of a full AuK checkpoint and strip ``prefix``. + + Released checkpoints wrap the backbone under ``transformer.`` alongside the + text-encoder layer-fusion parameters, which this module does not own. + """ + items = checkpoint.items() if isinstance(checkpoint, Mapping) else checkpoint + return {name[len(prefix) :]: tensor for name, tensor in items if name.startswith(prefix)} + + +@torch.no_grad() +def sample_latents( + dit: AuKTransformer, + *, + text: torch.Tensor, + c_mask: torch.Tensor | None, + ref: torch.Tensor, + ref_mask: torch.Tensor | None, + gen_frames: int, + nfe: int = 32, + cfg_strength: float = 1.0, + sway_sampling_coef: float | None = None, + t_grid: list[float] | None = None, + seed: int | None = None, + generator: torch.Generator | None = None, + latent_dim: int | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Integrate the flow from noise to audio latents with explicit Euler steps. + + Single-request only: the batch dimension carries the CFG branches, not + separate prompts, so no target padding mask is needed. + + Args: + dit: The velocity model. + text: Pre-encoded text hidden states ``[1, nt, text_hidden_dim]``. + c_mask: Text padding mask ``[1, nt]``. + ref: Reference prompt latents ``[1, np, latent_dim]``. + ref_mask: Reference padding mask ``[1, np]``. + gen_frames: Target latent frames to generate. + nfe: Euler steps, ignored when ``t_grid`` is given. + cfg_strength: Classifier-free guidance weight. Below ``1e-5`` the + uncond branch is skipped entirely. + sway_sampling_coef: Reshapes the uniform time grid towards ``t=0``. + t_grid: Explicit timesteps, overriding ``nfe`` and the sway reshape. + seed: Seeds the global RNG before drawing the noise. Ignored when + ``generator`` is given. + generator: Draws the initial noise from its own RNG, leaving the global + one untouched. + latent_dim: Latent channels; defaults to the model's. + device: Device for the noise; defaults to the reference's. + dtype: Dtype for the noise and the Euler accumulator; defaults to the + reference's, which is the reference implementation's rule. Pass + ``torch.float32`` explicitly when serving half-precision weights. + Drawing the noise at bf16 instead of fp32 moves a 32-step CFG + trajectory by 0.025 relative MSE, roughly 17 times what bf16 + weights themselves cost, because the draw consumes the generator + differently and starts the ODE somewhere else. Keeping it fp32 + costs a few tens of KB and no measurable time. + + Returns: + The target latents at ``t=1``, ``[1, gen_frames, latent_dim]``. + """ + device = torch.device(device) if device is not None else ref.device + dtype = dtype if dtype is not None else ref.dtype + latent_dim = latent_dim if latent_dim is not None else dit.latent_dim + + if generator is None: + if seed is not None: + torch.manual_seed(seed) + x = torch.randn(gen_frames, latent_dim, device=device, dtype=dtype).unsqueeze(0) + else: + x = ( + torch.randn(gen_frames, latent_dim, generator=generator, device=generator.device, dtype=dtype) + .to(device) + .unsqueeze(0) + ) + + if t_grid is not None: + t = torch.tensor(t_grid, device=device, dtype=torch.float32) + else: + t = torch.linspace(0, 1, nfe + 1, device=device, dtype=torch.float32) + if sway_sampling_coef is not None: + t = t + sway_sampling_coef * (torch.cos(math.pi / 2 * t) - 1 + t) + + guided = cfg_strength >= 1e-5 + if t.numel() < 2 or not bool(torch.all(t[1:] > t[:-1])): + raise ValueError( + f"AuK sampling needs a strictly increasing time grid with at least two points; got {t.tolist()}" + ) + try: + for i in range(t.shape[0] - 1): + if guided: + pred = dit( + x, + text, + t[i], + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + cfg_infer=True, + cache=True, + ) + v_cond, v_uncond = pred.chunk(2, dim=0) + v = v_cond + (v_cond - v_uncond) * cfg_strength + else: + v = dit(x, text, t[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask) + x = x + (t[i + 1] - t[i]) * v + finally: + # The cached text projections belong to this request only; a failed + # step must not leak them into the next one. + dit.clear_cache() + return x diff --git a/vllm_omni/diffusion/models/auk/auk_vae.py b/vllm_omni/diffusion/models/auk/auk_vae.py new file mode 100644 index 000000000..1fabeb327 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/auk_vae.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Adapted from the AuK reference release (MIT), which in turn adapts NVIDIA +# BigVGAN (MIT), jik876/hifi-gan (MIT), junjun3518/alias-free-torch +# (Apache-2.0), EdwardDixon/snake (MIT) and adefossez/julius (MIT). +"""BigVGAN-flow VAE for AuK: 24 kHz mono waveform <-> 50 Hz, 64-channel latents. + +Inference only. The reference module also carries a residual coupling flow, a KL +term and a discriminator; all three are training-only, so the released +``flow.*`` tensors are skipped (see :meth:`AuKVAE.load_weights`). + +Parameter and buffer names follow the reference release, so the published +``vae.safetensors`` loads without a key map. That is also why the convolutions +are built with the legacy ``weight_norm`` parametrization (``weight_g`` / +``weight_v``): :meth:`remove_weight_norm` folds them into plain weights once +the checkpoint is in. +""" + +from __future__ import annotations + +import inspect +import math +import warnings +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch import nn +from torch.nn.utils import remove_weight_norm as _fold_weight_norm +from torch.nn.utils import weight_norm as _apply_weight_norm +from vllm.logger import init_logger + +logger = init_logger(__name__) + +__all__ = ["AuKVAE"] + +# Checkpoint prefix owned by the training-only coupling flow. +TRAINING_ONLY_PREFIX = "flow." + + +def _weight_norm(module: nn.Module) -> nn.Module: + """Legacy ``weight_g``/``weight_v`` split, which is how the release stores these tensors.""" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*weight_norm.*deprecated.*", category=FutureWarning) + return _apply_weight_norm(module) + + +def _same_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +def _kaiser_sinc_filter(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Kaiser-windowed sinc low-pass, returned as a ``[1, 1, kernel_size]`` conv kernel.""" + half_size = kernel_size // 2 + attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21.0) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + if kernel_size % 2 == 0: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + return torch.zeros_like(time).view(1, 1, kernel_size) + taps = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + # Normalize to unit sum so a constant input passes through unscaled. + return (taps / taps.sum()).view(1, 1, kernel_size) + + +class Conv(nn.Conv1d): + """1-D convolution over ``[B, C, T]``; ``causal`` swaps centred padding for a left pad.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + dilation: int = 1, + bias: bool = True, + causal: bool = False, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0 if causal else _same_padding(kernel_size, dilation), + dilation=dilation, + bias=bias, + ) + self.causal = causal + self.left_padding = dilation * (kernel_size - 1) if causal else 0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.causal: + x = F.pad(x, (self.left_padding, 0)) + return super().forward(x) + + +class ConvTranspose(nn.ConvTranspose1d): + """Transposed 1-D convolution; ``causal`` drops the trailing ``stride`` samples instead of padding.""" + + def __init__( + self, in_channels: int, out_channels: int, kernel_size: int, stride: int, causal: bool = False + ) -> None: + if causal and kernel_size != 2 * stride: + raise ValueError(f"causal ConvTranspose needs kernel_size == 2 * stride, got {kernel_size} and {stride}") + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0 if causal else (kernel_size - stride) // 2, + ) + self.causal = causal + self.trim = stride + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = super().forward(x) + return x[:, :, : -self.trim] if self.causal else x + + +class NormConv(nn.Module): + """Weight-normalized conv held under ``layer``; the attribute name is load-bearing for the checkpoint.""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 1, stride: int = 1) -> None: + super().__init__() + self.layer = _weight_norm( + nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=(kernel_size - 1) // 2) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layer(x) + + +class ResStack(nn.Module): + """Stack of dilated residual conv pairs. The leaky slope is torch's 0.01 default, as upstream.""" + + def __init__(self, channels: int, kernel_size: int = 3, base: int = 3, nums: int = 4) -> None: + super().__init__() + self.layers = nn.ModuleList( + [ + nn.Sequential( + nn.LeakyReLU(), + _weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=base**i, padding=base**i)), + nn.LeakyReLU(), + _weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=1)), + ) + for i in range(nums) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + x = x + layer(x) + return x + + +class Encoder(nn.Module): + """Strided conv encoder over ``[B, 1, T]``, emitting ``[B, 2 * latent_dim, T / hop]`` mean/log-std pairs.""" + + def __init__( + self, + latent_dim: int, + channels: tuple[int, ...], + down_sample_factors: tuple[int, ...], + in_channels: int = 1, + proj_kernel_size: int = 3, + stack_kernel_size: int = 3, + stack_dilation_base: int = 2, + stacks: int = 6, + ) -> None: + super().__init__() + layers: list[nn.Module] = [ + NormConv(in_channels, channels[0], kernel_size=proj_kernel_size), + nn.LeakyReLU(0.2, True), + ] + for in_c, out_c, factor in zip(channels[:-1], channels[1:], down_sample_factors, strict=True): + layers += [ + NormConv(in_c, out_c, kernel_size=factor * 2, stride=factor), + ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks), + nn.LeakyReLU(0.2, True), + ] + layers.append(NormConv(channels[-1], latent_dim * 2, kernel_size=proj_kernel_size)) + self.generator = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.generator(x) + + +class SnakeBeta(nn.Module): + """``x + sin^2(alpha * x) / beta`` with per-channel learned alpha and beta (log-scale in AuK).""" + + def __init__(self, channels: int, alpha_logscale: bool = False) -> None: + super().__init__() + self.alpha_logscale = alpha_logscale + init = torch.zeros(channels) if alpha_logscale else torch.ones(channels) + self.alpha = nn.Parameter(init.clone()) + self.beta = nn.Parameter(init.clone()) + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) + + +class LowPass(nn.Module): + """Fixed FIR low-pass with replicate padding, applied per channel.""" + + def __init__( + self, + cutoff: float = 0.5, + half_width: float = 0.6, + stride: int = 1, + kernel_size: int = 12, + causal: bool = False, + ) -> None: + super().__init__() + if not 0.0 <= cutoff <= 0.5: + raise ValueError(f"cutoff must lie in [0, 0.5], got {cutoff}") + if causal: + self.pad_left, self.pad_right = kernel_size - 1, 0 + else: + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", _kaiser_sinc_filter(cutoff, half_width, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.size(1) + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) + + +class Upsample(nn.Module): + """Band-limited interpolation by ``ratio``. Always non-causal, matching how AuK builds it.""" + + def __init__(self, ratio: int = 2, kernel_size: int = 12) -> None: + super().__init__() + self.ratio = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2 + self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2 + self.register_buffer("filter", _kaiser_sinc_filter(0.5 / ratio, 0.6 / ratio, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.size(1) + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(channels, -1, -1), stride=self.ratio, groups=channels) + return x[..., self.pad_left : -self.pad_right] + + +class Downsample(nn.Module): + """Decimation by ``ratio`` behind a low-pass; ``lowpass`` holds the filter buffer.""" + + def __init__(self, ratio: int = 2, kernel_size: int = 12, causal: bool = False) -> None: + super().__init__() + self.lowpass = LowPass( + cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, causal=causal + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class AliasFreeActivation(nn.Module): + """Oversample, apply the pointwise activation, decimate, so snake harmonics do not fold back.""" + + def __init__( + self, + activation: nn.Module, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + causal: bool = False, + ) -> None: + super().__init__() + self.act = activation + self.upsample = Upsample(up_ratio, up_kernel_size) + self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.downsample(self.act(self.upsample(x))) + + +class AmpBlock(nn.Module): + """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs.""" + + def __init__( + self, + channels: int, + kernel_size: int, + dilations: tuple[int, ...], + snake_logscale: bool = True, + causal: bool = True, + act_causal: bool = False, + ) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=d, causal=causal)) for d in dilations] + ) + self.convs2 = nn.ModuleList( + [_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=1, causal=causal)) for _ in dilations] + ) + self.activations = nn.ModuleList( + [ + AliasFreeActivation(SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal) + for _ in range(len(self.convs1) + len(self.convs2)) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + acts1, acts2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): + x = conv2(act2(conv1(act1(x)))) + x + return x + + +class AuKVAE(nn.Module): + """AuK's audio codec: :meth:`encode` waveform to normalized latents, :meth:`decode` back. + + Build from the released ``model.vae.model_init_kwargs`` with :meth:`from_config`, then + :meth:`load_weights`. Latents are normalized by the checkpoint's global statistics, which is + the space the rectified-flow transformer works in. + """ + + def __init__( + self, + *, + upsample_rates: tuple[int, ...] = (5, 4, 3, 2, 2, 2), + upsample_kernel_sizes: tuple[int, ...] = (10, 8, 6, 4, 4, 4), + upsample_initial_channel: int = 1536, + resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11), + resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)), + downsample_rates: tuple[int, ...] = (2, 2, 2, 3, 4, 5), + downsample_channels: tuple[int, ...] = (12, 24, 48, 96, 192, 384, 768), + snake_logscale: bool = True, + latent_dim: int = 64, + use_vae: bool = True, + causal: bool = True, + act_causal: bool = True, + sample_rate: int = 24000, + ) -> None: + super().__init__() + if not use_vae: + raise ValueError("AuKVAE requires use_vae=True: the encoder head emits mean/log-std pairs") + self.latent_dim = latent_dim + self.sample_rate = sample_rate + self.hop_size = math.prod(downsample_rates) + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self._norm_folded = False + + self.register_buffer("global_mean", torch.zeros(latent_dim, dtype=torch.float32)) + self.register_buffer("global_log_std", torch.ones(latent_dim, dtype=torch.float32)) + + self.audio_encoder = Encoder( + latent_dim=latent_dim, + channels=tuple(downsample_channels), + down_sample_factors=tuple(downsample_rates), + ) + + self.conv_pre = _weight_norm(Conv(latent_dim, upsample_initial_channel, 7, 1, causal=False)) + self.ups = nn.ModuleList() + for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)): + # One-element inner list: the checkpoint keys are ``ups..0.*``. + self.ups.append( + nn.ModuleList( + [ + _weight_norm( + ConvTranspose( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel, + rate, + causal=causal, + ) + ) + ] + ) + ) + + self.resblocks = nn.ModuleList() + channels = upsample_initial_channel + for i in range(self.num_upsamples): + channels = upsample_initial_channel // (2 ** (i + 1)) + for kernel, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): + self.resblocks.append( + AmpBlock( + channels, + kernel, + tuple(dilations), + snake_logscale=snake_logscale, + causal=causal, + act_causal=act_causal, + ) + ) + + self.activation_post = AliasFreeActivation( + SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal + ) + self.conv_post = _weight_norm(Conv(channels, 1, 7, 1, causal=causal, bias=False)) + + @classmethod + def from_config(cls, cfg: dict) -> AuKVAE: + """Build from the released ``model_init_kwargs`` mapping, ignoring training-only entries.""" + accepted = set(inspect.signature(cls.__init__).parameters) - {"self"} + return cls(**{key: value for key, value in cfg.items() if key in accepted}) + + def load_weights( + self, path: str | Path, *, device: str | torch.device | None = None, fold_weight_norm: bool = True + ) -> tuple[list[str], list[str]]: + """Load a released ``vae.safetensors``, then fold weight norm away. Returns (missing, unexpected). + + ``device`` defaults to where the module already lives, which is the placement that matters: + the fold reduces ``weight_v`` to a norm, and a CPU reduction lands up to one ULP away from the + accelerator's. The encoder amplifies its activations roughly fiftyfold, so that one ULP grows + into a ~4e-4 relative drift by the latent head. Move the module first, then load. + """ + target = self.global_mean.device if device is None else torch.device(device) + missing, unexpected = self.load_state_dict(load_file(str(path), device=str(target)), strict=False) + missing, unexpected = list(missing), list(unexpected) + stray = [key for key in missing + unexpected if not key.startswith(TRAINING_ONLY_PREFIX)] + if stray: + raise ValueError(f"AuK VAE weights at {path} do not match this module: {sorted(stray)[:8]}") + logger.info( + "Loaded AuK VAE from %s (skipped %d training-only %s* tensors)", path, len(unexpected), TRAINING_ONLY_PREFIX + ) + if fold_weight_norm: + self.remove_weight_norm() + return missing, unexpected + + def remove_weight_norm(self) -> None: + """Fold ``weight_g``/``weight_v`` into plain weights, on whichever device the module is on. + + Idempotent, and meant to run once the checkpoint is in. See :meth:`load_weights` for why the + device this runs on is a numeric decision rather than a convenience. + """ + if self._norm_folded: + return + folded = 0 + for module in self.modules(): + if hasattr(module, "weight_g"): + _fold_weight_norm(module) + folded += 1 + self._norm_folded = True + logger.debug("Folded weight norm on %d AuK VAE convolutions", folded) + + def encode( + self, wav: torch.Tensor, *, sample: bool = False, generator: torch.Generator | None = None + ) -> torch.Tensor: + """Encode a mono ``[1, T]`` waveform at :attr:`sample_rate` into ``[1, T / hop_size, latent_dim]``. + + ``sample=False`` takes the posterior mean, which is what a reproducible pipeline wants; + ``sample=True`` draws ``mean + randn * exp(log_std)`` from ``generator``. Latents come back + normalized by the checkpoint's global statistics. + """ + if wav.dim() != 2 or wav.size(0) != 1: + raise ValueError(f"encode expects a mono [1, T] waveform, got {tuple(wav.shape)}") + # Statistics stay in fp32 even under an ambient autocast, as in the reference. + with torch.autocast(device_type=wav.device.type, enabled=False): + mean, log_std = self.audio_encoder(wav.unsqueeze(1)).chunk(2, dim=1) + if sample: + latents = mean + self._noise_like(mean, generator) * torch.exp(log_std) + else: + latents = mean + latents = latents.transpose(1, 2).float() + return (latents - self.global_mean.float()) / torch.sqrt(self.global_log_std.float()) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """Decode normalized ``[B, Np, latent_dim]`` latents into a ``[B, Np * hop_size]`` waveform in [-1, 1].""" + if latents.dim() != 3 or latents.size(-1) != self.latent_dim: + raise ValueError(f"decode expects [B, Np, {self.latent_dim}] latents, got {tuple(latents.shape)}") + x = latents.float() * torch.sqrt(self.global_log_std.float()) + self.global_mean.float() + x = self.conv_pre(x.permute(0, 2, 1)) + for i in range(self.num_upsamples): + for up in self.ups[i]: + x = up(x) + first = i * self.num_kernels + residual = self.resblocks[first](x) + for j in range(1, self.num_kernels): + residual = residual + self.resblocks[first + j](x) + x = residual / self.num_kernels + x = self.conv_post(self.activation_post(x)) + return torch.clamp(x, min=-1.0, max=1.0).squeeze(1) + + @staticmethod + def _noise_like(reference: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor: + """Draw standard normal noise shaped like ``reference``; ``generator`` decides which RNG stream.""" + device = reference.device if generator is None else generator.device + noise = torch.randn(reference.shape, generator=generator, dtype=reference.dtype, device=device) + return noise.to(reference.device) diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py new file mode 100644 index 000000000..6ad12b02d --- /dev/null +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""AuK audio-editing pipeline for vLLM-Omni. + +Turns a layer-fused text condition (produced by the encoder stage) plus an +optional source clip into a 24 kHz waveform: the source clip is encoded to VAE +latents, a rectified-flow DiT integrates the target latents conditioned on +both, and the BigVGAN-flow decoder renders the waveform. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterable +from typing import Any, ClassVar + +import numpy as np +import torch +import torchaudio +from safetensors import safe_open +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE +from vllm_omni.diffusion.models.interface import ( + SupportAudioInput, + SupportAudioOutput, + SupportsComponentDiscovery, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.model_extras.auk import resolve_gen_frames + +logger = init_logger(__name__) + +# Longest ref + target latent sequence the released checkpoints accept. +MAX_LATENT_FRAMES = 65536 + +# Weight-file layout: the DiT lives under this prefix; the layer-fusion +# parameters sit next to it but belong to the encoder stage. +_DIT_PREFIX = "transformer." +_FUSION_KEYS = frozenset({"layer_weights", "layer_scale"}) + +# The distilled variant only reproduces its training recipe at these settings. +_FLASH_NFE = 4 +_FLASH_CFG = 0.0 +# Decorrelates the VAE posterior draw from the initial latent for the same request seed. +_VAE_SEED_OFFSET = 0x5EED_0A0C + + +def get_auk_post_process_func(od_config: OmniDiffusionConfig): + """Create the post-processing function for AuK audio output. + + Tensor output types pass through; anything else becomes a numpy waveform. + The sample rate is not attached here: the output formatter reads + ``AuKPipeline.audio_sample_rate`` for audio-output pipelines. + """ + + del od_config # The conversion does not depend on the config. + + def post_process_func( + audio: torch.Tensor, + output_type: str = "np", + ): + if output_type in ("pt", "latent"): + return audio + return audio.cpu().float().numpy() + + return post_process_func + + +def _prompt_mapping(prompt: Any) -> dict[str, Any]: + """Return the prompt dict, or an empty mapping for a bare text prompt.""" + + if isinstance(prompt, dict): + return prompt + return {} + + +def _unwrap_single(value: Any) -> Any: + """Unwrap the one-element lists the serving layer sometimes produces.""" + + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + +def _split_audio(audio: Any, default_sample_rate: int) -> tuple[Any, int]: + """Split a multimodal audio item into raw samples and a sample rate.""" + + if isinstance(audio, (tuple, list)) and len(audio) == 2 and isinstance(audio[1], (int, float)): + return audio[0], int(audio[1]) + return audio, default_sample_rate + + +class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComponentDiscovery): + """Instruction-driven audio generation and editing with AuK. + + One request per forward: the rectified-flow ODE runs over the whole target + span and the batch dimension carries the CFG branches, so there is nothing + to share between requests yet. + + Args: + od_config: OmniDiffusion configuration. ``od_config.model`` must be an + assembled AuK directory (``config.json``, ``auk.safetensors``, + ``vae.safetensors``), as produced by + ``tools/prepare_auk_checkpoint.py``. + prefix: Unused; kept for the pipeline construction contract. + """ + + supports_request_batch = False + + # Picked up by ``supports_audio_output`` in the diffusion engine so the + # default stage metadata reports ``final_output_type="audio"`` and the + # ``multimodal_output`` payload includes the sample rate. + support_audio_output: ClassVar[bool] = True + support_audio_input: ClassVar[bool] = True + audio_sample_rate: ClassVar[int] = 24000 + + # The DiT cannot run without an encoder-stage text condition, which the + # engine's synthetic warmup request has no way to produce. + dummy_run_num_frames: ClassVar[int] = 0 + + _dit_modules: ClassVar[list[str]] = ["dit"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ) -> None: + super().__init__() + del prefix # Weights are not namespaced: one checkpoint, one pipeline. + self.od_config = od_config + self.device = get_local_device() + self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16 + + model_dir = od_config.model + if not model_dir or not os.path.isdir(model_dir): + raise ValueError( + f"AuK needs an assembled local checkpoint directory, got {model_dir!r}. " + "Build one with tools/prepare_auk_checkpoint.py." + ) + config = _read_config(model_dir) + + vae_config = dict(config["vae"]) + self.latent_dim = int(vae_config["latent_dim"]) + self.hop_size = int(vae_config["downsample_rate"]) + self.sample_rate = int(vae_config["target_sample_rate"]) + + self.variant = str(config.get("variant", "base")) + self.is_flash = self.variant == "flash" + self.flash_t_grid = [float(t) for t in config.get("flash_t_grid") or ()] + if self.is_flash and len(self.flash_t_grid) != _FLASH_NFE + 1: + raise ValueError(f"The flash variant needs a {_FLASH_NFE + 1}-point flash_t_grid, got {self.flash_t_grid}.") + defaults = dict(config.get("defaults") or {}) + self.default_nfe = int(defaults.get("nfe", 32)) + self.default_cfg = float(defaults.get("cfg", 2.0)) + self.default_sway = defaults.get("sway", -1.0) + self._flash_lock_logged = False + + # The codec is small and numerically sensitive: it stays in fp32. It + # must sit on the inference device BEFORE load_weights folds weight + # norm: a CPU fold lands one ULP off the accelerator's and the encoder + # amplifies that into a 1e-2 latent drift. + self.vae = AuKVAE.from_config(vae_config["model_init_kwargs"]).to(device=self.device, dtype=torch.float32) + self.vae.load_weights(os.path.join(model_dir, "vae.safetensors")) + self.vae = self.vae.eval() + self.vae.requires_grad_(False) + self._check_vae_geometry() + + # The transformer owns every key the checkpoint tool writes into `dit`, + # `attn_mask_enabled` included, so the section passes straight through. + self.dit = AuKTransformer(latent_dim=self.latent_dim, **config["dit"]) + self.dit = self.dit.to(dtype=self.dtype) + self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True) + self.dit = self.dit.to(device=self.device).eval() + self.dit.requires_grad_(False) + + logger.info( + "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d", + self.variant, + self.dtype, + self.latent_dim, + self.hop_size, + self.sample_rate, + ) + + # The assembled checkpoint is not a diffusers layout: __init__ reads + # auk.safetensors and vae.safetensors directly, so the loader has no + # component sources to stream and load_weights only reports what is loaded. + weights_sources: ClassVar[tuple] = () + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stray = [name for name, _ in weights] + if stray: + logger.warning("AuKPipeline ignores %d loader-provided weights (e.g. %s)", len(stray), stray[:3]) + return {name for name, _ in self.named_parameters()} + + def _check_vae_geometry(self) -> None: + """Fail fast when config.json and the VAE disagree on latent geometry.""" + + declared = { + "latent_dim": (self.latent_dim, int(self.vae.latent_dim)), + "downsample_rate": (self.hop_size, int(self.vae.hop_size)), + "target_sample_rate": (self.sample_rate, int(self.vae.sample_rate)), + } + bad = {name: pair for name, pair in declared.items() if pair[0] != pair[1]} + if bad: + raise ValueError(f"config.json and the AuK VAE disagree (config, vae): {bad}") + if self.sample_rate != self.audio_sample_rate: + raise ValueError( + f"AuK advertises {self.audio_sample_rate} Hz output but the checkpoint is {self.sample_rate} Hz." + ) + + def _resolve_generator(self, sampling_params: Any) -> torch.Generator | None: + """Per-request generator; the process-global RNG is never seeded.""" + + generator = sampling_params.generator + if isinstance(generator, list): + if len(generator) > 1: + logger.warning( + "AuKPipeline runs one request per forward; using the first of %d generators", len(generator) + ) + generator = generator[0] if generator else None + if generator is None and sampling_params.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(int(sampling_params.seed)) + return generator + + def _prepare_waveform(self, audio: Any) -> torch.Tensor: + """Return mono float32 samples ``[1, T]`` at the VAE sample rate. + + ``T`` is truncated to a whole number of latent frames. The encoder is + causal, so dropping a trailing partial frame leaves every emitted frame + unchanged. + """ + + data, sample_rate = _split_audio(audio, self.sample_rate) + if isinstance(data, np.ndarray): + wav = torch.from_numpy(np.ascontiguousarray(data)) + else: + wav = torch.as_tensor(data) + wav = wav.detach().to(device="cpu", dtype=torch.float32) + + if wav.ndim == 1: + wav = wav.unsqueeze(0) + elif wav.ndim == 2: + # Accept both [channels, samples] and soundfile's [samples, channels]; + # the channel axis is the short one (at most 8 wide). + if wav.shape[0] > wav.shape[1] and wav.shape[1] <= 8: + wav = wav.transpose(0, 1) + if wav.shape[0] > 1: + wav = wav.mean(dim=0, keepdim=True) + else: + raise ValueError(f"AuK audio input must be [samples] or [channels, samples], got {tuple(wav.shape)}.") + + if wav.shape[-1] == 0: + raise ValueError("AuK audio input is empty.") + if not torch.isfinite(wav).all(): + raise ValueError("AuK audio input contains NaN or Inf.") + + if sample_rate != self.sample_rate: + wav = torchaudio.functional.resample(wav, sample_rate, self.sample_rate) + + frames = wav.shape[-1] // self.hop_size + if frames < 1: + raise ValueError( + f"AuK audio input is shorter than one latent frame ({self.hop_size} samples at {self.sample_rate} Hz)." + ) + return wav[..., : frames * self.hop_size] + + def _encode_source( + self, + audio: Any, + *, + sample: bool, + generator: torch.Generator | None, + ) -> torch.Tensor: + """Encode the source clip to normalized latents ``[1, np, latent_dim]``.""" + + if audio is None: + return torch.zeros(1, 0, self.latent_dim, device=self.device, dtype=torch.float32) + wav = self._prepare_waveform(_unwrap_single(audio)).to(self.device) + return self.vae.encode(wav, sample=sample, generator=generator) + + def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int: + """Target latent length, capped by the remaining context.""" + + frames = resolve_gen_frames( + None if gen_seconds is None else float(gen_seconds), + ref_frames, + sample_rate=self.sample_rate, + hop=self.hop_size, + ) + budget = MAX_LATENT_FRAMES - ref_frames + if budget < 1: + raise ValueError( + f"The source clip already fills the {MAX_LATENT_FRAMES}-frame context ({ref_frames} frames)." + ) + if frames > budget: + logger.warning("AuK target of %d frames exceeds the remaining context; clamping to %d.", frames, budget) + frames = budget + return frames + + def _resolve_schedule( + self, + sampling_params: Any, + knobs: dict[str, Any], + ) -> tuple[int, float, float | None, list[float] | None]: + """Resolve (nfe, cfg, sway, t_grid), locking the distilled recipe on Flash. + + A knob present but ``None`` counts as unspecified: the stage input + processor fills every key it knows about, whether the caller set it or + not. + """ + + steps = sampling_params.num_inference_steps + nfe = self.default_nfe if steps is None else int(steps) + if nfe < 1: + raise ValueError(f"num_inference_steps must be >= 1 for AuK; got {steps}") + cfg = float(sampling_params.guidance_scale) if sampling_params.guidance_scale_provided else self.default_cfg + sway = knobs.get("sway") + sway = self.default_sway if sway is None else float(sway) + raw_grid = knobs.get("t_grid") + t_grid = [float(t) for t in raw_grid] if raw_grid else None + + if not self.is_flash: + return nfe, cfg, sway, t_grid + + locked = (_FLASH_NFE, _FLASH_CFG, None, self.flash_t_grid) + if (nfe, cfg, sway, t_grid) != locked and not self._flash_lock_logged: + self._flash_lock_logged = True + logger.warning( + "AuK-Flash ignores per-request sampling knobs (asked nfe=%s cfg=%s sway=%s t_grid=%s); " + "using the distilled recipe nfe=%d cfg=%s with the checkpoint time grid.", + nfe, + cfg, + sway, + t_grid, + _FLASH_NFE, + _FLASH_CFG, + ) + return locked + + def _dit_autocast(self): + """Match the reference, which runs the DiT under bf16 autocast.""" + + return torch.autocast( + device_type=self.device.type, + dtype=self.dtype, + enabled=self.dtype in (torch.bfloat16, torch.float16), + ) + + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + """Generate one waveform. + + Args: + req: Request batch holding exactly one request. The prompt carries + ``prompt_embeds`` (the fused text condition ``[nt, 2048]``), an + optional ``multi_modal_data["audio"]`` source clip, and + ``additional_information["auk"]`` with ``gen_seconds``, + ``sway``, ``t_grid`` and ``vae_sample``. ``num_inference_steps``, + ``guidance_scale`` and ``seed``/``generator`` come from the + sampling params. + + Returns: + One ``DiffusionOutput`` whose ``output`` is a float32 mono waveform + ``[T]`` at 24 kHz, or the normalized target latents + ``[1, gen_frames, latent_dim]`` when ``output_type`` is ``latent``. + """ + + assert req.num_reqs == 1, f"AuKPipeline runs one request per forward, got {req.num_reqs}." + prompt = _prompt_mapping(req.prompts[0]) + sampling_params = req.sampling_params + knobs = dict((prompt.get("additional_information") or {}).get("auk") or {}) + + text = _unwrap_single(prompt.get("prompt_embeds")) + if text is None: + raise ValueError("AuK stage 1 needs `prompt_embeds`, the fused text condition from the encoder stage.") + text = torch.as_tensor(text).to(device=self.device, dtype=self.dtype) + if text.ndim == 3 and text.shape[0] == 1: + text = text[0] + if text.ndim != 2: + raise ValueError(f"AuK `prompt_embeds` must be [nt, text_hidden_dim], got {tuple(text.shape)}.") + text = text.unsqueeze(0) + c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device) + + generator = self._resolve_generator(sampling_params) + audio = (prompt.get("multi_modal_data") or {}).get("audio") + output_type = sampling_params.output_type or "np" + + with torch.inference_mode(): + # The VAE posterior draw gets its own generator, seeded with a + # fixed offset from the request seed: the initial latent below then + # still equals the reference's fresh manual_seed draw, and the two + # noise streams are not copies of each other. + vae_generator = None + if knobs.get("vae_sample") and generator is not None: + vae_seed = (int(generator.initial_seed()) + _VAE_SEED_OFFSET) % (2**63) + vae_generator = torch.Generator(device=self.device).manual_seed(vae_seed) + ref = self._encode_source(audio, sample=bool(knobs.get("vae_sample")), generator=vae_generator) + ref_frames = ref.shape[1] + ref_mask = torch.ones(ref.shape[:2], dtype=torch.bool, device=self.device) + gen_frames = self._target_frames(knobs.get("gen_seconds"), ref_frames) + nfe, cfg, sway, t_grid = self._resolve_schedule(sampling_params, knobs) + + with self._dit_autocast(): + latents = sample_latents( + self.dit, + text=text, + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + gen_frames=gen_frames, + nfe=nfe, + cfg_strength=cfg, + sway_sampling_coef=sway, + t_grid=t_grid, + seed=None, + latent_dim=self.latent_dim, + device=self.device, + dtype=torch.float32, + generator=generator, + ) + + latents = latents.float() + if not torch.isfinite(latents).all(): + raise RuntimeError("AuK generated latents contain NaN or Inf.") + if output_type == "latent": + return [DiffusionOutput(output=latents.detach().cpu())] + + wav = self.vae.decode(latents) + + # One mono waveform per request; the formatter expects [T]. + wav = wav.detach().to(device="cpu", dtype=torch.float32).reshape(-1) + if not torch.isfinite(wav).all(): + raise RuntimeError("AuK generated audio contains NaN or Inf.") + return [DiffusionOutput(output=wav)] + + +def _read_config(model_dir: str) -> dict[str, Any]: + """Read the assembled AuK ``config.json``.""" + + path = os.path.join(model_dir, "config.json") + try: + with open(path) as handle: + config = json.load(handle) + except FileNotFoundError as exc: + raise ValueError(f"AuK checkpoint is missing config.json: {path}") from exc + except json.JSONDecodeError as exc: + logger.error("AuK config.json is not valid JSON: %s", path) + raise ValueError(f"AuK config.json is not valid JSON: {path}") from exc + + missing = [key for key in ("dit", "vae") if key not in config] + if missing: + raise ValueError(f"AuK config.json is missing required sections {missing}: {path}") + return config + + +def _read_dit_weights(model_dir: str, dtype: torch.dtype) -> dict[str, torch.Tensor]: + """Read the DiT state dict from ``auk.safetensors``, cast to ``dtype``. + + Tensors are cast one at a time so the fp32 checkpoint is never fully + resident alongside the model. + """ + + path = os.path.join(model_dir, "auk.safetensors") + if not os.path.isfile(path): + raise ValueError(f"AuK checkpoint is missing auk.safetensors: {path}") + + with safe_open(path, framework="pt", device="cpu") as checkpoint: + keys = list(checkpoint.keys()) + state = dit_state_dict(((key, checkpoint.get_tensor(key).to(dtype)) for key in keys), _DIT_PREFIX) + + if not state: + raise ValueError(f"AuK checkpoint has no '{_DIT_PREFIX}*' weights: {path}") + unexpected = sorted(key for key in keys if key not in _FUSION_KEYS and not key.startswith(_DIT_PREFIX)) + if unexpected: + logger.warning("Ignoring %d unrecognized keys in %s, e.g. %s", len(unexpected), path, unexpected[:5]) + return state diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 39fff7df2..26b052a3a 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -116,6 +116,11 @@ _DIFFUSION_MODELS = { "pipeline_minimax_h3", "MiniMaxH3Pipeline", ), + "AuKPipeline": ( + "auk", + "pipeline_auk", + "AuKPipeline", + ), "StableAudioPipeline": ( "stable_audio", "pipeline_stable_audio", @@ -567,6 +572,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", + "AuKPipeline": "get_auk_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", diff --git a/vllm_omni/model_executor/models/auk/__init__.py b/vllm_omni/model_executor/models/auk/__init__.py new file mode 100644 index 000000000..d77d6a64d --- /dev/null +++ b/vllm_omni/model_executor/models/auk/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from vllm_omni.model_executor.models.auk.auk import AuKForConditionalGeneration + +__all__ = ["AuKForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/auk/auk.py b/vllm_omni/model_executor/models/auk/auk.py new file mode 100644 index 000000000..ebb20711d --- /dev/null +++ b/vllm_omni/model_executor/models/auk/auk.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK (Tencent): instruction-driven speech generation and editing. + +Stage 0 (``model_stage="encoder"``) runs the frozen Qwen2.5-Omni thinker as a +prefill-only encoder. The outputs of all decoder layers are fused with learned +softmax weights into one ``[tokens, hidden]`` text condition, emitted through +``multimodal_outputs["hidden_states"]["output"]`` for the diffusion stage +(``vllm_omni/diffusion/models/auk``). The stage samples exactly one token +(``max_tokens=1``, forced to EOS); that token is discarded downstream. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping +from functools import cached_property +from itertools import islice + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open +from transformers import PretrainedConfig +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import Qwen2_5OmniForConditionalGeneration +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerForConditionalGeneration, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, +) +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights + +logger = init_logger(__name__) + +# File next to the thinker shards in the assembled checkpoint directory that +# carries the DiT (``transformer.*``) and the two layer-fusion tensors. +AUK_WEIGHTS_FILE = "auk.safetensors" +LAYER_WEIGHTS_KEY = "layer_weights" +LAYER_SCALE_KEY = "layer_scale" +# The fused text condition travels as ``multimodal_outputs["hidden_states"]["output"]`` +# (``HiddenStatesStruct.output``), the same stage-wire slot MiniMax H3 uses. +TEXT_COND_KEY = "output" +# ``F.layer_norm`` epsilon used by the upstream fusion (torch default). +_FUSION_LN_EPS = 1e-5 + + +class AuKProcessingInfo(Qwen2_5OmniThinkerProcessingInfo): + """Qwen2.5-Omni thinker processing restricted to AuK's single audio input. + + AuK conditions on text plus at most one reference or source clip; images + and video are not part of its contract, so they are not advertised (and the + thinker then builds no vision tower). + """ + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen2_5OmniThinkerMultiModalProcessor, + info=AuKProcessingInfo, + dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, +) +class AuKForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE): + """Encoder stage of AuK: Qwen2.5-Omni thinker with learned layer fusion. + + The forward pass re-implements the thinker language model's layer loop so + the per-layer outputs can be fused on the fly: vLLM's decoder layers carry + the residual stream separately, and ``hidden + residual`` after layer ``i`` + is the HF ``output_hidden_states[i + 1]`` entry, while the last entry is + the final-normed state. Only one accumulator lives at a time. + """ + + have_multimodal_outputs = True + # The fused condition is the payload; the runner's default last-hidden + # copy and prefix-cached hidden states are not needed (MiniMax H3 pattern). + omni_pooler_payload_include_hidden = False + requires_full_prefix_cached_hidden_states = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.has_preprocess = False + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + + self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "encoder" + if self.model_stage != "encoder": + raise ValueError( + f"AuKForConditionalGeneration only hosts the 'encoder' stage, got model_stage={self.model_stage!r}; " + "the DiT stage is the AuKPipeline diffusion model." + ) + + thinker_config = config.thinker_config + self.thinker_config = thinker_config + self.thinker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + hf_config=thinker_config, + architectures=["Qwen2_5OmniThinkerModel"], + ) + # The audio-only mm limits keep the thinker from building its vision + # tower; drop it explicitly as well, as the upstream implementation does, + # so no vision weights are loaded or resident. + if getattr(self.thinker, "visual", None) is not None: + del self.thinker.visual + self.thinker.visual = None + self.model = self.thinker + self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors + + num_layers = thinker_config.text_config.num_hidden_layers + # Filled by load_weights() from AUK_WEIGHTS_FILE; buffers so the vLLM + # loader does not expect them among the thinker shards. + self.register_buffer("layer_weights", torch.zeros(num_layers, dtype=torch.float32)) + self.register_buffer("layer_scale", torch.ones(1, dtype=torch.float32)) + self._fusion_loaded = False + + # -------------------- delegations to the thinker -------------------- + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + return Qwen2_5OmniThinkerForConditionalGeneration.get_placeholder_str(modality, i) + + def get_language_model(self) -> nn.Module: + return self.thinker.get_language_model() + + @cached_property + def sampler(self): + if hasattr(self.thinker, "sampler"): + return self.thinker.sampler + return Sampler() + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + return self.thinker.embed_input_ids( + input_ids=input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def embed_multimodal(self, **kwargs): + return self.thinker.embed_multimodal(**kwargs) + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec] | None = None, + *, + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float] | None = None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + # The omni runner filters kwargs by this signature (the extended + # Qwen2.5-Omni form); the unified Qwen2.5-Omni class implements the + # audio-aware position math without touching instance state. + return Qwen2_5OmniForConditionalGeneration.get_mrope_input_positions( + self, + input_tokens, + mm_features, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None: + """Emit EOS immediately; the useful result is the prefill text condition.""" + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + text_config = self.thinker_config.text_config + logits = hidden_states.new_full( + (hidden_states.shape[0], int(text_config.vocab_size)), + torch.finfo(hidden_states.dtype).min, + ) + logits[:, self._eos_token_id()] = 0 + return logits + + def _eos_token_id(self) -> int: + for cfg in (self.thinker_config.text_config, self.thinker_config, self.vllm_config.model_config.hf_config): + eos = getattr(cfg, "eos_token_id", None) + if isinstance(eos, (list, tuple)) and eos: + eos = eos[0] + if isinstance(eos, int): + return eos + # Qwen2.5-Omni tokenizers end turns with <|im_end|>. + return 151645 + + # -------------------- forward -------------------- + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> OmniOutput: + if intermediate_tensors is not None: + raise NotImplementedError("AuK encoder does not support pipeline parallelism") + language_model = self.thinker.language_model.model + num_layers = len(language_model.layers) + if language_model.start_layer != 0 or language_model.end_layer != num_layers: + raise NotImplementedError("AuK encoder needs all decoder layers on one rank") + + if inputs_embeds is None: + hidden_states = language_model.embed_input_ids(input_ids) + else: + hidden_states = inputs_embeds + residual: torch.Tensor | None = None + + weights = torch.softmax(self.layer_weights, dim=0) + fused: torch.Tensor | None = None + normed_shape = (hidden_states.shape[-1],) + layers = islice(language_model.layers, language_model.start_layer, language_model.end_layer) + for idx, layer in enumerate(layers): + hidden_states, residual = layer(positions, hidden_states, residual) + if idx + 1 < num_layers: + layer_output = hidden_states + residual + else: + # HF's last hidden_states entry is post final norm. + hidden_states, _ = language_model.norm(hidden_states, residual) + layer_output = hidden_states + normed = F.layer_norm(layer_output.float(), normed_shape, eps=_FUSION_LN_EPS) + contribution = normed * weights[idx] + fused = contribution if fused is None else fused + contribution + + assert fused is not None + text_cond = (fused * self.layer_scale).to(hidden_states.dtype) + # Stage-wire layout: [tokens, features] under hidden_states.output so + # the runner slices it per request (and concatenates across chunked + # prefill steps); the stage-1 processor reads the same path. + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={"hidden_states": {"output": text_cond}}, + ) + + # -------------------- weights -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + thinker_weights: list[tuple[str, torch.Tensor]] = [] + skipped = 0 + for name, tensor in weights: + if name.startswith("thinker."): + thinker_weights.append((name, tensor)) + elif name.startswith(("talker.", "token2wav.", "transformer.")): + skipped += 1 + else: + logger.warning("AuK encoder: ignoring unexpected weight %s", name) + if skipped: + logger.debug("AuK encoder: skipped %d non-thinker weights", skipped) + + loaded = self.thinker.load_weights(thinker_weights) + loaded = add_prefix_to_loaded_weights(loaded, "thinker") + self._load_fusion_params() + return loaded + + def _load_fusion_params(self) -> None: + model_dir = self.vllm_config.model_config.model + path = os.path.join(model_dir, AUK_WEIGHTS_FILE) + if not os.path.isfile(path): + raise FileNotFoundError( + f"AuK encoder: {AUK_WEIGHTS_FILE} not found in {model_dir!r}; assemble the checkpoint with " + "tools/prepare_auk_checkpoint.py" + ) + with safe_open(path, framework="pt", device="cpu") as f: + layer_weights = f.get_tensor(LAYER_WEIGHTS_KEY).float() + layer_scale = f.get_tensor(LAYER_SCALE_KEY).float().reshape(1) + if layer_weights.shape != self.layer_weights.shape: + raise ValueError( + f"AuK encoder: {LAYER_WEIGHTS_KEY} has {tuple(layer_weights.shape)} entries, " + f"thinker has {tuple(self.layer_weights.shape)} layers" + ) + self.layer_weights.copy_(layer_weights.to(self.layer_weights.device)) + self.layer_scale.copy_(layer_scale.to(self.layer_scale.device)) + self._fusion_loaded = True + logger.info( + "AuK encoder: loaded layer fusion (scale %.3f, top layer weight %.3f at layer %d)", + float(layer_scale[0]), + float(torch.softmax(layer_weights, 0).max()), + int(torch.softmax(layer_weights, 0).argmax()) + 1, + ) diff --git a/vllm_omni/model_executor/models/auk/pipeline.py b/vllm_omni/model_executor/models/auk/pipeline.py new file mode 100644 index 000000000..8b9d6855b --- /dev/null +++ b/vllm_omni/model_executor/models/auk/pipeline.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK pipeline topology (frozen). + +Stage 0: Encoder — frozen Qwen2.5-Omni thinker, prefill only, emits the + learned layer fusion as ``multimodal_outputs["hidden_states"]["output"]`` +Stage 1: DiT — rectified-flow transformer plus the BigVGAN-flow VAE, + emits a 24 kHz waveform +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.auk" + +AUK_PIPELINE = PipelineConfig( + model_type="auk", + default_deploy_config_name="auk.yaml", + model_arch="AuKForConditionalGeneration", + hf_architectures=("AuKForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="encoder", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=False, + owns_tokenizer=True, + requires_multimodal_data=True, + hf_config_name="thinker_config", + engine_output_type="latent", + model_arch="AuKForConditionalGeneration", + # No model_subdir: tools/prepare_auk_checkpoint.py assembles one + # flat directory, so the encoder finds the thinker shards, the + # tokenizer, the audio processor and the layer-fusion tensors in + # auk.safetensors all under the pipeline root. + # One prefill step produces the text condition; the sampled token + # is discarded, and detokenizing it would only cost latency. + sampling_constraints={"max_tokens": 1, "temperature": 0.0, "detokenize": False}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + requires_multimodal_data=True, + final_output=True, + final_output_type="audio", + model_arch="AuKPipeline", + custom_process_input_func=f"{_PROC}.encoder2dit", + omni_kv_config={"need_recv_cache": False}, + # Single replica, and the whole ODE runs inside one forward, so + # the stage stays in the orchestrator process. + inline_diffusion=True, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index b9373e5f2..febda68af 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -87,6 +87,11 @@ _OMNI_MODELS = { "cosyvoice3", "CosyVoice3Model", ), + "AuKForConditionalGeneration": ( + "auk", + "auk", + "AuKForConditionalGeneration", + ), "NemotronDenseForCausalLM": ( "audex", "audex_thinker", diff --git a/vllm_omni/model_executor/stage_input_processors/auk.py b/vllm_omni/model_executor/stage_input_processors/auk.py new file mode 100644 index 000000000..7c417aee1 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/auk.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stage input processor for AuK: thinker encoder -> diffusion transition.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +from vllm.inputs import TextPrompt +from vllm.logger import init_logger + +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + +# The encoder emits the fused condition as +# ``multimodal_outputs["hidden_states"]["output"]`` (``HiddenStatesStruct.output``, +# the same stage-wire slot MiniMax H3 uses). The payload flattener may deliver +# it under the dotted spelling instead, and a plain ``text_cond`` key is +# accepted so an encoder that emits the condition directly still works. +HIDDEN_STATES_KEY = "hidden_states" +TEXT_COND_KEY = "output" +_FALLBACK_KEYS = ("hidden_states.output", "text_cond") + +# Per-request knobs the prompt carries for stage 1. nfe/cfg/seed are not here: +# they ride the stage-1 OmniDiffusionSamplingParams instead. +_AUK_KNOBS = ("gen_seconds", "sway", "t_grid", "vae_sample") + + +def _as_dict(prompt: Any) -> dict[str, Any]: + """Normalize the many prompt spellings the orchestrator may pass.""" + if isinstance(prompt, list): + prompt = prompt[0] if prompt else {} + if prompt is None: + return {} + if isinstance(prompt, dict): + return prompt + if hasattr(prompt, "_asdict"): + return prompt._asdict() + if hasattr(prompt, "__dict__"): + return vars(prompt) + return {} + + +def _has_audio(mm_data: Any) -> bool: + return isinstance(mm_data, Mapping) and mm_data.get("audio") is not None + + +def _from_hidden_states(hidden: Any) -> Any: + """Read the fused condition out of the payload's hidden-states slot.""" + if isinstance(hidden, Mapping): + return hidden.get(TEXT_COND_KEY) + # A deserialized payload arrives as a msgspec struct rather than a dict. + return getattr(hidden, TEXT_COND_KEY, None) + + +def _extract_text_cond(ar_output: Any) -> Any: + """Read the fused text condition from the encoder output. + + ``multimodal_output`` is attached to the RequestOutput; the CompletionOutput + is checked as a fallback for vLLM versions that attach it there. + """ + for holder in (ar_output, *(getattr(ar_output, "outputs", None) or ())): + mm_output = getattr(holder, "multimodal_output", None) + if not isinstance(mm_output, Mapping): + continue + text_cond = _from_hidden_states(mm_output.get(HIDDEN_STATES_KEY)) + if text_cond is not None: + return text_cond + for key in _FALLBACK_KEYS: + if mm_output.get(key) is not None: + return mm_output[key] + return None + + +def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: + """Coerce the emitted condition to a single 2-D CPU tensor.""" + if isinstance(text_cond, list): + if not text_cond: + raise ValueError("AuK encoder emitted an empty text condition list") + parts = [_to_cpu_tensor(part) for part in text_cond] + return torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] + if isinstance(text_cond, torch.Tensor): + tensor = text_cond + else: + # numpy array or any array-like the connector round-tripped. + tensor = torch.as_tensor(text_cond) + tensor = tensor.detach().cpu() + if tensor.ndim != 2: + raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(tensor.shape)}") + return tensor + + +def encoder2dit( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | list | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> dict[str, Any] | None: + """Turn the AuK encoder stage's fused hidden states into a diffusion prompt. + + The source clip is not decoded here: the original ``multi_modal_data`` is + forwarded so stage 1 resamples and VAE-encodes it itself. It is forwarded + unconditionally because the editing and enhancement tasks cannot run + without it; ``requires_multimodal_data`` only records the stage's + declaration. + """ + del streaming_context, requires_multimodal_data + if not source_outputs: + return None + + ar_output = source_outputs[0] + text_cond = _extract_text_cond(ar_output) + if text_cond is None: + raise ValueError( + "AuK encoder stage produced no hidden_states.output multimodal payload; " + "stage 1 cannot run without the fused thinker hidden states" + ) + prompt_embeds = _to_cpu_tensor(text_cond) + + original = _as_dict(prompt) + original_knobs = original.get("additional_information") or {} + if isinstance(original_knobs, Mapping): + original_knobs = original_knobs.get("auk") or {} + if not isinstance(original_knobs, Mapping): + original_knobs = {} + + mm_data = original.get("multi_modal_data") + knobs: dict[str, Any] = {key: original_knobs.get(key) for key in _AUK_KNOBS} + declared_audio = original_knobs.get("has_audio") + knobs["has_audio"] = bool(declared_audio) if declared_audio is not None else _has_audio(mm_data) + + logger.debug( + "[encoder2dit] text_cond=%s dtype=%s has_audio=%s gen_seconds=%s", + tuple(prompt_embeds.shape), + prompt_embeds.dtype, + knobs["has_audio"], + knobs["gen_seconds"], + ) + + return { + "prompt": "", + "prompt_embeds": prompt_embeds, + "multi_modal_data": mm_data, + "additional_information": {"auk": knobs}, + } diff --git a/vllm_omni/model_extras/auk.py b/vllm_omni/model_extras/auk.py new file mode 100644 index 000000000..16d6af52a --- /dev/null +++ b/vllm_omni/model_extras/auk.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request helpers for AuK: prompt rendering, sampling params, duration policy. + +AuK's contract is one instruction plus an optional reference clip in, audio +out; the task (zero-shot TTS, instruct TTS, content or acoustic editing, +enhancement, separation) is carried by the instruction text alone, so there is +no per-task request surface. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Any + +from vllm import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +if TYPE_CHECKING: + import numpy as np + import torch + +# The upstream ChatML render of a single user turn through +# ``Qwen2_5OmniProcessor.apply_chat_template(..., add_generation_prompt=True)`` +# with the template's default system turn. +_CHATML = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "{instruction}{audio_part}<|im_end|>\n" + "<|im_start|>assistant\n" +) +_AUDIO_PART = "<|audio_bos|><|AUDIO|><|audio_eos|>" + +# Literal marker the reference appends to the instruction when no reference +# clip is supplied. The model was trained with it, so it is not decorative. +NO_PROMPT_AUDIO = "||" + +# Latent frame rate: 24 kHz waveform, VAE downsample 480 -> 50 frames/s. +SAMPLE_RATE = 24000 +HOP = 480 + +DEFAULT_NFE = 32 +DEFAULT_CFG = 2.0 +DEFAULT_SWAY = -1.0 + + +def auk_prompt( + instruction: str, + audio: tuple[np.ndarray | torch.Tensor, int] | None = None, + *, + gen_seconds: float | None = None, + sway: float = DEFAULT_SWAY, + t_grid: list[float] | None = None, + vae_sample: bool = False, +) -> dict[str, Any]: + """Build the stage-0 prompt for one AuK request. + + Args: + instruction: The task instruction; for text-only requests the + no-reference marker is appended if it is not already there. + audio: Reference or source clip as ``(waveform, sample_rate)``. + ``None`` selects the text-only (instruct TTS) path. + gen_seconds: Target duration. ``None`` reuses the source clip's + length and is an error for text-only requests. + sway: Sway coefficient for the Euler time grid. + t_grid: Explicit time grid, overriding ``sway``. + vae_sample: Draw the reference latent from the VAE posterior instead + of taking its mean. + + Returns: + A prompt dict for ``Omni.generate``. + """ + has_audio = audio is not None and not (isinstance(audio, (tuple, list)) and (len(audio) == 0 or audio[0] is None)) + if not has_audio and not instruction.endswith(NO_PROMPT_AUDIO): + instruction = instruction + NO_PROMPT_AUDIO + + prompt: dict[str, Any] = { + "prompt": _CHATML.format( + instruction=instruction, + audio_part=_AUDIO_PART if has_audio else "", + ), + "additional_information": { + "auk": { + "gen_seconds": gen_seconds, + "sway": sway, + "t_grid": list(t_grid) if t_grid is not None else None, + "vae_sample": vae_sample, + "has_audio": has_audio, + } + }, + } + if has_audio: + prompt["multi_modal_data"] = {"audio": audio} + return prompt + + +def auk_sampling_params( + *, + nfe: int = DEFAULT_NFE, + cfg: float = DEFAULT_CFG, + seed: int | None = 0, +) -> list[SamplingParams | OmniDiffusionSamplingParams]: + """Build the per-stage sampling params for one AuK request. + + Stage 0 is a prefill-only encoder, so its vLLM sampling is a no-op that + only has to satisfy the scheduler. Stage 1 carries the real knobs: ``nfe`` + Euler steps and CFG strength. AuK-Flash ignores both and pins its + distilled recipe. + """ + return [ + SamplingParams(max_tokens=1, temperature=0.0, detokenize=False, seed=seed), + OmniDiffusionSamplingParams(num_inference_steps=nfe, guidance_scale=cfg, seed=seed), + ] + + +def resolve_gen_frames( + gen_seconds: float | None, + ref_frames: int, + *, + sample_rate: int = SAMPLE_RATE, + hop: int = HOP, +) -> int: + """Resolve the target latent length in frames. + + An explicit duration wins; otherwise the target is as long as the source + clip. A text-only request with no duration has nothing to generate, which + is an error rather than a silent one-frame output. + """ + if gen_seconds is not None: + if gen_seconds <= 0: + raise ValueError(f"gen_seconds must be positive; got {gen_seconds}") + return max(1, math.ceil(gen_seconds * sample_rate / hop)) + if ref_frames > 0: + return ref_frames + raise ValueError("gen_seconds is required when the request carries no source audio") + + +__all__ = [ + "NO_PROMPT_AUDIO", + "auk_prompt", + "auk_sampling_params", + "resolve_gen_frames", +] diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 6beb91018..fdbfb4b60 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -34,11 +34,13 @@ _CLASS_TO_MODULE: dict[str, str] = { "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", + "AuKConfig": "vllm_omni.transformers_utils.configs.auk", "SenseNovaU1Config": "vllm_omni.transformers_utils.configs.sensenova_u1", } __all__ = [ "Audio8TTSConfig", + "AuKConfig", "Audio8TTSFastARConfig", "Audio8TTSSlowARConfig", "HiggsAudioV3Config", @@ -82,6 +84,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import auk as _auk # noqa: F401, E402 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/auk.py b/vllm_omni/transformers_utils/configs/auk.py new file mode 100644 index 000000000..d1e8fa805 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/auk.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Config for the AuK instruction-driven audio generation model. + +The upstream checkpoint ships ``config.yaml`` plus loose safetensors, so an +assembled directory is produced by ``tools/prepare_auk_checkpoint.py``: the +Qwen2.5-Omni-3B ``config.json`` with this ``model_type``, AuK's architecture, +and the five sections below added at the top level. + +Subclassing ``Qwen2_5OmniConfig`` is load-bearing, not stylistic. The encoder +stage reuses the in-tree Qwen2.5-Omni thinker processor, whose +``get_hf_config`` resolves ``self.ctx.get_hf_config(Qwen2_5OmniConfig)`` and +asserts the type, so a bare ``PretrainedConfig`` subclass would fail before +the first request. Inheriting also leaves the parent's nested +``thinker_config`` / ``talker_config`` / ``token2wav_config`` handling and its +``get_text_config`` in place, which is what the encoder stage reads. +""" + +from typing import Any + +from transformers import AutoConfig +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniConfig + +# The distilled AuK-Flash checkpoint pins a 4-step time grid recorded in its +# metadata. The values are exact and must never be rounded: the student was +# trained on them, and a rounded grid changes the output amplitude. +AUK_FLASH_T_GRID: tuple[float, ...] = ( + 0.0, + 0.07612049579620361, + 0.2928932309150696, + 0.6173166036605835, + 1.0, +) + +DEFAULT_DIT: dict[str, Any] = { + "dim": 1536, + "heads": 24, + "dim_head": 64, + "ff_mult": 2, + "text_hidden_dim": 2048, + "num_layers": 10, + "num_single_layers": 20, + "attn_mask_enabled": True, +} + +DEFAULT_VAE: dict[str, Any] = { + "latent_dim": 64, + "downsample_rate": 480, + "target_sample_rate": 24000, + "model_init_kwargs": { + "upsample_rates": [5, 4, 3, 2, 2, 2], + "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4], + "upsample_initial_channel": 1536, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "downsample_rates": [2, 2, 2, 3, 4, 5], + "downsample_channels": [12, 24, 48, 96, 192, 384, 768], + "snake_logscale": True, + "latent_dim": 64, + "use_vae": True, + "causal": True, + "flow_hidden_channels": 256, + "act_causal": True, + }, +} + +# Base AuK sampling recipe; AuK-Flash replaces nfe and cfg with its own. +DEFAULT_SAMPLING: dict[str, Any] = {"nfe": 32, "cfg": 2.0, "sway": -1.0} + +VARIANTS = ("base", "flash") + + +class AuKConfig(Qwen2_5OmniConfig): + """Assembled AuK checkpoint config: a Qwen2.5-Omni config plus DiT and VAE.""" + + model_type = "auk" + + def __init__( + self, + dit: dict[str, Any] | None = None, + vae: dict[str, Any] | None = None, + variant: str = "base", + flash_t_grid: list[float] | None = None, + defaults: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + if variant not in VARIANTS: + raise ValueError(f"AuK variant must be one of {VARIANTS}; got {variant!r}") + + # Assigned before the parent's init so nothing it runs during + # validation sees a half-built config. + self.dit = {**DEFAULT_DIT, **(dit or {})} + self.vae = {**DEFAULT_VAE, **(vae or {})} + self.variant = variant + self.flash_t_grid = list(flash_t_grid) if flash_t_grid is not None else list(AUK_FLASH_T_GRID) + self.defaults = {**DEFAULT_SAMPLING, **(defaults or {})} + + kwargs.setdefault("architectures", ["AuKForConditionalGeneration"]) + # The parent owns thinker_config / talker_config / token2wav_config and + # get_text_config; nothing here overrides them. + super().__init__(**kwargs) + + @property + def is_flash(self) -> bool: + """Whether this checkpoint is the distilled 4-step student.""" + return getattr(self, "variant", "base") == "flash" + + @property + def latent_dim(self) -> int: + return int(self.vae["latent_dim"]) + + @property + def downsample_rate(self) -> int: + return int(self.vae["downsample_rate"]) + + @property + def target_sample_rate(self) -> int: + return int(self.vae["target_sample_rate"]) + + +AutoConfig.register("auk", AuKConfig) + +__all__ = [ + "AUK_FLASH_T_GRID", + "DEFAULT_DIT", + "DEFAULT_SAMPLING", + "DEFAULT_VAE", + "VARIANTS", + "AuKConfig", +] From 6fcab74aa9b048ab16b2a3815484649ec1ea32f6 Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Fri, 11 Sep 2026 13:54:28 -0700 Subject: [PATCH 031/209] [XPU][Docker] Align XPU image and CI with vLLM v0.29.0 (#7441) Signed-off-by: Joshna-Medisetty --- .buildkite/intel/pipeline-intel.yml | 2 +- docker/Dockerfile.xpu | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/intel/pipeline-intel.yml b/.buildkite/intel/pipeline-intel.yml index d6134b028..1e2f86af2 100644 --- a/.buildkite/intel/pipeline-intel.yml +++ b/.buildkite/intel/pipeline-intel.yml @@ -10,7 +10,7 @@ steps: DOCKER_BUILDKIT: "1" # Buildkite will automatically replace this with the actual commit hash VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}" - VLLM_VERSION: "v0.28.0" + VLLM_VERSION: "v0.29.0" priority: 100 timeout_in_minutes: 60 soft_fail: false diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index acfa71588..018bc4316 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -1,4 +1,4 @@ -# Target vLLM version: v0.28.0. Keep VLLM_VERSION aligned with the target vLLM +# Target vLLM version: v0.29.0. Keep VLLM_VERSION aligned with the target vLLM # release on every rebase. # # This file layers vLLM-Omni on top of the published upstream XPU image, which @@ -18,7 +18,7 @@ # # .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback. ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu -ARG VLLM_VERSION=v0.28.0 +ARG VLLM_VERSION=v0.29.0 ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION} FROM ${VLLM_BASE} AS vllm-omni From 13b85c562248831a7c5f527cc08531dc72c0fbad Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Sat, 12 Sep 2026 00:59:49 +0200 Subject: [PATCH 032/209] [Bugfix] Add explicit error when using CFGP with distilled Cosmos3 models (#7427) Signed-off-by: Maciej Bala --- .../diffusion/parallelism/cfg_parallel.md | 6 +++ .../models/cosmos3/test_cosmos3_pipeline.py | 45 ++++++++++++++++--- .../models/cosmos3/pipeline_cosmos3.py | 35 ++++++++------- 3 files changed, 65 insertions(+), 21 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism/cfg_parallel.md b/docs/user_guide/diffusion/parallelism/cfg_parallel.md index 544754403..1e3b998d8 100644 --- a/docs/user_guide/diffusion/parallelism/cfg_parallel.md +++ b/docs/user_guide/diffusion/parallelism/cfg_parallel.md @@ -19,6 +19,12 @@ CFG-Parallel accelerates diffusion models by distributing positive and negative See supported models list in [Supported Models](../../diffusion_features.md#supported-models). +Distilled Cosmos3 checkpoints, including `nvidia/Cosmos3-Super-Image2Video-4Step`, +run without CFG and reject `--cfg-parallel-size` greater than 1 at startup. +Use `--ulysses-degree 2` for two-GPU inference. These checkpoints always use +`guidance_scale=1.0`; other explicitly requested values produce a warning, and +`negative_prompt` does not affect generation. + --- ## Quick Start diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 302409aeb..9892a72d7 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -9,6 +9,7 @@ import types from dataclasses import dataclass from types import SimpleNamespace from typing import Any +from unittest.mock import Mock import numpy as np import pytest @@ -376,23 +377,31 @@ def _capture_tokenize_calls(pipeline: Any) -> list[dict[str, Any]]: @pytest.mark.parametrize( - ("provided", "value", "default", "is_distilled", "expected"), + ("provided", "value", "default", "is_distilled", "expected", "expected_warning"), [ - (False, 1.0, 7.0, False, 7.0), - (True, 1.0, 7.0, False, 1.0), - (True, 4.5, 7.0, False, 4.5), - (False, 1.0, 7.0, True, 1.0), - (True, 4.5, 7.0, True, 1.0), + (False, 1.0, 7.0, False, 7.0, False), + (True, 1.0, 7.0, False, 1.0, False), + (True, 4.5, 7.0, False, 4.5, False), + (False, 1.0, 7.0, True, 1.0, False), + (True, 1.0, 7.0, True, 1.0, False), + (True, 4.5, 7.0, True, 1.0, True), + (True, 0.0, 7.0, True, 1.0, True), ], ) def test_resolve_guidance_scale( make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, provided: bool, value: float, default: float, is_distilled: bool, expected: float, + expected_warning: bool, ) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + warning_once = Mock() + monkeypatch.setattr(pipeline_cosmos3.logger, "warning_once", warning_once) pipeline = make_cosmos3_pipeline() pipeline.is_distilled_model = is_distilled sp = make_sampling_params( @@ -401,6 +410,12 @@ def test_resolve_guidance_scale( ) assert pipeline._resolve_guidance_scale(sp, default) == expected + if expected_warning: + warning_once.assert_called_once() + assert "overridden to 1.0" in warning_once.call_args.args[0] + assert "negative_prompt does not affect generation" in warning_once.call_args.args[0] + else: + warning_once.assert_not_called() def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None: @@ -735,6 +750,7 @@ def _make_od_config( custom_pipeline_args={}, model_config=model_config or {}, tf_model_config=tf_model_config, + parallel_config=SimpleNamespace(cfg_parallel_size=1, ulysses_degree=1), ) @@ -770,6 +786,7 @@ def test_pipeline_init_uses_flow_unipc_with_cosmos3_defaults(stub_real_pipeline_ assert pipeline._engine_init_flow_shift == 2.5 +@pytest.mark.parametrize("cfg_parallel_size,ulysses_degree", [(1, 1), (1, 2), (2, 1), (2, 2)]) @pytest.mark.parametrize( ("scheduler_class_name", "expected_distilled"), [ @@ -783,6 +800,8 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file( monkeypatch: pytest.MonkeyPatch, scheduler_class_name: str, expected_distilled: bool, + cfg_parallel_size: int, + ulysses_degree: int, ) -> None: import json @@ -831,6 +850,20 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file( od_config = _make_od_config(sound_gen=False) od_config.model = str(tmp_path) + od_config.parallel_config.cfg_parallel_size = cfg_parallel_size + od_config.parallel_config.ulysses_degree = ulysses_degree + if expected_distilled and cfg_parallel_size > 1: + monkeypatch.setattr( + pipeline_cosmos3.AutoTokenizer, + "from_pretrained", + lambda *args, **kwargs: pytest.fail("component loading must not start for distilled CFG parallelism"), + ) + with pytest.raises(ValueError, match="Set --cfg-parallel-size 1 and use --ulysses-degree"): + Cosmos3OmniDiffusersPipeline(od_config=od_config) + assert StubFlowMatchScheduler.from_config_calls == [] + assert StubFlowUniPCScheduler.from_config_calls == [] + return + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) assert pipeline.is_distilled_model is expected_distilled diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 241eae805..5977e0eea 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -948,6 +948,20 @@ class Cosmos3OmniDiffusersPipeline( model_path = od_config.model local_files_only = os.path.exists(model_path) + # Validate guidance parallelism from checkpoint metadata before loading + # components. Distilled checkpoints have only a conditional branch. + scheduler_config = FlowUniPCMultistepScheduler.load_config( + model_path, + subfolder="scheduler", + local_files_only=local_files_only, + ) + self.is_distilled_model = scheduler_config.get("_class_name") == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS + if self.is_distilled_model and od_config.parallel_config.cfg_parallel_size > 1: + raise ValueError( + "Distilled Cosmos3 checkpoints run without classifier-free guidance. " + "Set --cfg-parallel-size 1 and use --ulysses-degree for multi-GPU inference." + ) + # --- Tokenizer --- self.tokenizer = AutoTokenizer.from_pretrained( model_path, @@ -1010,20 +1024,7 @@ class Cosmos3OmniDiffusersPipeline( logger.info("Cosmos3: session state manager enabled (max_sessions=%d)", mm_max_sessions) # --- Scheduler --- - # Distilled model differs from regular one only by scheduler, - # distilled one uses FlowMatchEulerDiscreteScheduler, while - # regular should use FlowUniPCMultistepScheduler - - scheduler_config = FlowUniPCMultistepScheduler.load_config( - model_path, - subfolder="scheduler", - local_files_only=local_files_only, - ) - - scheduler_class_name = scheduler_config.get("_class_name") - - self.is_distilled_model = False - if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS: + if self.is_distilled_model: fixed_step_config = scheduler_config.get("fixed_step_sampler_config") if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde": raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.") @@ -1035,7 +1036,6 @@ class Cosmos3OmniDiffusersPipeline( stochastic_sampling=True, ) self._scheduler_init_t_list = list(t_list) - self.is_distilled_model = True else: # Preserve compatible solver settings from the checkpoint, but keep # the base shift neutral. The concrete request shift is applied when @@ -1811,6 +1811,11 @@ class Cosmos3OmniDiffusersPipeline( def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float: if self.is_distilled_model: + if sp.guidance_scale_provided and float(sp.guidance_scale) != 1.0: + logger.warning_once( + "Distilled Cosmos3 checkpoints run without classifier-free guidance. " + "The requested guidance_scale is overridden to 1.0; negative_prompt does not affect generation." + ) return 1.0 if sp.guidance_scale_provided: return float(sp.guidance_scale) From 220366db005a79e92877c50266358375d828c349 Mon Sep 17 00:00:00 2001 From: DanaerLee Date: Sat, 12 Sep 2026 11:31:05 +0800 Subject: [PATCH 033/209] [Perf][Diffusion] Run MammothModa2 DiT attention through the shared attention layer (#7094) Signed-off-by: MrlixiangWE Co-authored-by: Hongsheng Liu --- .../attention/test_attention_config.py | 170 +++++++++++++++- .../models/ltx2/test_ltx2_transformer.py | 1 + .../models/mammoth_moda2/__init__.py | 0 .../mammoth_moda2/test_dit_attention.py | 189 ++++++++++++++++++ .../mammoth_moda2/test_dit_attention_cuda.py | 151 ++++++++++++++ vllm_omni/diffusion/attention/layer.py | 31 ++- .../mammoth_moda2/mammothmoda2_dit_model.py | 148 +++++--------- 7 files changed, 582 insertions(+), 108 deletions(-) create mode 100644 tests/diffusion/models/mammoth_moda2/__init__.py create mode 100644 tests/diffusion/models/mammoth_moda2/test_dit_attention.py create mode 100644 tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 656161589..4a053b1bf 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -10,6 +10,7 @@ Tests cover: - AttentionMetadata.extra field """ +from dataclasses import replace from types import SimpleNamespace from unittest.mock import Mock @@ -17,7 +18,12 @@ import pytest import torch import vllm_omni.diffusion.attention.layer as layer_mod -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +import vllm_omni.diffusion.attention.selector as selector_mod +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, + PackedPaddingMetadata, + QueryRange, +) from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.config import ( @@ -477,6 +483,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig: @pytest.mark.parametrize("mindiesd_available", [False, True]) def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available): attention = Attention.__new__(Attention) + attention.allow_fp32_fallback = False attention._scheduler_paged_kv = True attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" @@ -519,6 +526,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig: @pytest.mark.parametrize("fa4_available", [False, True]) def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available): attention = Attention.__new__(Attention) + attention.allow_fp32_fallback = False attention._scheduler_paged_kv = True attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" @@ -561,6 +569,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig: def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch): attention = Attention.__new__(Attention) + attention.allow_fp32_fallback = False attention._scheduler_paged_kv = False attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" @@ -900,6 +909,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig: return sentinel fake_attention = SimpleNamespace( + allow_fp32_fallback=False, attention=SimpleNamespace( forward=_selected_forward, ), @@ -938,6 +948,164 @@ class TestAttentionInitUsesCurrentDiffusionConfig: Attention._run_ring_attention(fake_attention, query, query, query, metadata) +class TestOptInFloat32Fallback: + @staticmethod + def _make_attention(monkeypatch, *, backend=None, allow=True, kv_cache_dtype=None): + events, state = [], {} + context = object() + + class _SelectedImpl: + def __init__(self, **kwargs): + self.kwargs = kwargs + + @staticmethod + def supports_kv_cache_dtype(*args): + return True + + def forward(self, query, key, value, metadata): + events.append("selected") + state["kernel_call"] = (query, key, value, metadata) + state["kernel_output"] = query + 4 + return state["kernel_output"] + + class _FallbackImpl(_SelectedImpl): + def forward(self, query, key, value, metadata): + output = super().forward(query, key, value, metadata) + events[-1] = "fallback" + return output + + class _FlashBackend(FlashAttentionBackend): + @staticmethod + def get_impl_cls(): + return _SelectedImpl + + class _SDPABackend(layer_mod.SDPABackend): + @staticmethod + def get_impl_cls(): + return _SelectedImpl + + class _Strategy: + def pre_attention(self, query, key, value, metadata): + events.append("pre") + state["prepared"] = ( + query + 1, + key + 2, + value + 3, + replace(metadata) if metadata is not None else None, + ) + return (*state["prepared"], context) + + def post_attention(self, output, ctx): + events.append("post") + assert output is state["kernel_output"] + assert ctx is context + return output + 7 + + # Keep real role/config resolution and construction; replace hardware + # lookup and kernels. These are CPU dispatch tests, not SP collectives. + monkeypatch.setattr( + selector_mod, + "_cached_get_backend_cls", + lambda name, *args, **kwargs: _SDPABackend if name == "TORCH_SDPA" else _FlashBackend, + ) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FallbackImpl)) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: _Strategy()) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + config = OmniDiffusionConfig( + diffusion_attention_config={"default": backend} if backend else {}, + diffusion_kv_cache_dtype=kv_cache_dtype, + ) + with set_current_diffusion_config(config): + attention = Attention( + num_heads=4, + num_kv_heads=2, + head_size=8, + causal=False, + softmax_scale=0.25, + **({"allow_fp32_fallback": True} if allow else {}), + ) + return attention, events, state + + @pytest.mark.parametrize( + "backend,allow,dtype,is_cuda,expected", + [ + pytest.param(None, True, torch.float32, True, "fallback", id="auto-fp32"), + pytest.param(None, False, torch.float32, True, "selected", id="default-opt-out"), + pytest.param("FLASH_ATTN", True, torch.float32, True, "selected", id="explicit-flash"), + pytest.param("TORCH_SDPA", True, torch.float32, True, "selected", id="explicit-sdpa"), + pytest.param(None, True, torch.bfloat16, True, "selected", id="bf16"), + pytest.param(None, True, torch.float32, False, "selected", id="cpu-fp32"), + ], + ) + def test_public_forward_preserves_strategy_and_backend_choice( + self, monkeypatch, backend, allow, dtype, is_cuda, expected + ): + attention, events, state = self._make_attention(monkeypatch, backend=backend, allow=allow) + # Only the dispatch predicate is simulated; every tensor stays on CPU. + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: is_cuda)) + query = torch.ones(1, 2, 4, 8, dtype=dtype) + key = torch.full((1, 2, 2, 8), 2, dtype=dtype) + value = torch.full_like(key, 3) + metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False]])) + + output = attention(query, key, value, metadata) + + assert events == ["pre", expected, "post"] + assert all(got is prepared for got, prepared in zip(state["kernel_call"], state["prepared"])) + assert state["kernel_call"][3] is not metadata + assert state["kernel_call"][3].attn_mask is metadata.attn_mask + torch.testing.assert_close(output, query + 12) + assert attention.backend_explicit is (backend is not None) + impl = attention.sdpa_fallback if expected == "fallback" else attention.attention + assert impl.kwargs["num_kv_heads"] == 2 + assert impl.kwargs["softmax_scale"] == 0.25 + + def test_explicit_backend_failure_is_not_retried_with_sdpa(self, monkeypatch): + attention, events, _ = self._make_attention(monkeypatch, backend="FLASH_ATTN") + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + selected = Mock(side_effect=RuntimeError("selected backend rejects this dtype")) + monkeypatch.setattr(attention.attention, "forward", selected) + query = torch.ones(1, 2, 4, 8) + + with pytest.raises(RuntimeError, match="selected backend rejects this dtype"): + attention(query, query, query) + + selected.assert_called_once() + assert events == ["pre"] + + @pytest.mark.parametrize("contract", ["packed", "packed-padding", "piecewise", "quantized-kv"]) + def test_backend_metadata_is_not_silently_sent_to_dense_sdpa(self, monkeypatch, contract): + quant_dtype = "fp8" if contract == "quantized-kv" else None + attention, events, state = self._make_attention(monkeypatch, kv_cache_dtype=quant_dtype) + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + lengths = torch.tensor([0, 1, 2], dtype=torch.int32) + metadata = AttentionMetadata() + if contract == "packed": + metadata.extra = { + "cu_seqlens_q": lengths, + "cu_seqlens_k": lengths, + "max_seqlen_q": 1, + "max_seqlen_k": 1, + } + elif contract == "packed-padding": + metadata.packed_padding = PackedPaddingMetadata(1, 1, lengths[:2], lengths[:2]) + elif contract == "piecewise": + metadata.full_attn_spans = [[(0, 1)]] + metadata.query_ranges = (QueryRange(0, 2, 0),) + query = torch.ones(1, 2, 4, 8) + + attention(query, query, query, metadata) + + assert events == ["pre", "selected", "post"] + forwarded = state["kernel_call"][3] + if quant_dtype: + # Quantization is injected by the public forward from config. + assert forwarded.extra == {"kv_cache_dtype": "fp8"} + assert metadata.extra == {} + else: + assert forwarded is state["prepared"][3] + + class TestDiffusionKvCacheQuantization: @staticmethod def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True): diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index f12fbc466..679b7ca37 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -233,6 +233,7 @@ def _fake_ltx_parallel_attention(): layer = object.__new__(ltx2_transformer._LTX2ParallelAttention) nn.Module.__init__(layer) + layer.allow_fp32_fallback = False layer.attn_backend = Backend layer.attention = SimpleNamespace(forward=run("native")) layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa")) diff --git a/tests/diffusion/models/mammoth_moda2/__init__.py b/tests/diffusion/models/mammoth_moda2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py new file mode 100644 index 000000000..bab2fee7d --- /dev/null +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MammothModa2 DiT self-attention runs through the shared Omni attention layer. + +``_reference_attention`` is the arithmetic the processor used before it was +wired to the shared layer: SDPA over K/V replicated to the query head count, +padded rows zeroed. Both sides read the same module weights, so any drift in +the projections, the QK RMSNorm, the interleaved real RoPE, the softmax scale +or the padding mask shows up as a mismatch here. +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import ( + TransformerBlock, + apply_real_rotary_emb, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Small stand-in for the real 2520 / 21 / 7 (head_dim 120) layout. +DIM, HEADS = 48, 6 +SEQ, BATCH = 37, 2 + +# The shared layer resolves its backend from the diffusion config at construction; +# without one it takes the platform default, which on a CUDA host is FLASH_ATTN +# and cannot run CPU tensors. Pin SDPA the way a deployment would. +_SDPA_CONFIG = OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "TORCH_SDPA"}}) + + +def _block(kv_heads: int, modulation: bool = True) -> TransformerBlock: + torch.manual_seed(0) + with set_current_diffusion_config(_SDPA_CONFIG): + block = TransformerBlock( + DIM, + HEADS, + kv_heads, + multiple_of=8, + ffn_dim_multiplier=1.0, + norm_eps=1e-5, + modulation=modulation, + ) + return block.eval() + + +def _inputs(head_dim: int): + torch.manual_seed(1) + hidden = torch.randn(BATCH, SEQ, DIM) + mask = torch.ones(BATCH, SEQ, dtype=torch.bool) + mask[0, SEQ - 9 :] = False + mask[1, SEQ - 2 :] = False + angles = torch.rand(1, SEQ, head_dim) * 6.283 + rotary = (angles.cos(), angles.sin()) + temb = torch.randn(BATCH, min(DIM, 1024)) + return hidden, mask, rotary, temb + + +def _reference_attention(attn, hidden, mask, rotary): + """The pre-change processor, kept verbatim as the oracle.""" + batch, seq, _ = hidden.shape + query, key, value = attn.to_q(hidden), attn.to_k(hidden), attn.to_v(hidden) + head_dim = query.shape[-1] // attn.heads + kv_heads = key.shape[-1] // head_dim + query = attn.norm_q(query.view(batch, seq, attn.heads, head_dim)) + key = attn.norm_k(key.view(batch, seq, kv_heads, head_dim)) + value = value.view(batch, seq, kv_heads, head_dim) + query = apply_real_rotary_emb(query, rotary[0], rotary[1]) + key = apply_real_rotary_emb(key, rotary[0], rotary[1]) + query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + if kv_heads < attn.heads: + key = key.repeat_interleave(attn.heads // kv_heads, dim=1) + value = value.repeat_interleave(attn.heads // kv_heads, dim=1) + out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask.view(batch, 1, 1, seq), scale=attn.scale) + out = (out * mask[:, None, :, None]).transpose(1, 2).reshape(batch, seq, attn.heads * head_dim) + return attn.to_out[1](attn.to_out[0](out)) + + +@pytest.mark.parametrize("kv_heads", [2, HEADS], ids=["gqa_3to1", "mha"]) +def test_processor_matches_previous_arithmetic(kv_heads): + block = _block(kv_heads) + hidden, mask, rotary, _ = _inputs(block.head_dim) + with torch.no_grad(): + got = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary + ) + want = _reference_attention(block.attn, hidden, mask, rotary) + torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5) + assert torch.count_nonzero(got[~mask]) == 0, "padded rows must stay zero" + + +def test_no_mask_equals_all_valid_mask(): + block = _block(2) + hidden, _, rotary, _ = _inputs(block.head_dim) + full = torch.ones(BATCH, SEQ, dtype=torch.bool) + with torch.no_grad(): + unmasked = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=None, image_rotary_emb=rotary + ) + masked = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=full, image_rotary_emb=rotary + ) + torch.testing.assert_close(unmasked, masked, atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize("modulation", [True, False], ids=["noise_block", "context_refiner"]) +def test_block_forward_hands_native_kv_heads_to_shared_layer(monkeypatch, modulation): + """Through the block's real forward, on both of its branches (the modulated + noise/refiner blocks and the unmodulated context refiner): the shared layer + is what runs, and it receives K/V with their native head count -- no + replication before the call.""" + block = _block(2, modulation=modulation) + hidden, mask, rotary, temb = _inputs(block.head_dim) + if not modulation: + temb = None + seen = [] + original = block.attn.omni_attn.forward + + def spy(query, key, value, attn_metadata=None): + seen.append((tuple(query.shape), tuple(key.shape), tuple(value.shape), attn_metadata)) + return original(query, key, value, attn_metadata) + + monkeypatch.setattr(block.attn.omni_attn, "forward", spy) + with torch.no_grad(): + out = block(hidden, mask, rotary, temb) + assert out.shape == hidden.shape and torch.isfinite(out).all() + assert len(seen) == 1 + q_shape, k_shape, v_shape, metadata = seen[0] + assert q_shape == (BATCH, SEQ, HEADS, block.head_dim) + assert k_shape == v_shape == (BATCH, SEQ, 2, block.head_dim) + assert metadata is not None and metadata.attn_mask.dtype == torch.bool + assert torch.equal(metadata.attn_mask, mask) + + +def test_backend_selection_reaches_the_dit_layer(): + """The diffusion config selects the backend when the DiT layer is constructed.""" + assert _block(2).attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.TORCH_SDPA.get_class() + with set_current_diffusion_config( + OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "FLASH_ATTN"}}) + ): + block = TransformerBlock(DIM, HEADS, 2, multiple_of=8, ffn_dim_multiplier=1.0, norm_eps=1e-5) + assert block.attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.FLASH_ATTN.get_class() + + +def test_shared_layer_owns_no_parameters_and_keeps_checkpoint_keys(): + block = _block(2) + assert sum(p.numel() for p in block.attn.omni_attn.parameters()) == 0 + assert not list(block.attn.omni_attn.buffers()) + keys = set(block.state_dict()) + assert not any("omni_attn" in key for key in keys) + # The projections and QK norms the checkpoint provides are still where they were. + for expected in ( + "attn.to_q.weight", + "attn.to_k.weight", + "attn.to_v.weight", + "attn.to_out.0.weight", + "attn.norm_q.weight", + "attn.norm_k.weight", + ): + assert expected in keys, expected + + +@pytest.mark.parametrize("with_mask", [True, False], ids=["empty_mask", "no_mask"]) +def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask): + """CFG's unconditional branch carries zero text tokens (the pipeline's default + negative_prompt_embeds has no rows), so the context refiner attends over an + empty sequence. The flash-attention varlen fallback cannot take that + (arange step 0), so the processor must not reach the shared layer at all.""" + block = _block(kv_heads=2, modulation=False) + attn = block.attn + + def _never(*args, **kwargs): + raise AssertionError("shared layer called with an empty sequence") + + monkeypatch.setattr(attn.omni_attn, "forward", _never) + hidden = torch.randn(BATCH, 0, DIM) + mask = torch.ones(BATCH, 0, dtype=torch.bool) if with_mask else None + angles = torch.rand(1, 0, block.head_dim) + out = attn( + hidden_states=hidden, + encoder_hidden_states=hidden, + attention_mask=mask, + image_rotary_emb=(angles.cos(), angles.sin()), + ) + assert out.shape == (BATCH, 0, DIM) diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py new file mode 100644 index 000000000..62415953a --- /dev/null +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared +backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic +in bf16, with the padding mask the joint text+image sequence carries.""" + +import pytest +import torch + +import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock + +from .test_dit_attention import _reference_attention + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.advanced_model, + pytest.mark.cuda, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] + +# MammothModa2-Preview gen_dit_config: hidden 2520, 21 heads, 7 KV heads, head_dim 120. +DIM, HEADS, KV_HEADS = 2520, 21, 7 + + +@pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"]) +def test_real_shape_matches_previous_arithmetic_bf16(seq): + torch.manual_seed(0) + block = ( + TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5) + .cuda() + .to(torch.bfloat16) + .eval() + ) + hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(2, seq, dtype=torch.bool, device="cuda") + mask[0, seq - 300 :] = False + mask[1, seq - 17 :] = False + angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283 + rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)) + with torch.no_grad(): + got = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary + ) + want = _reference_attention(block.attn, hidden, mask, rotary) + assert torch.isfinite(got).all() + assert torch.count_nonzero(got[~mask]) == 0 + diff = (got.float() - want.float()).abs()[mask] + # bf16 kernels differ in accumulation order; measured on A800: max 1.0e-3, mean 5e-5. + assert diff.max().item() < 2e-2, diff.max().item() + assert diff.mean().item() < 1e-3, diff.mean().item() + + +def test_empty_text_stream_on_the_default_backend(): + """The recipe's text-to-image request with text_guidance_scale > 1 runs the + context refiner on a zero-token unconditional prompt. Before the guard this + raised ``RuntimeError: step must be nonzero`` from the FA varlen fallback.""" + torch.manual_seed(0) + block = ( + TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5, modulation=False) + .cuda() + .to(torch.bfloat16) + .eval() + ) + hidden = torch.randn(1, 0, DIM, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(1, 0, dtype=torch.bool, device="cuda") + angles = torch.rand(1, 0, block.head_dim, device="cuda") + with torch.no_grad(): + out = block.attn( + hidden_states=hidden, + encoder_hidden_states=hidden, + attention_mask=mask, + image_rotary_emb=(angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)), + ) + assert out.shape == (1, 0, DIM) + + +@pytest.mark.parametrize( + ("seq", "force_gqa_fallback"), + [(512, False), (77 + 4096, False), (32, True)], + ids=["short", "t2i_1024", "unsupported_native_gqa"], +) +def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_gqa_fallback): + """FP32 keeps its dtype and reaches shared SDPA's GQA compatibility check. + + Cover the long masked shape that otherwise falls back to quadratic math + attention, and force the unsupported-GQA branch independently of hardware. + """ + torch.manual_seed(0) + dt = torch.float32 + block = ( + TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5) + .cuda() + .to(dt) + .eval() + ) + hidden = torch.randn(1, seq, DIM, device="cuda", dtype=dt) + mask = torch.ones(1, seq, dtype=torch.bool, device="cuda") + mask[:, -13:] = False + angles = torch.rand(1, seq, block.head_dim, device="cuda") + rotary = (angles.cos().to(dt), angles.sin().to(dt)) + capabilities, calls, phases = [], [], [] + can_use_fused_gqa = sdpa_backend.can_sdpa_use_fused_gqa + sdpa = torch.nn.functional.scaled_dot_product_attention + strategy = block.attn.omni_attn._get_active_parallel_strategy() + pre_attention, post_attention = strategy.pre_attention, strategy.post_attention + + def prepare(*args): + phases.append("pre") + return pre_attention(*args) + + def restore(*args): + phases.append("post") + return post_attention(*args) + + def check_gqa(*args): + supported = False if force_gqa_fallback else can_use_fused_gqa(*args) + capabilities.append(supported) + return supported + + def record_sdpa(query, key, value, **kwargs): + phases.append("sdpa") + calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs)) + return sdpa(query, key, value, **kwargs) + + with torch.no_grad(), monkeypatch.context() as patch: + patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa) + patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa) + patch.setattr(strategy, "pre_attention", prepare) + patch.setattr(strategy, "post_attention", restore) + got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary) + + assert phases == ["pre", "sdpa", "post"], "FP32 must preserve the shared parallel dispatch" + assert len(capabilities) == len(calls) == 1, "FP32 must use shared SDPA's runtime GQA check" + q_shape, k_shape, v_shape, dtype, kwargs = calls[0] + expected_kv_heads = KV_HEADS if capabilities[0] else HEADS + assert q_shape == (1, HEADS, seq, block.head_dim) + assert k_shape == v_shape == (1, expected_kv_heads, seq, block.head_dim) + assert dtype == got.dtype == torch.float32 + assert kwargs["enable_gqa"] == capabilities[0] + assert kwargs["scale"] == block.attn.scale + assert kwargs["is_causal"] is False + assert torch.equal(kwargs["attn_mask"], mask[:, None, None, :]) + + with torch.no_grad(): + want = _reference_attention(block.attn, hidden, mask, rotary) + assert torch.isfinite(got).all() + assert torch.count_nonzero(got[~mask]) == 0 + diff = (got.float() - want.float()).abs()[mask] + assert diff.max().item() < 1e-3, diff.max().item() diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4b8ca8f49..4c588af0c 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -81,6 +81,8 @@ class Attention(nn.Module): # varlen attention with learned sink logits). The shared Attention # layer still owns parallel dispatch and compile boundaries. custom_attention: nn.Module | None = None, + # Preserve dense FP32 inference for models opting into CUDA auto fallback. + allow_fp32_fallback: bool = False, ): super().__init__() @@ -193,8 +195,8 @@ class Attention(nn.Module): role=role, backend_explicit=self.backend_explicit, ) - # Some model-specific paths call this helper directly. The main backend - # dispatch below never switches to it implicitly. + # Compatibility kernels run inside shared dispatch, between the + # parallel strategy's input preparation and output restoration. self.sdpa_fallback = SDPABackend.get_impl_cls()( num_heads=num_heads, head_size=head_size, @@ -221,6 +223,7 @@ class Attention(nn.Module): self.use_sync = use_sync self.causal = causal self.skip_sequence_parallel = skip_sequence_parallel + self.allow_fp32_fallback = allow_fp32_fallback self.use_ring = False self.ring_pg = None @@ -497,13 +500,33 @@ class Attention(nn.Module): self._assert_metadata_compatible(attn_metadata) + if ( + self.allow_fp32_fallback + and query.is_cuda + and query.dtype == torch.float32 + and query.ndim == 4 + and not self.backend_explicit + and self.attn_backend.get_name() == "FLASH_ATTN" + and ( + attn_metadata is None + or ( + attn_metadata.full_attn_spans is None + and attn_metadata.query_ranges is None + and attn_metadata.video_layout is None + and attn_metadata.packed_padding is None + and not attn_metadata.extra + ) + ) + ): + logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.") + return self.sdpa_fallback.forward(query, key, value, attn_metadata) + in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile # The startup KV-capacity profile needs tensor shapes, not a paged # attention result. If dense FLASH_ATTN deps are absent (NPU MindIE-SD # or CUDA CuTe FA4), SDPA provides that profile forward. Formal paged # requests never use this branch because their Worker adapter is active. - # Do not silently remap float32 (or other dtypes) to SDPA: a - # user-explicit backend must run or raise. + # A user-explicit backend must run or raise. if ( self._scheduler_paged_kv and self.paged_kv_cache_role is not None diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 000abe051..e2c086d4a 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -1,4 +1,5 @@ -import math +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import torch import torch.nn.functional as F @@ -6,27 +7,14 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.attention_processor import Attention from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_utils import ModelMixin -from einops import rearrange, repeat +from einops import rearrange from torch import nn from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from .rope_real import RotaryPosEmbedReal +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as OmniAttention -try: - from transformers.modeling_flash_attention_utils import ( # type: ignore - flash_attn_varlen_func, # pyright: ignore[reportAttributeAccessIssue] - is_flash_attn_available, - ) -except Exception: # pragma: no cover - best-effort compatibility - flash_attn_varlen_func = None # type: ignore[assignment] - - def is_flash_attn_available() -> bool: # type: ignore[override] - return False - - -from .rope_real import apply_real_rotary_emb - -_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None +from .rope_real import RotaryPosEmbedReal, apply_real_rotary_emb class LuminaRMSNormZero(nn.Module): @@ -278,9 +266,14 @@ class SimpleQFormerImageRefiner(nn.Module): class AttnProcessor: - def __init__(self) -> None: - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).") + """Self-attention for the MammothModa2 DiT through the shared Omni attention layer. + + The model-specific parts stay here: the projections, the per-head QK RMSNorm + and the interleaved real RoPE. The kernel, the padding mask and the + grouped-query heads are the shared backend's job, so K/V are handed over + with their native head count instead of being replicated to the query + head count first. + """ def __call__( self, @@ -289,118 +282,55 @@ class AttnProcessor: encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, image_rotary_emb: torch.Tensor | None = None, - base_sequence_length: int | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape - # Get Query-Key-Value Pair query = attn.to_q(hidden_states) key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) - query_dim = query.shape[-1] - inner_dim = key.shape[-1] - head_dim = query_dim // attn.heads + head_dim = query.shape[-1] // attn.heads + kv_heads = key.shape[-1] // head_dim dtype = query.dtype - # Get key-value heads - kv_heads = inner_dim // head_dim + if sequence_length == 0 or key.shape[1] == 0: + # The pipeline's default unconditional branch is a text stream with + # zero tokens (negative_prompt_embeds has no rows), so under CFG the + # context refiner attends over an empty sequence. SDPA returns an + # empty tensor for that; the flash-attention varlen fallback builds + # cu_seqlens with arange(step=seq_len) and rejects step 0. Nothing + # to attend to either way, so hand the projection an empty input. + empty = query.new_zeros(batch_size, sequence_length, attn.heads * head_dim) + return attn.to_out[1](attn.to_out[0](empty)) - # Reshape tensors for attention computation query = query.view(batch_size, -1, attn.heads, head_dim) key = key.view(batch_size, -1, kv_heads, head_dim) value = value.view(batch_size, -1, kv_heads, head_dim) - # Apply Query-Key normalization if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) - # Apply Rotary Position Embeddings if image_rotary_emb is not None: query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1]) key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1]) query, key = query.to(dtype), key.to(dtype) - # Calculate attention scale - if base_sequence_length is not None: - softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale - else: - softmax_scale = attn.scale - - if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda: - # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need - # the self-attention case (q/k/v share the same padding mask). + if attention_mask is not None: attention_mask = attention_mask.to(torch.bool) - seqlens = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() - max_seqlen = int(seqlens.max().item()) - cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) - query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices] - key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] - value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + hidden_states = attn.omni_attn(query, key, value, attn_metadata) - if kv_heads < attn.heads: - key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) - value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) + if attention_mask is not None: + # Padded rows carry nothing downstream; keep them at zero as the + # previous varlen path did, since neither backend does it for us. + hidden_states = hidden_states * attention_mask[:, :, None, None] - attn_output_unpad = flash_attn_varlen_func( - query_states, - key_states, - value_states, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - dropout_p=0.0, - causal=False, - softmax_scale=softmax_scale, - ) + hidden_states = hidden_states.reshape(batch_size, sequence_length, attn.heads * head_dim).type_as(query) - out = torch.zeros( - (batch_size * sequence_length, attn.heads, head_dim), - device=attn_output_unpad.device, - dtype=attn_output_unpad.dtype, - ) - out[indices] = attn_output_unpad - hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2) - hidden_states = hidden_states.type_as(query) - else: - # PyTorch SDPA path. - attn_mask = None - if attention_mask is not None: - attention_mask = attention_mask.to(torch.bool) - attn_mask = attention_mask.view(batch_size, 1, 1, -1) - - query = query.transpose(1, 2) # [B, H, S, D] - key = key.transpose(1, 2) # [B, H_kv, S, D] - value = value.transpose(1, 2) - - if kv_heads < attn.heads: - key = key.repeat_interleave(attn.heads // kv_heads, dim=1) - value = value.repeat_interleave(attn.heads // kv_heads, dim=1) - - hidden_states = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=attn_mask, - dropout_p=0.0, - is_causal=False, - scale=softmax_scale, - ) - - if attention_mask is not None: - # Keep padding tokens consistent with the flash-varlen path (zero output). - hidden_states = hidden_states * attention_mask[:, None, :, None] - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.type_as(query) - - # Apply output projection hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) @@ -441,6 +371,18 @@ class TransformerBlock(nn.Module): # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。 self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5) self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5) + # The kernel itself runs through the shared Omni attention layer: backend + # selection, padding-mask handling and native grouped-query heads live + # there. It owns no parameters, so checkpoint keys are unchanged. The + # diffusers Attention above is kept for its projections and QK norms. + self.attn.omni_attn = OmniAttention( + num_heads=num_attention_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.attn.scale, + num_kv_heads=num_kv_heads, + allow_fp32_fallback=True, + ) # Initialize feed-forward network self.feed_forward = LuminaFeedForward( From 4a7a8454ba5649af43a7b9d44b92733cc488c604 Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:00:15 +0800 Subject: [PATCH 034/209] [Bugfix] Give model CLI flags typed owners in the Omni config (#7390) Signed-off-by: Guangjian Co-authored-by: amy-why-3459 --- tests/config/test_omni_config.py | 50 ++++++++++++++++++++++++++++++++ vllm_omni/config/omni_config.py | 26 +++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index fdb2cd4b9..a8a2056b6 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -271,6 +271,44 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) +_MODEL_CLI_FLAGS = { + "served_model_name": "omni", + "allowed_local_media_path": "/data", + "allowed_media_domains": ["example.com"], + "max_logprobs": 7, + "logprobs_mode": "processed_logprobs", + "mm_processor_kwargs": {"max_slice_nums": 1}, + "mm_processor_cache_type": "shm", + "hf_token": "hf_test", + "hf_config_path": "/models/config", + "generation_config": "vllm", + "override_generation_config": {"temperature": 0.5}, + "enable_prompt_embeds": True, +} + + +@pytest.mark.parametrize("model_type", ["minicpmo_4_5", "qwen3_tts"]) +@pytest.mark.parametrize(("field", "value"), list(_MODEL_CLI_FLAGS.items()), ids=list(_MODEL_CLI_FLAGS)) +def test_from_pipeline_config_owns_model_cli_fields(model_type, field, value): + config = _from_pipeline_key(model_type, cli_overrides={field: value}) + + for stage in config.stage_configs: + assert getattr(stage.model_config, field) == value + + +def test_model_cli_fields_reach_typed_engine_args(monkeypatch): + from vllm_omni.engine import stage_init_utils + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + # Worker discovery requires hardware support; this test covers config transport. + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None) + config = _from_pipeline_key("minicpmo_4_5", cli_overrides=dict(_MODEL_CLI_FLAGS)) + + engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model") + + assert {field: engine_args.get(field) for field in _MODEL_CLI_FLAGS} == _MODEL_CLI_FLAGS + + @pytest.mark.parametrize("stage_id", [0, 2], ids=["ar", "generation"]) def test_llm_additional_config_roundtrip_and_isolation(stage_id, monkeypatch): from vllm_omni.engine import stage_init_utils @@ -723,6 +761,18 @@ def test_sub_config_fields_match_structured_scopes(): "model_subdir", "tokenizer_subdir", "requires_full_payload_input", + "served_model_name", + "allowed_local_media_path", + "allowed_media_domains", + "max_logprobs", + "logprobs_mode", + "mm_processor_kwargs", + "mm_processor_cache_type", + "hf_token", + "hf_config_path", + "generation_config", + "override_generation_config", + "enable_prompt_embeds", } vllm_load_fields = {f.name for f in fields(VllmLoadConfig)} assert issubclass(OmniStageLoadConfig, VllmLoadConfig) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index f0752077b..b58887099 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -168,6 +168,19 @@ class _ModelEngineOverrides(TypedDict, total=False): enable_broadcast_weight_load: bool num_weight_load_threads: int disable_autocast: bool + # Upstream ModelConfig inputs that users pass as global CLI flags. + served_model_name: str | list[str] + allowed_local_media_path: str + allowed_media_domains: list[str] + max_logprobs: int + logprobs_mode: str + mm_processor_kwargs: dict[str, Any] + mm_processor_cache_type: str + hf_token: bool | str + hf_config_path: str + generation_config: str + override_generation_config: dict[str, Any] + enable_prompt_embeds: bool class _LoadEngineOverrides(TypedDict, total=False): @@ -445,6 +458,19 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): model_subdir: str | None = None tokenizer_subdir: str | None = None requires_full_payload_input: bool = False + # Upstream ModelConfig inputs that users pass as global CLI flags. + served_model_name: str | list[str] | None = None + allowed_local_media_path: str | None = None + allowed_media_domains: list[str] | None = None + max_logprobs: int | None = None + logprobs_mode: str | None = None + mm_processor_kwargs: dict[str, Any] | None = None + mm_processor_cache_type: str | None = None + hf_token: bool | str | None = None + hf_config_path: str | None = None + generation_config: str | None = None + override_generation_config: dict[str, Any] | None = None + enable_prompt_embeds: bool | None = None @_enforce_keyword_only_init From 68294066e3df6cdd0eba65da0ff05c712564d32b Mon Sep 17 00:00:00 2001 From: longguo <107740309+abinggo@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:03:16 +0800 Subject: [PATCH 035/209] [Bugfix] Require a model for `vllm serve --omni` (fixes #4158) (#4167) Signed-off-by: abinggo <107740309+abinggo@users.noreply.github.com> --- tests/entrypoints/test_serve.py | 74 ++++++++++++++++++++++++++++++ vllm_omni/entrypoints/cli/serve.py | 44 ++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 083f96d9f..1f3b58170 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -44,6 +44,80 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: assert not explicit["async_chunk"] +def _parse_serve_args(argv: list[str]) -> TrackingNamespace: + """Parse a ``serve`` argv through the real Omni parser, returning the + TrackingNamespace (with ``explicit_keys``) that ``validate`` receives.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + return parser.parse_args(argv) + + +def test_omni_serve_requires_model_when_none_provided() -> None: + """Regression for https://github.com/vllm-project/vllm-omni/issues/4158: + ``vllm serve --omni`` with no model must fail fast instead of silently + falling back to vLLM's default model and crashing in the diffusion worker. + """ + args = _parse_serve_args(["serve", "--omni"]) + cmd = OmniServeCommand() + with pytest.raises(ValueError, match="requires an explicit model"): + cmd.validate(args) + + +@pytest.mark.parametrize( + "argv", + [ + # A deploy YAML supplied *without* an explicit model must still trip + # the guard. It carries per-stage engine args only -- the + # checkpoint is always threaded in from ``args.model`` -- so its mere + # presence does not establish that a model was provided. Without an + # explicit model ``args.model`` stays at vLLM's default and the launch + # reproduces the issue-4158 diffusion-registry crash. + ["serve", "--omni", "--deploy-config", "deploy.yaml"], + # An empty/whitespace model must not satisfy the guard either -- this is + # the ``--model "$MODEL"``-with-unset-var footgun. Otherwise the empty + # value flows downstream and crashes with the same confusing error. + ["serve", "", "--omni"], # empty positional + ["serve", "--omni", "--model", ""], # empty --model + ["serve", "--omni", "--model", " "], # whitespace-only --model + ], +) +def test_omni_serve_rejects_missing_or_empty_model(argv: list[str]) -> None: + """Regression for review feedback on PR #4167: a deploy YAML is not a + model source, so ``--deploy-config`` alone must not bypass the + require-a-model guard; neither may an empty/whitespace model value.""" + args = _parse_serve_args(argv) + cmd = OmniServeCommand() + with pytest.raises(ValueError, match="requires an explicit model"): + cmd.validate(args) + + +@pytest.mark.parametrize( + "argv", + [ + ["serve", "fake-model", "--omni"], # positional model + ["serve", "--omni", "--model", "fake-model"], # --model flag + # An explicit model with a deploy YAML is the legitimate multi-stage + # path and must pass the guard. + ["serve", "fake-model", "--omni", "--deploy-config", "deploy.yaml"], + ], +) +def test_omni_serve_accepts_explicit_model(argv: list[str], mocker: MockerFixture) -> None: + """A model supplied positionally or via ``--model`` -- with or without a + deploy YAML -- must NOT trip the require-a-model guard. Downstream + validation is mocked so only the new guard is exercised.""" + mocker.patch( + "vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", + return_value=False, + ) + mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args") + + args = _parse_serve_args(argv) + cmd = OmniServeCommand() + # Must not raise the "requires a model" error for any of these. + cmd.validate(args) + + def test_serve_parser_accepts_strategy_config() -> None: """``--strategy-config`` must parse onto the ``strategy_config`` dest and be forwarded as an explicit kwarg so the engine can overlay the strategy.""" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index f9297aa0c..0166cb4c1 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -148,6 +148,50 @@ class OmniServeCommand(CLISubcommand): if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None): raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set") + # Require an explicit model under --omni. ``args.model`` always carries + # vLLM's ModelConfig default (``Qwen/Qwen3-0.6B``), so an omit is silent: + # the default text LLM is routed into the diffusion stage and startup + # crashes deep in the diffusion worker with a confusing registry error. + # Fail fast instead. + # + # The model must come from the CLI -- positionally + # (``vllm serve --omni``) or via ``--model``. A deploy YAML + # (``--deploy-config``) is NOT a model source: it carries per-stage + # engine args only, and the checkpoint is always threaded in from + # ``args.model`` (see + # ``load_and_resolve_stage_configs``, which takes ``model`` as its first + # argument). Treating its mere presence as "model provided" let + # ``vllm serve --omni --deploy-config `` slip through with the + # default model and reproduce the very crash this guard prevents. See + # https://github.com/vllm-project/vllm-omni/issues/4158. + if getattr(args, "omni", False): + explicit_keys = getattr(args, "explicit_keys", None) or frozenset() + # Resolve the model the user actually supplied on the CLI. A + # positional ``model_tag`` takes precedence (``cmd`` later copies it + # onto ``args.model``); otherwise ``--model`` counts only when it was + # explicitly passed. An empty/whitespace value (e.g. ``--model + # "$MODEL"`` with ``MODEL`` unset) is treated as "not provided" so it + # fails here with a clear message instead of the same confusing + # downstream crash. + model_tag = getattr(args, "model_tag", None) + if model_tag is not None: + explicit_model = model_tag + elif "model" in explicit_keys: + explicit_model = getattr(args, "model", None) + else: + explicit_model = None + model_provided = explicit_model is not None and str(explicit_model).strip() != "" + if not model_provided: + raise ValueError( + "`vllm serve --omni` requires an explicit model. Pass it " + "positionally (`vllm serve --omni`) or via `--model`. " + "`--deploy-config` carries per-stage engine args only and " + "does not supply a model; without an " + "explicit model, vLLM's default (Qwen/Qwen3-0.6B) is selected " + "and routed into the diffusion stage, which fails with a " + "confusing diffusion-registry error." + ) + # --omni-replica-address is only consulted in run_headless(); reject it # on the head so a misconfigured launch fails loudly instead of being # silently ignored. From 9ef9fc442a8a85c8e6b42a46412ceffaec7e0999 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:10:37 +0800 Subject: [PATCH 036/209] [Bugfix] Send a downstream terminal chunk when a parked stage ends (#6889) Signed-off-by: psv666 <2693925048@qq.com> --- .../test_chunk_transfer_adapter.py | 350 ++++++++++++++++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 62 ++++ .../chunk_transfer_adapter.py | 120 +++++- 3 files changed, 519 insertions(+), 13 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 40d29e969..a80ded5f0 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -2526,6 +2526,356 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): assert order == ["adapter", "super"] +def _parked_sender_scheduler(adapter, session): + """A minimal OmniARScheduler holding one parked streaming session.""" + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + scheduler.requests = {session.request_id: session} + scheduler.running = [] + scheduler.waiting = DummyWaitingQueue() + scheduler.skipped_waiting = DummyWaitingQueue([session]) + scheduler.num_waiting_for_streaming_input = 1 + scheduler.finished_req_ids_dict = {} + + def _free_request(self, request, delay_free_blocks=False): + del delay_free_blocks + self.requests.pop(request.request_id, None) + return None, None + + scheduler._free_request = MethodType(_free_request, scheduler) + return scheduler + + +@pytest.mark.parametrize("send_first_chunk", [False, True]) +def test_parked_streaming_final_update_emits_downstream_terminal(build_adapter, send_first_chunk): + """#6670: a final update on a parked sender must terminate the receiver. + + Upstream turns "final update while parked in WAITING_FOR_STREAMING_REQ" + into a silent local abort. The downstream stage only learns a stream ended + from a chunk carrying ``finished``, so the abort used to leave it in + WAITING_FOR_CHUNK until the 600s input deadline. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + session = _req( + "req-parked-final", + RequestStatus.WAITING_FOR_STREAMING_REQ, + external_req_id="ext-parked-final", + ) + session.resumable = True + session._omni_segment_generation = 0 + # A segment stop queues its boundary before the final update arrives. + # The downstream request is already prewarmed even if the sender has not + # dequeued this first task, so it still needs a terminal in both cases. + adapter.save_async(None, session, is_segment_finished=True) + assert session.external_req_id not in adapter.put_req_chunk + if send_first_chunk: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.reset_mock() + + scheduler = _parked_sender_scheduler(adapter, session) + + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + OmniARScheduler.add_request(scheduler, final_update) + + # Local teardown happened... + assert session.request_id not in scheduler.requests + # ...but only after a terminal chunk was queued for the next stage. + assert session.status == RequestStatus.FINISHED_STOPPED + assert any(task["is_finished"] for task in adapter._pending_save_reqs) + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + connector.put.assert_called_once() + assert connector.put.call_args.kwargs["put_key"] == f"ext-parked-final_1_{int(send_first_chunk)}" + payload = connector.put.call_args.kwargs["data"] + assert bool(payload.meta.finished.item()) is True + assert bool(payload.meta.is_segment_finished.item()) is False + + +def test_queued_terminal_chunk_survives_sender_cleanup(build_adapter): + """#6670: ``cleanup_sender`` must not discard an already-queued terminal.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal") + request.resumable = False + + adapter.save_async(None, request, is_segment_finished=False) + assert len(adapter._pending_save_reqs) == 1 + + # The scheduler tears the request down before the save_loop drains. + adapter.cleanup_sender("ext-terminal") + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_called_once() + assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True + # The terminal put reclaims the generation it kept alive. + assert "ext-terminal" not in adapter._sender_tokens + + +def test_non_terminal_chunk_is_still_dropped_by_sender_cleanup(build_adapter): + """The #6670 fence is terminal-only: ordinary queued chunks still drop.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-mid", RequestStatus.RUNNING, external_req_id="ext-mid") + request.resumable = True + + adapter.save_async(None, request, is_segment_finished=False) + adapter.cleanup_sender("ext-mid") + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_not_called() + + +def test_final_update_on_non_sender_stage_keeps_upstream_abort(build_adapter): + """A final stage never put a chunk, so it must keep the upstream path.""" + adapter, _connector = build_adapter(stage_id=2, model_mode="generation") + session = _req( + "req-final-stage", + RequestStatus.WAITING_FOR_STREAMING_REQ, + external_req_id="ext-final-stage", + ) + session.resumable = True + scheduler = _parked_sender_scheduler(adapter, session) + + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + + with patch.object(VLLMScheduler, "add_request") as base_add: + OmniARScheduler.add_request(scheduler, final_update) + + base_add.assert_called_once() + assert not adapter._pending_save_reqs + + +def _queue_terminal( + adapter, + external_req_id: str, + req_id: str = "req-terminal", + segment_generation: int = 0, +) -> None: + """Queue one request-terminal chunk on the adapter's save queue. + + ``segment_generation`` must match the adapter's dedup watermark, which a + preceding segment stop advances; otherwise ``save_async`` drops the chunk + as a late duplicate. + """ + request = _req(req_id, RequestStatus.FINISHED_STOPPED, external_req_id=external_req_id) + request.resumable = False + request._omni_segment_generation = segment_generation + adapter.save_async(None, request, is_segment_finished=False) + + +def _assert_terminal_put(connector) -> None: + connector.put.assert_called_once() + assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True + + +def test_terminal_chunk_survives_cleanup_after_dequeue(build_adapter): + """#6670: the fence must hold across the connector handoff, not just the queue. + + ``_finish_parked_streaming_session`` queues the terminal and then calls + ``finish_requests``, so ``cleanup_sender`` can land *after* the save_loop + has already dequeued the terminal. Releasing the fence at dequeue left + that interleaving cancelling the chunk before ``put`` -- the same lost + terminal, and the same downstream WAITING_FOR_CHUNK hang. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-late-cleanup") + task = adapter._pending_save_reqs.popleft() + + real_send = adapter._send_single_request_for_generation + + def send_with_cleanup_in_flight(inner_task, sender_token=None): + # The scheduler thread reaches cleanup_sender here: the save_loop owns + # the terminal but has not handed it to the connector yet. + adapter.cleanup_sender("ext-late-cleanup") + return real_send(inner_task, sender_token) + + adapter._send_single_request_for_generation = send_with_cleanup_in_flight + adapter._send_single_request(task) + + _assert_terminal_put(connector) + # The generation the fence kept alive is reclaimed, not leaked. + assert "ext-late-cleanup" not in adapter._sender_tokens + assert "ext-late-cleanup" not in adapter.put_req_chunk + + +def test_terminal_chunk_survives_cleanup_racing_an_inflight_sibling(build_adapter): + """#6670: an in-flight ordinary chunk must not reclaim a fenced generation. + + ``_send_single_request``'s ``finally`` also tears a cancelled generation + down. If the stop and the abort both land while an earlier chunk is on the + wire, that teardown would unregister the token and the queued terminal + would then fail the identity check. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + running = _req("req-sibling", RequestStatus.RUNNING, external_req_id="ext-sibling") + running.resumable = True + adapter.save_async(None, running, is_segment_finished=True) + ordinary_task = adapter._pending_save_reqs.popleft() + + real_send = adapter._send_single_request_for_generation + + def send_then_stop_and_abort(inner_task, sender_token=None): + result = real_send(inner_task, sender_token) + # Terminal stop and abort both land while this chunk is on the wire. + # The segment stop above already armed the next segment's watermark. + _queue_terminal(adapter, "ext-sibling", req_id="req-sibling", segment_generation=1) + adapter.cleanup_sender("ext-sibling") + return result + + adapter._send_single_request_for_generation = send_then_stop_and_abort + adapter._send_single_request(ordinary_task) + adapter._send_single_request_for_generation = real_send + + assert len(adapter._pending_save_reqs) == 1 + connector.put.reset_mock() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + _assert_terminal_put(connector) + assert "ext-sibling" not in adapter._sender_tokens + + +@pytest.mark.parametrize("finish_inside_put", [False, True]) +def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_inside_put): + """A cleanup during put must not reuse a key already consumed downstream.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + receiver, receiver_connector = build_adapter(stage_id=2, model_mode="generation") + session = _req("req-sequence", RequestStatus.WAITING_FOR_STREAMING_REQ, external_req_id="ext-sequence") + session.resumable = True + session._omni_segment_generation = 0 + scheduler = _parked_sender_scheduler(adapter, session) + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + downstream = _req("req-downstream", RequestStatus.WAITING, external_req_id=session.external_req_id) + receiver.request_ids_mapping[downstream.request_id] = session.external_req_id + payloads = {} + + def receive_next(): + entry = _dequeue_load_entry(receiver, downstream) + return receiver._poll_single_request(entry) + + def put_then_finish(**kwargs): + payload = kwargs["data"] + payloads[kwargs["put_key"]] = { + "meta": { + "finished": bool(payload.meta.finished.item()), + "is_segment_finished": bool(payload.meta.is_segment_finished.item()), + } + } + if not payload.meta.finished.item(): + # Consume the ordinary boundary while put is still in flight. + assert receive_next() + if finish_inside_put: + scheduler.add_request(final_update) + return True, 1, {} + + def get_by_key(_from_stage, _to_stage, key): + payload = payloads.pop(key, None) + return (payload, 1) if payload is not None else None + + connector.put.side_effect = put_then_finish + receiver_connector.get.side_effect = get_by_key + adapter.save_async(None, session, is_segment_finished=True) + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + if not finish_inside_put: + scheduler.add_request(final_update) + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert [call.kwargs["put_key"] for call in connector.put.call_args_list] == ["ext-sequence_1_0", "ext-sequence_1_1"] + assert receive_next() + assert downstream.request_id in receiver.upstream_exhausted_requests + assert not payloads + assert session.external_req_id not in adapter._sender_tokens + + +def test_late_successful_put_cannot_advance_replacement_generation(build_adapter): + """The post-put guard must retain identity checks while relaxing cancellation.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-old", RequestStatus.RUNNING, external_req_id="ext-replaced") + request.resumable = True + adapter.save_async(None, request, is_segment_finished=True) + task = adapter._pending_save_reqs.popleft() + old_token = task["sender_token"] + replacement_token = type(old_token)() + + def replace_during_put(**kwargs): + # Model a completion that no longer owns its external id. The old + # terminal fence must never authorize writes to the successor's state. + with adapter._sender_state_lock: + old_token.cancelled = True + old_token.terminal_pending = True + adapter._sender_tokens[request.external_req_id] = replacement_token + adapter.put_req_chunk[request.external_req_id] = 4 + adapter.ramp_chunk_count[request.external_req_id] = 2 + return True, 1, {} + + connector.put.side_effect = replace_during_put + adapter._send_single_request(task) + assert adapter._sender_tokens[request.external_req_id] is replacement_token + assert adapter.put_req_chunk[request.external_req_id] == 4 + assert adapter.ramp_chunk_count[request.external_req_id] == 2 + + +def test_new_request_cannot_join_a_generation_with_a_queued_terminal(build_adapter): + """#6670: the fence must not readmit a request reusing the external id. + + ``save_async`` rejects chunks whose generation is ``cancelled``; that is + what stops a new request from inheriting a retiring generation's token and + chunk counters. Deferring the reclaim must not also defer that retirement. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-shared", req_id="req-old") + adapter.cleanup_sender("ext-shared") + + newcomer = _req("req-new", RequestStatus.RUNNING, external_req_id="ext-shared") + newcomer.resumable = True + adapter.save_async(None, newcomer, is_segment_finished=True) + + # Rejected loudly, so the scheduler fails the newcomer now instead of its + # chunks being silently discarded at send time. + assert len(adapter._pending_save_reqs) == 1 + assert adapter.collect_failed_send_request_ids() == {"req-new": "previous sender generation is still draining"} + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + _assert_terminal_put(connector) + assert "ext-shared" not in adapter._sender_tokens + + +def test_failed_terminal_put_runs_the_deferred_cleanup(build_adapter): + """#6670: a silent connector drop must not strand the deferred cleanup. + + ``connector.put`` returning False skips the success path that would have + called ``cleanup``. The reclaim ``cleanup_sender`` handed over is owed on + that path too, or the whole per-request sender state leaks for good. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-put-failed") + adapter.cleanup_sender("ext-put-failed") + connector.put.return_value = (False, 0, {}) + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert "ext-put-failed" not in adapter._sender_tokens + assert "ext-put-failed" not in adapter.put_req_chunk + + +def test_raising_terminal_put_runs_the_deferred_cleanup(build_adapter): + """#6670: the same reclaim is owed when the connector raises.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-put-raised") + adapter.cleanup_sender("ext-put-raised") + connector.put.side_effect = RuntimeError("connector down") + + # ``save_loop`` catches this and records a send failure. + with pytest.raises(RuntimeError): + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert "ext-put-raised" not in adapter._sender_tokens + assert "ext-put-raised" not in adapter.put_req_chunk + + def test_wire_round_trip_struct_to_dict_contract(): """Pin the wire contract: encoding ``OmniPayloadStruct`` and decoding it yields a dict equivalent to ``to_dict(struct)``. diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 9123b794d..f33359835 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -688,6 +688,68 @@ class OmniSchedulerMixin: engine_core_outputs[0] = output = EngineCoreOutputs() output.scheduler_stats = stats + def add_request(self, request: Request) -> None: + """Route a terminal streaming update that lands on a parked session. + + Upstream ``Scheduler.add_request`` turns "final update (``resumable`` + is False) arrives while the session sits in + ``WAITING_FOR_STREAMING_REQ``" into a local + ``finish_requests(FINISHED_ABORTED)``. That is fine for a stage whose + output the orchestrator forwards, but an async-chunk *sender* also + owes its downstream stage a terminal chunk: the receiver learns that a + stream ended only from a connector payload carrying ``finished``, and + an abort emits none. The downstream stage then parks in + ``WAITING_FOR_CHUNK`` until ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` (600s by + default), long after the client's own deadline -- the realtime + async-chunk hang in vllm-project/vllm-omni#6670. + + Whether this races is timing, not configuration: the stage parks only + when its own stop beats both the upstream terminal chunk and the + orchestrator's terminal update, a window of tens of milliseconds at + the end of the last segment. Emit the terminal chunk here so the + downstream stage terminates on the payload path either way. + """ + existing = self.requests.get(request.request_id) + if existing is not None and existing.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + adapter = None if getattr(request, "resumable", False) else self._adapter_owing_terminal(existing) + if adapter is not None: + self._finish_parked_streaming_session(existing, adapter) + return + super().add_request(request) + + def _adapter_owing_terminal(self, request: Request) -> OmniChunkTransferAdapter | None: + """The adapter that owes *request*'s downstream stage a terminal chunk. + + Sender ownership starts when the first chunk is queued, before the + background thread initializes its chunk counter. A prewarmed receiver + needs a terminal even if that first chunk has not been sent yet. + Final stages never queue sends, so they never match. + """ + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is None: + return None + external_req_id = getattr(request, "external_req_id", None) or request.request_id + return adapter if adapter.has_active_sender(external_req_id) else None + + def _finish_parked_streaming_session(self, request: Request, adapter: OmniChunkTransferAdapter) -> None: + """End a parked async-chunk session with a downstream terminal chunk.""" + request.resumable = False + # A segment stop already advanced the adapter's dedup watermark to + # ``generation + 1`` for the segment that never arrived. Claim that + # generation, exactly as ``_update_request_as_session`` would have, + # or ``save_async`` drops this chunk as a late duplicate. + request._omni_segment_generation = int(getattr(request, "_omni_segment_generation", 0) or 0) + 1 + # ``force_request_finished``: the terminal must be enqueued *before* + # ``finish_requests``, which skips requests that already report + # ``is_finished()``. + adapter.save_async( + None, + request, + is_segment_finished=False, + force_request_finished=True, + ) + self.finish_requests(request.request_id, RequestStatus.FINISHED_STOPPED) + def finish_requests( self, request_ids: str | Iterable[str] | None, diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index d71d55bf1..163363931 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -32,6 +32,16 @@ class _SenderGeneration: def __init__(self) -> None: self.in_flight = False self.cancelled = False + # A request-terminal chunk is queued but not yet on the wire. The + # downstream stage has no other way to learn the stream ended -- it + # polls the connector and parks in WAITING_FOR_CHUNK until a chunk + # carrying ``finished`` arrives -- so cleanup must not discard it. + # See vllm-project/vllm-omni#6670. + self.terminal_pending = False + # ``cleanup_sender`` asked to reclaim this generation while the fence + # was up. The terminal task owes that cleanup once it is done with the + # connector -- on every exit path, including a failed or raising put. + self.cleanup_deferred = False class _LoadEntry: @@ -321,6 +331,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._processor_accepts_step_tokens[processor] = accepts return accepts + def has_active_sender(self, external_req_id: str) -> bool: + """Whether this request owns a sender, including work still queued. + + ``save_async`` registers the generation before the background thread + can send its first chunk. Chunk counters are initialized later and + cannot determine whether a prewarmed receiver needs a terminal. + """ + with self._sender_state_lock: + sender_token = self._sender_tokens.get(external_req_id) + return sender_token is not None and not sender_token.cancelled + def save_async( self, multimodal_output: dict[str, Any] | None = None, @@ -329,6 +350,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): new_token_ids: Iterable[int] | None = None, confirmed_num_computed_tokens: int | None = None, segment_generation: int | None = None, + force_request_finished: bool = False, ): """Build and enqueue one chunk for asynchronous sending. @@ -352,8 +374,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): before a streaming transition can mutate ``request`` segment_generation: generation captured before a resumable stop can apply a queued update to the mutable request + force_request_finished: send this chunk as the request terminal + even though ``request`` is not in a finished status yet. The + scheduler needs this when it ends a *parked* streaming session: + the terminal has to be enqueued before ``finish_requests``, + because ``finish_requests`` skips requests that already report + ``is_finished()``. See ``_finish_parked_streaming_session``. """ - is_finished = request.is_finished() and not request.resumable + is_finished = force_request_finished or (request.is_finished() and not request.resumable) if not hasattr(self, "_segment_generation"): self._segment_generation = defaultdict(int) external_req_id = request.external_req_id @@ -414,6 +442,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._sender_tokens[external_req_id] = sender_token task["sender_token"] = sender_token self._pending_save_reqs.append(task) + if is_finished: + # Fence this generation against cleanup until the terminal + # chunk has been handed to the connector (#6670). + sender_token.terminal_pending = True if is_segment_finished: # The queued FIFO item now owns the old segment. Start the next # segment's deduplication watermark before the worker sends it. @@ -624,32 +656,71 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if sender_token is None: self._send_single_request_for_generation(task) return + is_terminal_task = bool(task.get("is_finished")) with self._sender_state_lock: is_current = self._sender_tokens.get(external_req_id) is sender_token - if is_current and not sender_token.cancelled: + # ``cancelled`` does not veto a terminal chunk. It is the + # downstream stage's only end-of-stream signal, so dropping it + # leaves that stage parked in WAITING_FOR_CHUNK until the input + # deadline -- the hang this fence exists to prevent (#6670). + # Losing ownership of the external id still vetoes: the state the + # chunk describes then belongs to a different generation. + if is_current and (is_terminal_task or not sender_token.cancelled): sender_token.in_flight = True else: is_current = False if not is_current: logger.debug("Discarding stale queued chunk for aborted request %s", external_req_id) + if is_terminal_task: + self._release_terminal_fence(external_req_id, sender_token) return try: self._send_single_request_for_generation(task, sender_token) finally: + if is_terminal_task: + # Drop the fence and run the cleanup it deferred, whether the + # put succeeded, reported failure, or raised. Until this runs, + # the block below must not reclaim the generation either. + self._release_terminal_fence(external_req_id, sender_token) with self._sender_state_lock: if self._sender_tokens.get(external_req_id) is sender_token: sender_token.in_flight = False - if sender_token.cancelled: + if sender_token.cancelled and not sender_token.terminal_pending: self._sender_tokens.pop(external_req_id, None) self._clear_sender_state_locked(external_req_id) + def _release_terminal_fence(self, external_req_id: str, sender_token: _SenderGeneration) -> None: + """Drop the terminal fence and run the cleanup it deferred. + + Called on every exit path of a terminal task. ``cleanup_sender`` + cannot reclaim a generation whose terminal chunk is still queued, so + it records the request on the token instead; without this the only + path that ever replayed it was a *successful* put, and a connector + failure leaked the whole per-request sender state (#6670). + """ + with self._sender_state_lock: + sender_token.terminal_pending = False + deferred = sender_token.cleanup_deferred + sender_token.cleanup_deferred = False + # A generation that no longer owns the external id was already + # reclaimed by whoever replaced it; cleaning now would clear the + # successor's state. + still_current = self._sender_tokens.get(external_req_id) is sender_token + if deferred and still_current: + self.cleanup_sender(external_req_id) + def _sender_generation_is_active( self, external_req_id: str, sender_token: _SenderGeneration, + *, + is_terminal: bool = False, ) -> bool: with self._sender_state_lock: - return self._sender_tokens.get(external_req_id) is sender_token and not sender_token.cancelled + if self._sender_tokens.get(external_req_id) is not sender_token: + return False + # Terminal chunks outlive cancellation -- see _send_single_request. + return is_terminal or not sender_token.cancelled def _send_single_request_for_generation( self, @@ -712,7 +783,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if payload_data.meta.is_segment_finished is None: payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool) - if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): + if sender_token is not None and not self._sender_generation_is_active( + external_req_id, sender_token, is_terminal=is_finished + ): logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id) return @@ -723,16 +796,25 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): data=payload_data, ) - if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): - # cleanup_sender() may cancel this generation while put() is - # blocked. Do not let that stale completion recreate state for a - # request whose cleanup is already in progress. - logger.debug("Ignoring completed put for cancelled request %s", external_req_id) - return + with self._sender_state_lock: + if sender_token is not None: + still_current = self._sender_tokens.get(external_req_id) is sender_token + retiring_without_terminal = sender_token.cancelled and not ( + is_finished or sender_token.terminal_pending + ) + if not still_current or retiring_without_terminal: + logger.debug("Ignoring completed put for retired request %s", external_req_id) + return + if success: + # A receiver may already have consumed this key while put was + # in flight. Even if cleanup retired this generation, its + # pending terminal must use the next key. Check ownership and + # advance counters under the same lock so a stale completion + # cannot modify a successor's state. + self.put_req_chunk[external_req_id] += 1 + self.ramp_chunk_count[external_req_id] += 1 if success: - self.put_req_chunk[external_req_id] += 1 - self.ramp_chunk_count[external_req_id] += 1 logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") # Sender uses struct attr access here; the receive path in # `_load_one_request` / `_update_request_payload` reads dict keys. @@ -862,6 +944,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if sender_token is None: self._clear_sender_state_locked(external_req_id) return + if sender_token.terminal_pending: + # A request-terminal chunk is still queued. Reclaiming the + # state now would drop it, and the downstream stage would sit + # in WAITING_FOR_CHUNK until VLLM_OMNI_INPUT_WAIT_TIMEOUT_S + # with no way to know the stream ended (#6670). Retire the + # generation anyway -- ``save_async`` rejects new chunks on + # ``cancelled``, which is what keeps a request reusing this + # external id from joining a generation that is going away -- + # and hand the reclaim to the terminal task. + sender_token.cancelled = True + sender_token.cleanup_deferred = True + return sender_token.cancelled = True if sender_token.in_flight: # The sender's finally block is the first point where no From b3ba67d71ea3466f1af76661275fe63a9f990f09 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Sat, 12 Sep 2026 12:14:42 +0800 Subject: [PATCH 037/209] [NPU] upgrade to v0.29.0 (#7433) Signed-off-by: Weiming Liao --- docker/Dockerfile.npu | 2 +- docker/Dockerfile.npu.ci | 2 +- .../installation/npu/npu.inc.md | 22 +++++++++---------- .../npu/worker/npu_ar_model_runner.py | 8 ++++++- .../npu/worker/npu_generation_model_runner.py | 21 +++++++++++++----- .../platforms/npu/worker/npu_model_runner.py | 21 +++++++++++++----- 6 files changed, 50 insertions(+), 26 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 0d91c0b32..f83cc512a 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_VERSION=v0.28.0 +ARG VLLM_ASCEND_VERSION=v0.29.0 ARG VLLM_ASCEND_DEVICE_SUFFIX= FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX} diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 0a150fdd8..b65b803d6 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_VERSION=v0.28.0 +ARG VLLM_ASCEND_VERSION=v0.29.0 ARG VLLM_ASCEND_DEVICE_SUFFIX= FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX} diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 9ff477522..21127caf2 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -10,25 +10,25 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built === "A2" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0 + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0 ``` === "A3" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a3 ``` === "A5" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a5 + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a5 ``` === "310P (Experimental)" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-310p + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-310p ``` ```bash @@ -53,7 +53,7 @@ docker run --rm \ # Inside the container, install vLLM-Omni from source cd /vllm-workspace -git clone -b v0.28.0 https://github.com/vllm-project/vllm-omni.git +git clone -b v0.29.0rc1 https://github.com/vllm-project/vllm-omni.git cd vllm-omni pip install -v -e . --no-build-isolation # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . @@ -71,10 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886 # --8<-- [start:installation-main] -You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.28.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU. +You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.29.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU. ```bash -# Inside the vLLM-Ascend v0.28.0 container, install vLLM-Omni from main +# Inside the vLLM-Ascend v0.29.0 container, install vLLM-Omni from main cd /vllm-workspace git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -103,7 +103,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "A2" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0 + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0 docker run --rm \ --name vllm-omni-a2 \ --shm-size=64g \ @@ -127,7 +127,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "A3" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a3 + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a3 docker run --rm \ --name vllm-omni-a3 \ --shm-size=64g \ @@ -151,7 +151,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "A5" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a5 + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a5 docker run --rm \ --name vllm-omni-a5 \ --shm-size=64g \ @@ -175,7 +175,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "310P (Experimental)" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-310p + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-310p docker run --rm \ --name vllm-omni-310p \ --shm-size=64g \ diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 734178e01..7bc6c21c6 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -334,6 +334,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if not isinstance(info, dict): return None val = info.get("omni_final_stage_id") + if val is None: + return None try: return int(val) except (TypeError, ValueError): @@ -512,7 +514,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. - self._dummy_run(1) + self._dummy_run(1, skip_gdn_state_update=True) kv_ids = self.kv_extracted_req_ids self.kv_extracted_req_ids = None @@ -1035,6 +1037,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex propose_draft_token_ids(sampler_output.sampled_token_ids) ( + num_nans_in_logits, + num_nans_device, logprobs_lists, valid_sampled_token_ids, prompt_logprobs_dict, @@ -1278,6 +1282,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, + num_nans_in_logits=num_nans_in_logits, ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids model_runner_output.routed_experts = routed_experts_lists @@ -1318,6 +1323,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, vocab_size=self.input_batch.vocab_size, + num_nans=num_nans_device, ) self.input_batch.set_async_sampled_token_ids( async_output.sampled_token_ids_cpu, diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index fc26882cd..59dfc88a3 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -166,7 +166,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. - self._dummy_run(1) + self._dummy_run(1, skip_gdn_state_update=True) if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) @@ -636,6 +636,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin profile_seq_lens: int | None = None, profile_cpp: bool = False, randomize_inputs: bool = False, + skip_gdn_state_update: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode() @@ -768,7 +769,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + if skip_gdn_state_update: + self.gdn_query_start_loc.np.fill(0) + else: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.gdn_query_start_loc.copy_to_gpu() if not profile_cpp: @@ -786,6 +790,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # rows as well so device-side metadata does not see stale block ids. self.input_batch.block_table.commit_block_table(num_reqs_padded) + # Invalidate slots before backends copy metadata for dummy execution. + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, @@ -796,11 +806,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, num_scheduled_tokens_np=num_scheduled_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + skip_gdn_state_update=skip_gdn_state_update, ) - if not is_graph_capturing: - for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): - blk_table = self.input_batch.block_table[kv_cache_gid] - blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 8f0369808..bd7c7dce4 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from functools import partial from typing import Any @@ -126,6 +126,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): profile_seq_lens: int | None = None, profile_cpp: bool = False, randomize_inputs: bool = False, + skip_gdn_state_update: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode() @@ -257,7 +258,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + if skip_gdn_state_update: + self.gdn_query_start_loc.np.fill(0) + else: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.gdn_query_start_loc.copy_to_gpu() if not profile_cpp: @@ -275,6 +279,12 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # rows as well so device-side metadata does not see stale block ids. self.input_batch.block_table.commit_block_table(num_reqs_padded) + # Invalidate slots before backends copy metadata for dummy execution. + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, @@ -285,11 +295,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, num_scheduled_tokens_np=num_scheduled_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + skip_gdn_state_update=skip_gdn_state_update, ) - if not is_graph_capturing: - for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): - blk_table = self.input_batch.block_table[kv_cache_gid] - blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, From bc0c9f4b45c45c59aa2f92471842e8c18ae403ca Mon Sep 17 00:00:00 2001 From: junpengw67-max Date: Sat, 12 Sep 2026 12:23:24 +0800 Subject: [PATCH 038/209] [Bugfix][Model][Lance] Support decoded video frames in video editing (#5128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 吴俊鹏 <248679769+junpengw67-max@users.noreply.github.com> Co-authored-by: 吴俊鹏 <248679769+junpengw67-max@users.noreply.github.com> --- .../models/lance/test_pipeline_lance.py | 171 ++++++++++++++++++ .../diffusion/models/lance/pipeline_lance.py | 3 + vllm_omni/model_executor/models/utils.py | 51 +++++- 3 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/models/lance/test_pipeline_lance.py diff --git a/tests/diffusion/models/lance/test_pipeline_lance.py b/tests/diffusion/models/lance/test_pipeline_lance.py new file mode 100644 index 000000000..32ff9bd10 --- /dev/null +++ b/tests/diffusion/models/lance/test_pipeline_lance.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +from PIL import Image + +from vllm_omni.diffusion.models.lance.lance_transformer import LanceBagel +from vllm_omni.diffusion.models.lance.pipeline_lance import LancePipeline +from vllm_omni.model_executor.models.utils import normalize_decoded_video_frames + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FramesWithMetadata(list[Image.Image]): + def __init__(self, frames: list[Image.Image], **metadata: object) -> None: + super().__init__(frames) + for name, value in metadata.items(): + setattr(self, name, value) + + +class StopAfterVideoPreprocessError(Exception): + pass + + +def test_normalize_decoded_video_frames_preserves_rgb_pixels_and_fps(): + red = Image.new("RGB", (3, 2), color=(255, 0, 0)) + rgba = Image.new("RGBA", (3, 2), color=(0, 255, 0, 64)) + rgba_bytes = rgba.tobytes() + + video, fps = normalize_decoded_video_frames(FramesWithMetadata([red, rgba], fps=8), default_fps=24.0) + + assert video.shape == (2, 2, 3, 3) + assert video.dtype == np.uint8 + assert video.flags.c_contiguous + np.testing.assert_array_equal(video[0, 0, 0], [255, 0, 0]) + np.testing.assert_array_equal(video[1, 0, 0], [0, 255, 0]) + assert fps == 8.0 + assert rgba.mode == "RGBA" + assert rgba.tobytes() == rgba_bytes + + +def test_normalize_decoded_video_frames_uses_frame_rate_when_fps_is_missing(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, frame_rate=12) + + _, fps = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert fps == 12.0 + + +def test_normalize_decoded_video_frames_prefers_valid_fps_over_frame_rate(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps=8, frame_rate=12) + + _, fps = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert fps == 8.0 + + +def test_normalize_decoded_video_frames_uses_default_fps_for_plain_list_and_tuple(): + frames = [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2))] + + assert normalize_decoded_video_frames(frames, default_fps=24.0)[1] == 24.0 + assert normalize_decoded_video_frames(tuple(frames), default_fps=24.0)[1] == 24.0 + + +def test_normalize_decoded_video_frames_accepts_numeric_string_fps(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps="12.5") + + _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert frame_rate == 12.5 + + +@pytest.mark.parametrize( + ("fps", "frame_rate", "expected"), + [ + (float("nan"), 12, 12.0), + (0, 15, 15.0), + (-1, 8, 8.0), + (float("inf"), 10, 10.0), + ("invalid", "6.5", 6.5), + ], +) +def test_normalize_decoded_video_frames_uses_valid_frame_rate_after_invalid_fps(fps, frame_rate, expected): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps=fps, frame_rate=frame_rate) + + _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert frame_rate == expected + + +def test_normalize_decoded_video_frames_uses_default_when_both_metadata_values_are_invalid(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps="invalid", frame_rate=0) + + _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert frame_rate == 24.0 + + +@pytest.mark.parametrize( + ("frames", "match"), + [ + ([], "empty decoded video frame sequence"), + ([Image.new("RGB", (2, 2)), object()], "must be a PIL.Image.Image"), + ( + [Image.new("RGB", (2, 2)), Image.new("RGB", (3, 2))], + "must have identical dimensions", + ), + ], +) +def test_normalize_decoded_video_frames_rejects_invalid_sequences(frames, match): + with pytest.raises(ValueError, match=match): + normalize_decoded_video_frames(frames, default_fps=24.0) + + +@pytest.mark.parametrize( + ("frames", "match"), + [ + ([Image.new("RGB", (0, 16))], r"index 0.*0x16"), + ([Image.new("RGB", (16, 0))], r"index 0.*16x0"), + ([Image.new("RGB", (0, 0))], r"index 0.*0x0"), + ( + [Image.new("RGB", (16, 16)), Image.new("RGB", (0, 16))], + r"index 1.*0x16", + ), + ], +) +def test_normalize_decoded_video_frames_rejects_zero_dimensions(frames, match): + with pytest.raises(ValueError, match=match): + normalize_decoded_video_frames(frames, default_fps=24.0) + + +@pytest.mark.parametrize( + ("video_input", "expected_fps"), + [ + (FramesWithMetadata([Image.new("RGB", (3, 2))] * 2, fps=8), 8.0), + (tuple([Image.new("RGBA", (3, 2))] * 2), 12.0), + ], +) +def test_lance_video_edit_normalizes_decoded_frame_sequences(monkeypatch, video_input, expected_fps): + captured = {} + + def capture_preprocessed_input(video, fps): + captured["video"] = video + captured["fps"] = fps + raise StopAfterVideoPreprocessError + + monkeypatch.setattr(LanceBagel, "_lance_video_preprocess", staticmethod(capture_preprocessed_input)) + request = SimpleNamespace( + prompts=[ + { + "modalities": ["video"], + "user_text": "edit the video", + "multi_modal_data": {"video": video_input}, + "extra_args": {"origin_fps": 12.0}, + } + ], + sampling_params=SimpleNamespace(extra_args={}, num_inference_steps=1), + ) + + with pytest.raises(StopAfterVideoPreprocessError): + object.__new__(LancePipeline)._forward_video_edit(request) + + video = captured["video"] + assert isinstance(video, np.ndarray) + assert video.shape == (2, 2, 3, 3) + assert video.dtype == np.uint8 + assert video.flags.c_contiguous + assert captured["fps"] == expected_fps diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index dc885ac44..f33a35b96 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -62,6 +62,7 @@ from vllm_omni.diffusion.models.bagel.pipeline_bagel import ( from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) +from vllm_omni.model_executor.models.utils import normalize_decoded_video_frames from .lance_transformer import ( LanceBagel, @@ -1454,6 +1455,8 @@ class LancePipeline(BagelPipeline): frames_bgr.append(_cv2.cvtColor(f, _cv2.COLOR_BGR2RGB)) cap.release() video_raw = _np.stack(frames_bgr, axis=0) + elif isinstance(video_input, (list, tuple)): + video_raw, origin_fps = normalize_decoded_video_frames(video_input, default_fps=origin_fps_default) elif isinstance(video_input, _np.ndarray): video_raw = video_input origin_fps = origin_fps_default diff --git a/vllm_omni/model_executor/models/utils.py b/vllm_omni/model_executor/models/utils.py index fd9ddd41b..73e9d5572 100644 --- a/vllm_omni/model_executor/models/utils.py +++ b/vllm_omni/model_executor/models/utils.py @@ -1,10 +1,59 @@ -from collections.abc import Callable +import math +from collections.abc import Callable, Sequence from contextlib import contextmanager +import numpy as np import torch +from PIL import Image from vllm.model_executor.models.utils import maybe_prefix +def normalize_decoded_video_frames( + video_input: Sequence[Image.Image], + *, + default_fps: float, +) -> tuple[np.ndarray, float]: + """Convert decoded image frames to a contiguous THWC uint8 video array.""" + if not video_input: + raise ValueError("video_edit received an empty decoded video frame sequence.") + + frame_rate = float(default_fps) + for fps in ( + getattr(video_input, "fps", None), + getattr(video_input, "frame_rate", None), + ): + try: + candidate = float(fps) + except (TypeError, ValueError): + continue + if math.isfinite(candidate) and candidate > 0: + frame_rate = candidate + break + + frames = [] + frame_size = None + for index, frame in enumerate(video_input): + if not isinstance(frame, Image.Image): + raise ValueError( + f"video_edit decoded video frame at index {index} must be a PIL.Image.Image, got {type(frame)}." + ) + width, height = frame.size + if width <= 0 or height <= 0: + raise ValueError( + f"video_edit decoded video frame at index {index} must have positive dimensions, got {width}x{height}." + ) + if frame_size is None: + frame_size = frame.size + elif frame.size != frame_size: + raise ValueError( + "video_edit decoded video frames must have identical dimensions; " + f"frame 0 is {frame_size}, frame {index} is {frame.size}." + ) + frames.append(np.asarray(frame.convert("RGB"), dtype=np.uint8)) + + return np.ascontiguousarray(np.stack(frames)), frame_rate + + @contextmanager def transformers_keys_to_ignore_compat(): """Make ``trust_remote_code`` weight loading robust to the transformers 5.9 From 28291df1795129ae03ce86a7b9b0a230166b5ac5 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:30:03 +0800 Subject: [PATCH 039/209] [Refactor][Diffusion] Remove model-specific names from LoRA and ModelOpt loader defaults (#5907) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/diffusion/lora/test_lora_manager.py | 4 +++- tests/diffusion/models/bagel/test_bagel_lora.py | 9 +++++++-- vllm_omni/diffusion/lora/manager.py | 17 ++++++++++------- .../checkpoint_adapters/modelopt.py | 9 ++++++++- .../diffusion/models/bagel/pipeline_bagel.py | 7 ++++++- .../models/z_image/z_image_transformer.py | 9 +++++++++ 6 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 10b023bf0..b2a0505f3 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -729,8 +729,10 @@ def test_lora_manager_discovers_bagel_component(monkeypatch): lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), ) - # Pipeline with a 'bagel' component (no 'transformer') + # Pipeline with a 'bagel' component (no 'transformer') declared via + # _lora_components, matching how BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() pipeline.bagel.language_model = torch.nn.Module() pipeline.bagel.language_model.qkv_proj = _FakeLinearBase() diff --git a/tests/diffusion/models/bagel/test_bagel_lora.py b/tests/diffusion/models/bagel/test_bagel_lora.py index c285758fe..1e73318ce 100644 --- a/tests/diffusion/models/bagel/test_bagel_lora.py +++ b/tests/diffusion/models/bagel/test_bagel_lora.py @@ -122,8 +122,10 @@ class TestStage1DiTLoRA: lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), ) - # Build pipeline with bagel component + # Build pipeline with bagel component declared via _lora_components, + # matching how BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() lm = torch.nn.Module() lm.stacked_params_mapping = [ @@ -194,8 +196,11 @@ class TestBagelLoRARoundTrip: monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA) - # Build pipeline with bagel.language_model.foo (simple non-packed layer) + # Build pipeline with bagel.language_model.foo (simple non-packed layer); + # the bagel component is declared via _lora_components, matching how + # BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() lm = torch.nn.Module() lm.foo = _FakeLinearBase() diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 11d193531..3a8230530 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -394,17 +394,20 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - # Default denoising components, declared DiT components, and any a - # pipeline opts into via ``_lora_components``. + # Components scanned for LoRA-capable layers: framework defaults, + # declared DiT components, and any a pipeline opts into via + # ``_lora_components``. The defaults only cover the generic diffusers + # naming convention: ``transformer`` (plus ``transformer_2`` for + # dual-DiT pipelines such as Wan2.2) and ``unet`` for SDXL-style + # pipelines. Model-specific attribute names must be declared by the + # pipeline itself via ``_dit_modules`` or ``_lora_components``. # - # NOTE: SDXL-style pipelines expose the denoiser as ``unet``. - # Without scanning this component, adapters can load/activate while - # effectively applying to zero layers, producing base-identical output. + # NOTE: if the denoiser component is not scanned here, adapters can + # load/activate while effectively applying to zero layers, producing + # base-identical output. default_components = ( "transformer", "transformer_2", - "dit", - "bagel", "unet", ) declared_components = tuple(getattr(self.pipeline, "_dit_modules", ()) or ()) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py index 68ec98d5a..f3c703b09 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py @@ -17,10 +17,17 @@ MODEL_OPT_SCALE_SUFFIXES = ( ".weight_scale_2", ".weight_scale_inv", ) +# Fused projections following the generic diffusers Attention naming +# convention. ``to_q``/``to_k``/``to_v`` and the joint-attention +# ``add_q_proj``/``add_k_proj``/``add_v_proj`` shard names are fused into +# ``to_qkv``/``add_kv_proj`` by most vllm-omni DiT implementations, many of +# which do not declare their own ``packed_modules_mapping``; keeping these +# two entries here avoids repeating them in every such model. Model-specific +# fusions (e.g. Z-Image's ``w13``) must be declared on the model class via +# ``packed_modules_mapping``, which is merged on top of these defaults. DEFAULT_PACKED_MODULES_MAPPING = { "to_qkv": ("to_q", "to_k", "to_v"), "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"), - "w13": ("w1", "w3"), } FP8_DTYPES = tuple( dtype diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index f27601148..d805fd3b8 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ BagelPipeline implementation for vLLM-Omni. """ @@ -316,6 +316,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf "bagel.connector", "bagel.vit_pos_embed", ] + # LoRA scanning target. ``_dit_modules`` holds the dotted path + # "language_model.model", which DiffusionLoRAManager cannot resolve as a + # pipeline attribute, so the top-level ``bagel`` module (LLM plus ViT, + # connector, and VAE projections) is declared explicitly here. + _lora_components: ClassVar[list[str]] = ["bagel"] supports_step_execution: ClassVar[bool] = True def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 1399eb0f0..f25c01fca 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -612,6 +612,15 @@ class ZImageTransformer2DModel(CachedTransformer): Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. """ + # Fused projections and their checkpoint shard names (mirrors + # stacked_params_mapping in load_weights). ``w13`` is the Z-Image FFN + # fusion of ``w1``/``w3``; declared here so quantized-checkpoint adapters + # can resolve shard keys without a framework-level default entry. + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + "w13": ["w1", "w3"], + } + _repeated_blocks = ["ZImageTransformerBlock"] _layerwise_offload_blocks_attrs = ["layers"] From 9edc1b4667d4b56d12effe1e5b300315520d1ae3 Mon Sep 17 00:00:00 2001 From: geray <48796550+gerayking@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:31:49 +0800 Subject: [PATCH 040/209] Optimize CosyVoice3 Stage1 flow batching (#4876) Signed-off-by: gerayking <399geray@gmail.com> Co-authored-by: Hongsheng Liu --- .../examples/online_serving/text_to_speech.md | 9 + tests/config/test_environment_variables.py | 6 +- .../cosyvoice3/test_cosyvoice3_components.py | 252 ++++++++++++++++++ .../test_cosyvoice3_model_helpers.py | 64 +++++ .../config/environment_variable_inventory.py | 2 + .../models/cosyvoice3_audio/cosyvoice3_dit.py | 173 +++++++----- .../models/cosyvoice3/code2wav_core/cfm.py | 176 ++++++------ .../models/cosyvoice3/cosyvoice3.py | 97 +++++-- .../models/cosyvoice3/cosyvoice3_code2wav.py | 213 ++++++++++++--- .../models/cosyvoice3/runtime.py | 31 +++ 10 files changed, 825 insertions(+), 198 deletions(-) create mode 100644 vllm_omni/model_executor/models/cosyvoice3/runtime.py diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 9a33dc858..2b4f617f4 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -120,6 +120,15 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path. +Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call: + +```bash +export COSYVOICE3_BATCH_FLOW=1 +vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code +``` + +Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead. + ### CLI client The client defaults to the official upstream zero-shot prompt, so it runs without extra flags: diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 7603945a0..cc37cd6b1 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 24, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 61, + EnvironmentVariableCategory.MODEL_SPECIFIC: 63, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -178,10 +178,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 37, + ModelEnvironmentVariableDisposition.PROMOTE: 38, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 13, + ModelEnvironmentVariableDisposition.INTERNALIZE: 14, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, } diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 4f01022f8..1b1df6367 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for CosyVoice3 components.""" +import types from contextlib import contextmanager from types import SimpleNamespace @@ -9,6 +10,7 @@ import pytest import torch import torch.nn as nn +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( @@ -29,6 +31,23 @@ def _force_torch_sdpa(): yield +def test_batch_flow_runtime_flags_are_opt_in(monkeypatch): + from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_enabled, + ) + + monkeypatch.delenv("COSYVOICE3_BATCH_FLOW", raising=False) + monkeypatch.delenv("COSYVOICE3_BATCH_FLOW_DEBUG", raising=False) + assert not cosyvoice3_batch_flow_enabled() + assert not cosyvoice3_batch_flow_debug() + + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW_DEBUG", "1") + assert cosyvoice3_batch_flow_enabled() + assert cosyvoice3_batch_flow_debug() + + @pytest.fixture def causal_hift(): return CausalHiFTGenerator( @@ -139,6 +158,39 @@ class TestDiTAttention: assert attention.to_k.out_features == 512 assert attention.to_v.out_features == 512 + @pytest.mark.core_model + @pytest.mark.cpu + def test_forward_passes_2d_mask_without_casting_qkv(self): + """The float32 fallback receives the per-row validity mask.""" + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention + + class CaptureAttention(nn.Module): + def __init__(self): + super().__init__() + self.last_metadata = None + self.last_dtype = None + + def forward(self, q, k, v, attn_metadata=None): + self.last_metadata = attn_metadata + self.last_dtype = q.dtype + return torch.zeros_like(q) + + attention = DiTAttention(dim=16, heads=2, dim_head=8, dropout=0.0) + capture = CaptureAttention() + attention.attn = capture + + x = torch.randn(2, 5, 16) + mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]]) + + out = attention(x, mask=mask) + + assert out.shape == x.shape + assert capture.last_metadata is not None + assert torch.equal(capture.last_metadata.attn_mask, mask) + assert capture.last_dtype == x.dtype + assert out.dtype == x.dtype + assert torch.allclose(out[~mask], torch.zeros_like(out[~mask])) + class TestDiTBlock: """Tests for DiTBlock.""" @@ -422,6 +474,52 @@ class TestCFM: assert reused_context is context assert reused_stream is stream + @pytest.mark.core_model + @pytest.mark.cpu + def test_causal_conditional_cfm_batches_cfg_estimator(self): + """Batched flow calls should invoke the CFG estimator with 2B rows.""" + from omegaconf import DictConfig + + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM + + class RecordingEstimator(nn.Module): + def __init__(self): + super().__init__() + self.shapes: list[tuple[int, ...]] = [] + + def forward(self, x, mask, mu, t, spks=None, cond=None): + self.shapes.append(tuple(x.shape)) + return torch.zeros_like(x) + + estimator = RecordingEstimator() + cfm_params = DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ) + cfm = CausalConditionalCFM( + in_channels=80, + cfm_params=cfm_params, + n_spks=1, + spk_emb_dim=80, + estimator=estimator, + ) + + batch, mel_dim, seq_len = 3, 80, 16 + mu = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond) + + assert out.shape == mu.shape + assert estimator.shapes == [(2 * batch, mel_dim, seq_len)] * 2 + class TestSDPAFallback: """Test SDPA fallback for float32 inputs.""" @@ -484,3 +582,157 @@ def test_code2wav_forward_finalizes_hift_tail(): assert out.shape == (1, 1, 8) assert model.hift.finalize_calls == [True] assert forward_mel_calls[0]["token_offset_tokens"] == 0 + + +def test_code2wav_streaming_batch_pads_codec_tokens_and_preserves_lengths(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + model = object.__new__(CosyVoice3Code2Wav) + model.flow_model = SimpleNamespace(pre_lookahead_len=1, token_mel_ratio=2) + forward_mel_calls = [] + + def fake_forward_mel(self, **kwargs): + forward_mel_calls.append(kwargs) + return torch.arange(2 * 80 * 8, dtype=torch.float32).reshape(2, 80, 8) + + def fake_stream_hift(self, feat, *, cache_state=None, finalize=False): + return feat, None + + model._forward_mel = types.MethodType(fake_forward_mel, model) + model._stream_hift_from_feat = types.MethodType(fake_stream_hift, model) + items = [ + { + "token": torch.ones(1, 3, dtype=torch.int32), + "prompt_token": torch.ones(1, 4, dtype=torch.int32), + "prompt_feat": torch.ones(1, 8, 80), + "embedding": torch.ones(1, 192), + "finalize": False, + }, + { + "token": torch.ones(1, 5, dtype=torch.int32), + "prompt_token": torch.ones(1, 4, dtype=torch.int32), + "prompt_feat": torch.ones(1, 8, 80), + "embedding": torch.ones(1, 192), + "finalize": False, + }, + ] + + results = model.forward_streaming_batch(items) + + assert len(forward_mel_calls) == 1 + call = forward_mel_calls[0] + assert call["token"].shape == (2, 5) + assert torch.equal(call["token_lens"], torch.tensor([3, 5], dtype=torch.int32)) + assert results[0][0].shape[-1] == 4 + assert results[1][0].shape[-1] == 8 + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch): + """A padded flow batch must match individual flow calls on valid mels.""" + from omegaconf import DictConfig + + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( + CausalConditionalCFM, + CausalMaskedDiffWithDiT, + ) + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + torch.manual_seed(0) + estimator = DiT( + dim=32, + depth=1, + heads=4, + dim_head=8, + dropout=0.0, + ff_mult=2, + mel_dim=80, + mu_dim=80, + spk_dim=80, + out_channels=80, + ) + decoder = CausalConditionalCFM( + in_channels=80, + cfm_params=DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ), + n_spks=1, + spk_emb_dim=80, + estimator=estimator, + ) + flow_model = ( + CausalMaskedDiffWithDiT( + input_size=80, + output_size=80, + spk_embed_dim=192, + vocab_size=64, + input_frame_rate=25, + only_mask_loss=True, + token_mel_ratio=2, + pre_lookahead_len=1, + pre_lookahead_layer=PreLookaheadLayer(in_channels=80, channels=80, pre_lookahead_len=1), + decoder=decoder, + ) + .to(device="cuda", dtype=torch.bfloat16) + .eval() + ) + + model = object.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = flow_model + + def return_mel(self, feat, *, cache_state=None, finalize=False): + return feat, None + + model._stream_hift_from_feat = types.MethodType(return_mel, model) + + original_randn = torch.randn + + def length_consistent_randn(*size, **kwargs): + shape = tuple(size[0]) if len(size) == 1 and isinstance(size[0], (tuple, list)) else tuple(size) + if len(shape) == 3 and shape[1] == 80: + device = kwargs.get("device") + dtype = kwargs.get("dtype", torch.float32) + channels = torch.arange(shape[1], device=device, dtype=torch.float32).view(1, -1, 1) + positions = torch.arange(shape[2], device=device, dtype=torch.float32).view(1, 1, -1) + noise = torch.sin(channels * 0.17 + positions * 0.31) + return noise.expand(shape[0], -1, -1).to(dtype=dtype).clone() + return original_randn(*size, **kwargs) + + monkeypatch.setattr(torch, "randn", length_consistent_randn) + common = { + "prompt_token": torch.tensor([[7, 8]], dtype=torch.int32), + "prompt_feat": torch.linspace(-0.5, 0.5, 4 * 80).reshape(1, 4, 80), + "embedding": torch.linspace(-1.0, 1.0, 192).reshape(1, 192), + "finalize": False, + } + items = [ + {**common, "token": torch.tensor([[1, 2, 3]], dtype=torch.int32)}, + {**common, "token": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32)}, + ] + + batched = model.forward_streaming_batch(items, n_timesteps=2) + individual = [ + model.forward_streaming( + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + n_timesteps=2, + ) + for item in items + ] + + for (batched_mel, _), (individual_mel, _) in zip(batched, individual): + assert batched_mel.shape == individual_mel.shape + rel_mean = (batched_mel - individual_mel).abs().mean() / individual_mel.abs().mean().clamp_min(1e-6) + assert rel_mean.item() < 0.05 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index e16927096..1d583de6c 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -42,6 +42,7 @@ class _DummyCode2Wav: self.outputs = list(outputs or []) self.forward_calls: list[dict[str, object]] = [] self.forward_streaming_calls: list[dict[str, object]] = [] + self.forward_streaming_batch_calls: list[list[dict[str, object]]] = [] def forward(self, **kwargs): self.forward_calls.append(kwargs) @@ -65,6 +66,22 @@ class _DummyCode2Wav: } return audio, new_state + def forward_streaming_batch(self, items, *, n_timesteps: int = 10): + self.forward_streaming_batch_calls.append(items) + return [ + self.forward_streaming( + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + cache_state=item.get("cache_state"), + n_timesteps=n_timesteps, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + for item in items + ] + def _make_code2wav_model( *, @@ -373,6 +390,53 @@ def test_forward_clears_streaming_cache_on_terminal_chunk(): assert "rid-stream" not in model._stream_vocoder_cache_by_req +def test_forward_batches_streaming_flow_items(monkeypatch): + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + model = _make_code2wav_model() + runtime_info = [ + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long), + "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32), + "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32), + }, + "meta": { + "req_id": ["rid-a"], + "stream_finished": torch.tensor(False), + "left_context_size": 0, + }, + }, + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long), + "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32), + "embedding": torch.tensor([[0.7, 0.8]], dtype=torch.float32), + }, + "meta": { + "req_id": ["rid-b"], + "stream_finished": torch.tensor(False), + "left_context_size": 1, + }, + }, + ] + + out = model.forward( + input_ids=torch.tensor([0, 1, 2, 0, 1, 2], dtype=torch.long), + positions=torch.arange(6, dtype=torch.long), + model_intermediate_buffer=runtime_info, + seq_token_counts=[3, 3], + ) + + assert len(out.multimodal_outputs["audio"]) == 2 + assert len(model.code2wav.forward_streaming_batch_calls) == 1 + batch_items = model.code2wav.forward_streaming_batch_calls[0] + assert [item["index"] for item in batch_items] == [0, 1] + assert torch.equal(batch_items[0]["token"], torch.tensor([[0, 1, 2]])) + assert batch_items[1]["token_offset_tokens"] == 1 + assert "rid-a" in model._stream_vocoder_cache_by_req + assert "rid-b" in model._stream_vocoder_cache_by_req + + def test_sample_uses_ras_rejection_for_recent_repetition(): model = _make_talker_model() metadata = _make_sampling_metadata(output_token_ids=[[1] * 10]) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index c1e380a22..4a4966d9c 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -137,6 +137,7 @@ _MODEL_PROMOTE = ( "AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH", "COSMOS3_SOUND_TOKENIZER_PATH", + "COSYVOICE3_BATCH_FLOW", "COSYVOICE3_ESTIMATOR_ONNX", "COSYVOICE3_TRT", "COSYVOICE3_TRT_CACHE", @@ -188,6 +189,7 @@ _MODEL_REQUEST_SCOPE = ( _MODEL_EXTERNAL: tuple[str, ...] = () _MODEL_INTERNALIZE = ( + "COSYVOICE3_BATCH_FLOW_DEBUG", "DOTS_TTS_BETA_TRACE", "DYNIN_S2U_VENDOR_ROOT", "DZ_PHASE_TIMING", diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py index 5ae75b256..d0538d79f 100644 --- a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow/DiT # Refactored to use vllm_omni diffusion infrastructure for optimized attention backends @@ -15,8 +15,10 @@ from torch import nn from vllm.logger import init_logger from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_batch_flow_profile logger = init_logger(__name__) @@ -52,6 +54,22 @@ def get_pos_embed_indices(start, length, max_pos, scale=1.0): return pos +def _normalize_attention_mask(mask: torch.Tensor | None, batch_size: int, seq_len: int) -> torch.Tensor | None: + if mask is None: + return None + if mask.dim() == 2: + attn_mask = mask + elif mask.dim() == 3 and mask.shape[1] == 1: + attn_mask = mask[:, 0] + elif mask.dim() == 4: + attn_mask = mask[:, 0, -1] + else: + return None + if attn_mask.shape != (batch_size, seq_len): + return None + return attn_mask.bool() + + class FeedForward(nn.Module): """Feed-forward network with GELU activation.""" @@ -114,42 +132,39 @@ class DiTAttention(nn.Module): ) -> torch.Tensor: batch_size, seq_len = x.shape[0], x.shape[1] - # Project to Q, K, V - query = self.to_q(x) - key = self.to_k(x) - value = self.to_v(x) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_qkv_projection"): + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) - # Apply rotary position embedding - if rope is not None: - freqs, xpos_scale = rope - q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) - query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) - key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_rope"): + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) - # Reshape for attention: (batch, seq, heads, head_dim) - query = query.view(batch_size, seq_len, self.heads, self.dim_head) - key = key.view(batch_size, seq_len, self.heads, self.dim_head) - value = value.view(batch_size, seq_len, self.heads, self.dim_head) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_backend"): + # Reshape for attention: (batch, seq, heads, head_dim) + query = query.view(batch_size, seq_len, self.heads, self.dim_head) + key = key.view(batch_size, seq_len, self.heads, self.dim_head) + value = value.view(batch_size, seq_len, self.heads, self.dim_head) - # Use diffusion attention backend - # The diffusion Attention layer expects (batch, seq, heads, head_dim) - out = self.attn(query, key, value, attn_metadata=None) + # The diffusion Attention layer expects (batch, seq, heads, head_dim) + attn_mask = _normalize_attention_mask(mask, batch_size, seq_len) + attn_metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None + out = self.attn(query, key, value, attn_metadata=attn_metadata) - # Some attention backends return a non-contiguous layout. - out = out.reshape(batch_size, seq_len, self.inner_dim) - out = out.to(query.dtype) - - # Output projection - out = self.to_out(out) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_output_projection"): + # Some attention backends return a non-contiguous layout. + out = out.reshape(batch_size, seq_len, self.inner_dim) + out = out.to(x.dtype) + out = self.to_out(out) # Apply mask if provided - if mask is not None: - if mask.dim() == 2: - mask = mask.unsqueeze(-1) - elif mask.dim() == 4: - # (batch, heads, seq, seq) -> use last dim - mask = mask[:, 0, -1].unsqueeze(-1) - out = out.masked_fill(~mask.bool(), 0.0) + attn_mask = _normalize_attention_mask(mask, batch_size, seq_len) + if attn_mask is not None: + out = out.masked_fill(~attn_mask.unsqueeze(-1), 0.0) return out @@ -172,18 +187,19 @@ class DiTBlock(nn.Module): self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") def forward(self, x, t, mask=None, rope=None): - # pre-norm & modulation for attention input - norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_adalayernorm_attention"): + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) - # attention - attn_output = self.attn(x=norm, mask=mask, rope=rope) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention"): + attn_output = self.attn(x=norm, mask=mask, rope=rope) - # process attention output for input x - x = x + gate_msa.unsqueeze(1) * attn_output + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_residual"): + x = x + gate_msa.unsqueeze(1) * attn_output - ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - ff_output = self.ff(ff_norm) - x = x + gate_mlp.unsqueeze(1) * ff_output + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_layernorm_ffn"): + ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(ff_norm) + x = x + gate_mlp.unsqueeze(1) * ff_output return x @@ -205,19 +221,20 @@ class CausalConvPositionEmbedding(nn.Module): ) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): - if mask is not None: - mask = mask[..., None] - x = x.masked_fill(~mask, 0.0) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_causal_conv_position_embedding"): + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) - x = x.permute(0, 2, 1) - x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) - x = self.conv1(x) - x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) - x = self.conv2(x) - out = x.permute(0, 2, 1) + x = x.permute(0, 2, 1) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv1(x) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv2(x) + out = x.permute(0, 2, 1) - if mask is not None: - out = out.masked_fill(~mask, 0.0) + if mask is not None: + out = out.masked_fill(~mask, 0.0) return out @@ -327,13 +344,15 @@ class InputEmbedding(nn.Module): self.conv_pos_embed = CausalConvPositionEmbedding(dim=out_dim) def forward(self, x, cond, text_embed, spks): - to_cat = [x, cond, text_embed] - if self.spk_dim > 0: - spks = repeat(spks, "b c -> b t c", t=x.shape[1]) - to_cat.append(spks) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_input_projection"): + to_cat = [x, cond, text_embed] + if self.spk_dim > 0: + spks = repeat(spks, "b c -> b t c", t=x.shape[1]) + to_cat.append(spks) - x = self.proj(torch.cat(to_cat, dim=-1)) - x = self.conv_pos_embed(x) + x + x = self.proj(torch.cat(to_cat, dim=-1)) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_input_causal_conv"): + x = self.conv_pos_embed(x) + x return x @@ -385,31 +404,37 @@ class DiT(nn.Module): self.num_decoding_left_chunks = num_decoding_left_chunks def forward(self, x, mask, mu, t, spks=None, cond=None): - x = x.transpose(1, 2) - mu = mu.transpose(1, 2) - cond = cond.transpose(1, 2) - spks = spks.unsqueeze(dim=1) - batch, seq_len = x.shape[0], x.shape[1] - if t.ndim == 0: - t = t.repeat(batch) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_prepare_inputs"): + x = x.transpose(1, 2) + mu = mu.transpose(1, 2) + cond = cond.transpose(1, 2) + spks = spks.unsqueeze(dim=1) + batch, seq_len = x.shape[0], x.shape[1] + if t.ndim == 0: + t = t.repeat(batch) - # t: conditioning time, c: context (text + masked cond audio), x: noised input audio - t = self.time_embed(t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_time_embedding"): + t = self.time_embed(t) x = self.input_embed(x, cond, mu, spks.squeeze(1)) - rope = self.rotary_embed.forward_from_seq_len(seq_len) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_rope_and_mask"): + rope = self.rotary_embed.forward_from_seq_len(seq_len) + attn_mask = mask[:, 0].bool() if mask.dim() == 3 and mask.shape[1] == 1 else mask.bool() if self.long_skip_connection is not None: residual = x - attn_mask = mask.bool().repeat(1, x.size(1), 1).unsqueeze(dim=1) - - for block in self.transformer_blocks: - x = block(x, t, mask=attn_mask.bool(), rope=rope) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_transformer_blocks"): + for block_index, block in enumerate(self.transformer_blocks): + with cosyvoice3_batch_flow_profile(f"cosyvoice3_dit_transformer_block_{block_index}"): + x = block(x, t, mask=attn_mask.bool(), rope=rope) if self.long_skip_connection is not None: - x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_long_skip_projection"): + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) - x = self.norm_out(x, t) - output = self.proj_out(x).transpose(1, 2) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_final_adalayernorm"): + x = self.norm_out(x, t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_final_projection"): + output = self.proj_out(x).transpose(1, 2) return output diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py index 2ec12a251..c86b4a305 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow """Conditional Flow Matching (CFM) classes for audio generation.""" @@ -11,6 +11,7 @@ from omegaconf import DictConfig from torch.nn import functional as F from vllm.logger import init_logger +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_batch_flow_profile from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask logger = init_logger(__name__) @@ -74,20 +75,23 @@ class ConditionalCFM(BASECFM): shape: (batch_size, n_feats, mel_timesteps) """ - z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature - cache_size = cache.shape[2] - # fix prompt and overlap part mu and z - if cache_size != 0: - z[:, :, :cache_size] = cache[:, :, :, 0] - mu[:, :, :cache_size] = cache[:, :, :, 1] - z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) - mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) - cache = torch.stack([z_cache, mu_cache], dim=-1) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_noise_cache"): + z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature + cache_size = cache.shape[2] + # fix prompt and overlap part mu and z + if cache_size != 0: + z[:, :, :cache_size] = cache[:, :, :, 0] + mu[:, :, :cache_size] = cache[:, :, :, 1] + z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) + mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) + cache = torch.stack([z_cache, mu_cache], dim=-1) - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == "cosine": - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_t_span"): + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + with cosyvoice3_batch_flow_profile(f"cosyvoice3_cfm_euler_{max(1, int(n_timesteps))}_steps"): + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache def solve_euler(self, x, t_span, mu, mask, spks, cond): """ @@ -111,26 +115,39 @@ class ConditionalCFM(BASECFM): # Do not use concat, it may cause memory format changed and trt infer with wrong results! # NOTE when flow run in amp mode, x.dtype is float32, which cause nan in trt fp16 - # inference, so set dtype=spks.dtype - x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) - mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=spks.dtype) - mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) - t_in = torch.zeros([2], device=x.device, dtype=spks.dtype) - spks_in = torch.zeros([2, 80], device=x.device, dtype=spks.dtype) - cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + # inference, so set dtype=spks.dtype. The batch is doubled for CFG: + # first B rows are conditioned, second B rows are unconditional. + batch_size = int(x.size(0)) + estimator_batch = 2 * batch_size + estimator_dtype = spks.dtype if spks is not None else x.dtype + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_allocate_2b"): + x_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) + mask_in = torch.zeros([estimator_batch, 1, x.size(2)], device=x.device, dtype=estimator_dtype) + mu_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) + t_in = torch.zeros([estimator_batch], device=x.device, dtype=estimator_dtype) + spks_in = torch.zeros([estimator_batch, 80], device=x.device, dtype=estimator_dtype) + cond_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) for step in range(1, len(t_span)): # Classifier-Free Guidance inference introduced in VoiceBox - x_in[:] = x - mask_in[:] = mask - mu_in[0] = mu - t_in[:] = t.unsqueeze(0) - spks_in[0] = spks - cond_in[0] = cond - dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) - dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0) - dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt - x = x + dt * dphi_dt - t = t + dt + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_prepare_2b"): + x_in[:batch_size] = x + x_in[batch_size:] = x + mask_in[:batch_size] = mask + mask_in[batch_size:] = mask + mu_in[:batch_size] = mu + t_in[:] = t + if spks is not None: + spks_in[:batch_size] = spks + if cond is not None: + cond_in[:batch_size] = cond + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_forward_estimator"): + dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_combine"): + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [batch_size, batch_size], dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_euler_update"): + x = x + dt * dphi_dt + t = t + dt sol.append(x) if step < len(t_span) - 1: dt = t_span[step + 1] - t @@ -158,12 +175,12 @@ class ConditionalCFM(BASECFM): spks_e = spks.to(io_dtype).contiguous() cond_e = cond.to(io_dtype).contiguous() out_e = torch.empty_like(x_e) - estimator.set_input_shape("x", (2, 80, x_e.size(2))) - estimator.set_input_shape("mask", (2, 1, x_e.size(2))) - estimator.set_input_shape("mu", (2, 80, x_e.size(2))) - estimator.set_input_shape("t", (2,)) - estimator.set_input_shape("spks", (2, 80)) - estimator.set_input_shape("cond", (2, 80, x_e.size(2))) + estimator.set_input_shape("x", tuple(x_e.shape)) + estimator.set_input_shape("mask", tuple(mask_e.shape)) + estimator.set_input_shape("mu", tuple(mu_e.shape)) + estimator.set_input_shape("t", tuple(t_e.shape)) + estimator.set_input_shape("spks", tuple(spks_e.shape)) + estimator.set_input_shape("cond", tuple(cond_e.shape)) data_ptrs = [ x_e.data_ptr(), mask_e.data_ptr(), @@ -211,22 +228,25 @@ class CausalConditionalCFM(ConditionalCFM): shape: (batch_size, n_feats, mel_timesteps) """ - z = ( - torch.randn( - (mu.size(0), mu.size(1), mu.size(2)), - device=mu.device, - dtype=mu.dtype, + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_noise_cache"): + z = ( + torch.randn( + (mu.size(0), mu.size(1), mu.size(2)), + device=mu.device, + dtype=mu.dtype, + ) + * temperature ) - * temperature - ) # fix prompt and overlap part mu and z - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_t_span"): + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == "cosine": - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None + with cosyvoice3_batch_flow_profile(f"cosyvoice3_cfm_euler_{max(1, int(n_timesteps))}_steps"): + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None class CausalMaskedDiffWithDiT(torch.nn.Module): @@ -299,36 +319,41 @@ class CausalMaskedDiffWithDiT(torch.nn.Module): finalize: bool = False, n_timesteps: int = 10, ): - assert token.shape[0] == 1 - # xvec projection + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_speaker_embedding"): + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) - embedding = F.normalize(embedding, dim=1) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_token_embedding_lookahead"): + # concat text and prompt_text + codec_token_len = token_len + token, total_token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + codec_token_len + mask = (~make_pad_mask(total_token_len, max_len=token.shape[1])).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + # text encode + if finalize is True: + h = self.pre_lookahead_layer(token) + else: + h = self.pre_lookahead_layer( + token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] + ) - embedding = self.spk_embed_affine_layer(embedding) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_repeat_to_mel_axis"): + h = h.repeat_interleave(self.token_mel_ratio, dim=1) - # concat text and prompt_text - token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len - mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding) - token = self.input_embedding(torch.clamp(token, min=0)) * mask - # text encode - if finalize is True: - h = self.pre_lookahead_layer(token) - else: - h = self.pre_lookahead_layer( - token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] - ) - h = h.repeat_interleave(self.token_mel_ratio, dim=1) + batch_size = int(token.shape[0]) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cond_prompt_mel"): + mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] - mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] + # get conditions + conds = torch.zeros([batch_size, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) - # get conditions - conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + lookahead = 0 if finalize else int(self.pre_lookahead_len) + valid_h_lens = torch.clamp(total_token_len.to(torch.long) - lookahead, min=0) + mel_lens = torch.clamp(valid_h_lens * int(self.token_mel_ratio), max=mel_len1 + mel_len2) + mask = (~make_pad_mask(mel_lens, max_len=mel_len1 + mel_len2)).to(h) - conds[:, :mel_len1] = prompt_feat - - conds = conds.transpose(1, 2) - - mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) feat, _ = self.decoder( mu=h.transpose(1, 2).contiguous(), mask=mask.unsqueeze(1), @@ -338,6 +363,7 @@ class CausalMaskedDiffWithDiT(torch.nn.Module): streaming=streaming, ) - feat = feat[:, :, mel_len1:] - assert feat.shape[2] == mel_len2 + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_crop_prompt_mel"): + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 return feat.float(), None diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 558e96b65..0ddc6dafa 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -39,6 +39,10 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_enabled, +) from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, @@ -1070,12 +1074,21 @@ class CosyVoice3Model( request_ids_list = self._split_request_ids(flat_ids, seq_token_counts) num_reqs = max(1, len(request_ids_list)) + debug_batch_flow = cosyvoice3_batch_flow_debug() + if debug_batch_flow: + logger.info( + "CosyVoice3 code2wav debug: forward num_reqs=%d seq_token_counts=%s flat_ids=%d", + num_reqs, + seq_token_counts, + int(flat_ids.numel()), + ) sample_rate = torch.tensor(int(self.config.sample_rate), dtype=torch.int32) empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device) audios: list[torch.Tensor] = [empty_audio] * num_reqs srs: list[torch.Tensor] = [sample_rate] * num_reqs if not isinstance(runtime_info, list): runtime_info = [] + streaming_flow_items: list[dict[str, object]] = [] for idx, req_ids in enumerate(request_ids_list): raw = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {} @@ -1142,23 +1155,21 @@ class CosyVoice3Model( with self._stream_audio_cache_lock: cache_state = self._stream_vocoder_cache_by_req.get(req_id) - tts_speech, new_cache_state = self.code2wav.forward_streaming( - token=token.unsqueeze(0), - prompt_token=speech_token[:1], - prompt_feat=speech_feat[:1], - embedding=embedding[:1], - cache_state=cache_state, - n_timesteps=10, - token_offset_tokens=token_offset, - finalize=stream_finished, + streaming_flow_items.append( + { + "index": idx, + "req_id": req_id, + "stream_finished": stream_finished, + "token": token.unsqueeze(0), + "prompt_token": speech_token[:1], + "prompt_feat": speech_feat[:1], + "embedding": embedding[:1], + "cache_state": cache_state, + "token_offset_tokens": token_offset, + "finalize": stream_finished, + } ) - - if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): - with self._stream_audio_cache_lock: - if new_cache_state is None or stream_finished: - self._stream_vocoder_cache_by_req.pop(req_id, None) - else: - self._stream_vocoder_cache_by_req[req_id] = new_cache_state + continue else: token_offset = max(0, meta.talker_prefill_offset or 0) if meta else 0 tts_speech = self.code2wav.forward( @@ -1174,6 +1185,60 @@ class CosyVoice3Model( audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished) + if streaming_flow_items: + if debug_batch_flow: + item_shapes = [ + ( + tuple(item["token"].shape), # type: ignore[union-attr] + tuple(item["prompt_token"].shape), # type: ignore[union-attr] + tuple(item["prompt_feat"].shape), # type: ignore[union-attr] + bool(item.get("finalize", False)), + ) + for item in streaming_flow_items + ] + logger.info( + "CosyVoice3 code2wav debug: streaming_items=%d item_shapes=%s", + len(streaming_flow_items), + item_shapes, + ) + if ( + len(streaming_flow_items) > 1 + and cosyvoice3_batch_flow_enabled() + and hasattr(self.code2wav, "forward_streaming_batch") + ): + streaming_results = self.code2wav.forward_streaming_batch( + streaming_flow_items, + n_timesteps=10, + ) + else: + streaming_results = [ + self.code2wav.forward_streaming( + token=item["token"], # type: ignore[arg-type] + prompt_token=item["prompt_token"], # type: ignore[arg-type] + prompt_feat=item["prompt_feat"], # type: ignore[arg-type] + embedding=item["embedding"], # type: ignore[arg-type] + cache_state=item.get("cache_state"), # type: ignore[arg-type] + n_timesteps=10, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + for item in streaming_flow_items + ] + + for item, (tts_speech, new_cache_state) in zip(streaming_flow_items, streaming_results): + idx = int(item["index"]) + req_id = item.get("req_id") + stream_finished = bool(item.get("stream_finished", False)) + if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): + with self._stream_audio_cache_lock: + if new_cache_state is None or stream_finished: + self._stream_vocoder_cache_by_req.pop(req_id, None) + else: + self._stream_vocoder_cache_by_req[req_id] = new_cache_state + + audio = tts_speech.reshape(-1).to(dtype=torch.float32) + audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished) + return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audios, "sr": srs}) else: raise ValueError(f"Unsupported model_stage: {self.model_stage}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 49e8e515b..ef706b6c8 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ CosyVoice3 Code2Wav Stage - Converts speech tokens to audio waveforms. @@ -11,6 +11,9 @@ This module contains the code2wav (token-to-waveform) stage which uses: from __future__ import annotations +from collections import Counter +from typing import cast + import numpy as np import torch import torch.nn as nn @@ -27,6 +30,10 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( CausalHiFTGenerator, ) from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer +from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_profile, +) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) @@ -159,6 +166,9 @@ class CosyVoice3Code2Wav(nn.Module): token_offset_tokens: int = 0, streaming: bool = True, finalize: bool = False, + token_lens: torch.Tensor | None = None, + prompt_token_lens: torch.Tensor | None = None, + prompt_feat_lens: torch.Tensor | None = None, ) -> torch.Tensor: """Generate mel features via the upstream flow-model inference path.""" flow_weight = next(self.flow_model.parameters()) @@ -169,9 +179,22 @@ class CosyVoice3Code2Wav(nn.Module): prompt_token = prompt_token.to(device=device, dtype=torch.int32) prompt_feat = prompt_feat.to(device=device, dtype=dtype) embedding = embedding.to(device=device, dtype=dtype) - token_len = torch.tensor([token.shape[1]], device=device, dtype=torch.int32) - prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32) - prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32) + batch_size = int(token.shape[0]) + token_len = ( + token_lens.to(device=device, dtype=torch.int32) + if token_lens is not None + else torch.full((batch_size,), token.shape[1], device=device, dtype=torch.int32) + ) + prompt_token_len = ( + prompt_token_lens.to(device=device, dtype=torch.int32) + if prompt_token_lens is not None + else torch.full((batch_size,), prompt_token.shape[1], device=device, dtype=torch.int32) + ) + prompt_feat_len = ( + prompt_feat_lens.to(device=device, dtype=torch.int32) + if prompt_feat_lens is not None + else torch.full((batch_size,), prompt_feat.shape[1], device=device, dtype=torch.int32) + ) feat, _ = self.flow_model.inference( token=token, @@ -192,37 +215,13 @@ class CosyVoice3Code2Wav(nn.Module): return feat - @torch.inference_mode() - def forward_streaming( + def _stream_hift_from_feat( self, - token: torch.Tensor, - prompt_token: torch.Tensor, - prompt_feat: torch.Tensor, - embedding: torch.Tensor, + feat: torch.Tensor, *, cache_state: dict[str, torch.Tensor] | None = None, - n_timesteps: int = 10, - token_offset_tokens: int = 0, finalize: bool = False, ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: - """Decode streaming audio using cumulative mel + emitted-speech offset. - - This mirrors upstream CosyVoice3 streaming semantics more closely than - waveform-domain overlap-add: keep a cumulative mel history per request, - re-run causal HiFT on the history, and emit only the newly grown speech - suffix. That preserves causal look-right handling without double - trimming or duplicated overlap at chunk boundaries. - """ - feat = self._forward_mel( - token=token, - prompt_token=prompt_token, - prompt_feat=prompt_feat, - embedding=embedding, - n_timesteps=n_timesteps, - token_offset_tokens=token_offset_tokens, - streaming=True, - finalize=finalize, - ) hift_weight = self.hift.m_source.l_linear.weight chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) @@ -257,6 +256,160 @@ class CosyVoice3Code2Wav(nn.Module): } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state + @torch.inference_mode() + def forward_streaming_batch( + self, + items: list[dict[str, object]], + *, + n_timesteps: int = 10, + ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: + """Batch the flow-matching mel path, then run HiFT per request. + + Items are grouped by prompt condition shape and finalization state. + Codec tokens may have different lengths; those are padded within the + group and passed to the flow as per-row token lengths. + """ + results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) + groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {} + for index, item in enumerate(items): + token = item["token"] + prompt_token = item["prompt_token"] + prompt_feat = item["prompt_feat"] + embedding = item["embedding"] + assert isinstance(token, torch.Tensor) + assert isinstance(prompt_token, torch.Tensor) + assert isinstance(prompt_feat, torch.Tensor) + assert isinstance(embedding, torch.Tensor) + key = ( + int(prompt_token.shape[1]), + int(prompt_feat.shape[1]), + int(embedding.shape[1]), + bool(item.get("finalize", False)), + ) + groups.setdefault(key, []).append((index, item)) + + if cosyvoice3_batch_flow_debug(): + group_summary = {key: len(group) for key, group in groups.items()} + group_size_distribution = Counter(group_summary.values()) + batchable_items = sum(size for size in group_summary.values() if size > 1) + logger.info( + "CosyVoice3 code2wav debug: forward_streaming_batch items=%d " + "groups=%s group_size_distribution=%s batchable_items=%d", + len(items), + group_summary, + dict(sorted(group_size_distribution.items())), + batchable_items, + ) + + for _key, group in groups.items(): + if len(group) == 1: + index, item = group[0] + result = self.forward_streaming( + token=item["token"], # type: ignore[arg-type] + prompt_token=item["prompt_token"], # type: ignore[arg-type] + prompt_feat=item["prompt_feat"], # type: ignore[arg-type] + embedding=item["embedding"], # type: ignore[arg-type] + cache_state=item.get("cache_state"), # type: ignore[arg-type] + n_timesteps=n_timesteps, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + results[index] = result + continue + + token_tensors = [item["token"] for _, item in group] + assert all(isinstance(token, torch.Tensor) for token in token_tensors) + token_lens = torch.tensor( + [int(token.shape[1]) for token in token_tensors], # type: ignore[union-attr] + dtype=torch.int32, + ) + max_token_len = int(token_lens.max().item()) + padded_tokens = [] + for token in token_tensors: + assert isinstance(token, torch.Tensor) + if int(token.shape[1]) == max_token_len: + padded_tokens.append(token) + else: + pad = torch.zeros( + (token.shape[0], max_token_len - int(token.shape[1])), + device=token.device, + dtype=token.dtype, + ) + padded_tokens.append(torch.cat([token, pad], dim=1)) + tokens = torch.cat(padded_tokens, dim=0) + prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) # type: ignore[list-item] + embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32) + prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32) + finalize = bool(group[0][1].get("finalize", False)) + + with cosyvoice3_batch_flow_profile(f"cosyvoice3_flow_batch_b{len(group)}_t{tokens.shape[1]}"): + feat = self._forward_mel( + token=tokens, + prompt_token=prompt_tokens, + prompt_feat=prompt_feats, + embedding=embeddings, + n_timesteps=n_timesteps, + token_offset_tokens=0, + streaming=True, + finalize=finalize, + token_lens=token_lens, + prompt_token_lens=prompt_token_lens, + prompt_feat_lens=prompt_feat_lens, + ) + + for row, (index, item) in enumerate(group): + trim_mel = max(0, int(item.get("token_offset_tokens", 0))) * int(self.token_mel_ratio) + valid_tokens = int(token_lens[row].item()) + if not finalize: + valid_tokens = max(0, valid_tokens - int(self.flow_model.pre_lookahead_len)) + valid_mel = valid_tokens * int(self.token_mel_ratio) + row_feat = feat[row : row + 1, :, :valid_mel] + if trim_mel > 0: + row_feat = row_feat[:, :, trim_mel:] + results[index] = self._stream_hift_from_feat( + row_feat, + cache_state=item.get("cache_state"), # type: ignore[arg-type] + finalize=finalize, + ) + + assert all(result is not None for result in results), "every streaming item must produce exactly one result" + return cast(list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]], results) + + @torch.inference_mode() + def forward_streaming( + self, + token: torch.Tensor, + prompt_token: torch.Tensor, + prompt_feat: torch.Tensor, + embedding: torch.Tensor, + *, + cache_state: dict[str, torch.Tensor] | None = None, + n_timesteps: int = 10, + token_offset_tokens: int = 0, + finalize: bool = False, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + """Decode streaming audio using cumulative mel + emitted-speech offset. + + This mirrors upstream CosyVoice3 streaming semantics more closely than + waveform-domain overlap-add: keep a cumulative mel history per request, + re-run causal HiFT on the history, and emit only the newly grown speech + suffix. That preserves causal look-right handling without double + trimming or duplicated overlap at chunk boundaries. + """ + feat = self._forward_mel( + token=token, + prompt_token=prompt_token, + prompt_feat=prompt_feat, + embedding=embedding, + n_timesteps=n_timesteps, + token_offset_tokens=token_offset_tokens, + streaming=True, + finalize=finalize, + ) + return self._stream_hift_from_feat(feat, cache_state=cache_state, finalize=finalize) + @torch.inference_mode() def forward( self, diff --git a/vllm_omni/model_executor/models/cosyvoice3/runtime.py b/vllm_omni/model_executor/models/cosyvoice3/runtime.py new file mode 100644 index 000000000..d94725148 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/runtime.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Runtime toggles shared by the CosyVoice3 stages.""" + +import os +from contextlib import nullcontext + +import torch + +_FALSE_ENV_VALUES = ("0", "false", "False", "") + + +def _env_flag(name: str, default: str) -> bool: + return os.environ.get(name, default) not in _FALSE_ENV_VALUES + + +def cosyvoice3_batch_flow_enabled() -> bool: + """Return whether cross-request Stage-1 flow batching is enabled.""" + return _env_flag("COSYVOICE3_BATCH_FLOW", "0") + + +def cosyvoice3_batch_flow_debug() -> bool: + """Return whether Stage-1 batching diagnostics are enabled.""" + return _env_flag("COSYVOICE3_BATCH_FLOW_DEBUG", "0") + + +def cosyvoice3_batch_flow_profile(name: str): + """Create a profiler scope only when batching diagnostics are enabled.""" + if cosyvoice3_batch_flow_debug(): + return torch.profiler.record_function(name) + return nullcontext() From 7e520f9cc47e6884520ea8f3f9ab94ab97b34780 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:58:11 +0800 Subject: [PATCH 041/209] [3/N] Encode streamed video on the worker with bounded batching (#7018) Signed-off-by: specture724 Co-authored-by: Claude Sonnet 5 Co-authored-by: Hongsheng Liu --- docs/serving/videos_api.md | 28 +++ .../minimax_h3/test_minimax_h3_contract.py | 163 ++++++++++++++++++ .../test_minimax_h3_step_execution.py | 17 +- .../test_serving_video_output_stream.py | 32 +++- .../openai_api/test_video_server.py | 54 ++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 163 ++++++++++++++++-- vllm_omni/diffusion/utils/media_utils.py | 35 +++- vllm_omni/entrypoints/openai/serving_video.py | 55 ++++-- .../openai/serving_video_output_stream.py | 10 +- 9 files changed, 517 insertions(+), 40 deletions(-) diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 09f023d27..c6e91df60 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -234,6 +234,34 @@ curl -X POST http://localhost:8091/v1/videos/sync \ -o output.mp4 ``` +## Output Encoding + +These `extra_params` control how the server turns decoded frames into MP4 bytes. + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| `preencode_mp4` | boolean | false | Encode the MP4 on the worker while the VAE is still decoding, instead of after the full video is materialized | +| `video_codec_options` | object | null | Encoder options passed through to the H.264 encoder, such as `{"preset": "ultrafast", "threads": "0"}` | + +With `preencode_mp4` enabled, each committed VAE chunk leaves the accelerator and +is encoded while later chunks are still decoding, so host transfer and CPU +encoding overlap the remaining decode instead of following it. The response is +unchanged: the same complete MP4, byte-for-byte equivalent frames. + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A small robot walking through a neon city" \ + -F 'extra_params={"preencode_mp4": true, "video_codec_options": {"preset": "ultrafast"}}' \ + -o output.mp4 +``` + +`preencode_mp4` applies to the complete-MP4 response paths only. The +`/v1/realtime/video` WebSocket endpoint rejects it, because that path already +overlaps encoding through its own incremental fragmented-MP4 encoder. + +Support is per model: MiniMax-H3 implements it, and other models ignore the flag +and take the full-decode path. + ## Storage Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 5470cd9c4..7868c255b 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -34,6 +34,169 @@ def _turbo_spec(filename: str) -> TurboSpec: return spec +def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.pushes = [] + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + self.pushes.append(np.array(frames, copy=True)) + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + for value in (0.0, 0.25, 0.5): + on_chunk(torch.full((1, 3, 3, 2, 2), value)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", + FakeEncoder, + ) + + output = pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + batch_frames=4, + ) + + assert output == b"mp4" + encoder = FakeEncoder.instances[-1] + assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] + assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) + assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) + assert encoder.pushes[1].shape == (3, 2, 2, 3) + + +def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): + """A client's encoder options must survive the worker-side pre-encode path.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + del frames + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + on_chunk(torch.zeros(1, 3, 1, 2, 2)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder) + + pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + video_codec_options={"preset": "ultrafast"}, + ) + + assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} + + +@pytest.mark.parametrize( + ("codec_extra", "expected"), + [ + ({}, {"preset": "ultrafast", "threads": "0"}), + ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), + ({"video_codec_options": {}}, {}), + ({"video_codec_options": None}, None), + ], +) +@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) +def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va"}) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline.device = torch.device("cpu") + pipeline.od_config = SimpleNamespace() + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._quality_policy = Mock() + pipeline._cache_dit_runtime = Mock() + pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode_to_mp4 = Mock(return_value=b"mp4") + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=50, + extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, + ) + batch = DiffusionRequestBatch( + [OmniDiffusionRequest(prompt="test", sampling_params=sampling, request_id="codec-defaults")] + ) + + pipeline.forward(batch) + + assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected + assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames + + +def test_video_codec_options_are_normalized_for_the_encoder(): + from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options + + assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { + "preset": "ultrafast", + "threads": "0", + } + assert normalize_video_codec_options(None) is None + with pytest.raises(ValueError, match="video_codec_options"): + normalize_video_codec_options("preset=ultrafast") + + def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 63c3a42f1..1b123929c 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -284,7 +284,8 @@ def test_mixed_step_batch_leaves_gated_attention_dense(): assert recorder.denoise_timestep is None -def test_prepare_encode_seeds_runner_visible_state(monkeypatch): +@pytest.mark.parametrize("batch_frames", [1, 33]) +def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8) @@ -293,6 +294,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch): context = { "height": 96, "width": 64, + "preencode_mp4": True, + "preencode_batch_frames": batch_frames, "latent_t": 2, "latent_h": 4, "latent_w": 6, @@ -335,6 +338,18 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch): assert state.extra[mod._STEP_BRANCH] is branch assert state.extra[mod._STEP_SHAPE]["height"] == 96 + pipeline.od_config = SimpleNamespace() + monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1))) + calls = [] + + def decode_to_mp4(*args, **kwargs): + calls.append(kwargs) + return b"mp4" + + monkeypatch.setattr(pipeline, "decode_to_mp4", decode_to_mp4) + assert pipeline.post_decode(state).output == (b"mp4", None) + assert calls[0]["batch_frames"] == batch_frames + def test_prepare_encode_rejects_request_mode_only_features(): """Multi-output and DLO have no representation in the step contract.""" diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index e57e3e279..291eada75 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for `/v1/realtime/video` WebSocket video output streaming.""" from __future__ import annotations @@ -7,12 +7,13 @@ from __future__ import annotations import asyncio import json from collections.abc import AsyncGenerator, Callable +from http import HTTPStatus from types import SimpleNamespace from unittest.mock import MagicMock import numpy as np import pytest -from fastapi import FastAPI, WebSocket +from fastapi import FastAPI, HTTPException, WebSocket from pytest_mock import MockerFixture from starlette.testclient import TestClient @@ -957,3 +958,30 @@ class TestStreamingVideoOutputPromptUpdate: "transition_chunks": 2, }, ) + + +class TestStreamingVideoOutputExtraParams: + """extra_params handling when building streaming generation inputs.""" + + @staticmethod + def _handler() -> OmniStreamingVideoOutputHandler: + handler = object.__new__(OmniStreamingVideoOutputHandler) + handler._engine_client = SimpleNamespace(default_sampling_params_list=[]) + return handler + + async def test_preencode_mp4_is_rejected_for_streaming_sessions(self): + """Worker-side pre-encoding emits one progressive MP4 the fMP4 encoder cannot consume.""" + request = VideoGenerationRequest(prompt="p", extra_params={"preencode_mp4": True}) + + with pytest.raises(HTTPException) as excinfo: + await self._handler()._build_prompt_and_sampling_params(request) + + assert excinfo.value.status_code == HTTPStatus.BAD_REQUEST.value + assert "preencode_mp4" in excinfo.value.detail + + async def test_other_extra_params_still_reach_the_engine(self): + request = VideoGenerationRequest(prompt="p", extra_params={"flow_shift": 3.0}) + + _, gen_params, _ = await self._handler()._build_prompt_and_sampling_params(request) + + assert gen_params.extra_args["flow_shift"] == 3.0 diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index b65b00e59..01332beec 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -134,6 +134,60 @@ def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixt handler.shutdown() +@pytest.mark.parametrize("batch_frames", [0, -1, True, 1.5, "17", None]) +def test_preencode_rejects_invalid_batch_frames_before_generation(batch_frames): + engine = FakeAsyncOmni() + handler = OmniOpenAIServingVideo.for_diffusion(engine, model_name="test-model") + request = VideoGenerationRequest( + prompt="test", extra_params={"preencode_mp4": True, "preencode_batch_frames": batch_frames} + ) + try: + with pytest.raises(HTTPException, match="preencode_batch_frames") as exc: + asyncio.run(handler.generate_video_bytes(request, "invalid-batch")) + assert exc.value.status_code == 400 + assert engine.captured_prompt is None + finally: + handler.shutdown() + + +def test_preencoded_video_bytes_preserve_metadata(mocker: MockerFixture): + from vllm_omni.entrypoints.openai.serving_video import VideoGenerationArtifacts + + handler = OmniOpenAIServingVideo.for_diffusion(FakeAsyncOmni(), model_name="test-model") + # Resolved frame count differs from anything the request asked for, so the + # metadata has to come from the encoded stream rather than request defaults. + preencoded = _make_test_video_bytes((32, 24), num_frames=7) + artifacts = VideoGenerationArtifacts( + videos=[preencoded], + audios=[None], + actions=[None], + audio_sample_rate=24000, + output_fps=24.0, + stage_durations={"decode": 0.5}, + peak_memory_mb=123.0, + metrics={"generation_time": 1.25}, + ) + mocker.patch.object(handler, "_run_and_extract", return_value=artifacts) + encoder = mocker.patch("vllm_omni.entrypoints.openai.serving_video._encode_video_bytes") + try: + result = asyncio.run(handler.generate_video_bytes(VideoGenerationRequest(prompt="test"), "preencoded")) + assert result == ( + preencoded, + {"decode": 0.5}, + 123.0, + None, + { + "fps": 24.0, + "num_frames": 7, + "duration_s": 7 / 24.0, + "metrics": {"generation_time": 1.25}, + }, + ) + encoder.assert_not_called() + finally: + handler.shutdown() + + def test_resolve_diffusion_od_config_falls_back_to_attribute(): od_config = SimpleNamespace(model_class_name="WanPipeline") handler = OmniOpenAIServingVideo.for_diffusion( diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index fb9f02b28..f80513d4e 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -64,6 +64,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.errors import OmniClientError, client_error_from_metadata from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -367,6 +368,19 @@ def _minimax_h3_post_process(output, output_type: str = "np"): if not isinstance(output, tuple) or len(output) != 2: return output video, audio = output + if isinstance(video, (bytes, bytearray, memoryview)): + video = [video] + if isinstance(video, list) and all(isinstance(item, (bytes, bytearray, memoryview)) for item in video): + encoded_videos = [bytes(item) for item in video] + else: + encoded_videos = None + if encoded_videos is not None: + return { + "video": encoded_videos, + "audio": [None] * len(encoded_videos), + "audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE, + "fps": MINIMAX_H3_FPS, + } if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4): # Float or channel-first frames would reach the muxer as a black or # banded video rather than as an error. @@ -2147,6 +2161,80 @@ class MiniMaxH3Pipeline( audio_t=audio_t, ) + def decode_to_mp4( + self, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + *, + height: int, + width: int, + max_pending: int = 2, + batch_frames: int = 17, + video_codec_options: dict[str, str] | None = None, + ) -> bytes: + """Decode and encode one output on the worker without full-video materialization. + + Audio is decoded first so the incremental mux session can attach its audio + stream before temporal video chunks arrive. The callback receives committed + float32 ``BCTHW`` frames, performs the requested-size crop and uint8 + conversion in the worker, then applies bounded backpressure to the encoder. + """ + from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder + + if batch_frames <= 0: + raise ValueError("batch_frames must be positive") + + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) + audio_np = audio.detach().float().cpu().numpy() + if audio_np.ndim == 3 and audio_np.shape[0] == 1: + audio_np = audio_np[0] + encoder = ChunkedMP4Encoder( + width=width, + height=height, + fps=MINIMAX_H3_FPS, + audio_waveform=audio_np, + audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, + max_pending=max_pending, + video_codec_options=video_codec_options, + ) + + pending_chunks: list[torch.Tensor] = [] + pending_frames = 0 + + def flush_pending() -> None: + nonlocal pending_frames + if not pending_chunks: + return + batched = torch.cat(pending_chunks, dim=1) + encoder.push(batched[0].cpu().numpy()) + pending_chunks.clear() + pending_frames = 0 + + def on_chunk(frames: torch.Tensor) -> None: + nonlocal pending_frames + prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width]) + if prepared.shape[0] != 1: + raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") + pending_chunks.append(prepared) + pending_frames += int(prepared.shape[1]) + if pending_frames >= batch_frames: + flush_pending() + + try: + with self._component_on_device(self.video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk) + flush_pending() + return encoder.finish() + except BaseException: + encoder.abort() + raise + def decode( self, video_latent: torch.Tensor, @@ -2230,6 +2318,11 @@ class MiniMaxH3Pipeline( quality = sampling.quality logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} + preencode_batch_frames = ( + normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17)) + if extra.get("preencode_mp4", False) + else 17 + ) turbo_spec = self._active_turbo_spec(sampling) has_native_lora = self._has_active_native_lora(sampling) task = self._resolve_task( @@ -2500,6 +2593,11 @@ class MiniMaxH3Pipeline( "audio_shift": audio_shift, "base_schedule": base_schedule, "num_outputs": num_outputs, + "preencode_mp4": bool(extra.get("preencode_mp4", False)), + "preencode_batch_frames": preencode_batch_frames, + "video_codec_options": normalize_video_codec_options( + extra.get("video_codec_options", {"preset": "ultrafast", "threads": "0"}) + ), } @staticmethod @@ -2526,16 +2624,33 @@ class MiniMaxH3Pipeline( audios = [] for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed}) - video, audio = self.decode( - video_latent, - audio_latent, - height=context["height"], - width=context["width"], - ) - videos.append(_prepare_minimax_h3_video_output(video)) - audios.append(audio) - video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) - audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) + if context["preencode_mp4"]: + videos.append( + self.decode_to_mp4( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + video_codec_options=context["video_codec_options"], + batch_frames=context["preencode_batch_frames"], + ) + ) + audios.append(None) + else: + video, audio = self.decode( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + ) + videos.append(_prepare_minimax_h3_video_output(video)) + audios.append(audio) + if videos and isinstance(videos[0], bytes): + video = videos[0] if len(videos) == 1 else videos + audio = None + else: + video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) + audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), @@ -2660,6 +2775,9 @@ class MiniMaxH3Pipeline( "latent_h": context["latent_h"], "latent_w": context["latent_w"], "audio_t": context["audio_t"], + "preencode_mp4": context.get("preencode_mp4", False), + "preencode_batch_frames": context.get("preencode_batch_frames", 17), + "video_codec_options": context.get("video_codec_options"), }, } ) @@ -2831,13 +2949,24 @@ class MiniMaxH3Pipeline( latent_w=shape["latent_w"], audio_t=shape["audio_t"], ) - video, audio = self.decode( - video_latent, - audio_latent, - height=shape["height"], - width=shape["width"], - ) - video = _prepare_minimax_h3_video_output(video) + if shape.get("preencode_mp4", False): + video = self.decode_to_mp4( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + video_codec_options=shape.get("video_codec_options"), + batch_frames=shape.get("preencode_batch_frames", 17), + ) + audio = None + else: + video, audio = self.decode( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + ) + video = _prepare_minimax_h3_video_output(video) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index 615df090f..99d20a0af 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -7,7 +7,7 @@ from __future__ import annotations import io import queue import threading -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from fractions import Fraction from typing import Any, cast @@ -17,6 +17,27 @@ import numpy as np _CHUNKED_MP4_DONE = object() +def normalize_preencode_batch_frames(value: Any) -> int: + """Validate the request's MP4 transfer/encoding batch threshold.""" + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("preencode_batch_frames must be a positive integer") + return value + + +def normalize_video_codec_options(value: Any) -> dict[str, str] | None: + """Coerce a request's ``video_codec_options`` into PyAV's str->str contract. + + The value reaches a worker straight from ``extra_params``, so reject a + non-mapping here rather than letting it fail inside the encoder after the + decode has already run. + """ + if value is None: + return None + if not isinstance(value, Mapping): + raise ValueError("video_codec_options must be a JSON object of encoder option names to values") + return {str(key): str(item) for key, item in value.items()} + + def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None: """Validate a ``(T, H, W, 3)`` uint8 RGB chunk against a session's frame size.""" if chunk.ndim != 4 or chunk.shape[-1] != 3: @@ -233,6 +254,18 @@ class FragmentedMP4Muxer: return chunk +def count_mp4_frames(video_bytes: bytes) -> int | None: + """Read an MP4's frame count from its sample table, without decoding it.""" + if not video_bytes: + return None + try: + with cast(Any, av.open(io.BytesIO(video_bytes), format="mp4")) as container: + frames = int(container.streams.video[0].frames) + except Exception: + return None + return frames or None + + def finalize_streaming_video_bytes( video_bytes: bytes, *, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index b3afe8b40..5ff055fa2 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -13,12 +13,14 @@ from http import HTTPStatus from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast +import pybase64 as base64 from fastapi import HTTPException from PIL import Image from vllm.engine.protocol import EngineClient from vllm.logger import init_logger from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.utils.media_utils import count_mp4_frames, normalize_preencode_batch_frames from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.videos import ( VideoAction, @@ -99,7 +101,12 @@ def _video_metadata_from_artifacts(artifacts: VideoGenerationArtifacts) -> dict[ if not artifacts.videos: return metadata - num_frames = count_video_frames(artifacts.videos[0]) + video = artifacts.videos[0] + # Pre-encoded outputs arrive as MP4 bytes, which carry no tensor shape. + if isinstance(video, (bytes, bytearray, memoryview)): + num_frames = count_mp4_frames(bytes(video)) + else: + num_frames = count_video_frames(video) if num_frames is not None and num_frames > 0: metadata["num_frames"] = num_frames if artifacts.output_fps > 0: @@ -402,6 +409,11 @@ class OmniOpenAIServingVideo: status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) + if request.extra_params.get("preencode_mp4") and "preencode_batch_frames" in request.extra_params: + try: + normalize_preencode_batch_frames(request.extra_params["preencode_batch_frames"]) + except ValueError as exc: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc # Merge extra_params into extra_args gen_params.extra_args.update(request.extra_params) @@ -479,26 +491,22 @@ class OmniOpenAIServingVideo: if "video_codec_options" in request.extra_params: video_codec_options = request.extra_params["video_codec_options"] + def encode_video_result(idx: int, video: Any) -> str: + if isinstance(video, bytes): + return base64.b64encode(video).decode("utf-8") + return encode_video_base64( + video, + fps=artifacts.output_fps, + audio=artifacts.audios[idx], + audio_sample_rate=artifacts.audio_sample_rate, + video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, + ) + _t_encode_start = time.perf_counter() video_data = [ VideoData( - b64_json=( - encode_video_base64( - video, - fps=artifacts.output_fps, - video_codec_options=video_codec_options, - frame_converter=self._video_frame_converter, - ) - if artifacts.audios[idx] is None - else encode_video_base64( - video, - fps=artifacts.output_fps, - audio=artifacts.audios[idx], - audio_sample_rate=artifacts.audio_sample_rate, - video_codec_options=video_codec_options, - frame_converter=self._video_frame_converter, - ) - ), + b64_json=encode_video_result(idx, video), action=artifacts.actions[idx], ) for idx, video in enumerate(artifacts.videos) @@ -549,6 +557,17 @@ class OmniOpenAIServingVideo: return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action, video_metadata _t_encode_start = time.perf_counter() + if isinstance(artifacts.videos[0], bytes): + video_bytes = artifacts.videos[0] + _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 + logger.info("Video response received pre-encoded MP4 bytes: %.2f ms", _t_encode_ms) + return ( + video_bytes, + artifacts.stage_durations, + artifacts.peak_memory_mb, + artifacts.actions[0], + video_metadata, + ) video_bytes = _encode_video_bytes( artifacts.videos[0], fps=artifacts.output_fps, diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index a9abad146..f8f4722ad 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """WebSocket handler for streaming generated video chunks. Protocol: @@ -556,6 +556,14 @@ class OmniStreamingVideoOutputHandler: status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) + # Worker-side pre-encoding produces one progressive MP4, which this + # path's incremental fMP4 encoder cannot consume; streaming already + # overlaps encoding through that encoder instead. + if request.extra_params.get("preencode_mp4"): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="preencode_mp4 is not supported for streaming video sessions.", + ) gen_params.extra_args.update(request.extra_params) self._apply_lora(request.lora, gen_params) From f7d9deb45ab56e6a2ccc1690279bd9e6bdefbfe3 Mon Sep 17 00:00:00 2001 From: KANG Qihan <3149604185@qq.com> Date: Sun, 13 Sep 2026 09:39:27 +0800 Subject: [PATCH 042/209] [Kernel][Boogu-Image] Fuse Q/K RMSNorm + interleaved RoPE via fused_qk_norm_rope (#6982) Signed-off-by: Qihan Kang --- docs/configuration/environment_variables.md | 1 + tests/config/test_environment_variables.py | 2 +- .../layers/test_fused_qk_norm_rope.py | 95 ++++- .../test_boogu_image_transformer.py | 153 +++++++- .../config/environment_variable_inventory.py | 1 + vllm_omni/diffusion/envs.py | 6 + .../diffusion/layers/fused_qk_norm_rope.py | 366 ++++++++++++++++-- .../boogu_image/boogu_image_transformer.py | 132 ++++++- 8 files changed, 703 insertions(+), 53 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 90b9ffae0..4ddc2939c 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -112,6 +112,7 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | +| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index cc37cd6b1..911cfbc80 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 24, + EnvironmentVariableCategory.PUBLIC_OMNI: 25, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 63, diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py index 40888556a..84edce0dc 100644 --- a/tests/diffusion/layers/test_fused_qk_norm_rope.py +++ b/tests/diffusion/layers/test_fused_qk_norm_rope.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch @@ -80,3 +80,96 @@ def test_fused_qk_norm_rope_matches_bf16_reference(seq_len): torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) + + +# --------------------------------------------------------------------------- +# General geometry (any even head_dim, here Boogu-Image's 120) in both +# pairing modes, against the module's own eager reference at the same +# tolerance as the MiniMax-H3 test above. +# --------------------------------------------------------------------------- + +_BOOGU_HEAD_DIM = 120 + + +def _boogu_inputs(strided: bool): + torch.manual_seed(11) + if strided: + # Slices of a wider head axis: the op must honour q/k strides + # (merged-QKV projections hand the op such views). + q = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :28] + k = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :7] + else: + q = torch.randn(4139, 28, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(4139, 7, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + q_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + freqs = torch.randn(4139, _BOOGU_HEAD_DIM // 2, device="cuda", dtype=torch.float32) + rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1) + return q, k, q_weight, k_weight, rope_table + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +@pytest.mark.parametrize("strided", [False, True]) +def test_fused_qk_norm_rope_interleaved(strided): + from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _eager_qk_norm_rope, + _launch_fused_qk_norm_rope, + fused_qk_norm_rope, + ) + + q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided) + expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM, True) + # Check the public op AND the launcher directly: the latter cannot fall + # back to eager, so a silent dispatch regression cannot go green here. + for actual in ( + fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True), + _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True), + ): + torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +def test_fused_qk_norm_rope_half_split_general_dim(): + """head_dim=120 off the MiniMax-H3 128 pin, half-split pairing. + + Exercises the generalised combined kernel's half-split mode via the + launcher. It is not routed in production (half-split traffic keeps the + untouched pre-existing per-tensor kernel and its 128/96 contract, so the + public op falls back to eager at this geometry), pending the + maintainers' call. + """ + from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _eager_qk_norm_rope, + _launch_fused_qk_norm_rope, + ) + + q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided=False) + expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM) + actual = _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=False) + torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02) + + +def test_fused_qk_norm_rope_min_tokens_resolution(monkeypatch): + """The op-level token-gate override: unset or blank -> caller's default; a + non-negative integer string overrides it (``0`` = always fuse); anything + else is rejected naming the variable.""" + from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope_min_tokens + + env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + monkeypatch.delenv(env, raising=False) + assert fused_qk_norm_rope_min_tokens(2048) == 2048 + for blank in ("", " "): + monkeypatch.setenv(env, blank) + assert fused_qk_norm_rope_min_tokens(2048) == 2048 + monkeypatch.setenv(env, "0") + assert fused_qk_norm_rope_min_tokens(2048) == 0 + monkeypatch.setenv(env, "4096") + assert fused_qk_norm_rope_min_tokens(2048) == 4096 + for bad in ("-1", "abc"): + monkeypatch.setenv(env, bad) + with pytest.raises(ValueError, match=env): + fused_qk_norm_rope_min_tokens(2048) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 43e2f12f4..79954f1db 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -26,7 +26,9 @@ NORM_EPS = 1e-5 def _init_distributed(request): """Initialize the minimal single-rank distributed environment required by the vLLM parallel linear layers (tensor-parallel group must exist).""" - if request.node.name.startswith("test_real_rotary_"): + if request.node.name.startswith( + ("test_real_rotary_", "test_packed_rope_table", "test_fused_qk_norm_rope", "test_fallback_with_no") + ): yield return @@ -600,3 +602,152 @@ def test_transformer_forward_ti2i_shape(): assert out.shape == (batch_size, model.out_channels, latent_h, latent_w) assert torch.isfinite(out).all() + + +# --------------------------------------------------------------------------- +# The three tests below are for the fused qk-norm+RoPE kernel's wiring into +# this transformer: the packed-table layout conversion, the fused path against +# the previous eager chain, and the fallback when no packed table is present. +# They need CUDA (the module-level marks are CPU; these add the cuda mark and +# skip on CPU-only runners) and real Boogu geometry rather than the mock dims +# above. +# --------------------------------------------------------------------------- + +_FUSED_HEAD_DIM = 120 +_FUSED_Q_HEADS = 28 +_FUSED_KV_HEADS = 7 +_FUSED_EPS = 1e-5 + + +def _fused_rotary_pair(tokens: int): + theta = torch.randn(1, tokens, _FUSED_HEAD_DIM // 2, device="cuda", dtype=torch.float32) + return ( + torch.cos(theta).repeat_interleave(2, dim=-1), + torch.sin(theta).repeat_interleave(2, dim=-1), + ) + + +def _fused_norms(): + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.layernorm import RMSNorm + + with set_current_vllm_config(VllmConfig()): + norm_q = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16) + norm_k = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16) + with torch.no_grad(): + norm_q.weight.copy_(torch.randn(_FUSED_HEAD_DIM)) + norm_k.weight.copy_(torch.randn(_FUSED_HEAD_DIM)) + return norm_q, norm_k + + +def _operand_ulp_bound(x, weight, cos, sin): + """One bf16 rounding of a normalized operand, propagated through RoPE.""" + import torch.nn.functional as F + + n = F.rms_norm(x, (_FUSED_HEAD_DIM,), weight, _FUSED_EPS).float().abs() + c = cos[..., ::2].float().abs().unsqueeze(2) + s = sin[..., ::2].float().abs().unsqueeze(2) + even_mag = n[..., ::2] * c + n[..., 1::2] * s + odd_mag = n[..., ::2] * s + n[..., 1::2] * c + return 2.0**-6 * torch.stack((even_mag, odd_mag), dim=-1).flatten(-2) + 1e-6 + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_packed_rope_table_layout(monkeypatch): + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _with_packed_rope_table, + ) + + monkeypatch.delenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", raising=False) + torch.manual_seed(0) + tokens = 2311 # above _FUSED_MIN_TOKENS + cos, sin = _fused_rotary_pair(tokens) + out = _with_packed_rope_table((cos, sin)) + assert len(out) == 3 + assert out[0] is cos and out[1] is sin + packed = out[2] + assert packed.shape == (tokens, _FUSED_HEAD_DIM) + assert packed.dtype == torch.float32 and packed.is_contiguous() + half = _FUSED_HEAD_DIM // 2 + assert torch.equal(packed[:, :half], cos[0, :, ::2]) + assert torch.equal(packed[:, half:], sin[0, :, ::2]) + # Below _FUSED_MIN_TOKENS the tuple passes through unchanged (host-bound + # regime: the fused path would cost more than it saves). + short = _with_packed_rope_table(_fused_rotary_pair(17)) + assert len(short) == 2 + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fused_qk_norm_rope_matches_eager_chain(): + """The fused path vs the previous eager chain at real Boogu shapes. + + Every element must sit within one bf16 ulp of a normalized operand + propagated through the rotation (the two paths legitimately differ by + single- vs double-rounding of the norm's weight multiply). + """ + from vllm.triton_utils import HAS_TRITON + + if not HAS_TRITON: + pytest.skip("Triton required") + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _qk_norm_rope, + _with_packed_rope_table, + ) + + torch.manual_seed(0) + norm_q, norm_k = _fused_norms() + q = torch.randn(1, 4139, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 4139, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + cos, sin = _fused_rotary_pair(4139) + with_table = _with_packed_rope_table((cos, sin)) + + fused = _qk_norm_rope(q, k, norm_q, norm_k, with_table, torch.bfloat16) + eager = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16) + for fused_t, eager_t, x, w in zip(fused, eager, (q, k), (norm_q.weight, norm_k.weight)): + bound = _operand_ulp_bound(x, w, cos, sin) + diff = (fused_t.float() - eager_t.float()).abs() + assert (diff <= bound).all(), "fused path beyond one operand ulp of the eager chain" + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fallback_with_no_packed_table(): + """Without a packed table the helper is bit-exact to the old chain.""" + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _qk_norm_rope, + apply_rotary_emb, + ) + + torch.manual_seed(0) + norm_q, norm_k = _fused_norms() + q = torch.randn(1, 64, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 64, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + cos, sin = _fused_rotary_pair(64) + + out_q, out_k = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16) + ref_q = apply_rotary_emb(norm_q(q), (cos, sin)).to(torch.bfloat16) + ref_k = apply_rotary_emb(norm_k(k), (cos, sin)).to(torch.bfloat16) + assert torch.equal(out_q, ref_q) + assert torch.equal(out_k, ref_k) + + +def test_packed_rope_table_env_override(monkeypatch): + """VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS steers the packing gate: 0 packs + even a 43-token table, a huge value leaves a 4139-token table untouched, + unset keeps the 2048 default routing.""" + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _with_packed_rope_table, + ) + + env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + short, long = _identity_rotary_emb(1, 43), _identity_rotary_emb(1, 4139) + + monkeypatch.setenv(env, "0") + assert len(_with_packed_rope_table(short)) == 3 + monkeypatch.setenv(env, str(10**9)) + assert _with_packed_rope_table(long) is long + monkeypatch.delenv(env) + assert len(_with_packed_rope_table(long)) == 3 + assert _with_packed_rope_table(short) is short diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 4a4966d9c..fe439457d 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -65,6 +65,7 @@ _PUBLIC_OMNI = ( "SPEAKER_SAMPLES_DIR", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", "VLLM_OMNI_EVENT_DRIVEN_ORCH", + "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "VLLM_OMNI_ORCH_MONITOR_PATH", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index 32e006f7e..4cbc0907c 100644 --- a/vllm_omni/diffusion/envs.py +++ b/vllm_omni/diffusion/envs.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: MASTER_PORT: int | None = None CUDA_HOME: str | None = None LOCAL_RANK: int = 0 + VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS: str | None = None environment_variables: dict[str, Callable[[], Any]] = { # ================== Runtime Env Vars ================== @@ -28,6 +29,11 @@ environment_variables: dict[str, Callable[[], Any]] = { # local rank of the process in the distributed setting, used to determine # the GPU device id "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # Minimum rotary-table span (B*S positions) at which consumers of the shared + # fused_qk_norm_rope op take the fused path instead of their eager chain; + # "0" = always fuse; unset = each consumer's own measured default. Raw + # string or None; validated by fused_qk_norm_rope_min_tokens(). + "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS": lambda: os.environ.get("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", None), } diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py index f7e27675c..9f6605726 100644 --- a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py +++ b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py @@ -1,18 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Q/K RMSNorm followed by packed non-interleaved RoPE. +"""Q/K RMSNorm followed by packed RoPE (non-interleaved or interleaved). The public contract is shared by diffusion attention implementations: * ``q`` and ``k`` are ``[tokens, heads, head_dim]``; * norm weights are one-dimensional ``[head_dim]`` tensors; * ``rope_table`` is ``[tokens, rotary_dim]`` and stores - ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. + ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``; + its dtype is either the activation dtype or float32; +* ``interleaved=False`` rotates half-split pairs ``(d, d + rotary_dim/2)`` + (MiniMax-H3); ``interleaved=True`` rotates adjacent pairs ``(2i, 2i + 1)`` + with ``theta_i`` (Boogu-Image, matching its ``apply_rotary_emb``). The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K -or rotary-product intermediates. Ascend composes its RMSNorm and rotary fused -primitives; unsupported inputs use the eager reference. +or rotary-product intermediates. The interleaved mode supports any even +``rotary_dim <= head_dim <= 256`` (``tl.arange`` padding to the next power of +two); the half-split mode keeps the pre-existing kernel and its MiniMax-H3 +geometry contract (``head_dim == 128``, ``rotary_dim == 96``). Ascend +composes its RMSNorm and rotary fused primitives on the MiniMax-H3 geometry; +unsupported inputs use the eager reference. """ from __future__ import annotations @@ -26,31 +34,73 @@ from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.utils.torch_utils import direct_register_custom_op +from vllm_omni.diffusion import envs from vllm_omni.platforms import current_omni_platform _FUSED_HEAD_DIM = 128 _FUSED_ROTARY_DIM = 96 +# Interleaved-mode geometry bounds (any even rotary_dim <= head_dim <= this). +_FUSED_MAX_HEAD_DIM = 256 _HEADS_PER_PROGRAM = 8 +# The combined interleaved kernel: small tiles + one warp measured fastest +# (latency-bound shapes; see the kernel docstring). +_INTERLEAVED_HEADS_PER_PROGRAM = 4 def _apply_rope_table( x: torch.Tensor, rope_table: torch.Tensor, rotary_dim: int, + *, + interleaved: bool = False, + dtype: torch.dtype | None = None, + output_dtype: torch.dtype | None = None, ) -> torch.Tensor: + """Eager RoPE application: the fallback and the unit-test reference. + + Shapes: ``x`` is the normalized ``[tokens, heads, head_dim]`` tensor; + ``rope_table`` is ``[tokens, rotary_dim]`` storing + ``[cos(theta) | sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. + + ``interleaved`` selects the pairing: half-split ``(d, d + rotary_dim/2)`` + sharing ``theta_d`` (MiniMax-H3) or adjacent pairs ``(2i, 2i + 1)`` + sharing ``theta_i`` (Boogu-Image). ``dtype`` is the precision the + rotation arithmetic runs in and ``output_dtype`` the precision the result + is cast to; both default to ``x``'s dtype (the historical half-split + behaviour). The interleaved caller passes ``dtype=torch.float32`` to + match the Triton kernel's composition — fp32 rotation of the + bf16-rounded normalized value, one final rounding — which is also how + vLLM's own ``fused_qk_norm_rope`` Triton kernel and Boogu's eager + ``apply_rotary_emb`` order their rounding. + """ + compute_dtype = x.dtype if dtype is None else dtype + out_dtype = x.dtype if output_dtype is None else output_dtype half = rotary_dim // 2 - cos = rope_table[..., :half].to(x.dtype).unsqueeze(1) - sin = rope_table[..., half:].to(x.dtype).unsqueeze(1) - first = x[..., :half] - second = x[..., half:rotary_dim] - return torch.cat( - ( - first * cos - second * sin, - second * cos + first * sin, - x[..., rotary_dim:], - ), - dim=-1, - ) + x_c = x.to(compute_dtype) + cos = rope_table[..., :half].to(compute_dtype).unsqueeze(1) + sin = rope_table[..., half:].to(compute_dtype).unsqueeze(1) + first = x_c[..., 0:rotary_dim:2] if interleaved else x_c[..., :half] + second = x_c[..., 1:rotary_dim:2] if interleaved else x_c[..., half:rotary_dim] + rotated_first = first * cos - second * sin + rotated_second = second * cos + first * sin + if interleaved: + # stack(..., dim=-1) pairs (new_even_i, new_odd_i) into + # [..., half, 2]; flatten(-2) lays the pairs back out as + # [e0, o0, e1, o1, ...] — back to their interleaved positions. + rotated = torch.stack( + (rotated_first, rotated_second), + dim=-1, + ).flatten(-2) + return torch.cat((rotated, x_c[..., rotary_dim:]), dim=-1).to(out_dtype) + else: + return torch.cat( + ( + rotated_first, + rotated_second, + x_c[..., rotary_dim:], + ), + dim=-1, + ).to(out_dtype) if HAS_TRITON: @@ -116,6 +166,135 @@ if HAS_TRITON: out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d tl.store(out_ptr + out_offsets, output, mask=mask) + @triton.jit + def _qk_norm_rope_kernel( + q_ptr, + k_ptr, + q_weight_ptr, + k_weight_ptr, + rope_table_ptr, + out_q_ptr, + out_k_ptr, + q_stride_t, + q_stride_h, + q_stride_d, + k_stride_t, + k_stride_h, + k_stride_d, + rope_stride_t, + out_q_stride_t, + out_q_stride_h, + out_q_stride_d, + out_k_stride_t, + out_k_stride_h, + out_k_stride_d, + num_q_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + head_dim: tl.constexpr, + padded_dim: tl.constexpr, # next power of 2 >= head_dim (tl.arange needs a power of 2) + rotary_half: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + q_head_groups: tl.constexpr, + interleaved: tl.constexpr, + ): + """Fused per-head RMSNorm + RoPE for Q and K in one launch. + + Grid axis 1 assigns the first ``q_head_groups`` programs of a token + to Q and the rest to K, so the small K grid hides inside Q's waves + instead of paying its own latency-bound launch. ``interleaved`` + selects the pairing, mirroring ``_apply_rope_table``: + + * ``True`` — adjacent pairs ``(2i, 2i+1)`` sharing ``theta_i``. + * ``False`` — half-split pairs ``(d, d + rotary_half)`` sharing + ``theta_d``. Production traffic still routes half-split inputs + through ``_rms_norm_rope_kernel``; this mode is test-covered and + awaits the maintainers' call before taking over that path. + + The modes differ only in the pair/frequency/selector index + expressions; the arithmetic is shared. Beyond the rotary width + (partial rotary or lane padding) the output is the normalized value + unchanged. + """ + token = tl.program_id(0) + head_group = tl.program_id(1) + dims = tl.arange(0, padded_dim) + dim_mask = dims < head_dim + if head_group < q_head_groups: + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + # Valid-domain mask: rows beyond the head count (when it is not a + # multiple of heads_per_program) and padded lanes beyond head_dim. + mask = (heads[:, None] < num_q_heads) & dim_mask[None, :] + in_ptr = q_ptr + weight_ptr = q_weight_ptr + out_ptr = out_q_ptr + in_stride_t, in_stride_h, in_stride_d = q_stride_t, q_stride_h, q_stride_d + out_stride_t, out_stride_h, out_stride_d = ( + out_q_stride_t, + out_q_stride_h, + out_q_stride_d, + ) + else: + heads = (head_group - q_head_groups) * heads_per_program + tl.arange(0, heads_per_program) + mask = (heads[:, None] < num_kv_heads) & dim_mask[None, :] + in_ptr = k_ptr + weight_ptr = k_weight_ptr + out_ptr = out_k_ptr + in_stride_t, in_stride_h, in_stride_d = k_stride_t, k_stride_h, k_stride_d + out_stride_t, out_stride_h, out_stride_d = ( + out_k_stride_t, + out_k_stride_h, + out_k_stride_d, + ) + + offsets = token * in_stride_t + heads[:, None] * in_stride_h + dims[None, :] * in_stride_d + x = tl.load(in_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + dims, mask=dim_mask, other=0.0).to(tl.float32) + # Padding lanes load 0 and do not perturb the sum; divide by the true + # head_dim, not the padded width. + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16) + + # The two pairings differ only in three index expressions; the pair + # reload, the table gather, the rotation formulas and the write-back + # are shared (the same structure as _apply_rope_table). + rotary_dim = rotary_half * 2 + if interleaved: + # Adjacent pairs (2i, 2i+1) sharing theta_i. + pair_dims = tl.where(dims < rotary_dim, dims ^ 1, dims) + freq_dims = tl.where(dims < rotary_dim, dims // 2, 0) + is_first = (dims % 2) == 0 + else: + # Half-split pairs (d, d + rotary_half) sharing theta_d. + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + freq_dims = tl.where( + dims < rotary_half, + dims, + tl.where(dims < rotary_dim, dims - rotary_half, 0), + ) + is_first = dims < rotary_half + pair_offsets = token * in_stride_t + heads[:, None] * in_stride_h + pair_dims[None, :] * in_stride_d + pair_x = tl.load(in_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_weight = tl.load(weight_ptr + pair_dims, mask=pair_dims < head_dim, other=0.0).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16) + table_offsets = token * rope_stride_t + freq_dims + cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32) + sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32) + first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin + second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin + output = tl.where( + dims < rotary_dim, + tl.where(is_first, first, second), + normalized.to(tl.float32), + ) + + out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d + tl.store(out_ptr + out_offsets, output, mask=mask) + def _eager_qk_norm_rope( q: torch.Tensor, @@ -126,12 +305,16 @@ def _eager_qk_norm_rope( eps: float, head_dim: int, rotary_dim: int, + interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: q_norm = F.rms_norm(q, (head_dim,), q_weight, eps) k_norm = F.rms_norm(k, (head_dim,), k_weight, eps) + # fp32 rotation for the interleaved (Boogu) semantics; the half-split + # (H3) reference keeps its historical x-dtype arithmetic via the defaults. + rope_dtype = torch.float32 if interleaved else None return ( - _apply_rope_table(q_norm, rope_table, rotary_dim), - _apply_rope_table(k_norm, rope_table, rotary_dim), + _apply_rope_table(q_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype), + _apply_rope_table(k_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype), ) @@ -206,17 +389,23 @@ def _fused_cuda_supported( k: torch.Tensor, head_dim: int, rotary_dim: int, + interleaved: bool = False, ) -> bool: - return ( + if not ( HAS_TRITON and current_platform.is_cuda() and q.is_cuda and k.is_cuda and q.dtype == torch.bfloat16 and k.dtype == torch.bfloat16 - and head_dim == _FUSED_HEAD_DIM - and rotary_dim == _FUSED_ROTARY_DIM - ) + ): + return False + if interleaved: + return rotary_dim % 2 == 0 and 2 <= rotary_dim <= head_dim <= _FUSED_MAX_HEAD_DIM + # Half-split traffic keeps the pre-existing per-tensor kernel and its + # exact MiniMax-H3 geometry contract (that kernel is untouched by the + # interleaved extension). + return head_dim == _FUSED_HEAD_DIM and rotary_dim == _FUSED_ROTARY_DIM def _fused_npu_supported( @@ -224,10 +413,12 @@ def _fused_npu_supported( k: torch.Tensor, head_dim: int, rotary_dim: int, + interleaved: bool = False, ) -> bool: """Return whether the MiniMax-H3 Ascend fused-op contract is satisfied.""" return ( - current_omni_platform.is_npu() + not interleaved + and current_omni_platform.is_npu() and q.device.type == "npu" and k.device.type == "npu" and q.dtype == torch.bfloat16 @@ -271,6 +462,69 @@ def _launch_fused_rms_norm_rope( return out +def _launch_fused_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + interleaved: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + tokens, q_heads, head_dim = q.shape + kv_heads = k.shape[1] + rotary_half = rope_table.shape[-1] // 2 + out_q = torch.empty(q.shape, dtype=q.dtype, device=q.device) + out_k = torch.empty(k.shape, dtype=k.dtype, device=k.device) + if tokens == 0: + return out_q, out_k + if interleaved: + heads_per_program = _INTERLEAVED_HEADS_PER_PROGRAM + num_warps, num_stages = 1, 2 + else: + # Match the historical per-tensor launch shape (8 heads per program, + # 8 warps) so the half-split arithmetic keeps its reduction tree. + heads_per_program = _HEADS_PER_PROGRAM + num_warps, num_stages = 8, 3 + q_head_groups = triton.cdiv(q_heads, heads_per_program) + kv_head_groups = triton.cdiv(kv_heads, heads_per_program) + grid = (tokens, q_head_groups + kv_head_groups) + _qk_norm_rope_kernel[grid]( + q, + k, + q_weight, + k_weight, + rope_table, + out_q, + out_k, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + rope_table.stride(0), + out_q.stride(0), + out_q.stride(1), + out_q.stride(2), + out_k.stride(0), + out_k.stride(1), + out_k.stride(2), + num_q_heads=q_heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + padded_dim=triton.next_power_of_2(head_dim), + rotary_half=rotary_half, + eps=eps, + heads_per_program=heads_per_program, + q_head_groups=q_head_groups, + interleaved=interleaved, + num_warps=num_warps, + num_stages=num_stages, + ) + return out_q, out_k + + def _fused_qk_norm_rope_impl( q: torch.Tensor, k: torch.Tensor, @@ -280,8 +534,9 @@ def _fused_qk_norm_rope_impl( eps: float, head_dim: int, rotary_dim: int, + interleaved: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - if _fused_npu_supported(q, k, head_dim, rotary_dim): + if _fused_npu_supported(q, k, head_dim, rotary_dim, interleaved): return _npu_qk_norm_rope( q, k, @@ -291,7 +546,7 @@ def _fused_qk_norm_rope_impl( eps, rotary_dim, ) - if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + if not _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved): return _eager_qk_norm_rope( q, k, @@ -301,7 +556,12 @@ def _fused_qk_norm_rope_impl( eps, head_dim, rotary_dim, + interleaved, ) + if interleaved: + return _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, eps, interleaved=True) + # Half-split production traffic keeps the historical per-tensor kernel; + # _qk_norm_rope_kernel's half-split mode awaits the maintainers' call. return ( _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps), _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps), @@ -317,8 +577,9 @@ def _fused_qk_norm_rope_fake( eps: float, head_dim: int, rotary_dim: int, + interleaved: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim + del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim, interleaved return torch.empty_like(q), torch.empty_like(k) @@ -333,6 +594,39 @@ if not hasattr(torch.ops.vllm_omni, "fused_qk_norm_rope"): ) +_MIN_TOKENS_ENV = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + + +def fused_qk_norm_rope_min_tokens(default: int) -> int: + """Token span below which a consumer should keep its eager chain. + + Taking the fused path costs a fixed amount of host time per call (the + custom-op dispatch into the Python impl and the Triton launcher), while the + GPU time it saves grows with the number of tokens the call covers. Below a + host- and stack-dependent crossover a call is host-bound and the fused path + can cost more end-to-end than it saves; above it the fused path wins. + Consumers pass the default they measured on their reference hardware + (Boogu-Image: 2048 tokens on one H200) and a deployment on other hardware + overrides it once, for every consumer, through + ``VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS`` (``0`` = always fuse). + + Resolved at call time — consumers call this once per forward, not per + attention site — so tests and operators can flip it without re-importing. + Unset or blank means the caller's default; anything else that is not a + non-negative integer raises ``ValueError``. + """ + raw = envs.VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS + if raw is None or not raw.strip(): + return default + try: + value = int(raw) + except ValueError: + value = -1 + if value < 0: + raise ValueError(f"{_MIN_TOKENS_ENV} must be a non-negative integer, got {raw!r}") + return value + + def fused_qk_norm_rope( q: torch.Tensor, k: torch.Tensor, @@ -343,8 +637,9 @@ def fused_qk_norm_rope( *, head_dim: int | None = None, rotary_dim: int | None = None, + interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Apply Q/K RMSNorm and packed non-interleaved RoPE.""" + """Apply Q/K RMSNorm and packed RoPE (half-split or adjacent pairs).""" if q.ndim != 3 or k.ndim != 3: raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}") if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]: @@ -364,15 +659,18 @@ def fused_qk_norm_rope( raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}") if q_weight.device != q.device or k_weight.device != q.device: raise ValueError("Q/K norm weights must be on the activation device") - if rope_table.device != q.device or rope_table.dtype != q.dtype: - raise ValueError("rope_table must have the same dtype and device as q/k") + if rope_table.device != q.device or rope_table.dtype not in ( + q.dtype, + torch.float32, + ): + raise ValueError("rope_table must be on q/k's device with their dtype or float32") if rope_table.shape != (q.shape[0], rotary_dim): raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}") q_weight = q_weight.contiguous() k_weight = k_weight.contiguous() rope_table = rope_table.contiguous() - if _fused_npu_supported(q, k, head_dim, rotary_dim): + if _fused_npu_supported(q, k, head_dim, rotary_dim, interleaved): return _npu_qk_norm_rope( q, k, @@ -382,7 +680,7 @@ def fused_qk_norm_rope( eps, rotary_dim, ) - if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + if not _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved): return _fused_qk_norm_rope_impl( q, k, @@ -392,6 +690,7 @@ def fused_qk_norm_rope( eps, head_dim, rotary_dim, + interleaved, ) return torch.ops.vllm_omni.fused_qk_norm_rope( q, @@ -402,7 +701,8 @@ def fused_qk_norm_rope( eps, head_dim, rotary_dim, + interleaved, ) -__all__ = ["fused_qk_norm_rope"] +__all__ = ["fused_qk_norm_rope", "fused_qk_norm_rope_min_tokens"] diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 66031ca3a..93d52b960 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py @@ -31,10 +31,17 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, + fused_qk_norm_rope_min_tokens, +) from vllm_omni.diffusion.models.utils import make_attention_mask from vllm_omni.platforms import current_omni_platform @@ -43,7 +50,10 @@ if TYPE_CHECKING: logger = init_logger(__name__) -RotaryEmbedding = tuple[torch.Tensor, torch.Tensor] +# (cos, sin) — optionally (cos, sin, packed_table) where packed_table is the +# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused +# qk-norm+RoPE op consumes; the eager path ignores the third element. +RotaryEmbedding = tuple[torch.Tensor, ...] RotaryFrequencyTables = list[RotaryEmbedding] @@ -56,9 +66,12 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tens Args: x: Query or key tensor of shape [B, S, H, D]. - rotary_emb: Cosine and sine tensors, each shaped [B, S, D]. + rotary_emb: Cosine and sine tensors, each shaped [B, S, D]. The tuple + may carry a third element (the packed table for the fused + qk-norm+RoPE path); it is ignored here, hence the indexed access + instead of tuple unpacking. """ - freqs_cos, freqs_sin = rotary_emb + freqs_cos, freqs_sin = rotary_emb[0], rotary_emb[1] freqs_cos = freqs_cos.unsqueeze(2) freqs_sin = freqs_sin.unsqueeze(2) @@ -73,6 +86,87 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tens return x_out.type_as(x) +# The fused qk-norm+RoPE path is always on where the fast path exists; on +# other platforms the eager chain below is the behaviour-identical fallback. +_FUSED_QK_NORM_ROPE = HAS_TRITON and current_platform.is_cuda() + +# Boogu's default for the fused-path token gate: fuse only when a rotary table +# spans at least this many positions (counted as B*S). Below the crossover a +# request is host-bound and the fused path's per-call Python launch (custom-op +# impl + Triton launcher) can cost more end-to-end than the kernel saves. +# Measured on one H200: on the pre-#6871 main, 512^2 (1.1k tokens) paid +# +191 ms/request and the crossover sat at ~2,350 tokens; on the current main +# the 512^2 family measures as noise (+/-25 ms, sign-inconsistent) while +# >=768^2 (2.3k tokens) gains. The crossover is host-dependent — re-benchmark +# on your hardware and override with VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS +# (0 = always fuse); see fused_qk_norm_rope_min_tokens(). +_FUSED_MIN_TOKENS = 2048 + + +def _with_packed_rope_table(rotary_emb: RotaryEmbedding) -> RotaryEmbedding: + """Append the packed fp32 table the fused qk-norm+RoPE op consumes. + + ``freqs_cos``/``freqs_sin`` are theta-repeat-interleaved ``[B, S, D]`` + (``cos[2i] == cos[2i+1] == cos(theta_i)``), so taking every even column + yields the theta-width halves and the packed table is + ``[B*S, D] = [cos(theta) | sin(theta)]``. Tables spanning fewer positions + (counted as ``B*S``) than the token gate — Boogu's ``_FUSED_MIN_TOKENS`` + default, overridable through ``VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS`` + (``0`` = always fuse) — are returned unchanged: without the third element + every consumer stays on the eager chain (this also keeps short tables like + the context refiner's off the fused path). + """ + freqs_cos, freqs_sin = rotary_emb[0], rotary_emb[1] + if freqs_cos.shape[0] * freqs_cos.shape[1] < fused_qk_norm_rope_min_tokens(_FUSED_MIN_TOKENS): + return rotary_emb + packed = torch.cat((freqs_cos[..., ::2], freqs_sin[..., ::2]), dim=-1) + return freqs_cos, freqs_sin, packed.reshape(-1, freqs_cos.shape[-1]).float() + + +def _qk_norm_rope( + query: torch.Tensor, + key: torch.Tensor, + norm_q, + norm_k, + rotary_emb: RotaryEmbedding | None, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-head Q/K RMSNorm + RoPE + cast, fused into one kernel launch when + a packed rope table is present; otherwise the original eager chain.""" + # Gate on the op's own fast-path predicate: when the fused CUDA kernel + # cannot run (dtype, device, or geometry outside even + # rotary_dim <= head_dim <= 256), the original chain below must be used + # bit-exactly rather than the op's internal eager fallback, whose rounding + # differs from vLLM's RMSNorm by one ulp on a fraction of elements. + if ( + rotary_emb is not None + and len(rotary_emb) > 2 + and rotary_emb[2] is not None + and _fused_cuda_supported(query, key, query.shape[-1], query.shape[-1], interleaved=True) + ): + batch, seq_len, num_heads, head_dim = query.shape + num_kv_heads = key.shape[2] + fused_q, fused_k = fused_qk_norm_rope( + query.view(batch * seq_len, num_heads, head_dim), + key.view(batch * seq_len, num_kv_heads, head_dim), + norm_q.weight, + norm_k.weight, + rotary_emb[2], + norm_q.variance_epsilon, + interleaved=True, + ) + return ( + fused_q.view(batch, seq_len, num_heads, head_dim), + fused_k.view(batch, seq_len, num_kv_heads, head_dim), + ) + query = norm_q(query) + key = norm_k(key) + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + return query.to(dtype), key.to(dtype) + + def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return F.silu(x.float(), inplace=False).to(x.dtype) * y @@ -559,13 +653,7 @@ class BooguImageSelfAttention(nn.Module): key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) value = value.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) - query = self.norm_q(query) - key = self.norm_k(key) - - if rotary_emb is not None: - query = apply_rotary_emb(query, rotary_emb) - key = apply_rotary_emb(key, rotary_emb) - query, key = query.to(dtype), key.to(dtype) + query, key = _qk_norm_rope(query, key, self.norm_q, self.norm_k, rotary_emb, dtype) attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None attn_output = self.attn(query, key, value, attn_metadata) @@ -711,13 +799,7 @@ class BooguImageJointAttention(nn.Module): key = key.view(batch_size, -1, self.num_local_kv_heads, self.head_dim) value = value.view(batch_size, -1, self.num_local_kv_heads, self.head_dim) - query = self.norm_q(query) - key = self.norm_k(key) - - if rotary_emb is not None: - query = apply_rotary_emb(query, rotary_emb) - key = apply_rotary_emb(key, rotary_emb) - query, key = query.to(dtype), key.to(dtype) + query, key = _qk_norm_rope(query, key, self.norm_q, self.norm_k, rotary_emb, dtype) attn_metadata = AttentionMetadata(attn_mask=joint_attention_mask) if joint_attention_mask is not None else None attn_output = self.attn(query, key, value, attn_metadata) @@ -1449,6 +1531,11 @@ class BooguImageTransformer2DModel(nn.Module): shift += ref_img_len idx += 1 + if _FUSED_QK_NORM_ROPE: + # The rebuilt per-reference-image tuple needs its own packed + # table; the forward-level tuples do not flow into this batch. + batch_ref_img_rotary_emb = _with_packed_rope_table(batch_ref_img_rotary_emb) + for layer in self.ref_image_refiner: batch_ref_image_hidden_states = layer( batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb @@ -1534,6 +1621,17 @@ class BooguImageTransformer2DModel(nn.Module): device, ) + if _FUSED_QK_NORM_ROPE: + # One packed [cos|sin] table per rotary embedding that reaches an + # attention, built once per forward; the tuples grow a third + # element that the fused path consumes and the eager path ignores. + # (ref_img_rotary_emb is packed inside img_patch_embed_and_refine, + # on the rebuilt per-reference-image batch tuple.) + context_rotary_emb = _with_packed_rope_table(context_rotary_emb) + noise_rotary_emb = _with_packed_rope_table(noise_rotary_emb) + rotary_emb = _with_packed_rope_table(rotary_emb) + combined_img_rotary_emb = _with_packed_rope_table(combined_img_rotary_emb) + # Context refinement. for layer in self.context_refiner: instruction_hidden_states = layer(instruction_hidden_states, instruction_attention_mask, context_rotary_emb) From 01a2f93256975c7ff9565c5414b0fe0225bc4765 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sun, 13 Sep 2026 10:05:56 +0800 Subject: [PATCH 043/209] [Bugfix][Frontend] Honor output_compression on the image generations route (#7447) Signed-off-by: hsliu_ustc Signed-off-by: amy-why-3459 Co-authored-by: amy-why-3459 --- recipes/Boogu/Boogu-Image.md | 33 ++++++++++++ .../openai_api/test_image_api_utils.py | 54 +++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 4 +- .../entrypoints/openai/protocol/images.py | 11 ++++ 4 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 tests/entrypoints/openai_api/test_image_api_utils.py diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 0353b9778..9a5ce0397 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -452,6 +452,39 @@ curl -s http://localhost:8091/v1/chat/completions \ - **Known limitations:** the same single-GPU limitations as Base and Edit apply; CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated. +## Output format and compression + +Image requests accept `output_format` (`png`, `jpeg`, or `webp`) and +`output_compression` (0-100). The default is `png` with the fastest, +least-compressed PNG encode, which is lossless but produces multi-megabyte +payloads. For latency-sensitive clients, `output_format: "jpeg"` removes +almost the entire response-encoding overhead: on a 1024x1024 Turbo request +the client-observed e2e latency drops from ~900 ms to ~766 ms and the +payload shrinks from ~3.1 MB to ~0.8 MB. Use `output_compression` to trade +encode time against payload size when PNG must stay lossless: lower values +(e.g. 1) compress harder and produce smaller payloads at proportionally +longer encode times. + +## Attention backend + +The diffusion attention backend defaults to `FLASH_ATTN` (FlashAttention-2). +On Hopper-class GPUs (SM90) pass `--diffusion-attention-backend CUDNN_ATTN` +for a measured ~12% engine-time reduction on Boogu-Image workloads: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Turbo \ + --omni \ + --diffusion-attention-backend CUDNN_ATTN \ + --port 8091 +``` + +Measured on a 143 GB SM90 card (1024x1024, 4 steps, 5-run mean): stage time +673 ms with `CUDNN_ATTN` vs 763 ms with the default. Outputs are numerically +different from FLASH_ATTN (different attention kernel) but visually +equivalent, deterministic per seed, and decode identically. On datacenter +Blackwell (SM100 / SM103) use `--diffusion-attention-backend TRTLLM_ATTN` +instead; it is rejected on SM90. + ## Performance validation The checked-in single-device A40 measurements remain the current Boogu diff --git a/tests/entrypoints/openai_api/test_image_api_utils.py b/tests/entrypoints/openai_api/test_image_api_utils.py new file mode 100644 index 000000000..31bc36c5e --- /dev/null +++ b/tests/entrypoints/openai_api/test_image_api_utils.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""L1 tests for the image response encoder's compression contract.""" + +import base64 +import io + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.openai.image_api_utils import ( + encode_image_base64_with_compression, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(scope="module") +def image(): + # Photographic-content proxy: noise-plus-gradient content compresses very + # differently across PNG compress levels, which is what the mapping under + # test controls. + import numpy as np + + rng = np.random.default_rng(0) + grad = np.linspace(0, 255, 256, dtype=np.float32) + arr = np.clip(grad[None, :, None] + rng.integers(-24, 25, (256, 256, 3)), 0, 255).astype(np.uint8) + return Image.fromarray(arr, mode="RGB") + + +def test_png_output_compression_maps_to_compress_level(image): + """output_compression 100 -> compress_level 0, 1 -> compress_level 9.""" + + def encode(compression): + buffer = io.BytesIO() + image.save(buffer, format="PNG", compress_level=max(0, min(9, 9 - compression // 11))) + return len(base64.b64encode(buffer.getvalue())) + + fast = encode_image_base64_with_compression(image, format="png", output_compression=100) + small = encode_image_base64_with_compression(image, format="png", output_compression=1) + + # Level-0 PNG barely compresses; the mapped level-9 encode is smaller. + assert len(small) < len(fast) + # The mapping must produce the same sizes as the equivalent compress_level. + assert len(fast) == encode(100) + assert len(small) == encode(1) + + +def test_jpeg_output_compression_maps_to_quality(image): + high = base64.b64decode(encode_image_base64_with_compression(image, format="jpeg", output_compression=100)) + low = base64.b64decode(encode_image_base64_with_compression(image, format="jpeg", output_compression=10)) + + assert len(high) > len(low) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 565362312..f4f22f95f 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1605,7 +1605,9 @@ def _build_image_generation_response( output_format = _choose_output_format(request.output_format or "png", None) image_data = [ ImageData( - b64_json=encode_image_base64_with_compression(image, format=output_format), + b64_json=encode_image_base64_with_compression( + image, format=output_format, output_compression=request.output_compression + ), revised_prompt=None, ) for image in images diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 0215e1ce0..b7800c13e 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -175,6 +175,17 @@ class ImageGenerationRequest(BaseModel): default=None, description="Output image format: 'png', 'jpeg', or 'webp'. Defaults to 'png'.", ) + output_compression: int = Field( + default=100, + ge=0, + le=100, + description=( + "Compression level 0-100. For 'jpeg'/'webp' this is the encoder " + "quality. For 'png' 100 keeps the fastest, least-compressed " + "encode and lower values trade encode time for smaller payloads " + "(100 -> compress_level 0, 1 -> compress_level 9)." + ), + ) return_stage_metrics: bool | None = Field( default=None, description="Return stage metrics for benchmark clients.", From 6753bbd18c9a1f45cfebcec3eacdd2f1e5ea859b Mon Sep 17 00:00:00 2001 From: Yuan Wu Date: Sun, 13 Sep 2026 16:35:02 +0800 Subject: [PATCH 044/209] [Core] Add NIXL omni connector (#6093) Signed-off-by: yuanwu Signed-off-by: Yuan Wu --- docs/.nav.yml | 1 + docs/configuration/environment_variables.md | 7 + .../design/feature/disaggregated_inference.md | 2 + .../feature/omni_connectors/nixl_connector.md | 135 +++ docs/design/index.md | 1 + tests/config/test_environment_variables.py | 2 +- .../sched/test_omni_scheduling_coordinator.py | 14 + .../test_chunk_transfer_adapter.py | 54 +- .../omni_connectors/test_nixl_connector.py | 891 ++++++++++++++ .../test_nixl_connector_native.py | 204 ++++ .../test_orchestrator_kv_sender_info.py | 159 ++- .../test_orchestrator_stage_input_bridge.py | 26 +- tests/worker/test_omni_connector_mixin.py | 192 +++ .../config/environment_variable_inventory.py | 2 + .../core/sched/omni_scheduling_coordinator.py | 2 + vllm_omni/core/sched/output.py | 1 + vllm_omni/distributed/__init__.py | 2 + .../distributed/omni_connectors/__init__.py | 6 + .../connectors/nixl_connector.py | 1044 +++++++++++++++++ .../distributed/omni_connectors/factory.py | 12 + .../omni_connectors/kv_transfer_manager.py | 11 +- .../omni_connector_payload_transport.py | 23 +- .../model_runner/omni_connector_runtime.py | 18 +- .../chunk_transfer_adapter.py | 28 +- .../omni_connectors/utils/config.py | 1 + .../omni_connectors/utils/initialization.py | 152 ++- .../omni_connectors/utils/kv_utils.py | 14 +- vllm_omni/engine/__init__.py | 5 + vllm_omni/engine/orchestrator.py | 22 + vllm_omni/engine/stage_client.py | 2 + vllm_omni/engine/stage_engine_core_client.py | 37 + vllm_omni/engine/stage_engine_core_proc.py | 1 + vllm_omni/engine/stage_init_utils.py | 27 +- 33 files changed, 3032 insertions(+), 66 deletions(-) create mode 100644 docs/design/feature/omni_connectors/nixl_connector.md create mode 100644 tests/distributed/omni_connectors/test_nixl_connector.py create mode 100644 tests/distributed/omni_connectors/test_nixl_connector_native.py create mode 100644 vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 8dac28d22..beb15bf47 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -136,6 +136,7 @@ nav: - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md - Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md - Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md + - NIXL Connector: design/feature/omni_connectors/nixl_connector.md - Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md - Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md - Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 4ddc2939c..3a37b52d2 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -90,6 +90,13 @@ depends on the installed kernels and model path. | `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental | | `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental | +### NIXL stage transfer + +| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | +| --- | --- | --- | --- | --- | +| `VLLM_OMNI_NIXL_LEASE_S` | Float seconds; default `3600` | Unclaimed producer payload expiry; connector construction | Nonempty environment value overrides `lease_seconds`; empty falls back to config/default. Invalid floats raise `ValueError`. Claimed READ allocations never expire by time. | Experimental | +| `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` | Float seconds; default `300` | Receiver READ wait; connector construction | Nonempty environment value overrides `transfer_timeout_s`; empty falls back to config/default. Invalid floats raise `ValueError`. Timeout does not cancel DMA or release active allocations. | Experimental | + ### Server storage Storage names use Pydantic's nested-settings delimiter (`__`). They are read diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index ba8c8f9fc..8b94325cb 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs: - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md) - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md) - [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md) +- [NixlConnector](omni_connectors/nixl_connector.md) - [YuanrongConnector](omni_connectors/yuanrong_connector.md) - [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md) @@ -25,6 +26,7 @@ Current connectors operate in D2H2D (device to host to device) mode. | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. | | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. | | Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. | +| Multi node (NIXL) | NixlConnector | NIXL `READ` over UCX; the only backend covering Intel XPU. | | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. | | Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. | diff --git a/docs/design/feature/omni_connectors/nixl_connector.md b/docs/design/feature/omni_connectors/nixl_connector.md new file mode 100644 index 000000000..7dfd73946 --- /dev/null +++ b/docs/design/feature/omni_connectors/nixl_connector.md @@ -0,0 +1,135 @@ +# NixlConnector + +## When to Use + +Multi-node or intra-node stage transfer over NIXL, which brokers RDMA/shared-memory +transports through a single API. Useful when the deployment already standardises on +NIXL for KV transfer, or when the target accelerator has no Mooncake/Mori backend -- +NIXL reaches Intel XPU through UCX's Level Zero support. + +## Mechanism + +Uses vLLM's `NixlWrapper` (`vllm.distributed.nixl_utils`) to register the producer's +tensors and let the consumer pull them with a NIXL `READ`. + +- Data Plane: NIXL agent-to-agent `READ`, GPU-to-GPU where the backend allows it. +- Control Plane: a ZMQ ROUTER socket atomically claims the source allocation before + returning descriptors. Forwarded `put()` metadata includes the generation and + endpoint; the consumer claims that generation before submitting any READ. + +Transfer metadata uses schema version 1. Tensor descriptors are grouped by NIXL +memory type (`DRAM`, `VRAM`, and so on), while each descriptor retains its global +tensor index so mixed CPU/accelerator payloads can be reconstructed in order. + +Payloads are not restricted to tensors. A single tensor or a list of tensors is +transferred as-is; a nested structure has its tensor leaves extracted and shipped +alongside a msgpack-encoded skeleton; anything else is msgpack-encoded into one +uint8 tensor. Mapping and sequence structure is preserved. `msgspec.Struct` values, +including `OmniPayloadStruct`, are normalised to mappings so their tensor leaves use +the direct NIXL path rather than CPU object serialisation. + +Each NIXL agent uses strict thread synchronization. A consumer performs one short +metadata query per scheduler poll (10ms by default), which prevents an unavailable +producer from blocking unrelated requests. The regular transfer-completion handshake +keeps the configured transfer timeout. + +## Installation + +CUDA hosts can install the published wheel: + +```bash +pip install nixl +``` + +For Intel XPU, use `docker/Dockerfile.xpu`. Its upstream vLLM XPU base image +provides a matching torch-xpu, UCX with Level Zero support, and NIXL runtime. + +## Configuration + +```yaml +connectors: + nixl_connector: + name: NixlConnector + extra: + host: "auto" + zmq_port: 50061 + backends: ["UCX"] + +stages: + - stage_id: 0 + output_connectors: + to_stage_1: nixl_connector + + - stage_id: 1 + input_connectors: + from_stage_0: nixl_connector +``` + +Parameters: + +- `host`: address the producer binds its handshake socket to (`"auto"` to detect). +- `zmq_port`: handshake port. If a producer omits it, an ephemeral port is opened + for ownership claims and returned in `put()` metadata. The value in the deploy YAML is a + base port. Purpose, replica, rank, and producer-stage offsets are added centrally, + so colocated connectors and tensor-parallel workers receive distinct endpoints + without one config entry per rank. +- `sender_host` / `sender_zmq_port`: consumer-side override naming the producer's + handshake endpoint. Only needed when the consumer cannot learn it from metadata. + In replicated deployments, the producer endpoint is carried with each request and + selected at `get()` time; it is not mutable connector-wide state. + Explicit sender ports are preserved, not recomputed using the receiver replica. + Intermediate NIXL stages retain their incoming receiver configuration and a + separate outgoing bind configuration, so one worker can receive and publish. +- `backends`: NIXL backends to register memory with. Defaults to `["UCX"]`. +- `receive_device`: forces where received tensors land. By default the consumer + keeps the producer's device *type* but uses its own current device of that type. +- `memory_type`: overrides the NIXL memory type, otherwise `DRAM` for CPU tensors + and `VRAM` for accelerator tensors. +- `lease_seconds`: how long a `put()` payload stays registered while waiting to be + read (default 3600). The consumer reports completion, so this only bounds payloads + nobody has claimed. An internal reaper releases only unclaimed registrations. + `VLLM_OMNI_NIXL_LEASE_S` overrides it. +- `transfer_timeout_s`: how long a `get()` waits for its `READ` to complete + (default 300). NIXL 1.3 has no transfer cancellation API, so a timed-out transfer's + buffers and registrations remain owned by the connector until NIXL reports a + terminal state; `close()` waits for that state before releasing them. + `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` overrides it. + +### Source ownership and failure limits + +Claims and expiry use the same lock. Completion identifies both a payload +generation and a unique READ claim; duplicate or stale ACKs cannot release a +different reader or replacement payload. Replacing a claimed key is rejected. +Terminal errors also release claims; timeout, unknown status, and active sibling +transfers retain the complete destination bundle and source claim until terminal. +Zero-element tensors retain their shape/dtype/skeleton positions but never create +DMA descriptors or registrations. + +A lost metadata response, lost completion ACK, or abandoned consumer can retain a +claim indefinitely. NIXL 1.3 cannot prove remote cancellation, so neither TTL nor +`cleanup()` frees those allocations. Producer `close()` rejects new work and +retains its agent, listener and claimed allocations; a later `close()` finishes +teardown after claims drain. Permanently abandoned claims remain until process +exit. This favors memory safety over bounded shutdown/memory usage. Only trusted +peers may access this unauthenticated control plane. Both endpoints must use the +claim-aware protocol; rolling compatibility with older GET_META clients is not +provided. Legacy externally-owned metadata without a generation is accepted only +under the caller's lifetime guarantee, not managed as a leased allocation here. + +## Validation + +The native CUDA smoke test requires NIXL and at least two visible CUDA devices: + +```bash +UCX_RCACHE_MAX_UNRELEASED=1024 pytest -o addopts='' -q -s \ + tests/distributed/omni_connectors/test_nixl_connector_native.py +``` + +It starts independent producer and consumer processes, loads the real UCX backend, +performs the ZMQ `GET_META` handshake and a native NIXL `READ` from GPU 0 to GPU 1, +checks mixed CPU/GPU tensor leaves from an `OmniPayloadStruct`, sends `XFER_DONE`, +and verifies that producer registrations are released. This validates same-host, +cross-GPU operation. Inter-node UCX transport is supported by the design but was not +validated by this test. + +For more details, refer to the [NIXL repository](https://github.com/ai-dynamo/nixl). diff --git a/docs/design/index.md b/docs/design/index.md index 6481aa49b..f70571cb5 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -36,6 +36,7 @@ implementation contract; it is not, by itself, a general support claim. - [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md) - [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md) - [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md) +- [NIXL Connector](feature/omni_connectors/nixl_connector.md) - [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md) - [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md) - [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md) diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 911cfbc80..d958bbf18 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 25, + EnvironmentVariableCategory.PUBLIC_OMNI: 27, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 63, diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index ae1b3924e..f37aac29f 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -58,6 +58,7 @@ def _make_request(req_id: str, status: str = "waiting") -> SimpleNamespace: num_computed_tokens=0, _all_token_ids=[], _output_token_ids=[], + payload_sender_info=None, ) @@ -270,6 +271,19 @@ class TestWaitingForInputTransition(unittest.TestCase): self.assertEqual(len(coord.pending_input_registrations), 1) self.assertEqual(coord.pending_input_registrations[0].request_id, "r1") + def test_pending_input_registration_carries_payload_sender_info(self): + coord = OmniSchedulingCoordinator(stage_id=1) + req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT) + req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + waiting = MockQueue([req]) + + coord.process_pending_full_payload_inputs(waiting, set()) + + self.assertEqual( + coord.pending_input_registrations[0].payload_sender_info, + {"host": "10.0.0.1", "zmq_port": 50051}, + ) + def test_idle_cycles_retain_received_marker_before_request_appears(self): coord = OmniSchedulingCoordinator(stage_id=1) coord._full_payload_input_received.add("late") diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index a80ded5f0..92c2a7343 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -4,7 +4,7 @@ import threading from collections import deque from types import MethodType, SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import Mock, call, patch import pytest import torch @@ -543,7 +543,7 @@ def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, ex assert connector == "ok" assert isinstance(captured["spec"], ConnectorSpec) assert captured["spec"].name == expected_name - assert captured["spec"].extra == expected_extra + assert captured["spec"].extra == {**expected_extra, "stage_id": 0} def test_load_poll(build_adapter): @@ -570,8 +570,51 @@ def test_load_poll(build_adapter): assert request.additional_information == payload assert adapter.get_req_chunk["req-1"] == 1 assert "req-1" not in adapter._finished_load_reqs - assert "req-1" in adapter.upstream_exhausted_requests - assert "req-1" not in adapter._pending_load_reqs + + +def test_load_poll_uses_request_scoped_payload_sender_endpoint(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1") + request.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + connector.get.return_value = None + + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) + + connector.get.assert_called_once_with( + "1", + "2", + "req-1_1_0", + {"source_host": "10.0.0.1", "source_port": 50051}, + ) + + +def test_load_poll_keeps_concurrent_payload_sender_endpoints_distinct(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + first = _req("req-1", RequestStatus.WAITING) + first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + second = _req("req-2", RequestStatus.WAITING) + second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051} + connector.get.return_value = None + + second_entry = _dequeue_load_entry(adapter, second) + first_entry = _dequeue_load_entry(adapter, first) + adapter._poll_single_request(second_entry) + adapter._poll_single_request(first_entry) + + assert connector.get.call_args_list == [ + call( + "1", + "2", + "req-2_1_0", + {"source_host": "10.0.0.2", "source_port": 51051}, + ), + call( + "1", + "2", + "req-1_1_0", + {"source_host": "10.0.0.1", "source_port": 50051}, + ), + ] def test_load_async_does_not_requeue_registered_inflight_request(build_adapter): @@ -2770,7 +2813,8 @@ def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_insi scheduler.add_request(final_update) return True, 1, {} - def get_by_key(_from_stage, _to_stage, key): + def get_by_key(_from_stage, _to_stage, key, metadata=None): + assert metadata is None payload = payloads.pop(key, None) return (payload, 1) if payload is not None else None diff --git a/tests/distributed/omni_connectors/test_nixl_connector.py b/tests/distributed/omni_connectors/test_nixl_connector.py new file mode 100644 index 000000000..5b2161eee --- /dev/null +++ b/tests/distributed/omni_connectors/test_nixl_connector.py @@ -0,0 +1,891 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +""" +Unit tests for the NixlConnector metadata handshake. + +NIXL itself is stubbed out, so these tests exercise the ZMQ control plane and +the payload normalisation rather than any actual RDMA transfer. +""" + +import ctypes +import sys +import time +import types + +import msgspec +import pytest +import torch +import zmq + +from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct + +pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model] + +PORT = 47431 + + +def _wait_for_metadata(consumer, key, metadata=None, *, timeout=2.0): + """Retry outside the connector's deliberately bounded metadata query.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + resolved = consumer._resolve_metadata(key, metadata) + if resolved is not None: + return resolved + time.sleep(0.01) + pytest.fail(f"Timed out waiting for metadata for {key!r}") + + +@pytest.mark.parametrize("misses", [0, 2]) +def test_wait_for_metadata_retries_until_success(monkeypatch, misses): + expected = {"claim_id": "claim"} + forwarded = {"source_host": "127.0.0.1", "source_port": PORT} + calls = [] + + def resolve(key, metadata): + calls.append((key, metadata)) + return None if len(calls) <= misses else expected + + ticks = iter(range(10)) + monkeypatch.setattr(time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(time, "sleep", lambda _: None) + consumer = types.SimpleNamespace(_resolve_metadata=resolve) + assert _wait_for_metadata(consumer, "retry", forwarded, timeout=5) is expected + assert calls == [("retry", forwarded)] * (misses + 1) + + +def test_wait_for_metadata_stops_at_deadline(monkeypatch): + calls = [] + ticks = iter([0.0, 0.0, 1.0, 2.0]) + monkeypatch.setattr(time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(time, "sleep", lambda _: None) + consumer = types.SimpleNamespace(_resolve_metadata=lambda *args: calls.append(args)) + with pytest.raises(pytest.fail.Exception, match="Timed out waiting for metadata for 'missing'"): + _wait_for_metadata(consumer, "missing") + assert calls == [("missing", None)] * 2 + + +@pytest.mark.parametrize("async_chunk", [False, True]) +def test_three_stage_incoming_and_outgoing_endpoints(nixl_connector_cls, async_chunk): + from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig + from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + config = OmniTransferConfig( + connectors={ + ("0", "1"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT}), + ("1", "2"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT + 10}), + } + ) + connectors = [] + try: + for stage in range(3): + spec = get_stage_connector_spec(config, stage, async_chunk) + resolved = resolve_connector_spec(ConnectorSpec(**spec), stage_id=stage, role=spec["extra"]["role"]) + connectors.append(nixl_connector_cls(resolved.extra)) + first, middle, last = connectors + assert middle._sender_zmq_port == PORT + assert middle._serving_handshake + assert middle._zmq_port == PORT + 11 + first.put("0", "1", "incoming", torch.ones(1)) + incoming = _wait_for_metadata(middle, "incoming") + assert incoming["sender_zmq_port"] == PORT + assert incoming["generation"] == first._published["incoming"]["generation"] + middle.put("1", "2", "outgoing", torch.ones(1)) + outgoing = _wait_for_metadata(last, "outgoing") + assert outgoing["sender_zmq_port"] == PORT + 11 + assert outgoing["generation"] == middle._published["outgoing"]["generation"] + finally: + for connector in reversed(connectors): + for pending in connector._pending.values(): + pending.claims.clear() # No DMA in this metadata-only probe. + connector.close() + + +@pytest.mark.parametrize("name", ["NixlConnector", "MooncakeTransferEngineConnector"]) +def test_explicit_sender_port_is_not_derived(name): + from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec + + spec = ConnectorSpec(name=name, extra={"sender_zmq_port": 55000, "from_stage": 0}) + resolved = resolve_connector_spec(spec, stage_id=1, role="receiver", replica_id=3, local_rank=2) + assert resolved.extra["sender_zmq_port"] == 55000 + assert spec.extra == {"sender_zmq_port": 55000, "from_stage": 0} + spec.extra["zmq_port"] = "${UNUSED_NIXL_BASE_PORT}" + assert resolve_connector_spec(spec, stage_id=1, role="receiver").extra["sender_zmq_port"] == 55000 + + +@pytest.mark.parametrize("need_recv_cache", [False, True]) +@pytest.mark.parametrize("update_before_init", [False, True]) +def test_manager_receiver_does_not_bind_shared_producer_port(producer, need_recv_cache, update_before_init): + """Payload-only and KV receivers must dial, not bind, the incoming edge.""" + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ) + + manager = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={ + "type": "NixlConnector", + "role": "receiver", + "host": "127.0.0.1", + "zmq_port": PORT, + "backends": ["UCX"], + }, + from_stage="0", + to_stage="1", + stage_id=1, + need_recv_cache=need_recv_cache, + ) + ) + sender_info = {"host": producer.host, "zmq_port": producer._zmq_port} + if update_before_init: + manager.update_sender_info(sender_info) + receiver = manager.connector + try: + assert receiver is not None, "Receiver incorrectly tried to bind the producer's occupied port" + assert receiver._zmq_port is None + assert not receiver._serving_handshake + if not update_before_init: + assert receiver._sender_zmq_port is None + manager.update_sender_info(sender_info) + assert receiver._sender_zmq_port == producer._zmq_port + assert manager.config.connector_config["zmq_port"] == PORT # Do not mutate the shared specification. + success, _, _ = producer.put("0", "1", "manager-receive", torch.ones(1)) + assert success + metadata = _wait_for_metadata(receiver, "manager-receive") + assert metadata["generation"] == producer._published["manager-receive"]["generation"] + receiver._notify_transfer_done("manager-receive", metadata) + assert not producer._pending + finally: + if receiver is not None: + receiver.close() + + +@pytest.mark.parametrize("need_recv_cache", [False, True]) +def test_manager_receiver_preserves_explicit_standalone_sender(producer, need_recv_cache): + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ) + + manager = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={ + "type": "NixlConnector", + "role": "receiver", + "zmq_port": PORT, + "sender_host": producer.host, + "sender_zmq_port": producer._zmq_port, + }, + need_recv_cache=need_recv_cache, + ) + ) + receiver = manager.connector + try: + assert receiver is not None + assert receiver._zmq_port is None and not receiver._serving_handshake + assert (receiver._sender_host, receiver._sender_zmq_port) == (producer.host, producer._zmq_port) + success, _, _ = producer.put("0", "1", "standalone-receive", torch.ones(1)) + assert success + metadata = _wait_for_metadata(receiver, "standalone-receive") + assert metadata["generation"] == producer._published["standalone-receive"]["generation"] + receiver._notify_transfer_done("standalone-receive", metadata) + assert not producer._pending + finally: + if receiver is not None: + receiver.close() + + +def test_middle_stage_advertises_its_outgoing_replica_endpoint(): + from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient + + client = object.__new__(StageEngineCoreClient) + client.stage_id = 1 + client.replica_id = 3 + client.vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace( + stage_connector_config={ + "name": "NixlConnector", + "extra": { + "role": "receiver", + "host": "10.0.0.1", + "zmq_port": 47000, + "from_stage": 0, + "outgoing": {"host": "10.0.0.2", "zmq_port": 48000, "from_stage": 1}, + }, + } + ) + ) + assert client._build_payload_sender_info() == {"host": "10.0.0.2", "zmq_port": 51073} + + +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.usefixtures("reliable_claim_queries") +def test_claimed_source_survives_expiry_and_cleanup(producer, consumer, direct): + _, _, metadata = producer.put("0", "1", "claimed", torch.ones(1)) + resolved = consumer._resolve_metadata("claimed", metadata if direct else None) + pending = producer._pending["claimed"] + pending.deadline = 0 + producer._cleanup_expired_pending() + assert producer._pending.get("claimed") is pending + producer.cleanup("claimed") + assert producer._pending.get("claimed") is pending + assert producer._agent.registered + consumer._notify_transfer_done("claimed", resolved) + consumer._notify_transfer_done("claimed", resolved) + assert producer._pending == {} + assert producer._agent.registered == [] + + +@pytest.mark.usefixtures("reliable_claim_queries") +def test_claims_are_generation_scoped_and_duplicate_ack_cannot_release_sibling(producer, consumer): + _, _, original = producer.put("0", "1", "owners", torch.ones(1)) + first = consumer._resolve_metadata("owners", original) + second = consumer._resolve_metadata("owners", original) + assert first["claim_id"] != second["claim_id"] + assert not producer.put("0", "1", "owners", torch.zeros(1))[0] + consumer._notify_transfer_done("owners", first) + consumer._notify_transfer_done("owners", first) + assert len(producer._pending["owners"].claims) == 1 + consumer._notify_transfer_done("owners", second) + assert not producer._pending + producer.put("0", "1", "owners", torch.zeros(1)) + consumer._notify_transfer_done("owners", second) + assert "owners" in producer._pending + assert consumer._resolve_metadata("owners", original) is None + + +@pytest.mark.parametrize("get_metadata", [None, {"schema_version": 1, "tensor_specs": [], "descriptor_groups": []}]) +@pytest.mark.usefixtures("reliable_claim_queries") +def test_abandoned_claim_survives_close_and_late_completion(producer, consumer, get_metadata): + _, _, metadata = producer.put("0", "1", "abandoned", torch.ones(1)) + claimed = consumer._resolve_metadata("abandoned", metadata) + producer._pending["abandoned"].deadline = 0 + producer.close() + assert producer._closing and not producer._closed + assert producer._agent.registered + assert producer._listener_thread.is_alive() + with pytest.raises(RuntimeError, match="closed"): + producer.put("0", "1", "new", torch.ones(1)) + metrics = dict(producer._metrics) + with pytest.raises(RuntimeError, match="Cannot get data: NixlConnector is closed"): + producer.get("0", "1", "new", get_metadata) + assert producer._metrics == metrics + assert producer._pending["abandoned"].claims == {claimed["claim_id"]} + assert consumer._resolve_metadata("abandoned", metadata) is None + consumer._notify_transfer_done("abandoned", claimed) + assert not producer._pending + assert not producer._agent.registered + assert producer._listener_thread.is_alive() + producer.close() + assert not producer._agent.registered + assert producer._closed + + +@pytest.fixture +def copying_native_agent(consumer, monkeypatch): + """Strict fake descriptors and actual CPU copies; not native NIXL evidence.""" + calls = [] + + def descriptors(regions, memory_type): + assert regions and all(region[0] > 0 and region[1] > 0 for region in regions) + return regions + + def prepare(operation, local, local_ids, remote, remote_ids): + assert operation == "READ" + assert local_ids == remote_ids == list(range(len(local))) + return list(zip(local, remote, strict=True)) + + def transfer(pairs): + calls.append(pairs) + for destination, source in pairs: + assert destination[1] == source[1] > 0 + ctypes.memmove(destination[0], source[0], source[1]) + + agent = consumer._agent + monkeypatch.setattr(agent, "add_remote_agent", lambda metadata: "producer", raising=False) + monkeypatch.setattr(agent, "get_xfer_descs", descriptors, raising=False) + monkeypatch.setattr(agent, "prep_xfer_dlist", lambda agent, descs: descs, raising=False) + monkeypatch.setattr(agent, "make_prepped_xfer", prepare, raising=False) + monkeypatch.setattr(agent, "transfer", transfer, raising=False) + monkeypatch.setattr(agent, "check_xfer_state", lambda handle: "DONE", raising=False) + for method in ("release_xfer_handle", "release_dlist_handle", "remove_remote_agent"): + monkeypatch.setattr(agent, method, lambda handle: None, raising=False) + return calls + + +@pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"]) +def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case): + empty = torch.empty((2, 0, 3), dtype=torch.float64) + other = torch.empty((0,), dtype=torch.int64) + scalar = torch.tensor(7) + vector = torch.arange(6, dtype=torch.float32) + payload = { + "empty": empty, + "all_empty": [empty, other], + "mixed": [empty, scalar, other, vector], + "scalar": scalar, + "structured_empty": {"leaves": (empty, [other]), "label": "empty"}, + "structured_mixed": {"leaves": (empty, [scalar, other, vector]), "label": "mixed"}, + }[case] + ok, size, metadata = producer.put("0", "1", "empty-slots", payload) + assert ok + indices = [index for group in metadata["descriptor_groups"] for index in group["tensor_indices"]] + assert sorted(indices) == [i for i, spec in enumerate(metadata["tensor_specs"]) if spec["size"]] + actual, received_size = consumer.get("0", "1", "empty-slots", metadata) + if isinstance(payload, dict): + assert actual["label"] == payload["label"] + torch.testing.assert_close(actual["leaves"], payload["leaves"], rtol=0, atol=0) + else: + torch.testing.assert_close(actual, payload, rtol=0, atol=0) + assert size == received_size + if case in ("empty", "all_empty"): + assert not copying_native_agent + assert not producer._pending and not producer._agent.registered + assert not consumer._agent.registered + + +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.parametrize("outcome", ["done", "error", "timeout", "unknown"]) +def test_read_ownership_through_terminal_and_deferred_paths( + producer, consumer, copying_native_agent, monkeypatch, direct, outcome +): + import threading + + _, _, metadata = producer.put("0", "1", "active-read", torch.ones(2)) + entered, proceed = threading.Event(), threading.Event() + state = ["PROC"] + monkeypatch.setattr(consumer._agent, "check_xfer_state", lambda handle: state[0]) + + def wait(handle, key): + entered.set() + assert proceed.wait(5) + if outcome == "done": + state[0] = "DONE" + elif outcome == "error": + state[0] = "ERR" + raise RuntimeError("terminal error") + else: + if outcome == "unknown": + state[0] = "UNKNOWN" + raise TimeoutError("still owned") + + monkeypatch.setattr(consumer, "_wait_for_transfer", wait) + worker = threading.Thread(target=consumer.get, args=("0", "1", "active-read", metadata if direct else None)) + worker.start() + try: + assert entered.wait(5) + pending = producer._pending["active-read"] + pending.deadline = 0 + producer._cleanup_expired_pending() + producer.cleanup("active-read") + assert producer._pending["active-read"] is pending + assert producer._agent.registered + finally: + proceed.set() + worker.join(5) + assert not worker.is_alive() + if outcome in ("timeout", "unknown"): + assert producer._agent.registered + assert consumer._deferred_transfers + consumer._reap_deferred_transfers() + assert producer._agent.registered + state[0] = "DONE" + deadline = time.monotonic() + 5 + while producer._pending and time.monotonic() < deadline: + time.sleep(0.01) + assert not producer._pending + assert not producer._agent.registered + + +class _FakeNixlAgent: + def __init__(self, name, config=None): + self.name = name + self.config = config + self.registered = [] + + def get_reg_descs(self, regions, memory_type): + return ("reg", tuple(regions), memory_type) + + def register_memory(self, descs, backends=None): + self.registered.append(descs) + + def deregister_memory(self, descs): + if descs in self.registered: + self.registered.remove(descs) + + def get_agent_metadata(self): + return f"agent-metadata-{self.name}".encode() + + +@pytest.fixture +def nixl_connector_cls(monkeypatch): + """Import NixlConnector with vLLM's optional NIXL dependency stubbed out.""" + stub = types.ModuleType("vllm.distributed.nixl_utils") + stub.NixlWrapper = _FakeNixlAgent # type: ignore[attr-defined] + config_module = types.ModuleType("test_nixl_config_api") + config_sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="strict") + config_module.nixl_thread_sync_t = config_sync_type # type: ignore[attr-defined] + + def agent_config(**kwargs): + return kwargs + + agent_config.__module__ = config_module.__name__ + config_module.nixl_agent_config = agent_config # type: ignore[attr-defined] + stub.nixl_agent_config = agent_config # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "vllm.distributed.nixl_utils", stub) + monkeypatch.setitem(sys.modules, config_module.__name__, config_module) + sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="wrong-public-enum") + nixl_module = types.ModuleType("nixl") + nixl_module.nixl_thread_sync_t = sync_type # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "nixl", nixl_module) + + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector + + return NixlConnector + + +@pytest.fixture +def producer(nixl_connector_cls): + connector = nixl_connector_cls({"role": "sender", "host": "127.0.0.1", "zmq_port": PORT}) + yield connector + # These control-plane tests do not submit DMA. Explicitly complete claims + # made by metadata-only probes before tearing down the shared port. + for key, pending in list(connector._pending.items()): + for claim in list(pending.claims): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _XFER_DONE_MSG + + connector._handle_handshake_message( + _XFER_DONE_MSG + + msgspec.msgpack.encode({"key": key, "generation": pending.generation, "claim_id": claim}) + ) + connector.close() + + +@pytest.fixture +def consumer(nixl_connector_cls): + connector = nixl_connector_cls({"role": "receiver", "sender_host": "127.0.0.1", "sender_zmq_port": PORT}) + yield connector + connector.close() + + +@pytest.fixture +def reliable_claim_queries(producer, consumer, monkeypatch): + """Ownership probes need exact claims, not claims retained after lost replies. + + Exercise the real resolver, wire encoding and producer handler synchronously; + discovery tests separately cover bounded queries over real ZMQ sockets. + """ + import uuid + + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _GET_META_MSG, _META_NOT_FOUND + + def query(key, host, port, *, generation=None): + assert (host, port) == (producer.host, producer._zmq_port) + request = {"key": key, "generation": generation, "claim_id": uuid.uuid4().hex} + reply = producer._handle_handshake_message(_GET_META_MSG + msgspec.msgpack.encode(request)) + return None if reply == _META_NOT_FOUND else msgspec.msgpack.decode(reply) + + monkeypatch.setattr(consumer, "_query_metadata_at", query) + + +def test_put_publishes_its_handshake_endpoint(producer): + ok, size, metadata = producer.put("0", "1", "req-0", torch.arange(8, dtype=torch.float32)) + + assert ok is True + assert size == 32 + assert metadata["sender_host"] == "127.0.0.1" + assert metadata["sender_zmq_port"] == PORT + + +def test_agent_uses_strict_thread_synchronization(producer): + assert producer._agent.config["sync_mode"] == "strict" + + +def test_handshake_serves_metadata_when_caller_has_none(producer, consumer): + _, _, published = producer.put("0", "1", "req-1", torch.arange(4, dtype=torch.float32)) + + resolved = _wait_for_metadata(consumer, "req-1") + + # msgpack has no tuple type, so region descriptors arrive as lists; get() + # re-tuples them before handing them to NIXL. + assert resolved.pop("claim_id") + assert resolved == msgspec.msgpack.decode(msgspec.msgpack.encode(published)) + assert [tuple(region) for group in resolved["descriptor_groups"] for region in group["regions"]] == [ + region for group in published["descriptor_groups"] for region in group["regions"] + ] + + +def test_legacy_metadata_without_generation_skips_the_handshake(consumer): + """Externally owned metadata without a generation needs no ownership claim.""" + direct = {"schema_version": 1, "kind": "tensors", "tensor_specs": []} + + assert consumer._resolve_metadata("req-2", direct) is direct + + +def test_source_endpoint_metadata_overrides_configured_sender(consumer, monkeypatch): + queried = [] + expected = {"schema_version": 1, "kind": "tensors", "tensor_specs": []} + + def query(get_key, host, port): + queried.append((get_key, host, port)) + return expected + + monkeypatch.setattr(consumer, "_query_metadata_at", query) + + resolved = consumer._resolve_metadata( + "req-tp4-rank3", + {"source_host": "10.0.0.3", "source_port": PORT + 3 * 16}, + ) + + assert resolved is expected + assert queried == [("req-tp4-rank3", "10.0.0.3", PORT + 3 * 16)] + + +def test_unknown_key_is_queried_once(producer, consumer, monkeypatch): + socket = consumer._get_req_socket(f"tcp://127.0.0.1:{PORT}") + send_count = 0 + original_send = socket.send + + def count_send(message): + nonlocal send_count + send_count += 1 + return original_send(message) + + monkeypatch.setattr(socket, "send", count_send) + + assert consumer._resolve_metadata("never-published", None) is None + assert send_count == 1 + assert socket.getsockopt(zmq.RCVTIMEO) == 10 + + +@pytest.mark.usefixtures("reliable_claim_queries") +def test_transfer_done_releases_the_producer_buffer(producer, consumer): + _, _, metadata = producer.put("0", "1", "req-3", torch.arange(4, dtype=torch.float32)) + assert producer._pending and producer._agent.registered + + metadata = consumer._resolve_metadata("req-3", metadata) + consumer._notify_transfer_done("req-3", metadata) + + assert producer._pending == {} + assert producer._published == {} + assert producer._agent.registered == [] + + +def test_direct_metadata_has_ephemeral_ownership_endpoint(nixl_connector_cls): + connector = nixl_connector_cls({}) + try: + assert connector._zmq_ctx is not None + assert connector._listener_thread is not None + _, _, metadata = connector.put("0", "1", "req-4", torch.zeros(2)) + assert metadata["sender_host"] + assert metadata["sender_zmq_port"] > 0 + finally: + connector.close() + + +def test_idle_producer_lease_expires_without_another_put(nixl_connector_cls): + connector = nixl_connector_cls({"role": "sender", "lease_seconds": 0.01}) + try: + connector.put("0", "1", "req-expire", torch.zeros(2)) + + deadline = time.monotonic() + 1.0 + while connector._pending and time.monotonic() < deadline: + time.sleep(0.01) + + assert connector._pending == {} + assert connector._published == {} + assert connector._agent.registered == [] + assert connector._registered_descs == [] + finally: + connector.close() + + +def test_close_stops_lease_reaper(nixl_connector_cls): + connector = nixl_connector_cls({"role": "sender"}) + lease_thread = connector._lease_thread + + connector.close() + + assert lease_thread is not None + assert not lease_thread.is_alive() + assert connector._lease_thread is None + + +def test_deferred_transfer_is_retained_while_active(nixl_connector_cls): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer + + connector = nixl_connector_cls({"role": "receiver"}) + released = [] + connector._agent.check_xfer_state = lambda handle: "PROC" + connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle)) + connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle)) + connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent)) + transfer = _DeferredTransfer( + tensors=[torch.zeros(1)], + registrations=["registration"], + dlists=["local", "remote"], + handles=["transfer"], + remote_agent="producer", + ) + connector._defer_transfer(transfer) + try: + connector._reap_deferred_transfers() + + assert released == [] + assert transfer in connector._deferred_transfers + assert transfer.tensors + finally: + connector._agent.check_xfer_state = lambda handle: "DONE" + connector.close() + + +def test_deferred_transfer_releases_exactly_once_after_done(nixl_connector_cls): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer + + connector = nixl_connector_cls({"role": "receiver"}) + released = [] + connector._agent.check_xfer_state = lambda handle: "DONE" + connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle)) + connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle)) + connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent)) + connector._agent.deregister_memory = lambda descs: released.append(("registration", descs)) + transfer = _DeferredTransfer( + tensors=[torch.zeros(1)], + registrations=["registration"], + dlists=["local", "remote"], + handles=["transfer"], + remote_agent="producer", + ) + connector._defer_transfer(transfer) + try: + connector._reap_deferred_transfers() + connector._reap_deferred_transfers() + + assert released == [ + ("handle", "transfer"), + ("dlist", "local"), + ("dlist", "remote"), + ("agent", "producer"), + ("registration", "registration"), + ] + assert connector._deferred_transfers == [] + assert transfer.tensors == [] + finally: + connector.close() + + +def test_active_sibling_transfer_requires_deferred_ownership(nixl_connector_cls): + connector = nixl_connector_cls({"role": "receiver"}) + connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle] + try: + assert connector._transfers_may_be_active(["failed", "active"]) is True + assert connector._transfers_may_be_active(["failed"]) is False + finally: + connector.close() + + +def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_connector_cls): + connector = nixl_connector_cls({"role": "receiver"}) + connector._agent.add_remote_agent = lambda metadata: "producer" + connector._agent.get_xfer_descs = lambda regions, memory_type: (regions, memory_type) + connector._agent.prep_xfer_dlist = lambda agent, descs: (agent, tuple(descs[0])) + handles = iter(["failed", "active"]) + connector._agent.make_prepped_xfer = lambda *_args: next(handles) + connector._agent.transfer = lambda handle: None + connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle] + released = [] + connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle)) + connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle)) + connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent)) + connector._agent.deregister_memory = lambda descs: released.append(("registration", descs)) + metadata = { + "schema_version": 1, + "kind": "tensors", + "agent_metadata": b"producer-metadata", + "tensor_specs": [ + {"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4}, + {"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4}, + ], + "descriptor_groups": [ + {"memory_type": "DRAM", "tensor_indices": [0], "regions": [(1, 4, 0, "")]}, + {"memory_type": "DRAM", "tensor_indices": [1], "regions": [(2, 4, 0, "")]}, + ], + "size": 8, + } + + try: + assert connector.get("0", "1", "req-active-sibling", metadata) is None + assert released == [] + assert len(connector._deferred_transfers) == 1 + deferred = connector._deferred_transfers[0] + assert deferred.handles == ["failed", "active"] + assert len(deferred.registrations) == 2 + assert len(deferred.dlists) == 4 + assert deferred.remote_agent == "producer" + assert len(deferred.tensors) == 2 + finally: + connector._agent.check_xfer_state = lambda handle: "DONE" + connector.close() + + +@pytest.mark.parametrize( + "payload,expected_kind", + [ + (torch.zeros(4), "tensors"), + ([torch.zeros(2), torch.ones(3)], "tensors"), + ({"hidden": torch.zeros(4), "meta": {"token_role_ids": [1, 2]}}, "structured"), + ({"prompt": "a cat", "steps": 8}, "object"), + ], +) +def test_payload_kinds_round_trip_through_the_handshake(producer, consumer, payload, expected_kind): + _, _, published = producer.put("0", "1", "req-kind", payload) + assert published["kind"] == expected_kind + + resolved = _wait_for_metadata(consumer, "req-kind") + + assert resolved["kind"] == expected_kind + assert resolved["tensor_specs"] == published["tensor_specs"] + + +def test_structured_payload_groups_descriptors_by_memory_type(producer, monkeypatch): + monkeypatch.setattr( + producer, + "_resolve_memory_type", + lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM", + ) + + _, _, metadata = producer.put( + "0", + "1", + "req-mixed", + {"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)}, + ) + + assert metadata["schema_version"] == 1 + assert metadata["descriptor_groups"] == [ + { + "memory_type": "DRAM", + "tensor_indices": [0], + "regions": metadata["descriptor_groups"][0]["regions"], + }, + { + "memory_type": "VRAM", + "tensor_indices": [1], + "regions": metadata["descriptor_groups"][1]["regions"], + }, + ] + assert len(producer._agent.registered) == 2 + + +def test_omni_payload_struct_uses_structured_tensor_path(producer): + payload = OmniPayloadStruct( + hidden_states=HiddenStatesStruct(output=torch.ones(4, dtype=torch.float32)), + meta=MetaStruct(left_context_size=2), + ) + + _, _, metadata = producer.put("0", "1", "req-struct", payload) + skeleton, tensors = producer._extract_tensor_leaves(payload) + + assert metadata["kind"] == "structured" + assert len(tensors) == 1 + assert torch.equal(tensors[0], payload.hidden_states.output) + assert skeleton["meta"]["left_context_size"] == 2 + + +def test_partial_group_registration_failure_rolls_back(producer, monkeypatch): + monkeypatch.setattr( + producer, + "_resolve_memory_type", + lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM", + ) + original_register = producer._agent.register_memory + + def fail_second_group(descs, backends=None): + if descs[2] == "VRAM": + raise RuntimeError("registration failed") + original_register(descs, backends=backends) + + monkeypatch.setattr(producer._agent, "register_memory", fail_second_group) + + success, _, metadata = producer.put( + "0", + "1", + "req-partial", + {"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)}, + ) + + assert success is False + assert metadata is None + assert producer._agent.registered == [] + assert producer._registered_descs == [] + assert "req-partial" not in producer._pending + + +def test_reusing_put_key_releases_previous_registration(producer): + producer.put("0", "1", "req-reuse", torch.zeros(2)) + previous_registration = producer._agent.registered[0] + + producer.put("0", "1", "req-reuse", torch.ones(2)) + + assert previous_registration not in producer._agent.registered + assert len(producer._agent.registered) == 1 + assert len(producer._registered_descs) == 1 + + +def test_expired_snapshot_cannot_claim_replacement(producer): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _PendingPayload + + old = _PendingPayload([torch.zeros(1)], ["old"], 0.0) + replacement = _PendingPayload([torch.ones(1)], ["new"], time.monotonic() + 60) + producer._pending["req-race"] = replacement + + assert producer._take_pending("req-race", expected=old) is None + assert producer._pending["req-race"] is replacement + + +@pytest.mark.parametrize( + "indices", + [ + [0, 0], + [0], + [0, 2], + [-1, 0], + ], +) +def test_descriptor_groups_require_exact_tensor_index_partition(nixl_connector_cls, indices): + metadata = { + "schema_version": 1, + "descriptor_groups": [ + { + "memory_type": "DRAM", + "tensor_indices": indices, + "regions": [(1, 4, 0, "")] * len(indices), + } + ], + } + + with pytest.raises(RuntimeError, match="exact partition"): + nixl_connector_cls._validated_descriptor_groups( + metadata, [{"shape": [1], "dtype": "torch.float32", "size": 4}] * 2 + ) + + +def test_receive_device_ignores_the_producer_index(nixl_connector_cls): + """A producer on cuda:3 must not dictate the consumer's card.""" + connector = nixl_connector_cls({}) + try: + assert connector._resolve_receive_device("cpu") == torch.device("cpu") + assert connector._resolve_receive_device(None) == torch.device("cpu") + finally: + connector.close() + + +def test_receive_device_config_wins(nixl_connector_cls): + connector = nixl_connector_cls({"receive_device": "cpu"}) + try: + assert connector._resolve_receive_device("cuda:3") == torch.device("cpu") + finally: + connector.close() diff --git a/tests/distributed/omni_connectors/test_nixl_connector_native.py b/tests/distributed/omni_connectors/test_nixl_connector_native.py new file mode 100644 index 000000000..e0936a02d --- /dev/null +++ b/tests/distributed/omni_connectors/test_nixl_connector_native.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import multiprocessing as mp +import socket +import time +from typing import Any + +import pytest +import torch + +from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.parallel] + + +def _native_nixl_available() -> bool: + try: + from vllm.distributed.nixl_utils import NixlWrapper + except ImportError: + return False + return NixlWrapper is not None and torch.cuda.is_available() and torch.accelerator.device_count() >= 2 + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _producer(port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any) -> None: + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector + + current_omni_platform.set_device(0) + connector = NixlConnector( + { + "role": "sender", + "host": "127.0.0.1", + "zmq_port": port, + "agent_name": "native-smoke-producer", + } + ) + try: + payload = OmniPayloadStruct( + hidden_states=HiddenStatesStruct(output=torch.arange(8, dtype=torch.float32, device="cuda:0")), + meta=MetaStruct(token_role_ids=torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu")), + kv_metadata={"empty": torch.empty((2, 0, 3), device="cuda:0")}, + request_id="native-smoke", + ) + success, size, metadata = connector.put("0", "1", "native-smoke", payload) + result.put(("put", success, size, metadata["kind"] if metadata else None)) + wire.put(metadata) + ready.set() + if not claimed.wait(timeout=30): + raise TimeoutError("consumer did not claim source") + pending = connector._pending["native-smoke"] + pending.deadline = 0 + connector._cleanup_expired_pending() + connector.cleanup("native-smoke") + assert connector._pending["native-smoke"] is pending + assert pending.claims and connector._registered_descs + result.put(("source_retained", len(pending.claims))) + proceed.set() + if not consumed.wait(timeout=30): + raise TimeoutError("consumer did not complete the native NIXL transfer") + deadline = time.monotonic() + 10 + while connector._pending and time.monotonic() < deadline: + time.sleep(0.01) + result.put(("cleanup", len(connector._pending), len(connector._registered_descs))) + except Exception as error: + result.put(("producer_error", repr(error))) + raise + finally: + connector.close() + + +def _consumer( + port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any, direct: bool +) -> None: + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ) + + if not ready.wait(timeout=30): + raise TimeoutError("producer did not publish native NIXL metadata") + current_omni_platform.set_device(1) + manager = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={ + "type": "NixlConnector", + "role": "receiver", + "host": "127.0.0.1", + "zmq_port": port, # Shared deployment edge: already bound by producer. + "backends": ["UCX"], + "receive_device": "cuda", + "agent_name": "native-smoke-consumer", + }, + from_stage="0", + to_stage="1", + stage_id=1, + need_recv_cache=direct, # Exercise both payload-only and KV receiver setup. + ) + ) + manager.update_sender_info({"host": "127.0.0.1", "zmq_port": port}) + connector = manager.connector + assert connector is not None, "Receiver must not bind the producer's occupied port" + assert connector._zmq_port is None and not connector._serving_handshake + assert connector._backends == ["UCX"] + try: + original_wait = connector._wait_for_transfer + states = [] + + def wait(handle, key): + states.append(connector._agent.check_xfer_state(handle)) + claimed.set() + if not proceed.wait(timeout=30): + raise TimeoutError("producer did not verify active ownership") + original_wait(handle, key) + + connector._wait_for_transfer = wait + metadata = wire.get(timeout=5) + received = None + deadline = time.monotonic() + 30 + while received is None and time.monotonic() < deadline: + received = connector.get( + "0", + "1", + "native-smoke", + metadata if direct else {"source_host": "127.0.0.1", "source_port": port}, + ) + if received is None: + raise TimeoutError("native NIXL transfer did not complete") + payload, size = received + hidden = payload["hidden_states"]["output"] + token_roles = payload["meta"]["token_role_ids"] + empty = payload["kv_metadata"]["empty"] + assert empty.shape == (2, 0, 3) and str(empty.device) == "cuda:1" + result.put(("native_states", states)) + result.put( + ( + "get", + size, + hidden.cpu().tolist(), + str(hidden.dtype), + list(hidden.shape), + str(hidden.device), + token_roles.cpu().tolist(), + str(token_roles.dtype), + list(token_roles.shape), + str(token_roles.device), + ) + ) + consumed.set() + except Exception as error: + result.put(("consumer_error", repr(error))) + raise + finally: + connector.close() + + +@pytest.mark.skipif(not _native_nixl_available(), reason="requires NIXL and at least two CUDA devices") +@pytest.mark.parametrize("direct", [False, True]) +def test_native_two_process_structured_mixed_device_transfer(direct): + context = mp.get_context("spawn") + ready = context.Event() + consumed = context.Event() + claimed = context.Event() + proceed = context.Event() + wire = context.Queue() + result = context.Queue() + port = _free_port() + producer = context.Process(target=_producer, args=(port, ready, consumed, result, claimed, proceed, wire)) + consumer = context.Process(target=_consumer, args=(port, ready, consumed, result, claimed, proceed, wire, direct)) + + producer.start() + consumer.start() + producer.join(timeout=60) + consumer.join(timeout=60) + if producer.is_alive(): + producer.terminate() + if consumer.is_alive(): + consumer.terminate() + + records = [result.get(timeout=5) for _ in range(5)] + print("Native ownership records:", records) + assert producer.exitcode == 0, records + assert consumer.exitcode == 0, records + assert records[0][0] == "put" + assert records[0][1] is True + assert records[0][2] >= 56 + assert records[0][3] == "structured" + get_record = next(record for record in records if record[0] == "get") + assert get_record[1] == records[0][2] + assert get_record[2:6] == ( + list(range(8)), + "torch.float32", + [8], + "cuda:1", + ) + assert get_record[6:] == ([1, 2, 3], "torch.int64", [3], "cuda:1") + assert next(record for record in records if record[0] == "cleanup") == ("cleanup", 0, 0) + assert next(record for record in records if record[0] == "source_retained") == ("source_retained", 1) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index e3f58c129..ca2a26b70 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -1,13 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio +from collections.abc import Callable from types import SimpleNamespace +from typing import Any import pytest from vllm import SamplingParams +from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.messages import OutputMessage from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState -from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.stage_engine_core_client import ( + DPLBStageEngineCoreClient, + StageEngineCoreClient, +) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -28,7 +37,7 @@ class _DummySenderStage: class _DummyDiffusionStage: stage_type = "diffusion" final_output = True - custom_process_input_func = None + custom_process_input_func: Callable[..., Any] | None = None def __init__(self, engine_input_source=None): self.engine_input_source = engine_input_source or [0] @@ -54,6 +63,97 @@ def _build_sender_pool(stage_id: int, sender_info: dict[str, object]) -> StagePo ) +@pytest.mark.parametrize( + ("client_class", "base_client_class"), + [ + (StageEngineCoreClient, AsyncMPClient), + (DPLBStageEngineCoreClient, DPLBAsyncMPClient), + ], +) +def test_stage_engine_core_client_builds_payload_sender_info_after_base_init( + monkeypatch, client_class, base_client_class +): + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + omni_kv_config=None, + hf_config=None, + stage_connector_config=None, + ) + ) + metadata = SimpleNamespace( + stage_id=0, + replica_id=0, + stage_type="llm", + model_stage="main", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source=[], + final_output=False, + final_output_type=None, + default_sampling_params=None, + prompt_transform_func=None, + prompt_expand_func=None, + custom_process_input_func=None, + ) + + def fake_base_init(self, config, *_args, **_kwargs): + self.vllm_config = config + self.resources = SimpleNamespace(engine_dead=False) + + monkeypatch.setattr(base_client_class, "__init__", fake_base_init) + + client = client_class(vllm_config, object, metadata=metadata) + + assert client.vllm_config is vllm_config + assert client.get_payload_sender_info() is None + + +@pytest.mark.parametrize("outgoing", [False, True]) +@pytest.mark.parametrize("base_port", [None, 48000]) +def test_payload_sender_endpoint_matches_resolver_with_unequal_replicas(outgoing, base_port): + from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec + + edge = {"host": "10.0.0.2", "from_stage": 1} + if base_port is not None: + edge["zmq_port"] = base_port + extra = ( + { + "role": "receiver", + "host": "10.0.0.1", + "zmq_port": 47000, + "from_stage": 0, + "outgoing": edge, + } + if outgoing + else {"role": "sender", **edge} + ) + client = object.__new__(StageEngineCoreClient) + client.stage_id = 1 + client.replica_id = 3 + client.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(stage_connector_config={"name": "NixlConnector", "extra": extra}) + ) + producer = resolve_connector_spec( + ConnectorSpec(name="NixlConnector", extra=extra), + stage_id=1, + role=extra["role"], + replica_id=client.replica_id, + ) + consumer = resolve_connector_spec( + ConnectorSpec(name="NixlConnector", extra=edge), stage_id=2, role="receiver", replica_id=7 + ) + + sender_info = client._build_payload_sender_info() + + assert sender_info == { + "host": "10.0.0.2", + "zmq_port": (50051 if base_port is None else base_port) + 3 * 1024 + 1, + } + assert sender_info["zmq_port"] == producer.extra["zmq_port"] + assert sender_info["zmq_port"] != consumer.extra["sender_zmq_port"] + + def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address(): client = object.__new__(StageEngineCoreClient) client.stage_id = 0 @@ -259,3 +359,58 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): 0: {"host": "10.0.0.3", "zmq_port": 50151}, } assert req_state.stage_submit_ts[1] > 0 + + +def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas(): + orchestrator = object.__new__(Orchestrator) + endpoints = {"a": {"host": "10.0.0.2", "zmq_port": 52099}, "b": {"host": "10.0.0.3", "zmq_port": 54147}} + source = SimpleNamespace( + get_bound_client=lambda key: SimpleNamespace(get_payload_sender_info=lambda: endpoints[key]), + get_bound_replica_id=lambda key: {"a": 2, "b": 4}[key], + ) + submitted = {} + + async def submit(key, state, request, **kwargs): + await asyncio.sleep(0) + submitted[key] = request + + target = SimpleNamespace( + stage_type="llm", + stage_client=SimpleNamespace(engine_input_source=[0]), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=None, max_model_len=64)), + submit_initial=submit, + get_bound_replica_id=lambda key: {"a": 7, "b": 1}[key], + ) + orchestrator.stage_pools = [source, target] + orchestrator._stage_receives_async_chunks = lambda stage: True + orchestrator._record_duplex_stage_submission = lambda *args: None + orchestrator._emit_tx_edge = lambda **kwargs: None + + async def dispatch(callback, **kwargs): + await callback() + return True + + orchestrator._dispatch_or_fail_request = dispatch + + async def run(): + await asyncio.gather( + *[ + orchestrator._prewarm_async_chunk_stages( + key, + SimpleNamespace(prompt_token_ids=[1]), + OrchestratorRequestState( + request_id=key, + prompt={}, + sampling_params_list=[SamplingParams(), SamplingParams()], + final_stage_id=1, + ), + ) + for key in endpoints + ] + ) + + asyncio.run(run()) + from vllm_omni.engine import OmniEngineCoreRequest + + assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values()) + assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 68836e2a4..69df71231 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -104,9 +104,13 @@ class FakeInputProcessor: class FakePrewarmPool: + """Single-replica pool double that binds requests only on submission.""" + stage_type = "llm" def __init__(self, role: str) -> None: + self.stage_client = SimpleNamespace() + self._bound_request_ids: set[str] = set() self.stage_vllm_config = SimpleNamespace( model_config=SimpleNamespace( max_model_len=64, @@ -115,12 +119,16 @@ class FakePrewarmPool: ) self.submitted: list[Any] = [] - async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): + async def submit_initial(self, request_id, _req_state, request, prompt_text=None): self.submitted.append(request) + self._bound_request_ids.add(request_id) return 0 - def get_bound_replica_id(self, _request_id): - return 0 + def get_bound_replica_id(self, request_id): + return 0 if request_id in self._bound_request_ids else None + + def get_bound_client(self, request_id): + return self.stage_client if self.get_bound_replica_id(request_id) is not None else None def _duplex_stage_port_submission(): @@ -261,11 +269,14 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: @pytest.mark.asyncio -async def test_async_prewarm_skips_outgoing_only_stage() -> None: +@pytest.mark.parametrize("payload_sender_info", [None, {"host": "10.0.0.2", "zmq_port": 52099}]) +async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> None: orchestrator = object.__new__(Orchestrator) stage0 = FakePrewarmPool("sender") stage1 = FakePrewarmPool("sender") stage2 = FakePrewarmPool("receiver") + if payload_sender_info is not None: + stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None orchestrator._record_duplex_stage_submission = MagicMock() @@ -285,7 +296,14 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None: assert prewarmed is True assert stage1.submitted == [] + assert stage1.get_bound_client("req-prewarm") is None assert len(stage2.submitted) == 1 + assert stage2.get_bound_client("req-prewarm") is stage2.stage_client + assert stage2.submitted[0].payload_sender_info == payload_sender_info + if payload_sender_info is not None: + stage1.stage_client.get_payload_sender_info.assert_called_once_with() + assert stage2.submitted[0].external_req_id == "req-prewarm" + assert stage2.submitted[0].resumable is True assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts orchestrator._record_duplex_stage_submission.assert_called_once_with( diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 8366e177a..603ca65ff 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -10,6 +10,7 @@ from __future__ import annotations import time import unittest +from copy import deepcopy from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock, patch @@ -20,6 +21,10 @@ import torch from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( OmniKVTransferManager, ) +from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + resolve_connector_spec, +) from vllm_omni.outputs import OmniConnectorOutput from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, @@ -564,6 +569,55 @@ class TestChunkStreamCompletedGuard(unittest.TestCase): host.shutdown_omni_connectors() + def test_poll_uses_request_scoped_payload_sender_endpoint(self): + host = self._make_host(stage_id=1) + host._omni_connector.get = MagicMock(return_value=None) + req = _make_request("req-1", "ext-req-1") + req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + + host.register_chunk_recv(req) + host._poll_single_request("req-1") + + host._omni_connector.get.assert_called_once_with( + "0", + "1", + "ext-req-1_0_0", + {"source_host": "10.0.0.1", "source_port": 50051}, + ) + host.shutdown_omni_connectors() + + def test_concurrent_requests_keep_distinct_payload_sender_endpoints(self): + host = self._make_host(stage_id=1) + host._omni_connector.get = MagicMock(return_value=None) + first = _make_request("req-1", "ext-req-1") + first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + second = _make_request("req-2", "ext-req-2") + second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051} + + host.register_chunk_recv(first) + host.register_chunk_recv(second) + host._poll_single_request("req-2") + host._poll_single_request("req-1") + + self.assertEqual( + host._omni_connector.get.call_args_list, + [ + unittest.mock.call( + "0", + "1", + "ext-req-2_0_0", + {"source_host": "10.0.0.2", "source_port": 51051}, + ), + unittest.mock.call( + "0", + "1", + "ext-req-1_0_0", + {"source_host": "10.0.0.1", "source_port": 50051}, + ), + ], + ) + host.shutdown_omni_connectors() + def test_finish_sentinel_populates_completed_set(self): """Receiving is_finished=True adds to _chunk_stream_completed.""" host = self._make_host(stage_id=1) @@ -1381,6 +1435,144 @@ class TestConnectorConfigValidation(unittest.TestCase): host.init_omni_connectors(model_config=model_config) +class TestRankAwareHandshakePort(unittest.TestCase): + """A sender must bind the port the orchestrator advertises for its rank.""" + + @staticmethod + def _resolve(name, extra, local_rank=0): + spec = resolve_connector_spec( + ConnectorSpec(name=name, extra=extra), + stage_id=int(extra.get("stage_id", 0)), + role=extra.get("role"), + local_rank=local_rank, + ) + return spec.extra + + def test_sender_port_matches_the_orchestrator_formula(self): + for local_rank in (0, 1, 3): + with self.subTest(local_rank=local_rank): + extra = {"zmq_port": 50071, "stage_id": 2, "role": "sender"} + resolved = self._resolve("NixlConnector", extra, local_rank=local_rank) + self.assertEqual(resolved["zmq_port"], 50071 + 2 + local_rank * 16) + self.assertEqual(extra["zmq_port"], 50071, "the source config must not be mutated") + + def test_receiver_gets_sender_endpoint_without_a_bind_port(self): + receiver = self._resolve( + "NixlConnector", + {"zmq_port": 50071, "stage_id": 3, "from_stage": 2, "role": "receiver"}, + local_rank=1, + ) + self.assertNotIn("zmq_port", receiver) + self.assertEqual(receiver["sender_zmq_port"], 50071 + 2 + 16) + + def test_handshake_less_connector_keeps_its_config(self): + shm = self._resolve("SharedMemoryConnector", {"zmq_port": 50071, "role": "sender"}, local_rank=1) + self.assertEqual(shm["zmq_port"], 50071) + + def test_unresolvable_port_is_rejected(self): + extra = {"zmq_port": "${MISSING_PORT_VAR}", "stage_id": 0, "role": "sender"} + with self.assertRaises(ValueError): + self._resolve("NixlConnector", extra) + + +@pytest.mark.parametrize( + ("extra", "expected"), + [ + ( + {"role": "sender", "from_stage": 2, "zmq_port": 50071}, + {"zmq_port": 52153}, + ), + ( + {"role": "receiver", "from_stage": 1, "zmq_port": 50071, "host": "producer"}, + {"sender_zmq_port": 52152, "sender_host": "producer"}, + ), + ( + { + "role": "receiver", + "from_stage": 1, + "zmq_port": 50071, + "sender_zmq_port": 49000, + "sender_host": "bound-producer", + "outgoing": {"from_stage": 2, "zmq_port": 51071, "host": "local-worker"}, + }, + {"sender_zmq_port": 49000, "sender_host": "bound-producer", "zmq_port": 53153, "host": "local-worker"}, + ), + ], +) +def test_split_runtime_factory_resolves_nixl_endpoints(extra, expected): + """The runner must pass the resolved spec, not the raw config, to the factory.""" + from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_runtime as runtime + + original = deepcopy(extra) + config = SimpleNamespace(stage_id=2, stage_connector_config={"name": "NixlConnector", "extra": extra}) + with ( + patch.object(runtime, "get_local_tp_rank", return_value=2, create=True), + patch.object(runtime, "get_omni_replica_id", return_value=2, create=True), + patch.object(runtime.OmniConnectorFactory, "create_connector") as factory, + ): + assert MixinHost._create_connector(config) is factory.return_value + + spec = factory.call_args.args[0] + assert spec.name == "NixlConnector" + assert spec.extra["stage_id"] == 2 + for key, value in expected.items(): + assert spec.extra[key] == value + if extra["role"] == "receiver" and "outgoing" not in extra: + assert "zmq_port" not in spec.extra + assert extra == original + + +@pytest.mark.parametrize("async_chunk", [False, True]) +def test_split_payload_endpoint_retry_cleanup_and_reuse(async_chunk): + host = MixinHost() + host.init_omni_connectors(_make_model_config(async_chunk=async_chunk)) + host._stage_id = 1 + host._omni_connector = MagicMock() + host._omni_connector.get.return_value = None + first = _make_request("req-1", "ext-req-1") + first.payload_sender_info = {"host": "producer-a", "zmq_port": "50051"} + second = _make_request("req-2", "ext-req-2") + second.payload_sender_info = {"host": "producer-b", "zmq_port": 51051} + + try: + host.register_chunk_recv(first) + host.register_chunk_recv(second) + for req_id in ("req-1", "req-2", "req-1"): + assert not host._poll_single_request(req_id) + assert host._omni_connector.get.call_args_list == [ + unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}), + unittest.mock.call("0", "1", "ext-req-2_0_0", {"source_host": "producer-b", "source_port": 51051}), + unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}), + ] + host.cleanup_finished_request("req-1") + assert "req-1" not in host._pending_load_reqs + assert host._pending_load_reqs["req-2"] is second + host.register_chunk_recv(_make_request("req-1", "reused")) + assert not host._poll_single_request("req-1") + host._omni_connector.get.assert_called_with("0", "1", "reused_0_0") + finally: + host.shutdown_omni_connectors() + + +@pytest.mark.parametrize("method", ["_recv_full_payload_result", "_recv_async_chunk_result"]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_split_payload_sender_metadata_respects_tp_leader(method, rank): + host = MixinHost() + connector = MagicMock() + metadata = {"source_host": "producer", "source_port": 50051} + with ( + patch.object(host, "_get_local_tp_group", return_value=_FakeTPGroup(world_size=2, rank_in_group=rank)), + patch.object(host, "is_data_transfer_rank", return_value=rank == 0), + ): + result = getattr(host, method)(connector, "0", "1", "req_0_0", metadata) + if rank == 0: + connector.get.assert_called_once_with("0", "1", "req_0_0", metadata) + assert result is connector.get.return_value + else: + connector.get.assert_not_called() + assert result is None + + class _FailingConnector: """Connector whose put() fails a configurable number of times.""" diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index fe439457d..fd7617144 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -67,6 +67,8 @@ _PUBLIC_OMNI = ( "VLLM_OMNI_EVENT_DRIVEN_ORCH", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", + "VLLM_OMNI_NIXL_LEASE_S", + "VLLM_OMNI_NIXL_XFER_TIMEOUT_S", "VLLM_OMNI_ORCH_MONITOR_PATH", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_TTL", diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 93f40daa0..3d516f5e3 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -121,6 +121,7 @@ class OmniSchedulingCoordinator: OmniChunkRecvHandle( request_id=request.request_id, external_req_id=getattr(request, "external_req_id", None), + payload_sender_info=getattr(request, "payload_sender_info", None), ) ) elif request.status == RequestStatus.WAITING_FOR_INPUT: @@ -134,6 +135,7 @@ class OmniSchedulingCoordinator: OmniChunkRecvHandle( request_id=request.request_id, external_req_id=getattr(request, "external_req_id", None), + payload_sender_info=getattr(request, "payload_sender_info", None), ) ) if to_remove: diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index b73fbd6aa..deaf44df8 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -105,6 +105,7 @@ class OmniChunkRecvHandle: request_id: str external_req_id: str | None = None + payload_sender_info: dict[str, object] | None = None @dataclass diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index 3506f02f6..40374cd6d 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -7,6 +7,7 @@ from .omni_connectors import ( MooncakeStoreConnector, MooncakeTransferEngineConnector, MoriTransferEngineConnector, + NixlConnector, OmniConnectorBase, OmniConnectorFactory, OmniTransferConfig, @@ -27,6 +28,7 @@ __all__ = [ "MooncakeStoreConnector", "MooncakeTransferEngineConnector", "MoriTransferEngineConnector", + "NixlConnector", "SharedMemoryConnector", "YuanrongConnector", "YuanrongTransferEngineConnector", diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index 0cee596bb..e46142b2a 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -22,6 +22,11 @@ try: from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector except ImportError: MoriTransferEngineConnector = None # RDMA deps (msgspec/zmq/mori) not installed + +try: + from .connectors.nixl_connector import NixlConnector +except ImportError: + NixlConnector = None # NIXL deps not installed from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec, OmniTransferConfig from .utils.initialization import ( @@ -50,6 +55,7 @@ __all__ = [ "MooncakeStoreConnector", "MooncakeTransferEngineConnector", "MoriTransferEngineConnector", + "NixlConnector", "SharedMemoryConnector", "YuanrongConnector", "YuanrongTransferEngineConnector", diff --git a/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py b/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py new file mode 100644 index 000000000..7dda744f6 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py @@ -0,0 +1,1044 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +import os +import socket +import sys +import threading +import time +import uuid +from dataclasses import dataclass, field +from typing import Any + +import msgspec +import torch +import zmq + +from ..utils.logging import get_connector_logger +from .base import OmniConnectorBase + +logger = get_connector_logger(__name__) + +_SCHEMA_VERSION = 1 +_KIND_TENSORS = "tensors" +_KIND_STRUCTURED = "structured" +_KIND_OBJECT = "object" +_INIT_AGENT = "NIXL_INIT_AGENT" +_TENSOR_MARKER = "__nixl_tensor_index__" +_TUPLE_MARKER = "__nixl_tuple__" + +# Handshake control-plane messages. Consumers claim connector-managed payloads +# before READ, even with out-of-band metadata, and acknowledge completion here. +_GET_META_MSG = b"nixl_get_meta" +_XFER_DONE_MSG = b"nixl_xfer_done" +_META_NOT_FOUND = b"nixl_meta_not_found" +_ACK = b"nixl_ack" + +# NIXL has no remote READ cancellation primitive. A closed producer with an +# abandoned claim must keep both its agent and allocations alive until process +# exit; dropping the last Python reference is not a safe cancellation policy. +_RETAINED_PRODUCERS: set[Any] = set() + + +@dataclass +class _DeferredTransfer: + tensors: list[torch.Tensor] + registrations: list[Any] = field(default_factory=list) + dlists: list[Any] = field(default_factory=list) + handles: list[Any] = field(default_factory=list) + remote_agent: Any = None + source_key: str | None = None + source_metadata: dict[str, Any] | None = None + + +@dataclass +class _PendingPayload: + tensors: list[torch.Tensor] + registrations: list[Any] + deadline: float + generation: str = field(default_factory=lambda: uuid.uuid4().hex) + claims: set[str] = field(default_factory=set) + + +class NixlConnector(OmniConnectorBase): + """OmniConnector backed by vLLM's native NIXL wrapper. + + This connector intentionally depends on vLLM's optional NIXL integration + (``vllm.distributed.nixl_utils``). It transfers raw tensor payloads + directly through NIXL READ operations. Non-tensor Python + payloads are serialized with OmniSerializer, packed into a uint8 CPU tensor, + and moved through the same NIXL path. + + A NIXL READ needs the producer's agent metadata and memory descriptors, + which ``put`` returns to its caller. Consumers use the ZMQ control plane to + claim connector-managed payloads before READ and acknowledge completion, + including when metadata is forwarded directly. Producers without a fixed + ``zmq_port`` use an ephemeral listener advertised in that metadata. Callers + passing ``metadata=None`` need a configured or request-specific producer + endpoint to fetch and claim metadata by key. + """ + + supports_raw_data: bool = True + + def __init__(self, config: dict[str, Any]): + self.config = dict(config or {}) + self.stage_id = int(self.config.get("stage_id", 0)) + self._closed = True + self._registered_descs: list[Any] = [] + self._pending: dict[str, _PendingPayload] = {} + self._published: dict[str, dict[str, Any]] = {} + self._state_lock = threading.RLock() + self._remote_agents: list[str] = [] + self._metrics: dict[str, int] = { + "puts": 0, + "gets": 0, + "errors": 0, + "bytes_transferred": 0, + } + + from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config + + if NixlWrapper is None: + raise RuntimeError("NIXL is not available. Install the optional nixl/rixl package to use NixlConnector.") + + backends = self.config.get("backends", ["UCX"]) + self._backends = list(backends) if isinstance(backends, (list, tuple)) else [str(backends)] + if nixl_agent_config is None: + agent_config = None + else: + from nixl import nixl_thread_sync_t + + config_module = sys.modules.get(nixl_agent_config.__module__) + sync_type = getattr(config_module, "nixl_thread_sync_t", nixl_thread_sync_t) + sync_mode = sync_type.NIXL_THREAD_SYNC_STRICT + non_ucx_backends = [backend for backend in self._backends if backend != "UCX"] + if non_ucx_backends: + agent_config = nixl_agent_config( + backends=self._backends, + capture_telemetry=False, + sync_mode=sync_mode, + ) + else: + num_threads = int(self.config.get("num_threads", 4)) + agent_config = nixl_agent_config( + num_threads=num_threads, + capture_telemetry=False, + sync_mode=sync_mode, + ) + + self._agent = NixlWrapper(str(self.config.get("agent_name", uuid.uuid4())), agent_config) + self._receive_device = self._parse_device(self.config.get("receive_device")) + self._default_memory_type = self.config.get("memory_type") + # The lease bounds how long a ``put`` payload stays registered while it + # waits to be read. A consumer stage may legitimately queue for a long + # time (e.g. video diffusion with concurrency above the replica count), + # so the default must be far larger than any plausible queueing delay. + # The environment variables intentionally take precedence over the + # deployment YAML, which ships a fixed 300s lease that is too short for + # long-queueing pipelines and cannot be overridden per deployment. + self._lease_seconds = float( + os.environ.get("VLLM_OMNI_NIXL_LEASE_S") or self.config.get("lease_seconds", 3600.0) + ) + self._transfer_timeout_s = float( + os.environ.get("VLLM_OMNI_NIXL_XFER_TIMEOUT_S") or self.config.get("transfer_timeout_s", 300.0) + ) + self._poll_interval_s = float(self.config.get("poll_interval_s", 0.001)) + self._closed = False + self._init_handshake() + self._lease_wakeup = threading.Event() + self._lease_thread: threading.Thread | None = None + if self._serving_handshake: + self._lease_thread = threading.Thread( + target=self._lease_reaper_loop, + name="nixl-lease-reaper", + daemon=True, + ) + self._lease_thread.start() + self._transfer_wakeup = threading.Event() + self._deferred_transfers: list[_DeferredTransfer] = [] + self._transfer_thread: threading.Thread | None = None + if self._role != "sender": + self._transfer_thread = threading.Thread( + target=self._transfer_reaper_loop, + name="nixl-transfer-reaper", + daemon=True, + ) + self._transfer_thread.start() + + def _init_handshake(self) -> None: + """Set up the ZMQ metadata and payload-ownership control plane. + + Producers listen on a configured or ephemeral port, including for + directly forwarded metadata. Receive-only connectors dial the producer; + intermediate stages can also listen for their outgoing payloads. + """ + role = self.config.get("role") + self._role = str(role).lower() if role else None + # Even out-of-band metadata needs an atomic claim before READ. An + # ephemeral listener supplies that ownership channel when no fixed + # data-plane metadata port was configured. + self._zmq_port = self.config.get("zmq_port", 0 if self._role != "receiver" else None) + self._sender_host = self.config.get("sender_host") + self._sender_zmq_port = self.config.get("sender_zmq_port") + self._handshake_timeout_ms = int(self.config.get("handshake_timeout_ms", 5000)) + self._metadata_query_timeout_ms = int( + self.config.get("metadata_query_timeout_ms", min(self._handshake_timeout_ms, 10)) + ) + self._handshake_max_wait_s = float(self.config.get("handshake_max_wait_s", 60.0)) + self._handshake_retry_s = float(self.config.get("handshake_retry_s", 0.05)) + self._zmq_ctx: zmq.Context | None = None + self._req_local = threading.local() + self._listener_thread: threading.Thread | None = None + self._listener_ready = threading.Event() + self._stop_event = threading.Event() + self._bind_error: BaseException | None = None + self.host: str | None = None + self._serving_handshake = False + + self._handshake_enabled = self._zmq_port is not None or self._sender_host is not None + if not self._handshake_enabled: + return + + self._zmq_ctx = zmq.Context() + # A receiver only dials out, so it never needs a port of its own. + if self._zmq_port is None: + return + + host_value = str(self.config.get("host", "auto")) + self.host = ( + self._get_local_ip() + if host_value.lower() == "auto" or host_value in {"", "*", "0.0.0.0", "::"} + else host_value + ) + self._zmq_port = int(self._zmq_port) + self._listener_thread = threading.Thread( + target=self._handshake_listener_loop, name="nixl-handshake", daemon=True + ) + self._listener_thread.start() + self._listener_ready.wait(timeout=5.0) + if self._bind_error is not None: + raise RuntimeError( + f"NixlConnector failed to bind handshake socket on {self.host}:{self._zmq_port}" + ) from self._bind_error + self._serving_handshake = True + logger.info("NixlConnector handshake listener bound on %s:%s", self.host, self._zmq_port) + + def put( + self, + from_stage: str, + to_stage: str, + put_key: str, + data: Any, + ) -> tuple[bool, int, dict[str, Any] | None]: + if self._closed or getattr(self, "_closing", False): + raise RuntimeError("Cannot put data: NixlConnector is closed") + + try: + self._cleanup_expired_pending() + kind = _KIND_TENSORS + if self._is_tensor_payload(data): + tensors, tensor_specs = self._normalize_tensor_payload(data) + elif self._contains_tensor(data): + skeleton, payload_tensors = self._extract_tensor_leaves(data) + header = self.serialize_obj(skeleton) + header_tensor = torch.frombuffer(header, dtype=torch.uint8).clone().contiguous() + tensors, tensor_specs = self._normalize_tensor_payload([header_tensor, *payload_tensors]) + kind = _KIND_STRUCTURED + else: + payload = self.serialize_obj(data) + tensor = torch.frombuffer(payload, dtype=torch.uint8).clone().contiguous() + tensors, tensor_specs = self._normalize_tensor_payload(tensor) + kind = _KIND_OBJECT + + grouped_tensors: dict[str, list[tuple[int, torch.Tensor]]] = {} + for tensor_index, tensor in enumerate(tensors): + # Preserve spec/skeleton slots, but never register empty DMA regions. + if tensor.numel() == 0: + continue + memory_type = self._resolve_memory_type(tensor) + grouped_tensors.setdefault(memory_type, []).append((tensor_index, tensor)) + + descriptor_groups = [] + registered_descs = [] + try: + for memory_type, indexed_tensors in grouped_tensors.items(): + tensor_indices = [tensor_index for tensor_index, _ in indexed_tensors] + regions = [self._tensor_region(tensor) for _, tensor in indexed_tensors] + reg_descs = self._agent.get_reg_descs(regions, memory_type) + self._agent.register_memory(reg_descs, backends=self._backends) + registered_descs.append(reg_descs) + descriptor_groups.append( + { + "memory_type": memory_type, + "tensor_indices": tensor_indices, + "regions": regions, + } + ) + except Exception: + for reg_descs in registered_descs: + self._safe_call(self._agent.deregister_memory, reg_descs) + raise + with self._state_lock: + self._registered_descs.extend(registered_descs) + + size = sum(spec["size"] for spec in tensor_specs) + metadata = { + "schema_version": _SCHEMA_VERSION, + "kind": kind, + "agent_metadata": self._agent.get_agent_metadata(), + "descriptor_groups": descriptor_groups, + "tensor_specs": tensor_specs, + "size": size, + } + if self._serving_handshake: + metadata["sender_host"] = self.host + metadata["sender_zmq_port"] = self._zmq_port + with self._state_lock: + previous = self._pending.get(put_key) + if previous is not None and previous.claims: + for descs in registered_descs: + self._safe_call(self._agent.deregister_memory, descs) + self._registered_descs.remove(descs) + raise RuntimeError(f"Cannot replace claimed NIXL payload {put_key}") + previous = self._take_pending(put_key) + if previous is not None: + self._release_pending(previous) + pending = _PendingPayload( + tensors=tensors, + registrations=registered_descs, + deadline=time.monotonic() + self._lease_seconds, + ) + metadata["generation"] = pending.generation + self._pending[put_key] = pending + if self._serving_handshake: + self._published[put_key] = metadata + self._lease_wakeup.set() + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += size + logger.debug("NixlConnector put %s->%s key=%s size=%d", from_stage, to_stage, put_key, size) + return True, size, metadata + except Exception: + self._metrics["errors"] += 1 + logger.error("NixlConnector put failed for %s", put_key, exc_info=True) + return False, 0, None + + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata: dict[str, Any] | None = None, + ) -> tuple[Any, int] | None: + if self._closed or getattr(self, "_closing", False): + raise RuntimeError("Cannot get data: NixlConnector is closed") + + remote_agent = None + local_reg_descs_list = [] + dlist_handles = [] + xfer_handles = [] + source_metadata = None + try: + metadata = self._resolve_metadata(get_key, metadata) + if not isinstance(metadata, dict) or metadata.get("schema_version") != _SCHEMA_VERSION: + logger.error("NixlConnector get has invalid metadata for %s", get_key) + return None + + source_metadata = metadata + + tensor_specs = metadata.get("tensor_specs") + if not isinstance(tensor_specs, list): + raise RuntimeError(f"Invalid NIXL metadata for {get_key}: missing tensor_specs") + descriptor_groups = self._validated_descriptor_groups(metadata, tensor_specs) + + local_tensors = [self._allocate_tensor_from_spec(spec, metadata.get("kind")) for spec in tensor_specs] + if descriptor_groups: + remote_agent = self._agent.add_remote_agent(metadata["agent_metadata"]) + self._remote_agents.append(remote_agent) + for descriptor_group in descriptor_groups: + remote_memory_type = descriptor_group["memory_type"] + indexed_regions = list( + zip(descriptor_group["tensor_indices"], descriptor_group["regions"], strict=True) + ) + local_groups: dict[str, list[tuple[int, Any]]] = {} + for tensor_index, remote_region in indexed_regions: + local_memory_type = self._resolve_memory_type(local_tensors[tensor_index]) + local_groups.setdefault(local_memory_type, []).append((tensor_index, remote_region)) + + for local_memory_type, entries in local_groups.items(): + group_tensors = [local_tensors[tensor_index] for tensor_index, _ in entries] + local_regions = [self._tensor_region(tensor) for tensor in group_tensors] + local_reg_descs = self._agent.get_reg_descs(local_regions, local_memory_type) + self._agent.register_memory(local_reg_descs, backends=self._backends) + local_reg_descs_list.append(local_reg_descs) + + remote_regions = [tuple(region)[:3] for _, region in entries] + remote_descs = self._agent.get_xfer_descs(remote_regions, remote_memory_type) + local_descs = self._agent.get_xfer_descs( + [tuple(region)[:3] for region in local_regions], local_memory_type + ) + remote_dlist = self._agent.prep_xfer_dlist(remote_agent, remote_descs) + local_dlist = self._agent.prep_xfer_dlist(_INIT_AGENT, local_descs) + dlist_handles.extend([local_dlist, remote_dlist]) + + desc_ids = list(range(len(entries))) + xfer_handle = self._agent.make_prepped_xfer( + "READ", + local_dlist, + desc_ids, + remote_dlist, + desc_ids, + ) + xfer_handles.append(xfer_handle) + self._agent.transfer(xfer_handle) + + for xfer_handle in xfer_handles: + self._wait_for_transfer(xfer_handle, get_key) + for xfer_handle in xfer_handles: + self._agent.release_xfer_handle(xfer_handle) + xfer_handles.clear() + + size = int(metadata.get("size", sum(spec.get("size", 0) for spec in tensor_specs))) + if metadata.get("kind") == _KIND_OBJECT: + raw = local_tensors[0].detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + payload = self.deserialize_obj(raw) + elif metadata.get("kind") == _KIND_STRUCTURED: + raw = local_tensors[0].detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + skeleton = self.deserialize_obj(raw) + payload = self._restore_tensor_leaves(skeleton, local_tensors[1:]) + else: + payload = local_tensors[0] if len(local_tensors) == 1 else local_tensors + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += size + logger.debug("NixlConnector get %s->%s key=%s size=%d", from_stage, to_stage, get_key, size) + return payload, size + except Exception: + if xfer_handles and self._transfers_may_be_active(xfer_handles): + self._defer_transfer( + _DeferredTransfer( + tensors=local_tensors, + registrations=local_reg_descs_list, + dlists=dlist_handles, + handles=xfer_handles, + remote_agent=remote_agent, + source_key=get_key, + source_metadata=source_metadata, + ) + ) + local_tensors = [] + local_reg_descs_list = [] + dlist_handles = [] + xfer_handles = [] + remote_agent = None + source_metadata = None + self._metrics["errors"] += 1 + logger.error("NixlConnector get failed for %s", get_key, exc_info=True) + return None + finally: + for xfer_handle in xfer_handles: + self._safe_call(self._agent.release_xfer_handle, xfer_handle) + for dlist_handle in dlist_handles: + self._safe_call(self._agent.release_dlist_handle, dlist_handle) + if remote_agent is not None: + self._safe_call(self._agent.remove_remote_agent, remote_agent) + if remote_agent in self._remote_agents: + self._remote_agents.remove(remote_agent) + for local_reg_descs in local_reg_descs_list: + self._safe_call(self._agent.deregister_memory, local_reg_descs) + if source_metadata is not None: + self._notify_transfer_done(get_key, source_metadata) + + def cleanup(self, request_id: str) -> None: + pending = self._take_pending(request_id) + if pending is None: + return + self._release_pending(pending) + + def _release_pending(self, pending: _PendingPayload) -> None: + for reg_descs in pending.registrations: + self._safe_call(self._agent.deregister_memory, reg_descs) + with self._state_lock: + if reg_descs in self._registered_descs: + self._registered_descs.remove(reg_descs) + + def _take_pending(self, request_id: str, expected: _PendingPayload | None = None) -> _PendingPayload | None: + with self._state_lock: + pending = self._pending.get(request_id) + if pending is None or (expected is not None and pending is not expected): + return None + if pending.claims: + return None + self._published.pop(request_id, None) + return self._pending.pop(request_id) + + def close(self) -> None: + if self._closed: + return + self._closed = True + # Stop accepting puts/gets, but keep serving completion ACKs for outstanding + # remote READs. A later close() call can finish teardown once drained. + for request_id in list(self._pending): + self.cleanup(request_id) + with self._state_lock: + if any(pending.claims for pending in self._pending.values()): + _RETAINED_PRODUCERS.add(self) + self._closed = False + self._closing = True + logger.warning("NIXL close deferred: remote READ claims still own source allocations") + return + _RETAINED_PRODUCERS.discard(self) + self._stop_event.set() + self._lease_wakeup.set() + self._transfer_wakeup.set() + if self._listener_thread is not None: + self._listener_thread.join(timeout=5.0) + self._listener_thread = None + if self._lease_thread is not None: + self._lease_thread.join(timeout=5.0) + self._lease_thread = None + if self._transfer_thread is not None: + self._transfer_thread.join(timeout=5.0) + self._transfer_thread = None + while self._deferred_transfers: + self._reap_deferred_transfers() + if self._deferred_transfers: + time.sleep(max(self._poll_interval_s, 0.01)) + if self._zmq_ctx is not None: + # destroy() rather than term(): REQ sockets live in thread-local + # caches this thread cannot reach, and term() blocks until every + # socket in the context is closed. + self._safe_call(self._zmq_ctx.destroy, 0) + self._zmq_ctx = None + for request_id in list(self._pending): + self.cleanup(request_id) + deferred_agents = { + transfer.remote_agent for transfer in self._deferred_transfers if transfer.remote_agent is not None + } + for agent_name in [agent for agent in self._remote_agents if agent not in deferred_agents]: + self._safe_call(self._agent.remove_remote_agent, agent_name) + self._remote_agents.remove(agent_name) + for reg_descs in list(self._registered_descs): + self._safe_call(self._agent.deregister_memory, reg_descs) + self._registered_descs.clear() + + def health(self) -> dict[str, Any]: + return { + "status": "unhealthy" if self._closed else "healthy", + "pending_requests": len(self._pending), + **self._metrics, + } + + def get_connection_info(self) -> dict[str, Any]: + """Endpoint a consumer needs to reach this producer's handshake socket.""" + return {"host": self.host, "zmq_port": self._zmq_port} + + def update_sender_info(self, sender_host: str, sender_zmq_port: int) -> None: + """Register the producer handshake endpoint on the consumer side. + + Used when the endpoint is only known after both stages have started and + therefore cannot be baked into the connector config. + """ + self._sender_host = sender_host + self._sender_zmq_port = int(sender_zmq_port) + if self._zmq_ctx is None: + self._zmq_ctx = zmq.Context() + self._handshake_enabled = True + + @staticmethod + def _metadata_endpoint(metadata: dict[str, Any] | None) -> tuple[str, int] | None: + if not isinstance(metadata, dict): + return None + host = metadata.get("sender_host") + port = metadata.get("sender_zmq_port") + if not host or not port: + host = metadata.get("source_host") + port = metadata.get("source_port") + if not host or not port: + return None + return str(host), int(port) + + def _resolve_metadata(self, get_key: str, metadata: dict[str, Any] | None) -> dict[str, Any] | None: + """Return complete NIXL transfer metadata for ``get_key``. + + Connector-managed metadata, including directly forwarded ``put`` results, + requires a handshake claim at the supplied or configured producer endpoint. + Only legacy externally owned metadata without a generation bypasses it. + """ + direct = isinstance(metadata, dict) and metadata.get("schema_version") == _SCHEMA_VERSION + if direct and "generation" not in metadata: + # Legacy externally owned metadata has no connector-managed lease. + return metadata + + endpoint = self._metadata_endpoint(metadata) + if endpoint is None and self._sender_host and self._sender_zmq_port: + endpoint = (str(self._sender_host), int(self._sender_zmq_port)) + if endpoint is None: + logger.error( + "NixlConnector get(%s) received no usable metadata and no producer handshake " + "endpoint is configured. Set sender_host/sender_zmq_port on the consumer, " + "pass source_host/source_port for the producer rank, or forward the metadata " + "returned by put().", + get_key, + ) + return None + if self._zmq_ctx is None: + with self._state_lock: + if self._zmq_ctx is None: + self._zmq_ctx = zmq.Context() + + if direct: + return self._query_metadata_at(get_key, *endpoint, generation=metadata["generation"]) + return self._query_metadata_at(get_key, *endpoint) + + def _query_metadata_at( + self, get_key: str, host: str, port: int, *, generation: str | None = None + ) -> dict[str, Any] | None: + """Fetch transfer metadata for ``get_key`` from a producer's ROUTER socket. + + Each call performs one bounded query so a missing key cannot starve + other requests in the shared receive loop. The caller retries later. + """ + zmq_addr = f"tcp://{host}:{port}" + request = _GET_META_MSG + msgspec.msgpack.encode( + {"key": get_key, "generation": generation, "claim_id": uuid.uuid4().hex} + ) + sock = self._get_req_socket(zmq_addr, self._metadata_query_timeout_ms) + try: + sock.send(request) + reply = sock.recv() + except Exception: + self._invalidate_req_socket(zmq_addr) + logger.debug("NixlConnector handshake query to %s failed for %s", zmq_addr, get_key, exc_info=True) + return None + if reply == _META_NOT_FOUND: + return None + return msgspec.msgpack.decode(reply) + + def _notify_transfer_done(self, get_key: str, metadata: dict[str, Any]) -> None: + """Tell the producer its buffer is drained so it can deregister now. + + Without this the producer would hold the registration until the lease + expires, which for long-lived stages means unbounded growth. + """ + endpoint = self._metadata_endpoint(metadata) + if endpoint is None or self._zmq_ctx is None: + return + host, port = endpoint + zmq_addr = f"tcp://{host}:{port}" + sock = self._get_req_socket(zmq_addr) + try: + sock.send( + _XFER_DONE_MSG + + msgspec.msgpack.encode( + {"key": get_key, "generation": metadata.get("generation"), "claim_id": metadata.get("claim_id")} + ) + ) + sock.recv() + except Exception: + self._invalidate_req_socket(zmq_addr) + logger.debug("NixlConnector failed to notify completion for %s", get_key, exc_info=True) + + def _handshake_listener_loop(self) -> None: + router = self._zmq_ctx.socket(zmq.ROUTER) + try: + if self._zmq_port == 0: + self._zmq_port = router.bind_to_random_port(f"tcp://{self.host}") + else: + router.bind(f"tcp://{self.host}:{self._zmq_port}") + except zmq.ZMQError as exc: + logger.error("NixlConnector handshake bind failed on %s:%s: %s", self.host, self._zmq_port, exc) + self._bind_error = exc + self._listener_ready.set() + router.close(linger=0) + return + self._listener_ready.set() + + poller = zmq.Poller() + poller.register(router, zmq.POLLIN) + try: + while not self._stop_event.is_set(): + try: + if not dict(poller.poll(500)): + self._cleanup_expired_pending() + continue + identity, _, payload = router.recv_multipart() + router.send_multipart([identity, b"", self._handle_handshake_message(payload)]) + except zmq.ContextTerminated: + break + except Exception: + logger.debug("NixlConnector handshake listener error", exc_info=True) + finally: + self._safe_call(router.close, 0) + + def _handle_handshake_message(self, payload: bytes) -> bytes: + if payload.startswith(_GET_META_MSG): + request = msgspec.msgpack.decode(payload[len(_GET_META_MSG) :]) + key = request.get("key") + claim = request.get("claim_id") + with self._state_lock: + pending = self._pending.get(key) + metadata = self._published.get(key) + if ( + pending is None + or metadata is None + or not claim + or request.get("generation") not in (None, pending.generation) + or (not pending.claims and time.monotonic() >= pending.deadline) + or self._closed + or getattr(self, "_closing", False) + ): + return _META_NOT_FOUND + # Publish ownership before descriptors can leave this lock. + # Failed/lost replies retain the claim conservatively. + reply = msgspec.msgpack.encode({**metadata, "claim_id": claim}) + pending.claims.add(claim) + return reply + if payload.startswith(_XFER_DONE_MSG): + request = msgspec.msgpack.decode(payload[len(_XFER_DONE_MSG) :]) + key = request.get("key") + with self._state_lock: + pending = self._pending.get(key) + if ( + pending is not None + and request.get("generation") == pending.generation + and request.get("claim_id") in pending.claims + ): + pending.claims.remove(request["claim_id"]) + self.cleanup(key) + return _ACK + logger.warning("NixlConnector handshake received an unknown message") + return _META_NOT_FOUND + + def _get_req_socket(self, zmq_addr: str, timeout_ms: int | None = None) -> zmq.Socket: + """Return a thread-local REQ socket so concurrent calls never interleave.""" + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache is None: + cache = {} + self._req_local.cache = cache + sock = cache.get(zmq_addr) + if sock is None: + sock = self._zmq_ctx.socket(zmq.REQ) + sock.connect(zmq_addr) + cache[zmq_addr] = sock + timeout_ms = self._handshake_timeout_ms if timeout_ms is None else timeout_ms + sock.setsockopt(zmq.SNDTIMEO, timeout_ms) + sock.setsockopt(zmq.RCVTIMEO, timeout_ms) + return sock + + def _invalidate_req_socket(self, zmq_addr: str) -> None: + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache is None: + return + sock = cache.pop(zmq_addr, None) + if sock is not None: + self._safe_call(sock.close, 0) + + @staticmethod + def _get_local_ip() -> str: + """Resolve the externally routable local address for the handshake bind.""" + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("8.8.8.8", 80)) + return probe.getsockname()[0] + except Exception: + try: + return socket.gethostbyname(socket.gethostname()) + except Exception: + logger.warning("NixlConnector could not resolve a local IP; binding loopback") + return "127.0.0.1" + + def _wait_for_transfer(self, handle: int, request_id: str) -> None: + deadline = time.monotonic() + self._transfer_timeout_s + while True: + state = self._agent.check_xfer_state(handle) + if state == "DONE": + return + if state != "PROC": + raise RuntimeError(f"NIXL transfer for {request_id} failed with state={state}") + if time.monotonic() >= deadline: + raise TimeoutError(f"NIXL transfer for {request_id} timed out") + time.sleep(self._poll_interval_s) + + def _transfers_may_be_active(self, handles: list[Any]) -> bool: + try: + return any(self._agent.check_xfer_state(handle) not in {"DONE", "ERR"} for handle in handles) + except Exception: + return True + + def _cleanup_expired_pending(self) -> None: + now = time.monotonic() + with self._state_lock: + expired = [ + (key, pending, sum(t.numel() * t.element_size() for t in pending.tensors)) + for key, pending in self._pending.items() + if not pending.claims and now >= pending.deadline + ] + for request_id, pending, size in expired: + logger.warning( + "NixlConnector lease expired for request %s after %.0fs; its %d bytes are " + "unclaimed and can safely be reclaimed.", + request_id, + self._lease_seconds, + size, + ) + claimed = self._take_pending(request_id, expected=pending) + if claimed is not None: + self._release_pending(claimed) + + def _lease_reaper_loop(self) -> None: + while not self._stop_event.is_set(): + now = time.monotonic() + with self._state_lock: + next_deadline = min( + (pending.deadline for pending in self._pending.values() if not pending.claims), default=None + ) + if next_deadline is None: + timeout = None + else: + timeout = max(0.0, next_deadline - now) + self._lease_wakeup.wait(timeout=timeout) + self._lease_wakeup.clear() + if not self._stop_event.is_set(): + self._cleanup_expired_pending() + + def _defer_transfer(self, transfer: _DeferredTransfer) -> None: + with self._state_lock: + self._deferred_transfers.append(transfer) + self._transfer_wakeup.set() + + def _transfer_reaper_loop(self) -> None: + while not self._stop_event.is_set(): + self._transfer_wakeup.wait(timeout=max(self._poll_interval_s, 0.01)) + self._transfer_wakeup.clear() + self._reap_deferred_transfers() + + def _reap_deferred_transfers(self) -> None: + with self._state_lock: + transfers = list(self._deferred_transfers) + for transfer in transfers: + try: + states = [self._agent.check_xfer_state(handle) for handle in transfer.handles] + except Exception: + logger.debug("Failed to poll deferred NIXL transfer", exc_info=True) + continue + if any(state not in {"DONE", "ERR"} for state in states): + continue + self._release_deferred_transfer(transfer) + if ( + not transfer.handles + and not transfer.dlists + and not transfer.registrations + and transfer.remote_agent is None + ): + transfer.tensors.clear() + if transfer.source_metadata is not None: + self._notify_transfer_done(transfer.source_key, transfer.source_metadata) + transfer.source_metadata = None + with self._state_lock: + if transfer in self._deferred_transfers: + self._deferred_transfers.remove(transfer) + + def _release_deferred_transfer(self, transfer: _DeferredTransfer) -> None: + transfer.handles = self._release_owned_resources(self._agent.release_xfer_handle, transfer.handles) + transfer.dlists = self._release_owned_resources(self._agent.release_dlist_handle, transfer.dlists) + if transfer.remote_agent is not None: + try: + self._agent.remove_remote_agent(transfer.remote_agent) + except Exception: + logger.debug("Failed to remove deferred NIXL remote agent", exc_info=True) + else: + if transfer.remote_agent in self._remote_agents: + self._remote_agents.remove(transfer.remote_agent) + transfer.remote_agent = None + transfer.registrations = self._release_owned_resources(self._agent.deregister_memory, transfer.registrations) + + @staticmethod + def _release_owned_resources(release: Any, resources: list[Any]) -> list[Any]: + remaining = [] + for resource in resources: + try: + release(resource) + except Exception: + logger.debug("Failed to release deferred NIXL resource", exc_info=True) + remaining.append(resource) + return remaining + + def _resolve_memory_type(self, tensor: torch.Tensor) -> str: + if self._default_memory_type is not None: + return str(self._default_memory_type) + if tensor.device.type == "cpu": + return "DRAM" + return "VRAM" + + @staticmethod + def _validated_descriptor_groups( + metadata: dict[str, Any], tensor_specs: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + # Derive emptiness from shape/dtype rather than the advertised size. + sizes = [] + for spec in tensor_specs: + if not isinstance(spec, dict): + raise RuntimeError("Invalid NIXL metadata: tensor spec must be a mapping") + shape = spec.get("shape") + if not isinstance(shape, list) or any(type(dim) is not int or dim < 0 for dim in shape): + raise RuntimeError("Invalid NIXL metadata: invalid tensor shape") + dtype = getattr(torch, str(spec.get("dtype", "")).removeprefix("torch."), None) + if not isinstance(dtype, torch.dtype): + raise RuntimeError("Invalid NIXL metadata: invalid tensor dtype") + size = math.prod(shape) * dtype.itemsize + if type(spec.get("size")) is not int or spec["size"] != size: + raise RuntimeError("Invalid NIXL metadata: tensor size does not match shape/dtype") + sizes.append(size) + groups = metadata.get("descriptor_groups") + if not isinstance(groups, list): + raise RuntimeError("Invalid NIXL metadata: missing descriptor_groups") + + seen_indices = [] + for group in groups: + if not isinstance(group, dict): + raise RuntimeError("Invalid NIXL metadata: descriptor group must be a mapping") + indices = group.get("tensor_indices") + regions = group.get("regions") + if not isinstance(indices, list) or not isinstance(regions, list) or len(indices) != len(regions): + raise RuntimeError("Invalid NIXL metadata: tensor_indices and regions must have equal lengths") + if not indices: + raise RuntimeError("Invalid NIXL metadata: empty descriptor group") + if not group.get("memory_type"): + raise RuntimeError("Invalid NIXL metadata: descriptor group is missing memory_type") + if any(type(index) is not int or not 0 <= index < len(sizes) for index in indices): + raise RuntimeError("Invalid NIXL metadata: tensor indices must form an exact partition") + for index, region in zip(indices, regions, strict=True): + if ( + not isinstance(region, (list, tuple)) + or len(region) not in (3, 4) + or type(region[0]) is not int + or region[0] <= 0 + or type(region[1]) is not int + or region[1] <= 0 + or region[1] != sizes[index] + or type(region[2]) is not int + or region[2] < 0 + ): + raise RuntimeError("Invalid NIXL metadata: invalid tensor region or byte size") + seen_indices.extend(indices) + + if sorted(seen_indices) != [index for index, size in enumerate(sizes) if size > 0]: + raise RuntimeError("Invalid NIXL metadata: tensor indices must form an exact partition") + return groups + + @staticmethod + def _tensor_region(tensor: torch.Tensor) -> tuple[int, int, int, str]: + device_id = max(tensor.get_device(), 0) if tensor.device.type != "cpu" else 0 + return (tensor.data_ptr(), tensor.numel() * tensor.element_size(), device_id, "") + + @staticmethod + def _is_tensor_payload(payload: Any) -> bool: + return isinstance(payload, torch.Tensor) or ( + isinstance(payload, (list, tuple)) + and bool(payload) + and all(isinstance(item, torch.Tensor) for item in payload) + ) + + @classmethod + def _contains_tensor(cls, payload: Any) -> bool: + if isinstance(payload, torch.Tensor): + return True + if isinstance(payload, msgspec.Struct): + payload = msgspec.structs.asdict(payload) + if isinstance(payload, dict): + return any(cls._contains_tensor(value) for value in payload.values()) + if isinstance(payload, (list, tuple)): + return any(cls._contains_tensor(value) for value in payload) + return False + + @classmethod + def _extract_tensor_leaves(cls, payload: Any) -> tuple[Any, list[torch.Tensor]]: + tensors: list[torch.Tensor] = [] + + def visit(value: Any) -> Any: + if isinstance(value, torch.Tensor): + index = len(tensors) + tensors.append(value) + return {_TENSOR_MARKER: index} + if isinstance(value, msgspec.Struct): + value = msgspec.structs.asdict(value) + if isinstance(value, dict): + return {key: visit(item) for key, item in value.items()} + if isinstance(value, list): + return [visit(item) for item in value] + if isinstance(value, tuple): + return {_TUPLE_MARKER: [visit(item) for item in value]} + return value + + return visit(payload), tensors + + @classmethod + def _restore_tensor_leaves(cls, skeleton: Any, tensors: list[torch.Tensor]) -> Any: + if isinstance(skeleton, dict): + if set(skeleton) == {_TENSOR_MARKER}: + return tensors[int(skeleton[_TENSOR_MARKER])] + if set(skeleton) == {_TUPLE_MARKER}: + return tuple(cls._restore_tensor_leaves(item, tensors) for item in skeleton[_TUPLE_MARKER]) + return {key: cls._restore_tensor_leaves(value, tensors) for key, value in skeleton.items()} + if isinstance(skeleton, list): + return [cls._restore_tensor_leaves(item, tensors) for item in skeleton] + return skeleton + + @staticmethod + def _normalize_tensor_payload(payload: Any) -> tuple[list[torch.Tensor], list[dict[str, Any]]]: + tensors = [payload] if isinstance(payload, torch.Tensor) else list(payload) + normalized: list[torch.Tensor] = [] + specs: list[dict[str, Any]] = [] + for tensor in tensors: + contiguous = tensor.detach().contiguous() + normalized.append(contiguous) + specs.append( + { + "shape": list(contiguous.shape), + "dtype": str(contiguous.dtype), + "device": str(contiguous.device), + "size": contiguous.numel() * contiguous.element_size(), + } + ) + return normalized, specs + + def _allocate_tensor_from_spec(self, spec: dict[str, Any], kind: str | None) -> torch.Tensor: + shape = spec.get("shape") + if not isinstance(shape, list): + raise RuntimeError(f"Invalid NIXL tensor shape: {shape!r}") + dtype_name = str(spec.get("dtype", "")).removeprefix("torch.") + dtype = getattr(torch, dtype_name, None) + if dtype is None: + raise RuntimeError(f"Unsupported NIXL tensor dtype: {spec.get('dtype')!r}") + device = torch.device("cpu") if kind == _KIND_OBJECT else self._resolve_receive_device(spec.get("device")) + return torch.empty(tuple(int(dim) for dim in shape), dtype=dtype, device=device) + + def _resolve_receive_device(self, spec_device: Any) -> torch.device: + if self._receive_device is not None: + return self._receive_device + device = self._parse_device(spec_device) + if device is None or device.type == "cpu": + return torch.device("cpu") + # The producer's device index is meaningless in this process, so keep + # only its type and land the buffer on this stage's own card. + backend = getattr(torch, device.type, None) + index = backend.current_device() if hasattr(backend, "current_device") else 0 + return torch.device(device.type, index) + + @staticmethod + def _parse_device(device_like: Any) -> torch.device | None: + if device_like is None: + return None + try: + return torch.device(device_like) + except Exception as exc: + raise RuntimeError(f"Invalid NIXL receive device: {device_like!r}") from exc + + @staticmethod + def _safe_call(func: Any, *args: Any) -> None: + try: + func(*args) + except Exception: + logger.debug("Ignoring NIXL cleanup failure", exc_info=True) diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index c49da1c6c..66dbca079 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -125,6 +125,17 @@ def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnec return MoriTransferEngineConnector(config) +def _create_nixl_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from .connectors.nixl_connector import NixlConnector + except ImportError: + import sys + + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from omni_connectors.connectors.nixl_connector import NixlConnector + return NixlConnector(config) + + # Register connectors OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector) OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector) @@ -132,5 +143,6 @@ OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_con OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector) OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector) OmniConnectorFactory.register_connector("MoriTransferEngineConnector", _create_mori_transfer_engine_connector) +OmniConnectorFactory.register_connector("NixlConnector", _create_nixl_connector) # Backward-compatible aliases – will be removed in the future OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector) diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 6b83349f2..6ab0b36a2 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -564,6 +564,13 @@ class OmniKVTransferManager: # the wrong process. Explicit sender_* YAML values # are still preserved for standalone connector use. + if c_type == "NixlConnector" and c_extra.get("role") == "receiver": + # This manager receives on the shared incoming edge; + # its YAML zmq_port belongs to the producer. NIXL treats + # an explicit zmq_port as a local listener, even for a + # receiver. Keep only the request-scoped sender_* endpoint. + c_extra.pop("zmq_port", None) + logger.info( "Initializing OmniConnector type=%s role=%s", c_type, @@ -922,7 +929,9 @@ class OmniKVTransferManager: The base host/port are also stored so that the receive path can construct per-rank metadata for heterogeneous TP scenarios. """ - if not self.config.need_recv_cache: + # A stage can sit on the receiving end of a payload-only edge, where the + # sender endpoint still has to be applied even though no KV is expected. + if not self.config.need_recv_cache and (self.config.connector_config or {}).get("role") != "receiver": return actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 25515cfaf..2343315b9 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -182,14 +182,19 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage: str, to_stage: str, connector_get_key: str, + metadata: dict[str, Any] | None = None, ) -> Any: """Receive one ordinary non-KV stage payload on the local leader rank only.""" tp_group = self._get_local_tp_group() if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: - return connector.get(from_stage, to_stage, connector_get_key) + if metadata is None: + return connector.get(from_stage, to_stage, connector_get_key) + return connector.get(from_stage, to_stage, connector_get_key, metadata) if not self.is_data_transfer_rank(): return None - return connector.get(from_stage, to_stage, connector_get_key) + if metadata is None: + return connector.get(from_stage, to_stage, connector_get_key) + return connector.get(from_stage, to_stage, connector_get_key, metadata) def _recv_full_payload_result( self, @@ -197,6 +202,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage: str, to_stage: str, connector_get_key: str, + metadata: dict[str, Any] | None = None, ) -> Any: """Receive one full-payload transfer on the local leader rank only.""" return self._recv_ordinary_stage_result( @@ -204,6 +210,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage, to_stage, connector_get_key, + metadata, ) def _recv_async_chunk_result( @@ -212,6 +219,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage: str, to_stage: str, connector_get_key: str, + metadata: dict[str, Any] | None = None, ) -> Any: """Receive one ordinary async chunk on the local leader rank only.""" return self._recv_ordinary_stage_result( @@ -219,6 +227,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage, to_stage, connector_get_key, + metadata, ) @staticmethod @@ -1027,6 +1036,14 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): chunk_id = self._get_req_chunk[req_id] external_req_id = self._request_ids_mapping.get(req_id, req_id) connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" + request = self._pending_load_reqs.get(req_id) + sender_info = getattr(request, "payload_sender_info", None) + metadata = None + if isinstance(sender_info, dict): + host = sender_info.get("host") + port = sender_info.get("zmq_port") + if host and port: + metadata = {"source_host": str(host), "source_port": int(port)} if self._async_chunk: result = self._recv_async_chunk_result( @@ -1034,6 +1051,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): str(target_stage_id), str(self._stage_id), connector_get_key, + metadata, ) else: result = self._recv_full_payload_result( @@ -1041,6 +1059,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): str(target_stage_id), str(self._stage_id), connector_get_key, + metadata, ) if result is None: diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 275a2c0af..95a6cac22 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -19,6 +19,8 @@ from vllm_omni.distributed.omni_connectors.utils.config import ( ConnectorSpec, get_stage_connector_role, ) +from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec +from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") @@ -186,7 +188,11 @@ class _OmniConnectorRuntimeMixin: getattr(topology, "target_tp_size", None), getattr(topology, "local_rank", None), ) - if all(isinstance(value, int) for value in effective_mapping): + if ( + isinstance(effective_mapping[0], int) + and isinstance(effective_mapping[1], int) + and isinstance(effective_mapping[2], int) + ): rank_cfg = { "from_tp": effective_mapping[0], "to_tp": effective_mapping[1], @@ -521,7 +527,15 @@ class _OmniConnectorRuntimeMixin: elif not isinstance(extra, dict): raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}") - spec = ConnectorSpec(name=name, extra=extra) + stage_id = int(getattr(model_config, "stage_id", extra.get("stage_id", 0))) + role = extra.get("role") + spec = resolve_connector_spec( + ConnectorSpec(name=name, extra=extra), + stage_id=stage_id, + role=str(role) if role is not None else None, + local_rank=get_local_tp_rank(), + replica_id=get_omni_replica_id(), + ) try: return OmniConnectorFactory.create_connector(spec) except Exception as exc: diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 163363931..b801622f9 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -20,6 +20,8 @@ from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_p from ..adapter import construct_next_stage_streaming_input_prompt from ..factory import OmniConnectorFactory from ..utils.config import ConnectorSpec, stage_receives_chunks +from ..utils.initialization import resolve_connector_spec +from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id from ..utils.logging import get_connector_logger from .base import OmniTransferAdapterBase @@ -47,10 +49,20 @@ class _SenderGeneration: class _LoadEntry: """Identify one receiver registration across queue and I/O boundaries.""" - __slots__ = ("request",) + __slots__ = ("request", "source_metadata") def __init__(self, request: Request) -> None: self.request = request + sender_info = getattr(request, "payload_sender_info", None) + self.source_metadata = None + if isinstance(sender_info, dict): + host = sender_info.get("host") + port = sender_info.get("zmq_port") + if host and port: + self.source_metadata = { + "source_host": str(host), + "source_port": int(port), + } @property def request_id(self) -> str: @@ -280,11 +292,20 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): "extra": getattr(connector_config, "extra", {}), } - connector_specs = ConnectorSpec( + connector_spec = ConnectorSpec( name=connector_config.get("name", "SharedMemoryConnector"), extra=connector_config.get("extra", {}), ) - return OmniConnectorFactory.create_connector(connector_specs) + stage_id = int(getattr(model_config, "stage_id", connector_spec.extra.get("stage_id", 0))) + role = connector_spec.extra.get("role") + resolved_spec = resolve_connector_spec( + connector_spec, + stage_id=stage_id, + role=str(role) if role is not None else None, + local_rank=get_local_tp_rank(), + replica_id=get_omni_replica_id(), + ) + return OmniConnectorFactory.create_connector(resolved_spec) def load_async(self, request: Request): """Register a request for asynchronous chunk retrieval. @@ -484,6 +505,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): str(target_stage_id), str(stage_id), connector_get_key, + entry.source_metadata, ) except Exception as e: logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}") diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index 1a6d395c2..176a8bf69 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -13,6 +13,7 @@ TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( "MooncakeTransferEngineConnector", "MoriTransferEngineConnector", "YuanrongTransferEngineConnector", + "NixlConnector", } ) diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index bbbda44e3..d0139b59a 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -4,9 +4,10 @@ """Utilities for OmniConnector configuration and validation.""" import json +import os import warnings from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from ..factory import OmniConnectorFactory from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig @@ -38,6 +39,91 @@ KV_RANK_PORT_STRIDE = 16 # comfortably sized block per replica for TP-rank and stage offsets. KV_REPLICA_PORT_STRIDE = 1024 +REQUEST_FORWARDING_PORT_OFFSET = 0 +ORCHESTRATOR_PORT_OFFSET = 200 + +ConnectorPurpose = Literal["request_forwarding", "kv_transfer", "orchestrator"] + + +def compute_connector_zmq_port( + base_port: int, + *, + purpose: ConnectorPurpose, + from_stage: int, + local_rank: int = 0, + replica_id: int = 0, +) -> int: + """Return a collision-free connector endpoint for an edge and worker.""" + purpose_offsets = { + "request_forwarding": REQUEST_FORWARDING_PORT_OFFSET, + "kv_transfer": KV_TRANSFER_PORT_OFFSET, + "orchestrator": ORCHESTRATOR_PORT_OFFSET, + } + if purpose not in purpose_offsets: + raise ValueError(f"Unsupported connector purpose: {purpose}") + if not 0 <= from_stage < KV_RANK_PORT_STRIDE: + raise ValueError(f"from_stage must be in [0, {KV_RANK_PORT_STRIDE}), got {from_stage}") + if local_rank < 0: + raise ValueError(f"local_rank must be non-negative, got {local_rank}") + if replica_id < 0: + raise ValueError(f"replica_id must be non-negative, got {replica_id}") + + port = ( + base_port + + purpose_offsets[purpose] + + replica_id * KV_REPLICA_PORT_STRIDE + + local_rank * KV_RANK_PORT_STRIDE + + from_stage + ) + if not 1 <= port <= 65535: + raise ValueError(f"Resolved connector port must be in [1, 65535], got {port}") + return port + + +def resolve_connector_spec( + spec: ConnectorSpec, + *, + stage_id: int, + role: str | None, + purpose: ConnectorPurpose = "request_forwarding", + local_rank: int = 0, + replica_id: int = 0, +) -> ConnectorSpec: + """Copy and resolve a connector spec for one stage worker.""" + extra = dict(spec.extra or {}) + extra["stage_id"] = stage_id + if role is not None: + extra["role"] = role + if spec.name not in TRANSFER_ENGINE_CONNECTOR_NAMES: + return ConnectorSpec(name=spec.name, extra=extra) + + resolved_port = None + if role != "receiver" or "sender_zmq_port" not in extra: + resolved_port = compute_connector_zmq_port( + expand_env_int(extra.get("zmq_port", 50051), "zmq_port"), + purpose=purpose, + from_stage=int(extra.get("from_stage", stage_id)), + local_rank=local_rank, + replica_id=replica_id, + ) + if role == "receiver": + extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) + extra.setdefault("sender_zmq_port", resolved_port) + extra.pop("zmq_port", None) + outgoing = extra.get("outgoing") + if spec.name == "NixlConnector" and isinstance(outgoing, dict): + extra["zmq_port"] = compute_connector_zmq_port( + expand_env_int(outgoing.get("zmq_port", 50051), "zmq_port"), + purpose=purpose, + from_stage=int(outgoing.get("from_stage", stage_id)), + local_rank=local_rank, + replica_id=replica_id, + ) + extra["host"] = outgoing.get("host", "auto") + else: + extra["zmq_port"] = resolved_port + return ConnectorSpec(name=spec.name, extra=extra) + def initialize_connectors_from_config( config_path: str | Path | None = None, @@ -82,52 +168,40 @@ def create_connectors_from_config( Returns: A dictionary of connectors. """ - purpose_port_offsets = { - "request_forwarding": 0, - "kv_transfer": KV_TRANSFER_PORT_OFFSET, - } - port_offset = purpose_port_offsets.get(purpose, 0) - orchestrator_port_offset = 200 - connectors = {} for edge_key, connector_spec in connectors_config.items(): from_stage, to_stage = edge_key try: if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES: extra = dict(connector_spec.extra) if connector_spec.extra else {} - base_port = expand_env_int(extra.get("zmq_port", 50051), "zmq_port") try: - stage_offset = int(from_stage) - except (TypeError, ValueError): - stage_offset = 0 - - local_rank = 0 if str(caller_stage_id) == "orchestrator" else get_connector_local_rank() - if str(caller_stage_id) == "orchestrator": - rank0_port = base_port + orchestrator_port_offset + stage_offset - else: - rank0_port = base_port + port_offset + stage_offset - adjusted_port = rank0_port + local_rank * KV_RANK_PORT_STRIDE - extra["zmq_port"] = adjusted_port + from_stage_id = int(from_stage) + except (TypeError, ValueError) as exc: + raise ValueError(f"Transfer-engine stage id must be an integer, got {from_stage!r}") from exc + is_orchestrator = str(caller_stage_id) == "orchestrator" + local_rank = 0 if is_orchestrator else get_connector_local_rank() + resolved_purpose: ConnectorPurpose = "orchestrator" if is_orchestrator else purpose # type: ignore[assignment] if is_sender is not None: - extra["role"] = "sender" if is_sender else "receiver" - if not is_sender: - extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) - extra.setdefault("sender_zmq_port", adjusted_port) - elif caller_stage_id is not None: - caller_str = str(caller_stage_id) - if caller_str == from_stage: - extra["role"] = "sender" - elif caller_str == to_stage: - extra["role"] = "receiver" - extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) - extra.setdefault("sender_zmq_port", adjusted_port) - else: - extra["role"] = "sender" + role = "sender" if is_sender else "receiver" + elif caller_stage_id is not None and str(caller_stage_id) == to_stage: + role = "receiver" + elif caller_stage_id is not None and str(caller_stage_id) == from_stage: + role = "sender" else: - extra["role"] = extra.get("role", "auto") - - connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=connector_spec.name, extra=extra)) + role = str(extra.get("role", "auto")) + extra["from_stage"] = from_stage_id + extra["to_stage"] = int(to_stage) + replica_id = max(int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")), 0) + resolved_spec = resolve_connector_spec( + ConnectorSpec(name=connector_spec.name, extra=extra), + stage_id=int(caller_stage_id) if str(caller_stage_id).isdigit() else from_stage_id, + role=role, + purpose=resolved_purpose, + local_rank=local_rank, + replica_id=replica_id, + ) + connector = OmniConnectorFactory.create_connector(resolved_spec) else: connector = OmniConnectorFactory.create_connector(connector_spec) connectors[edge_key] = connector @@ -174,12 +248,16 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None, # Incoming edge → this stage is the receiver extra = dict(spec.extra) if spec.extra else {} extra.setdefault("role", "receiver") + extra["from_stage"] = int(from_stage) + extra["to_stage"] = int(to_stage) stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": extra}} elif from_stage == target_stage and target_stage == "0": # Outgoing edge for stage 0 — included for async_chunk spec # extraction (omni_stage.py), NOT for connector instantiation. extra = dict(spec.extra) if spec.extra else {} extra.setdefault("role", "sender") + extra["from_stage"] = int(from_stage) + extra["to_stage"] = int(to_stage) stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": extra}} return stage_connectors_config diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py index b7f51460c..4be32ba7e 100644 --- a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -16,7 +16,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.logger import init_logger -from .initialization import KV_RANK_PORT_STRIDE, KV_REPLICA_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET +from .initialization import compute_connector_zmq_port logger = init_logger(__name__) @@ -123,12 +123,12 @@ def kv_zmq_port( """ replica = get_omni_replica_id() if replica_id is None else max(int(replica_id), 0) - return ( - base_port - + KV_TRANSFER_PORT_OFFSET - + replica * KV_REPLICA_PORT_STRIDE - + local_rank * KV_RANK_PORT_STRIDE - + from_stage + return compute_connector_zmq_port( + base_port, + purpose="kv_transfer", + from_stage=from_stage, + local_rank=local_rank, + replica_id=replica, ) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index e72ce9fed..a96c8027f 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -79,6 +79,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): # GPUModelRunner.model_intermediate_buffer instead of using the deprecated # additional_information request transport. model_intermediate_buffer: dict[str, Any] | None = None + payload_sender_info: dict[str, Any] | None = None @classmethod def from_request( @@ -88,6 +89,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): prompt_embeds: torch.Tensor | None = None, additional_information: AdditionalInformationPayload | None = None, model_intermediate_buffer: dict[str, Any] | None = None, + payload_sender_info: dict[str, Any] | None = None, ) -> "OmniEngineCoreRequest": """Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides.""" @@ -97,6 +99,8 @@ class OmniEngineCoreRequest(EngineCoreRequest): additional_information = getattr(request, "additional_information", None) if model_intermediate_buffer is None: model_intermediate_buffer = getattr(request, "model_intermediate_buffer", None) + if payload_sender_info is None: + payload_sender_info = getattr(request, "payload_sender_info", None) return cls( request_id=request.request_id, @@ -121,6 +125,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): abort_immediately=request.abort_immediately, additional_information=additional_information, model_intermediate_buffer=model_intermediate_buffer, + payload_sender_info=payload_sender_info, ) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 299e1dd07..498a9dc89 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -2148,6 +2148,7 @@ class Orchestrator: *, mm_features: list | None = None, resumable: bool = False, + payload_sender_info: dict[str, Any] | None = None, ) -> Any: next_pool = self.stage_pools[next_stage_id] if self._next_stage_input_is_tokens(next_input): @@ -2160,6 +2161,7 @@ class Orchestrator: resumable=resumable, ) request.external_req_id = request.request_id + request.payload_sender_info = payload_sender_info return request processor = self._get_stage_input_processor(next_stage_id) @@ -2181,6 +2183,7 @@ class Orchestrator: ) request = self._upgrade_processed_stage_request(request, next_input) request.external_req_id = req_id + request.payload_sender_info = payload_sender_info return request @staticmethod @@ -2829,6 +2832,7 @@ class Orchestrator: params=params, mm_features=mm_features, resumable=next_stage_resumable, + payload_sender_info=self._build_payload_sender_info(src_stage_id, request_id=req_id), ) if already_submitted: @@ -2967,6 +2971,7 @@ class Orchestrator: model_config=next_pool.stage_vllm_config.model_config, resumable=downstream_resumable, ) + request.payload_sender_info = self._build_payload_sender_info(next_stage_id - 1, request_id=request_id) request.external_req_id = request.request_id submitted = await self._dispatch_or_fail_request( lambda: next_pool.submit_initial( @@ -3043,6 +3048,23 @@ class Orchestrator: return sender_infos or None + def _build_payload_sender_info( + self, + sender_stage_id: int, + *, + request_id: str | None = None, + ) -> dict[str, Any] | None: + if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools): + return None + sender_pool = self.stage_pools[sender_stage_id] + sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None + if sender_stage is None: + sender_stage = sender_pool.stage_client + get_sender_info = getattr(sender_stage, "get_payload_sender_info", None) + if not callable(get_sender_info): + return None + return get_sender_info() + # ---- Shutdown / lifecycle ---- def _shutdown_stages(self) -> None: diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 2a3c26177..7ee949146 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -94,6 +94,8 @@ class StagePoolLLMClient(StagePoolClient, Protocol): kv_transfer_port_offset: int = ..., ) -> dict[str, Any] | None: ... + def get_payload_sender_info(self) -> dict[str, Any] | None: ... + class StagePoolDiffusionClient(StagePoolClient, Protocol): """Pool-facing API for diffusion stages.""" diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 87baaed5f..5fdd74912 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Stage Engine Core Client for vLLM-Omni multi-stage runtime. @@ -210,6 +213,7 @@ class StageEngineCoreClientBase(StageClientBase): ) raise + self._payload_sender_info = self._build_payload_sender_info() self._initialize_kv_sender_endpoint() logger.info( @@ -291,6 +295,39 @@ class StageEngineCoreClientBase(StageClientBase): return None return connector_config + def _build_payload_sender_info(self) -> dict[str, Any] | None: + model_config = getattr(self.vllm_config, "model_config", None) + connector_config = getattr(model_config, "stage_connector_config", None) + if not isinstance(connector_config, dict): + return None + extra = connector_config.get("extra") + if not isinstance(extra, dict): + return None + if isinstance(extra.get("outgoing"), dict): + extra = extra["outgoing"] + elif extra.get("role") != "sender": + return None + base_port = extra.get("zmq_port", 50051) + if base_port is None: + return None + sender_host = self._resolve_sender_host_from_config(extra) + if sender_host is None: + return None + from vllm_omni.distributed.omni_connectors.utils.initialization import compute_connector_zmq_port + + return { + "host": sender_host, + "zmq_port": compute_connector_zmq_port( + int(os.path.expandvars(str(base_port))), + purpose="request_forwarding", + from_stage=int(extra.get("from_stage", self.stage_id)), + replica_id=self.replica_id, + ), + } + + def get_payload_sender_info(self) -> dict[str, Any] | None: + return dict(self._payload_sender_info) if self._payload_sender_info is not None else None + def _resolve_sender_host_from_config(self, connector_config: dict[str, Any]) -> str | None: host = connector_config.get("sender_host") or connector_config.get("host") if host in {None, "", "auto", "*", "0.0.0.0", "::"}: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index e483e1bd4..a52f58925 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -88,6 +88,7 @@ class StageEngineCoreProc(EngineCoreProc): scheduler_request, current_wave = super().preprocess_add_request(request) scheduler_request.additional_information = request.additional_information scheduler_request.external_req_id = getattr(request, "external_req_id", request.request_id) + scheduler_request.payload_sender_info = request.payload_sender_info return scheduler_request, current_wave @staticmethod diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index de97f7aa3..9668dfc99 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -50,6 +50,9 @@ from vllm_omni.config.omni_config import ( ) from vllm_omni.config.stage_config import StageType from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.distributed.omni_connectors.utils.config import ( + TRANSFER_ENGINE_CONNECTOR_NAMES, +) from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs @@ -1882,17 +1885,37 @@ def get_stage_connector_spec( stage_connectors_cfg = get_stage_connector_config(omni_transfer_config, stage_id) for cfg in stage_connectors_cfg.values(): - return dict(cfg.get("spec", {})) + connector_spec = dict(cfg.get("spec", {})) + if connector_spec.get("name") == "NixlConnector": + # A middle worker uses one connector for both directions. Preserve + # the incoming edge and carry the outgoing bind config separately. + for (source, target), outgoing in getattr(omni_transfer_config, "connectors", {}).items(): + if source == str(stage_id): + if outgoing.name != "NixlConnector": + raise ValueError("A NIXL middle stage requires NIXL on its outgoing edge") + extra = dict(connector_spec.get("extra", {})) + extra["outgoing"] = {**(outgoing.extra or {}), "from_stage": int(source), "to_stage": int(target)} + connector_spec["extra"] = extra + break + if connector_spec.get("name") not in TRANSFER_ENGINE_CONNECTOR_NAMES: + extra = dict(connector_spec.get("extra", {})) + extra.pop("from_stage", None) + extra.pop("to_stage", None) + connector_spec["extra"] = extra + return connector_spec # A producer does not consume connector data itself. Keep its connector # for both async-chunk and terminal full-payload sends, but mark it # sender-only so the scheduler does not park orchestrator-provided inputs # waiting for an upstream payload. target_stage = str(stage_id) - for (from_stage, _to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items(): + for (from_stage, to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items(): if from_stage == target_stage: extra = dict(spec.extra or {}) extra.setdefault("role", "sender") + if spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES: + extra["from_stage"] = int(from_stage) + extra["to_stage"] = int(to_stage) return {"name": spec.name, "extra": extra} return {} From 311e714d34be1c80714833abc754635aa3c1b59b Mon Sep 17 00:00:00 2001 From: linzhenpl07 <60651557+linzhenpl07@users.noreply.github.com> Date: Sun, 13 Sep 2026 18:17:11 +0800 Subject: [PATCH 045/209] [AR-Diffusion] Add session-owned streaming VAE decode (#6533) Signed-off-by: linzhenpl07 Co-authored-by: Claude Sonnet 5 --- .../ar_diffusion_streaming_decode.py | 349 ++++++++++++++++ .../ar_diffusion/test_streaming_decode.py | 361 +++++++++++++++++ .../test_pipeline_lingbot_world.py | 379 +++++++++++++++++- .../models/lingbot_world/pipeline.py | 175 +++++++- .../ar_diffusion/streaming_decode.py | 272 +++++++++++++ 5 files changed, 1520 insertions(+), 16 deletions(-) create mode 100644 examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py create mode 100644 tests/diffusion/ar_diffusion/test_streaming_decode.py create mode 100644 vllm_omni/experimental/ar_diffusion/streaming_decode.py diff --git a/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py new file mode 100644 index 000000000..2aa326f6d --- /dev/null +++ b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stream LingBot-World 2.0 pixels from one stepwise request. + +One ``AsyncOmni.generate()`` drives the whole rollout: each AR block leaves +``post_decode()`` as one video chunk, decoded through that session's own +temporal VAE cache rather than as an isolated clip. That is what this script +shows end to end -- chunk *N + 1* continues chunk *N*, so the streamed timeline +carries the same frame count as an offline whole-clip decode of the same +latents instead of losing each block's opening frame to a restarted decoder. +Only the session's first latent frame expands to a single raw frame; on the +default checkpoint, whose block is three latent frames at a temporal factor of +four, that is 9 frames for the opening chunk and 12 for every later one. + +The decode lives inside the pipeline, keyed by ``request_id``, so nothing here +owns a VAE: this file is a client of the served path, and the same chunks reach +a WebSocket client of ``/v1/realtime/video`` unchanged. + +Examples:: + + python examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py \ + --image scene.png --prompt "a lit hallway" \ + --action-script actions.json --output-dir /tmp/lingbot-stream + +``--action-script`` is a JSON list with one camera action list per AR block, +for example ``[[["w"], ["w"], ["w"]], [["a"], [], []]]`` for a checkpoint that +generates three latent frames per block. The frames per block, the temporal +compression and the number of blocks that fit the condition horizon are read +off the checkpoint rather than assumed, so the script can be repointed with +``--model``. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import os +import time +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + # Only for type hints -- the real imports stay deferred inside run() so + # --help and the offline validation helpers don't require a CUDA-enabled + # vLLM install. + import numpy as np + +_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" +# Mirrors _MAX_RAW_FRAMES in vllm_omni/diffusion/models/lingbot_world/pipeline.py: +# the image-condition horizon is the pipeline's constant, not the checkpoint's. +_MAX_RAW_FRAMES = 117 + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Stream decoded LingBot-World 2.0 video chunks from a single stepwise request.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.") + parser.add_argument("--image", required=True, help="Initial RGB image.") + parser.add_argument("--prompt", required=True, help="Scene prompt for the rollout.") + parser.add_argument( + "--action-script", + required=True, + help=( + "JSON file holding one camera action list per AR block. Both the per-block frame count " + "and the number of blocks that fit are read off the checkpoint, not assumed." + ), + ) + parser.add_argument("--output-dir", required=True, help="Directory for decoded frames and metadata.") + parser.add_argument("--request-id", default="lingbot-world-stream", help="Request id; also the session id.") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--gpu-memory-fraction", type=float, default=0.6) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--enforce-eager", action="store_true") + return parser.parse_args(argv) + + +def _load_action_script(path: Path) -> list[list[list[str]]]: + """Parse the per-block camera action script, without assuming its geometry. + + The shape is the one ``sampling_params.extra_args["camera_action_script"]`` + takes: one entry per AR block, each entry one per-frame action list. How + many frames a block holds is a property of the checkpoint, so it is checked + in :func:`_validate_against_checkpoint` once the configs have been read. + """ + try: + value = json.loads(path.read_text()) + except json.JSONDecodeError as exc: + raise ValueError(f"--action-script is not valid JSON: {exc.msg}.") from None + if not isinstance(value, list) or not value: + raise ValueError("--action-script must be a non-empty JSON list of per-block action lists.") + script: list[list[list[str]]] = [] + for block_index, block in enumerate(value): + if not isinstance(block, list) or not block: + raise ValueError(f"--action-script block {block_index} must be a non-empty list of frame action lists.") + frames: list[list[str]] = [] + for frame in block: + if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame): + raise ValueError(f"--action-script block {block_index} frames must contain only action strings.") + frames.append(list(frame)) + script.append(frames) + return script + + +def _model_config_json(model: str, relative_path: str) -> dict[str, Any]: + """Read one config file out of a local checkpoint or a Hub repo.""" + if os.path.isdir(model): + return json.loads((Path(model) / relative_path).read_text()) + from huggingface_hub import hf_hub_download + + return json.loads(Path(hf_hub_download(repo_id=model, filename=relative_path)).read_text()) + + +def _checkpoint_geometry(model: str) -> tuple[int, int]: + """Read ``(frames_per_block, temporal_compression)`` off the checkpoint. + + The pipeline takes both from the model -- + ``transformer.config.num_frames_per_block`` and + ``vae.config.scale_factor_temporal`` -- so hardcoding 3 and 4 here would + silently build a wrong action script and an inconsistent ``num_frames`` the + moment ``--model`` points at a checkpoint with different geometry. + """ + transformer_config = _model_config_json(model, "transformer/config.json") + vae_config = _model_config_json(model, "vae/config.json") + if "num_frames_per_block" not in transformer_config: + raise ValueError( + f"{model}'s transformer/config.json declares no num_frames_per_block, " + "so this checkpoint does not generate AR blocks and cannot be streamed this way." + ) + frames_per_block = int(transformer_config["num_frames_per_block"]) + temporal_compression = int(vae_config.get("scale_factor_temporal", 4)) + if frames_per_block <= 0 or temporal_compression <= 0: + raise ValueError( + "the checkpoint reports a non-positive block geometry " + f"(num_frames_per_block={frames_per_block}, scale_factor_temporal={temporal_compression})." + ) + return frames_per_block, temporal_compression + + +def _validate_against_checkpoint( + script: list[list[list[str]]], *, frames_per_block: int, temporal_compression: int +) -> int: + """Check the script against the checkpoint and return the frame count to request.""" + for block_index, block in enumerate(script): + if len(block) != frames_per_block: + raise ValueError( + f"--action-script block {block_index} holds {len(block)} frame action lists, but this " + f"checkpoint generates {frames_per_block} latent frames per AR block." + ) + max_chunks = ((_MAX_RAW_FRAMES - 1) // temporal_compression + 1) // frames_per_block + if len(script) > max_chunks: + raise ValueError( + f"--action-script holds {len(script)} blocks, but this checkpoint's " + f"{_MAX_RAW_FRAMES}-frame image-condition horizon fits at most {max_chunks}." + ) + return (len(script) * frames_per_block - 1) * temporal_compression + 1 + + +def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]: + image = Path(args.image).expanduser().resolve() + action_script = Path(args.action_script).expanduser().resolve() + output_dir = Path(args.output_dir).expanduser().resolve() + if not image.is_file(): + raise ValueError("--image must point to an existing file.") + if not action_script.is_file(): + raise ValueError("--action-script must point to an existing JSON file.") + if not args.prompt.strip(): + raise ValueError("--prompt must contain non-whitespace text.") + if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16: + raise ValueError("--height and --width must be positive multiples of 16.") + if args.tensor_parallel_size <= 0: + raise ValueError("--tensor-parallel-size must be positive.") + if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1: + raise ValueError("--gpu-memory-fraction must be in (0, 1].") + if not args.request_id.strip(): + raise ValueError("--request-id must contain non-whitespace text.") + return image, action_script, output_dir + + +def _as_frames(value: Any) -> np.ndarray: + """Normalize one streamed chunk to ``[T, H, W, 3]`` uint8. + + ``output_type="np"`` reaches the client as a float array in ``[0, 1]``, + with or without the leading batch axis depending on how the formatter + unwrapped it; a nested list is the ``"pil"`` shape. Anything else means + the chunk is not pixels -- most likely the request still asked for + latents -- and is worth failing on rather than saving. + """ + import numpy as np + import PIL.Image + import torch + + if isinstance(value, list): + if len(value) == 1 and isinstance(value[0], list): + value = value[0] + if value and all(isinstance(frame, PIL.Image.Image) for frame in value): + return np.stack([np.asarray(frame.convert("RGB")) for frame in value]) + if isinstance(value, torch.Tensor): + value = value.detach().float().cpu().numpy() + if not isinstance(value, np.ndarray): + raise RuntimeError(f"Expected decoded pixels from each AR block, got {type(value).__name__}.") + if value.ndim == 5 and value.shape[0] == 1: + value = value[0] + if value.ndim != 4 or value.shape[-1] != 3: + raise RuntimeError(f"Expected one [T, H, W, 3] pixel chunk per AR block, got shape {tuple(value.shape)}.") + if value.dtype != np.uint8: + value = (np.clip(value, 0.0, 1.0) * 255.0).round().astype(np.uint8) + return value + + +def _chunk_metadata(output: Any) -> dict[str, Any]: + multimodal = getattr(output, "multimodal_output", None) or {} + metadata = multimodal.get("metadata") if isinstance(multimodal, dict) else None + ar_diffusion = metadata.get("ar_diffusion") if isinstance(metadata, dict) else None + if not isinstance(ar_diffusion, dict): + raise RuntimeError("Streamed chunk is missing its ar_diffusion metadata envelope.") + return dict(ar_diffusion) + + +async def run(argv: Sequence[str] | None = None) -> Path: + args = parse_args(argv) + image, action_script_path, output_dir = _validate_args(args) + camera_action_script = _load_action_script(action_script_path) + output_dir.mkdir(parents=True, exist_ok=True) + + # Imports stay below pure input validation so --help and helper tests do not + # require a CUDA-enabled vLLM installation. + import numpy as np + + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + frames_per_block, temporal_compression = _checkpoint_geometry(args.model) + num_chunks = len(camera_action_script) + num_frames = _validate_against_checkpoint( + camera_action_script, + frames_per_block=frames_per_block, + temporal_compression=temporal_compression, + ) + + engine = AsyncOmni( + model=args.model, + engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine", + enforce_eager=args.enforce_eager, + tensor_parallel_size=args.tensor_parallel_size, + max_num_seqs=1, + # One request, many chunks: the stepwise contract from #6844. Streaming + # output is what turns each post_decode() into a delivered chunk. + step_execution=True, + diffusion_streaming_output=True, + model_config={ + "ar_diffusion_height": args.height, + "ar_diffusion_width": args.width, + "ar_diffusion_kv_config": { + "gpu_memory_fraction": args.gpu_memory_fraction, + "warmup_cudagraph": True, + }, + }, + ) + sampling = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_frames=num_frames, + num_inference_steps=4, + max_sequence_length=512, + seed=args.seed, + # Pixels, not latents: this is what routes the chunk through the + # session's streaming VAE decode instead of handing back a tensor. + output_type="np", + extra_args={"flow_shift": 5.0, "camera_action_script": camera_action_script}, + ) + prompt = {"prompt": args.prompt, "multi_modal_data": {"image": str(image)}} + + measurements: list[dict[str, Any]] = [] + t_start = time.perf_counter() + first_frame_at: float | None = None + chunk_started = t_start + try: + async for output in engine.generate(prompt, sampling, request_id=args.request_id): + images = getattr(output, "images", None) + if not images: + continue + if len(images) != 1: + raise RuntimeError("Expected exactly one decoded chunk per AR block.") + frames = _as_frames(images[0]) + now = time.perf_counter() + if first_frame_at is None: + first_frame_at = now - t_start + chunk_index = len(measurements) + metadata = _chunk_metadata(output) + + np.save(output_dir / f"chunk_{chunk_index:03d}.npy", frames) + (output_dir / f"chunk_{chunk_index:03d}.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n" + ) + measurements.append( + { + "chunk_index": chunk_index, + "chunk_latency_seconds": now - chunk_started, + "frames": int(frames.shape[0]), + "frame_shape": list(frames.shape), + "metadata": metadata, + } + ) + chunk_started = now + print(json.dumps(measurements[-1], sort_keys=True), flush=True) + finally: + engine.shutdown() + + delivered_frames = sum(int(chunk["frames"]) for chunk in measurements) + summary_path = output_dir / "summary.json" + summary_path.write_text( + json.dumps( + { + "request_id": args.request_id, + "chunks_delivered": len(measurements), + "chunks_expected": num_chunks, + "frames_delivered": delivered_frames, + "frames_expected": num_frames, + # The point of session-owned decode: a streamed rollout carries + # the same frame timeline as an offline decode of the same + # latents. Per-chunk decode would deliver fewer, because every + # block would re-expand its own opening frame. + "matches_offline_frame_timeline": delivered_frames == num_frames, + "time_to_first_frame_seconds": first_frame_at, + "total_seconds": time.perf_counter() - t_start, + "chunks": measurements, + }, + indent=2, + sort_keys=True, + ) + + "\n" + ) + return summary_path + + +def main(argv: Sequence[str] | None = None) -> Path: + return asyncio.run(run(argv)) + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/ar_diffusion/test_streaming_decode.py b/tests/diffusion/ar_diffusion/test_streaming_decode.py new file mode 100644 index 000000000..52889bfbd --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_streaming_decode.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Streaming VAE decode against the real Wan autoencoder, on CPU. + +The autoencoder is built from a config rather than a checkpoint -- weights are +random, only shapes and the temporal cache protocol matter -- so these run +without a device or a download. + +The load-bearing test is equivalence: decoding a session chunk by chunk must +produce exactly what decoding the whole clip produces. Everything else about +streaming is worthless if that does not hold. +""" + +from __future__ import annotations + +from contextlib import contextmanager + +import pytest +import torch + +from vllm_omni.experimental.ar_diffusion.streaming_decode import ( + StreamingDecodeState, + SupportsStreamingDecode, + WanStreamingDecoder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Small enough to decode on CPU in a test, same architecture as the shipped +# checkpoint: base_dim 96 / dim_mult [1,2,4,4] would be far too slow, so the +# widths are scaled down while the causal structure is unchanged. +VAE_CONFIG = { + "base_dim": 8, + "z_dim": 4, + "dim_mult": [1, 2], + "num_res_blocks": 1, + "temperal_downsample": [True], + "attn_scales": [], + "dropout": 0.0, +} +LATENT_H = LATENT_W = 4 + + +@pytest.fixture(scope="module") +def vae(): + diffusers = pytest.importorskip("diffusers") + torch.manual_seed(0) + model = diffusers.AutoencoderKLWan(**VAE_CONFIG).eval() + for parameter in model.parameters(): + parameter.requires_grad_(False) + return model + + +@pytest.fixture(scope="module") +def decoder(vae): + return WanStreamingDecoder(vae) + + +def _latent(num_frames: int, *, seed: int = 0) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn(1, VAE_CONFIG["z_dim"], num_frames, LATENT_H, LATENT_W, generator=generator) + + +@torch.no_grad() +def _decode_whole(vae, latent: torch.Tensor) -> torch.Tensor: + """The non-streaming reference: one call, cache cleared at both ends.""" + return vae._decode(latent, return_dict=False)[0] + + +# -------------------------------------------------------------------------- +# Equivalence: the property that makes streaming worth doing at all +# -------------------------------------------------------------------------- + + +def _stream(decoder, latent: torch.Tensor, chunk_sizes, *, session_id: str = "s") -> torch.Tensor: + state = decoder.new_decode_state(session_id) + pieces, offset = [], 0 + for size in chunk_sizes: + pieces.append(decoder.decode_chunk(latent[:, :, offset : offset + size], state)) + offset += size + return torch.cat(pieces, dim=2) + + +@torch.no_grad() +@pytest.mark.parametrize("chunk_sizes", [(3, 3, 3), (1, 1, 1, 1, 1), (2, 4, 3)]) +def test_chunking_is_neutral(vae, decoder, chunk_sizes) -> None: + """Splitting a session into chunks must change nothing about its output. + + This is the property streaming actually introduces, so it is the one + asserted exactly. The control is the same session decoded in a single + call through the same code path -- not ``_decode`` -- because ``_decode`` + differs in a way that has nothing to do with chunking; see + ``test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact``. + """ + total = sum(chunk_sizes) + latent = _latent(total) + reference = _stream(decoder, latent, (total,), session_id="ref") + streamed = _stream(decoder, latent, chunk_sizes) + + assert streamed.shape == reference.shape + torch.testing.assert_close(streamed, reference, rtol=0, atol=0) + + +@torch.no_grad() +def test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact(vae, decoder) -> None: + """Explain, and bound, the gap against diffusers' ``_decode``. + + ``_decode`` applies ``post_quant_conv`` to the whole latent at once; the + tiled path in this repo, and this decoder, apply it per frame. It is a + 1x1x1 convolution, so the two are mathematically identical -- but they can + select different kernels, and on bf16 CUDA that difference is real and is + amplified by the causal convolutions downstream. + + Asserting equality with ``_decode`` would therefore pin a kernel-selection + coincidence rather than a property of streaming. Assert instead that the + gap is entirely explained by where ``post_quant_conv`` runs. + """ + latent = _latent(6) + per_frame_pq = torch.cat([vae.post_quant_conv(latent[:, :, k : k + 1]) for k in range(latent.shape[2])], dim=2) + whole_pq = vae.post_quant_conv(latent) + + streamed = _stream(decoder, latent, (3, 3)) + one_shot = _decode_whole(vae, latent) + + # Same inputs to the decoder => same output, whatever the backend does. + if torch.equal(whole_pq, per_frame_pq): + torch.testing.assert_close(streamed, one_shot, rtol=0, atol=0) + else: # pragma: no cover - backend dependent + pytest.skip("backend batches post_quant_conv differently; the gap is that, not chunking") + + +@torch.no_grad() +def test_a_restarted_stream_does_not_match_a_continued_one(decoder) -> None: + """Negative control: the equivalence test above must be discriminating. + + Clearing state between chunks -- what the non-streaming path does on every + call -- has to change the pixels, otherwise the temporal cache carries + nothing and the assertion above holds for a write-only cache too. + + Restarting also changes the frame count, because a fresh state expands its + opening latent frame to one raw frame. Comparing totals would therefore + pass on the frame count alone and never look at a pixel, so this compares + the tail both runs agree on: the frames from the chunk's later latents, + which are the full temporal factor wide on either side. + """ + latent = _latent(6) + state = decoder.new_decode_state("continued") + decoder.decode_chunk(latent[:, :, :3], state) + continued = decoder.decode_chunk(latent[:, :, 3:], state) + restarted = decoder.decode_chunk(latent[:, :, 3:], decoder.new_decode_state("restarted")) + + tail = min(continued.shape[2], restarted.shape[2]) - 1 + assert tail > 0, "the comparison needs at least one fully expanded frame on both sides" + continued_tail, restarted_tail = continued[:, :, -tail:], restarted[:, :, -tail:] + assert continued_tail.shape == restarted_tail.shape, "the tails must be comparable for this to mean anything" + assert not torch.equal(continued_tail, restarted_tail), ( + "restarting the decoder produced the same pixels as continuing it, so the carried cache is not being read" + ) + + +# -------------------------------------------------------------------------- +# Cross-session isolation: the failure that timing metrics cannot see +# -------------------------------------------------------------------------- + + +@torch.no_grad() +@pytest.mark.parametrize("seed_a, seed_b", [(1, 2), (11, 12)]) +def test_interleaved_sessions_keep_independent_temporal_context(vae, decoder, seed_a: int, seed_b: int) -> None: + """Two sessions ticking alternately must each match their solo decode. + + A session that reads another's temporal cache produces a perfectly + plausible but wrong video, and every latency metric stays green, so this + has to be asserted against a recorded solo run rather than inferred. The + control is a solo stream rather than ``_decode``, so it differs from the + run under test only in the interleaving. + """ + latent_a = _latent(6, seed=seed_a) + latent_b = _latent(6, seed=seed_b) + solo_a = _stream(decoder, latent_a, (6,), session_id="solo-a") + solo_b = _stream(decoder, latent_b, (6,), session_id="solo-b") + assert not torch.equal(solo_a, solo_b), "negative control: the two sessions must differ" + + state_a = decoder.new_decode_state("a") + state_b = decoder.new_decode_state("b") + out_a, out_b = [], [] + for start in (0, 3): + window = slice(start, start + 3) + out_a.append(decoder.decode_chunk(latent_a[:, :, window], state_a)) + out_b.append(decoder.decode_chunk(latent_b[:, :, window], state_b)) + + torch.testing.assert_close(torch.cat(out_a, dim=2), solo_a, rtol=0, atol=0) + torch.testing.assert_close(torch.cat(out_b, dim=2), solo_b, rtol=0, atol=0) + + +@torch.no_grad() +def test_the_module_cache_is_untouched_so_sessions_can_share_one_decoder(vae, decoder) -> None: + vae.clear_cache() + before = [entry for entry in vae._feat_map] + decoder.decode_chunk(_latent(3), decoder.new_decode_state("s")) + assert vae._feat_map == before, "streaming decode must not write the module-owned cache" + + +# -------------------------------------------------------------------------- +# Causal geometry and boundedness +# -------------------------------------------------------------------------- + + +@torch.no_grad() +def test_opening_chunk_is_shorter_and_later_chunks_are_full_length(vae, decoder) -> None: + """(n - 1) * factor + 1 for the opening chunk, n * factor after it.""" + factor = 2 ** sum(VAE_CONFIG["temperal_downsample"]) + state = decoder.new_decode_state("s") + first = decoder.decode_chunk(_latent(3), state) + second = decoder.decode_chunk(_latent(3, seed=5), state) + + assert first.shape[2] == (3 - 1) * factor + 1 + assert second.shape[2] == 3 * factor + assert state.chunks_decoded == 2 + assert state.frames_decoded == 6 + + +@torch.no_grad() +def test_resident_state_does_not_grow_with_session_length(vae, decoder) -> None: + """Each causal convolution retains at most CACHE_T temporal slices.""" + state = decoder.new_decode_state("s") + decoder.decode_chunk(_latent(3), state) + baseline = state.nbytes() + assert baseline > 0 + + for step in range(8): + decoder.decode_chunk(_latent(3, seed=step + 10), state) + assert state.nbytes() == baseline + + assert sum(state.nbytes_by_device().values()) == baseline + + +@torch.no_grad() +def test_release_returns_the_session_to_its_pre_stream_state(vae, decoder) -> None: + state = decoder.new_decode_state("s") + decoder.decode_chunk(_latent(3), state) + assert state.started and state.nbytes() > 0 + + decoder.release(state) + assert state.nbytes() == 0 + assert not state.started + assert state.chunks_decoded == 0 + + # A released session restarts rather than resumes, so its first chunk is + # short again -- the same semantics reset() has for AR KV. + factor = 2 ** sum(VAE_CONFIG["temperal_downsample"]) + assert decoder.decode_chunk(_latent(3), state).shape[2] == (3 - 1) * factor + 1 + + +# -------------------------------------------------------------------------- +# Contract +# -------------------------------------------------------------------------- + + +def test_the_vaes_execution_context_is_entered_for_every_frame(vae) -> None: + """Decode must run inside the autoencoder's own execution context. + + Every public compute entry point on ``OmniAutoencoderKLWan`` enters + ``_execution_context()``, which is where autocast is established for + fp16/bf16 parameters -- and on NPU the only place it is established at all. + Driving ``post_quant_conv`` and ``decoder`` directly bypasses it unless this + decoder enters it itself, which would compute streamed chunks under a + different numeric context from the whole-clip decode they are compared to. + """ + entered: list[str] = [] + + class _ContextualVAE: + """The parts of the contract this decoder uses, plus the hook.""" + + config = vae.config + dtype = torch.float32 + + def __init__(self) -> None: + self.decoder = vae.decoder + self.post_quant_conv = vae.post_quant_conv + self._cached_conv_counts = vae._cached_conv_counts + + @contextmanager + def _execution_context(self): + entered.append("enter") + try: + yield + finally: + entered.append("exit") + + decoder = WanStreamingDecoder(_ContextualVAE()) + with torch.no_grad(): + decoder.decode_chunk(_latent(3), decoder.new_decode_state("s")) + + # One context for the call, not one per frame, and it closes. + assert entered == ["enter", "exit"] + + +def test_a_plain_diffusers_autoencoder_needs_no_execution_context(decoder) -> None: + """The hook is optional: a VAE without one still decodes.""" + with torch.no_grad(): + out = decoder.decode_chunk(_latent(1), decoder.new_decode_state("s")) + assert out.shape[2] == 1 + + +def test_wan_decoder_satisfies_the_protocol(decoder) -> None: + assert isinstance(decoder, SupportsStreamingDecode) + + +def test_state_from_another_decoder_is_rejected(decoder) -> None: + foreign = StreamingDecodeState(session_id="s", feat_map=[None]) + with pytest.raises(ValueError, match="does not belong to this decoder"): + decoder.decode_chunk(_latent(1), foreign) + + +def test_a_non_wan_module_is_rejected() -> None: + with pytest.raises(TypeError, match="Wan-family autoencoder"): + WanStreamingDecoder(object()) + + +def test_empty_and_malformed_latents_are_rejected(decoder) -> None: + state = decoder.new_decode_state("s") + with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"): + decoder.decode_chunk(torch.zeros(1, 4, 4, 4), state) + with pytest.raises(ValueError, match="at least one frame"): + decoder.decode_chunk(torch.zeros(1, 4, 0, LATENT_H, LATENT_W), state) + + +def test_a_batch_is_refused_rather_than_threaded_through_one_cache(decoder) -> None: + """One state is one session, so a batch has no correct interpretation here. + + ``AutoencoderKLWan.decode`` handles a batch by decoding each sample with + its own cleared cache (``use_slicing``). This decoder carries exactly one + temporal context, so a batch would mix samples into it and still return + correctly shaped pixels. + """ + state = decoder.new_decode_state("s") + batched = torch.zeros(2, VAE_CONFIG["z_dim"], 3, LATENT_H, LATENT_W) + with pytest.raises(ValueError, match="one sample per session state"): + decoder.decode_chunk(batched, state) + assert state.frames_decoded == 0, "a refused call must not advance the session" + + +def test_session_id_must_be_meaningful(decoder) -> None: + with pytest.raises(ValueError, match="session_id"): + decoder.new_decode_state(" ") + + +def test_declared_state_bytes_scales_with_area_and_dtype(vae) -> None: + # A local decoder, not the module-scoped fixture: the measured constant + # belongs to this test, and leaking it would change what a reordered run + # of the default-None case sees. Passing it here is also the constructor + # argument's intended use. + # Measured for the shipped checkpoint: 37832 KiB of fp32 cache at 64x64. + decoder = WanStreamingDecoder(vae, bytes_per_pixel_fp32=37832 * 1024 / (64 * 64)) + fp32 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.float32) + bf16 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.bfloat16) + half_area = decoder.declared_state_bytes(height=240, width=832, dtype=torch.float32) + assert bf16 == fp32 // 2 + assert half_area == fp32 // 2 + assert fp32 / 2**20 == pytest.approx(3602.0, abs=5.0) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index a58eee935..86d8e70b8 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -170,10 +170,49 @@ class _RecordingTransformer(nn.Module): return {name for name, _ in self.loaded_weights} +class _StubCausalDecoder: + """Stand-in for the Wan causal decoder's cache protocol. + + Reproduces only what streaming depends on: each call walks ``feat_idx`` + across one cache slot per causal convolution, and the session's opening + frame (``first_chunk``) expands to a single raw frame while every later + frame expands by the temporal factor. + """ + + NUM_CONVS = 3 + + def __init__(self) -> None: + self.first_chunk_flags: list[bool] = [] + # Identity of the cache each call wrote through, so an interleaved run + # can assert that a session never advanced another session's context. + self.cache_ids: list[int] = [] + + def __call__(self, x: torch.Tensor, *, feat_cache, feat_idx, first_chunk: bool = False): + self.first_chunk_flags.append(bool(first_chunk)) + self.cache_ids.append(id(feat_cache)) + for _ in range(self.NUM_CONVS): + index = feat_idx[0] + feat_cache[index] = x.detach().clone() + feat_idx[0] = index + 1 + num_frames = 1 if first_chunk else 4 + return x.new_zeros(x.shape[0], 3, num_frames, x.shape[-2] * 8, x.shape[-1] * 8) + + class _StubVAE(_FakePretrained): dtype = torch.float32 - def __init__(self): + def __init__(self, *, streaming: bool = True): + if streaming: + self.decoder = _StubCausalDecoder() + # Recorded, not discarded: on the streaming branch this is the only + # place the tensor handed to the decoder can be observed, and it is + # what carries the checkpoint's latent rescale. + self.post_quant_inputs: list[torch.Tensor] = [] + self.post_quant_conv = self._record_post_quant + self._cached_conv_counts = {"decoder": _StubCausalDecoder.NUM_CONVS, "encoder": 2} + # Module-owned cache the shared VAE keeps for whole-clip decode; + # streaming must never write through it. + self._feat_map = ["module-owned"] self.config = SimpleNamespace( z_dim=16, scale_factor_temporal=4, @@ -185,6 +224,10 @@ class _StubVAE(_FakePretrained): self.decode_inputs: list[torch.Tensor] = [] self.on_decode = None + def _record_post_quant(self, latent: torch.Tensor) -> torch.Tensor: + self.post_quant_inputs.append(latent.detach().clone()) + return latent + def encode(self, video: torch.Tensor): self.encode_inputs.append(video.detach().clone()) latent_frames = (video.shape[2] - 1) // 4 + 1 @@ -455,7 +498,41 @@ def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None: assert spec.sink_frames == 3 assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)] assert spec.cross_attention_lengths == {"text": 512} - assert spec.model_owned_state_bytes_per_session == 9_600 + # The image condition (9,600 B) plus the streaming decoder's own per-session + # temporal cache. The cache dwarfs the condition at any real resolution, and + # the runner subtracts this field from the KV budget before sizing its pools, + # so leaving the cache out of it over-commits device memory. + assert spec.model_owned_state_bytes_per_session == 9_600 + 2_421_248 + + +def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None: + """The declaration scales with output area, and is the condition alone when + the VAE cannot stream.""" + module = _load_pipeline_module() + + def spec_pipeline(): + pipeline = _pipeline(module) + # The spec reads the constructed attention's head count; the head count + # plays no part in model-owned bytes, so any value will do here. + pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1) + return pipeline + + small = spec_pipeline().ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session + + wide = spec_pipeline() + wide._ar_width = wide._ar_width * 2 + widened = wide.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session + # Twice the area: the decoder cache doubles, so the total grows by far more + # than the image condition alone could account for. + assert widened - small > 2_000_000 + + # A VAE without the causal-cache contract cannot stream, so there is no + # decoder cache to declare and the figure falls back to the condition. + bare = spec_pipeline() + for attribute in ("decoder", "post_quant_conv", "_cached_conv_counts"): + if hasattr(bare.vae, attribute): + delattr(bare.vae, attribute) + assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 9_600 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None: @@ -1852,6 +1929,27 @@ def _run_stepwise(pipeline, state): return outputs +def _stepwise_chunks(pipeline, state, ar_state): + """Yield one chunk at a time, binding AR state around each invocation. + + The runner binds runner-owned KV for the duration of one stepwise + invocation and releases it again, so a generator that suspends at each + chunk boundary is what lets two requests be driven the way the scheduler + drives them at session_capacity > 1: A-chunk0, B-chunk0, A-chunk1, ... + """ + with pipeline.bind_ar_diffusion_state(state.request_id, ar_state): + pipeline.prepare_encode(state) + while not state.request_denoise_completed: + output = None + with pipeline.bind_ar_diffusion_state(state.request_id, ar_state): + noise = pipeline.denoise_step(None, states=[state]) + pipeline.step_scheduler(state, noise) + if state.chunk_denoise_completed: + output = pipeline.post_decode(state) + if output is not None: + yield output + + def test_pipeline_declares_step_execution_support() -> None: module = _load_pipeline_module() pipeline = _pipeline(module) @@ -2039,12 +2137,11 @@ def test_stepwise_emits_latents_without_touching_the_vae() -> None: assert pipeline.vae.decode_inputs == [] -def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None: - """A non-latent request must stream pixels under the primary "video" key.""" +def _capture_video_processor(monkeypatch) -> list[tuple[tuple[int, ...], str]]: + """Record what reaches postprocess_video, and name each chunk's output.""" import diffusers.video_processor as video_processor_module - module = _load_pipeline_module() - processed = [] + processed: list[tuple[tuple[int, ...], str]] = [] class VideoProcessor: def __init__(self, *, vae_scale_factor): @@ -2055,9 +2152,21 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non return f"frames-{len(processed)}" monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor) + return processed + + +def _streaming_pipeline(module): pipeline = _pipeline(module, transformer=_RecordingTransformer()) pipeline._ar_height = 16 pipeline._ar_width = 16 + return pipeline + + +def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None: + """A non-latent request must stream pixels under the primary "video" key.""" + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) state = _stepwise_state(num_frames=21) state.sampling.output_type = "np" fake = _FakeARState(state.request_id) @@ -2065,14 +2174,268 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non with pipeline.bind_ar_diffusion_state(state.request_id, fake): outputs = _run_stepwise(pipeline, state) - # One decode per AR block, each seeing only that block's three latent frames. - assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2 assert [output_type for _, output_type in processed] == ["np", "np"] assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}] # Identity metadata survives the switch to pixel output. assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1] +def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> None: + """Chunk N + 1 must continue chunk N instead of restarting the decoder. + + Only the session's opening latent frame collapses to a single raw frame, + so the streamed timeline carries the same frame count an offline decode of + the same latents would: 9 frames, then 12. + """ + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + _run_stepwise(pipeline, state) + + assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)] + # One opening frame for the session, not one per chunk. + assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5 + # The whole-clip decode path is not involved, and the module-owned cache + # the shared VAE keeps for it is untouched. + assert pipeline.vae.decode_inputs == [] + assert pipeline.vae._feat_map == ["module-owned"] + + +def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None: + """Decoder state is keyed by request id and released with the AR session.""" + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(request_id="req-stream", num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + _run_stepwise(pipeline, state) + + decode_state = pipeline._streaming_decode_states["req-stream"] + assert decode_state.session_id == "req-stream" + assert decode_state.frames_decoded == 6 + assert decode_state.nbytes() > 0 + + pipeline.close_ar_diffusion_session("req-stream") + assert "req-stream" not in pipeline._streaming_decode_states + # Release drops the cache itself, not just the pipeline's reference to it. + assert decode_state.nbytes() == 0 + assert decode_state.frames_decoded == 0 + + +def test_streaming_decode_state_is_dropped_on_session_reset(monkeypatch) -> None: + """A reset session must not resume the temporal context it just abandoned.""" + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(request_id="req-reset", num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + _run_stepwise(pipeline, state) + + pipeline.reset_ar_diffusion_session("req-reset") + assert pipeline._streaming_decode_states == {} + # Releasing a session that never decoded is a no-op, not an error. + pipeline.reset_ar_diffusion_session("req-reset") + + +def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> None: + """Two sessions ticking alternately must not share one temporal cache. + + Running one rollout to completion before the other starts would only show + that the dict is keyed by request id. The failure worth catching is chunk N + of one session advancing the other's cache mid-rollout, which needs the + interleaving the scheduler actually produces. + """ + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + states = [] + for request_id in ("req-a", "req-b"): + state = _stepwise_state(request_id=request_id, num_frames=21) + state.sampling.output_type = "np" + states.append(state) + runs = [_stepwise_chunks(pipeline, state, _FakeARState(state.request_id)) for state in states] + + chunks: list[tuple[str, int]] = [] + while runs: + for run in list(runs): + try: + output = next(run) + except StopIteration: + runs.remove(run) + continue + metadata = output.output["metadata"]["ar_diffusion"] + chunks.append((metadata["session_id"], metadata["chunk_index"])) + + # A-chunk0, B-chunk0, A-chunk1, B-chunk1: the sessions really are interleaved. + assert chunks == [("req-a", 0), ("req-b", 0), ("req-a", 1), ("req-b", 1)] + + caches = pipeline._streaming_decode_states + assert set(caches) == {"req-a", "req-b"} + assert caches["req-a"].feat_map is not caches["req-b"].feat_map + # Every frame of a session was decoded through that session's own cache, + # three frames per chunk, and each session opened exactly one rollout. + cache_a, cache_b = (id(caches["req-a"].feat_map), id(caches["req-b"].feat_map)) + assert pipeline.vae.decoder.cache_ids == [cache_a] * 3 + [cache_b] * 3 + [cache_a] * 3 + [cache_b] * 3 + assert pipeline.vae.decoder.first_chunk_flags == ( + [True, False, False] + [True, False, False] + [False] * 3 + [False] * 3 + ) + assert [shape for shape, _ in processed] == [ + (1, 3, 9, 16, 16), + (1, 3, 9, 16, 16), + (1, 3, 12, 16, 16), + (1, 3, 12, 16, 16), + ] + + +def test_streaming_decode_failure_drops_the_half_advanced_cache(monkeypatch) -> None: + """A failed chunk must not leave a cache a later chunk would continue from.""" + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(request_id="req-boom", num_frames=21) + state.sampling.output_type = "np" + + def explode(*args, **kwargs): + raise RuntimeError("decoder blew up") + + pipeline.vae.decoder = explode + + with pytest.raises(RuntimeError, match="decoder blew up"): + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + _run_stepwise(pipeline, state) + + assert pipeline._streaming_decode_states == {} + + +def _enable_vae_tiling(pipeline, *, tile_sample_min: int) -> None: + """Configure the stub VAE the way registry.py configures the real one.""" + pipeline.vae.use_tiling = True + pipeline.vae.spatial_compression_ratio = 8 + pipeline.vae.tile_sample_min_height = tile_sample_min + pipeline.vae.tile_sample_min_width = tile_sample_min + + +def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None: + """The streaming branch must invert the checkpoint's latent statistics. + + Nothing downstream can catch a miss here: the pixels stay finite and + correctly shaped, only wrongly scaled. Request mode and the fallback path + observe the rescale through ``vae.decode``'s recorded input, which streaming + never reaches, so this asserts on what ``post_quant_conv`` received. + """ + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + while not state.chunk_denoise_completed: + pipeline.step_scheduler(state, pipeline.denoise_step(None, states=[state])) + model_space = state.latents.clone() + pipeline.post_decode(state) + + shape = (1, -1, 1, 1, 1) + latent_mean = torch.as_tensor(pipeline.vae.config.latents_mean, dtype=model_space.dtype).view(*shape) + latent_std = torch.as_tensor(pipeline.vae.config.latents_std, dtype=model_space.dtype).view(*shape) + expected = (model_space * latent_std + latent_mean).to(dtype=pipeline.vae.dtype) + + recorded = torch.cat(pipeline.vae.post_quant_inputs, dim=2) + assert recorded.shape == expected.shape, "one recorded frame per latent frame in the chunk" + torch.testing.assert_close(recorded, expected, rtol=0, atol=0) + # And the rescale is not a no-op on this fixture, so the assertion bites. + assert not torch.equal(expected, model_space.to(dtype=pipeline.vae.dtype)) + + +def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None: + """A streamed rollout must report decode time, not zero. + + The profiler wraps ``vae.decode``, which streaming never calls, so without + its own stage the one stage this path changes would be invisible -- and a + perf regression in the streaming decoder would be too. + """ + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _pipeline( + module, + transformer=_RecordingTransformer(), + od_config=_od_config( + enable_diffusion_pipeline_profiler=True, + ), + ) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + outputs = _run_stepwise(pipeline, state) + + stage = f"{type(pipeline).__name__}._streaming_decode_chunk" + durations = pipeline.stage_durations + assert stage in durations, f"streamed decode is unprofiled; stages seen: {sorted(durations)}" + assert durations[stage] > 0.0 + # The stage reaches the consumer on every streamed chunk, not just at the end. + assert all(stage in (output.stage_durations or {}) for output in outputs) + + +def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> None: + """Tiled decode drives its own cache, so a tiling shape keeps per-chunk decode.""" + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + # A latent of 2x2 exceeds a one-latent-cell tile, so vae.decode would tile. + _enable_vae_tiling(pipeline, tile_sample_min=8) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + outputs = _run_stepwise(pipeline, state) + + # One whole-clip decode per AR block, each seeing only that block's frames. + assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2 + assert pipeline.vae.decoder.first_chunk_flags == [] + assert pipeline._streaming_decode_states == {} + # Nine frames per chunk instead of 9 then 12: the frame loss being removed. + assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)] + assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}] + + +def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monkeypatch) -> None: + """``use_tiling`` alone does not tile, so such a shape must still stream. + + ``AutoencoderKLWan._decode`` tiles only when the latent exceeds a tile. A + gate on the flag alone would give up streaming -- and a frame per block -- + on a configuration that would never have tiled. + """ + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + # A 2x2 latent is within a tile this size, so vae.decode would not tile. + _enable_vae_tiling(pipeline, tile_sample_min=256) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + _run_stepwise(pipeline, state) + + assert pipeline.vae.decode_inputs == [] + assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5 + assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)] + + def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None: module = _load_pipeline_module() resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry( diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index d2dfde3d4..6e32844fc 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -18,6 +18,7 @@ import torch.nn.functional as F from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel +from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -62,11 +63,26 @@ from vllm_omni.experimental.ar_diffusion.capability import ( ARDiffusionKVBranchSpec, ARDiffusionKVCacheSpec, ) +from vllm_omni.experimental.ar_diffusion.streaming_decode import ( + StreamingDecodeState, + WanStreamingDecoder, +) from vllm_omni.experimental.ar_diffusion.tick_protocol import ( ARDiffusionChunkMetadata, ARDiffusionTickRequest, ) +logger = init_logger(__name__) + +# Resident streaming-decoder cache per session, as bytes per output pixel in +# fp32: measured once on the reference checkpoint at 37832 KiB for a 64x64 +# output. Every cached tensor is [B, C, <= CACHE_T, H/8, W/8] per causal +# convolution, so the total scales with height * width and never with session +# length; tests/diffusion/ar_diffusion/test_streaming_decode.py pins that +# scaling, and at 832x480 in bf16 this yields the 1801 MiB the decoder was +# measured to hold. +_STREAMING_DECODE_BYTES_PER_PIXEL_FP32 = 37832 * 1024 / (64 * 64) + if TYPE_CHECKING: from diffusers.video_processor import VideoProcessor from tqdm.std import tqdm as TqdmProgressBar @@ -552,10 +568,20 @@ class LingBotWorldCausalDMDPipeline( self._ar_width = int(model_config.get("ar_diffusion_width", 832)) self._ar_diffusion_kv_state: ARDiffusionKVState | None = None self._ar_sessions: dict[str, _LingBotARSessionState] = {} + # One temporal decoder cache per stepwise session, keyed the way the + # runner keys AR sessions (session_id == request_id) so both are + # released together rather than through two independent lifecycles. + self._streaming_decode_states: dict[str, StreamingDecodeState] = {} + self._cached_streaming_decoder: WanStreamingDecoder | None = None + self._streaming_decode_unsupported = False self.setup_diffusion_pipeline_profiler( profiler_targets=[ "vae.encode", "vae.decode", + # Streaming decode never reaches vae.decode, so it needs its + # own stage or a streamed rollout reports no decode time at + # all -- the one stage this path changes. + "_streaming_decode_chunk", "_generate_block", "text_encoder.forward", "tokenizer.forward", @@ -608,7 +634,32 @@ class LingBotWorldCausalDMDPipeline( _MAX_SEQUENCE_LENGTH, ), ), - model_owned_state_bytes_per_session=condition_bytes_per_session, + model_owned_state_bytes_per_session=( + condition_bytes_per_session + self._streaming_decode_bytes_per_session() + ), + ) + + def _streaming_decode_bytes_per_session(self) -> int: + """Resident streaming-decoder bytes one session can hold, for admission. + + ``model_owned_state_bytes_per_session`` is what the runner subtracts + from the KV budget before sizing the block pools, and this cache is + three orders of magnitude larger than the image condition that used to + be the only entry: leaving it out lets the profiler report a + configuration as fitting and then OOM partway through the second + session's rollout. + + Declared for every session rather than only for pixel-output ones: the + field is documented as a conservative upper bound, it is fixed at load + time, and whether a request asks for pixels is not known then. + """ + decoder = self._streaming_decoder() + if decoder is None: + return 0 + return decoder.declared_state_bytes( + height=self._ar_height, + width=self._ar_width, + dtype=self.vae.dtype, ) @contextmanager @@ -629,9 +680,11 @@ class LingBotWorldCausalDMDPipeline( def reset_ar_diffusion_session(self, session_id: str) -> None: self._ar_sessions.pop(session_id, None) + self._release_streaming_decode_state(session_id) def close_ar_diffusion_session(self, session_id: str) -> None: self._ar_sessions.pop(session_id, None) + self._release_streaming_decode_state(session_id) def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs: if req.num_reqs != 1 or len(req.prompts) != 1: @@ -1279,15 +1332,90 @@ class LingBotWorldCausalDMDPipeline( stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) + def _vae_would_tile(self, latents: torch.Tensor) -> bool: + """Whether ``vae.decode`` would take its tiled path for this shape. + + Tiled decode splits a frame spatially and drives the module's own + ``feat_cache`` across those tiles, so a per-session cache threaded + through it is a separate design and is not attempted here. The test + mirrors ``AutoencoderKLWan._decode``: the flag alone does not tile, the + latent also has to exceed a tile, so a sub-threshold resolution with + tiling enabled still streams. A VAE that sets the flag but reports no + tile size is treated as tiling, because the safe answer when the + threshold is unknown is the one that keeps ``vae.decode`` in charge. + """ + if not bool(getattr(self.vae, "use_tiling", False)): + return False + compression = int(getattr(self.vae, "spatial_compression_ratio", None) or self.vae_scale_factor_spatial) + tile_latent_min_height = int(getattr(self.vae, "tile_sample_min_height", None) or 0) // compression + tile_latent_min_width = int(getattr(self.vae, "tile_sample_min_width", None) or 0) // compression + return bool(latents.shape[-2] > tile_latent_min_height or latents.shape[-1] > tile_latent_min_width) + + def _streaming_decoder(self) -> WanStreamingDecoder | None: + """The session-owned streaming decoder, or ``None`` on a VAE that cannot stream. + + Only the decoder's own capability is decided here, and the decision is + cached including the negative. Whether a particular *shape* would have + been tiled is a per-call question and lives in :meth:`_vae_would_tile`. + """ + if self._streaming_decode_unsupported: + return None + if self._cached_streaming_decoder is not None: + return self._cached_streaming_decoder + try: + decoder = WanStreamingDecoder( + self.vae, + bytes_per_pixel_fp32=_STREAMING_DECODE_BYTES_PER_PIXEL_FP32, + ) + # Materialize the cache geometry now: a decoder that cannot report + # it would otherwise fail on the first chunk instead. + _ = decoder.num_causal_convs + except TypeError as exc: + logger.warning( + "LingBot streaming decode is unavailable (%s); AR blocks will be decoded independently, " + "which drops each block's opening frame and restarts temporal context at every boundary.", + exc, + ) + self._streaming_decode_unsupported = True + return None + self._cached_streaming_decoder = decoder + return decoder + + def _streaming_decode_state(self, decoder: WanStreamingDecoder, session_id: str) -> StreamingDecodeState: + state = self._streaming_decode_states.get(session_id) + if state is None: + state = decoder.new_decode_state(session_id) + self._streaming_decode_states[session_id] = state + return state + + def _release_streaming_decode_state(self, session_id: str) -> None: + """Drop one session's decoder cache. Safe for a session that never decoded. + + The state is released directly rather than through + ``WanStreamingDecoder.release``, which is the same one-line call: the + decoder's method exists for a caller holding only the + ``SupportsStreamingDecode`` protocol, and this pipeline holds the state. + """ + state = self._streaming_decode_states.pop(session_id, None) + if state is not None: + state.release() + def _decode_chunk_to_pixels( - self, latents: torch.Tensor, *, output_type: str + self, latents: torch.Tensor, *, output_type: str, session_id: str | None = None ) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]: """Decode one AR block so streaming consumers receive pixels, not latents. - Blocks are decoded independently, which keeps the shared VAE stateless - and lets a rollout be served without a per-session temporal cache. That - is the same shape Helios streams today; cross-block ``feat_cache`` - continuity is a separate decision and is not made here. + With a ``session_id`` the block is decoded through that session's own + temporal cache, so chunk *N + 1* continues chunk *N*: the causal decoder + expands the session's opening latent frame to one raw frame and every + later one to the full temporal factor, which is the frame timeline an + offline whole-clip decode of the same latents produces. The cache lives + on the session rather than on the shared VAE module, so concurrent + sessions cannot overwrite one another's temporal context. + + Without a ``session_id`` -- or on a VAE that cannot stream -- blocks are + decoded independently: every block re-expands its own opening frame, so + the rollout loses frames and continuity at each block boundary. The registered post-process hook returns a ``{"payload", "metadata"}`` envelope untouched, so the pixel conversion it would do for a bare @@ -1295,9 +1423,34 @@ class LingBotWorldCausalDMDPipeline( """ latent_mean, latent_std = self._vae_latent_stats(latents) vae_latents = (latents * latent_std + latent_mean).to(dtype=self.vae.dtype) - video = self.vae.decode(vae_latents, return_dict=False)[0] + decoder = None + if session_id is not None and not self._vae_would_tile(vae_latents): + decoder = self._streaming_decoder() + if decoder is None: + video = self.vae.decode(vae_latents, return_dict=False)[0] + else: + video = self._streaming_decode_chunk(decoder, vae_latents, cast(str, session_id)) return self._video_processor().postprocess_video(video, output_type=output_type) + def _streaming_decode_chunk( + self, decoder: WanStreamingDecoder, vae_latents: torch.Tensor, session_id: str + ) -> torch.Tensor: + """Decode one chunk through ``session_id``'s temporal cache. + + A separate method so the profiler has a stage to wrap: streaming never + calls ``vae.decode``, and without this a profiled rollout would report + no decode time for the one stage this path changes. + """ + state = self._streaming_decode_state(decoder, session_id) + try: + return decoder.decode_chunk(vae_latents, state) + except Exception: + # A half-advanced cache cannot be resumed: the session's next chunk + # would continue from a temporal context that was never completed, + # so drop it here rather than at close. + self._release_streaming_decode_state(session_id) + raise + def _video_processor(self) -> VideoProcessor: processor = getattr(self, "_cached_video_processor", None) if processor is None: @@ -1566,7 +1719,13 @@ class LingBotWorldCausalDMDPipeline( if inputs.output_type == "latent": payload: dict[str, Any] = {"latents": latents} else: - payload = {"video": self._decode_chunk_to_pixels(latents, output_type=inputs.output_type)} + payload = { + "video": self._decode_chunk_to_pixels( + latents, + output_type=inputs.output_type, + session_id=state.request_id, + ) + } output = { "payload": payload, "metadata": { diff --git a/vllm_omni/experimental/ar_diffusion/streaming_decode.py b/vllm_omni/experimental/ar_diffusion/streaming_decode.py new file mode 100644 index 000000000..e4316c83c --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/streaming_decode.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Session-owned streaming VAE decode for realtime AR-Diffusion. + +A causal video decoder already decodes one latent frame at a time, carrying a +bounded temporal cache between frames. What it does not do today is carry that +cache *between calls*: every decode call clears it, so a chunk can only be +decoded as part of a whole clip. + +This module moves the cache from the decoder module to the session, which is +the entire difference between "decode a clip" and "stream a session": + +* the cache becomes :class:`StreamingDecodeState`, threaded in and out by the + caller, so two concurrent sessions cannot overwrite each other's temporal + context -- a failure that produces a plausible but wrong video while every + timing metric stays green; +* ``first_chunk`` becomes a property of the session's first frame rather than + of the call, which is what makes chunk ``N + 1`` continue chunk ``N``. + +Resident state is bounded by construction: each causal convolution retains at +most ``CACHE_T`` temporal slices, so the cache is a function of resolution and +never of session length. It is not free -- see :meth:`StreamingDecodeState.nbytes`, +which is the quantity session admission has to account for. + +Measured on the reference checkpoint at 832x480 in bf16, on one RTX PRO 6000 +with torch 2.11: chunking is neutral -- a session decoded in four chunks is +bit-identical to the same session decoded in one call -- and resident decoder +state holds at 1801 MiB across every chunk. What a session gains is structural +rather than a speedup: a chunk becomes decodable on its own, so pixels leave on +the block that produced them instead of after the last one, and the session's +opening frame is expanded once rather than once per block. + +Note what "neutral" is measured against. ``AutoencoderKLWan._decode`` applies +``post_quant_conv`` to the whole latent at once, while the tiled path in this +repo -- and this decoder -- apply it per frame. It is a 1x1x1 convolution, so +the two are mathematically identical, but they can select different kernels; +on bf16 CUDA that seeds a 0.004 difference which the causal convolutions +downstream amplify to 0.047. Comparing against ``_decode`` therefore measures +where ``post_quant_conv`` runs, not anything about streaming, and the result +moves with the torch version. The control that isolates chunking is the same +session decoded in a single call through this same path. + +The consumer is the stepwise execution path: one ``generate()`` is one +session, ``post_decode`` emits one chunk per AR block, and the pipeline holds +one state per ``request_id``, releasing it when the runner closes or resets +that AR session -- so decoder state and KV state share a single lifetime. + +Only ``torch`` is imported here so the contract can be exercised without a +device, a checkpoint, or the distributed VAE stack. +""" + +from __future__ import annotations + +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +import torch + + +@dataclass +class StreamingDecodeState: + """One session's temporal decoder cache. + + ``feat_map`` holds one entry per causal convolution. Entries are tensors, + ``None`` before that convolution has run, or a non-tensor sentinel the + decoder uses for its first upsample step, so byte accounting must skip + anything that is not a tensor. + + The cache has **no recompute source**: it is the compressed tail of every + frame decoded so far and cannot be rebuilt from the prompt or the latents. + Releasing it ends the stream; there is no partial eviction. + """ + + session_id: str + feat_map: list[Any] + conv_idx: list[int] = field(default_factory=lambda: [0]) + frames_decoded: int = 0 + chunks_decoded: int = 0 + + def nbytes(self) -> int: + """Resident bytes of the cache, skipping non-tensor entries.""" + return sum(t.numel() * t.element_size() for t in self.feat_map if torch.is_tensor(t)) + + def nbytes_by_device(self) -> dict[str, int]: + """Resident bytes grouped by device, for admission accounting.""" + totals: dict[str, int] = {} + for tensor in self.feat_map: + if torch.is_tensor(tensor): + key = str(tensor.device) + totals[key] = totals.get(key, 0) + tensor.numel() * tensor.element_size() + return totals + + @property + def started(self) -> bool: + return self.frames_decoded > 0 + + def release(self) -> None: + """Drop the cache and return the session to its pre-stream state.""" + self.feat_map = [None] * len(self.feat_map) + self.conv_idx = [0] + self.frames_decoded = 0 + self.chunks_decoded = 0 + + +@runtime_checkable +class SupportsStreamingDecode(Protocol): + """A decoder that can emit a session's chunks incrementally. + + The lifecycle mirrors the AR session's: ``new_decode_state`` at session + creation, ``decode_chunk`` per committed chunk, and ``release`` on reset or + close, so decoder state and KV state are released together rather than + through two independent lifecycles. + """ + + def new_decode_state(self, session_id: str) -> StreamingDecodeState: + """Create empty decoder state for a new session.""" + ... + + def decode_chunk(self, latent: torch.Tensor, state: StreamingDecodeState) -> torch.Tensor: + """Decode one committed chunk, advancing ``state`` in place.""" + ... + + def declared_state_bytes(self, *, height: int, width: int, dtype: torch.dtype) -> int: + """Upper bound on resident decoder bytes per session at this shape.""" + ... + + def release(self, state: StreamingDecodeState) -> None: + """Drop ``state``'s cache and return it to its pre-stream state. + + Declared here, not just implemented on :class:`WanStreamingDecoder`, + so a caller programming against this protocol -- a session runtime + releasing decoder state alongside KV state on close or failure -- can + rely on it being present without depending on the concrete class. + """ + ... + + +class WanStreamingDecoder: + """Streaming decode over a Wan-family causal autoencoder. + + Wraps any module exposing ``decoder``, ``post_quant_conv`` and the cached + causal-convolution counts. It never touches the module's own ``_feat_map``, + so a module shared by several sessions stays stateless between calls. + """ + + def __init__(self, vae: Any, *, bytes_per_pixel_fp32: float | None = None) -> None: + for attribute in ("decoder", "post_quant_conv"): + if not hasattr(vae, attribute): + raise TypeError(f"Streaming decode requires a Wan-family autoencoder exposing {attribute!r}.") + self._vae = vae + # A one-off measured constant for this checkpoint; see declared_state_bytes. + # Optional at construction because most callers never need admission + # accounting, but a constructor argument (rather than requiring an + # external attribute assignment after the fact) means the value can't + # be silently unset for a caller that does. + self._bytes_per_pixel_fp32 = bytes_per_pixel_fp32 + + def _execution_context(self) -> AbstractContextManager[None]: + """The autoencoder's own execution context, when it exposes one. + + Every public compute entry point on ``OmniAutoencoderKLWan`` -- + ``encode``, ``decode``, ``decode_with_chunks``, ``tile_exec`` -- runs + inside ``_execution_context()``, which is where autocast is entered for + fp16/bf16 parameters (and on NPU is the only place + ``torch.npu.amp.autocast`` is entered at all). This decoder drives + ``post_quant_conv`` and ``decoder`` directly, so it has to enter the + same context or streamed chunks compute under a different numeric and + dispatch context from the whole-clip decode they are compared against. + + A plain ``diffusers`` autoencoder has no such hook, and then there is + nothing to enter. + """ + context = getattr(self._vae, "_execution_context", None) + if callable(context): + return context() + return nullcontext() + + @property + def num_causal_convs(self) -> int: + counts = getattr(self._vae, "_cached_conv_counts", None) + if not isinstance(counts, dict) or "decoder" not in counts: + raise TypeError("The autoencoder does not expose cached decoder convolution counts.") + return int(counts["decoder"]) + + def new_decode_state(self, session_id: str) -> StreamingDecodeState: + if not isinstance(session_id, str) or not session_id.strip(): + raise ValueError("session_id must be a non-empty string.") + return StreamingDecodeState(session_id=session_id, feat_map=[None] * self.num_causal_convs) + + def decode_chunk(self, latent: torch.Tensor, state: StreamingDecodeState) -> torch.Tensor: + """Decode ``latent`` as the continuation of ``state``'s session. + + ``latent`` is ``[1, C, T, H, W]`` in latent space, already rescaled by + the pipeline's latent statistics. Frames are decoded one at a time, the + same loop the non-streaming path runs, with the cache carried across + calls instead of cleared. The returned tensor is ``[1, 3, T', H', W']`` + where ``T'`` is smaller for a session's first chunk: the opening latent + frame expands to a single raw frame and every later one to the full + temporal factor. + """ + if latent.ndim != 5: + raise ValueError(f"latent must be [B, C, T, H, W]; got shape {tuple(latent.shape)}.") + if latent.shape[0] != 1: + # One state is one session's temporal context, so a batch would + # thread several samples through a single cache. ``decode()`` + # handles a batch by decoding each sample with its own cleared + # cache (``use_slicing``); there is no equivalent here, and the + # streaming contract is one session per state, so refuse rather + # than silently mix samples. + raise ValueError( + "streaming decode takes one sample per session state; " + f"got batch size {latent.shape[0]}. Decode each sample through its own state." + ) + if len(state.feat_map) != self.num_causal_convs: + raise ValueError( + "Decoder state does not belong to this decoder: " + f"{len(state.feat_map)} cache slots against {self.num_causal_convs} causal convolutions." + ) + num_frames = int(latent.shape[2]) + if num_frames == 0: + raise ValueError("latent must carry at least one frame.") + + decoded_frames = [] + with self._execution_context(): + for index in range(num_frames): + state.conv_idx[0] = 0 + frame = self._vae.post_quant_conv(latent[:, :, index : index + 1]) + # first_chunk marks the session's opening frame, not the call's: + # that is what makes chunk N + 1 continue chunk N rather than + # restart the causal expansion. + decoded_frames.append( + self._vae.decoder( + frame, + feat_cache=state.feat_map, + feat_idx=state.conv_idx, + first_chunk=(state.frames_decoded == 0), + ) + ) + state.frames_decoded += 1 + + out = torch.cat(decoded_frames, dim=2) + patch_size = getattr(getattr(self._vae, "config", None), "patch_size", None) + if patch_size is not None: + from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify + + out = unpatchify(out, patch_size=patch_size) + out = torch.clamp(out, min=-1.0, max=1.0) + state.chunks_decoded += 1 + return out + + def release(self, state: StreamingDecodeState) -> None: + state.release() + + def declared_state_bytes(self, *, height: int, width: int, dtype: torch.dtype) -> int: + """Resident decoder bytes per session at this output shape. + + Every cached tensor is ``[B, C_l, <= CACHE_T, H_l, W_l]`` with ``H_l`` + and ``W_l`` fixed fractions of the output, so the total scales exactly + with ``height * width``. Callers that need the constant should measure + it once for their checkpoint and pass it as ``bytes_per_pixel_fp32`` to + the constructor; this helper scales that measured constant so admission + does not have to re-measure per resolution. + """ + constant = self._bytes_per_pixel_fp32 + if constant is None: + raise RuntimeError( + "declared_state_bytes needs a measured bytes-per-pixel constant for this checkpoint; " + "pass bytes_per_pixel_fp32= to WanStreamingDecoder() from a one-off measurement." + ) + element_ratio = torch.empty((), dtype=dtype).element_size() / 4 + return int(constant * height * width * element_ratio) From e284d907b5554038bef6ab861bdb78b0b99dd5ea Mon Sep 17 00:00:00 2001 From: Qi Jia Date: Sun, 13 Sep 2026 20:33:10 +0800 Subject: [PATCH 046/209] [Perf][Frontend] Chunk in-memory image file responses (#7459) Signed-off-by: QI JIA Co-authored-by: QI JIA Co-authored-by: Cursor --- .../benchmark_image_file_response.py | 202 ++++++++++++++++++ .../openai_api/test_image_protocol.py | 75 +++++++ .../entrypoints/openai/protocol/images.py | 14 +- 3 files changed, 289 insertions(+), 2 deletions(-) create mode 100644 tests/benchmarks/benchmark_image_file_response.py create mode 100644 tests/entrypoints/openai_api/test_image_protocol.py diff --git a/tests/benchmarks/benchmark_image_file_response.py b/tests/benchmarks/benchmark_image_file_response.py new file mode 100644 index 000000000..acd8967a6 --- /dev/null +++ b/tests/benchmarks/benchmark_image_file_response.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark newline-split and fixed-size chunking of image file responses. + +Both variants stream the exact payload `ImageGenerationResponse.stream_response()` +produces, so the measurement isolates the response layer: PNG encoding and ZIP +building are shared by the variants and excluded from the timings. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import io +import json +import statistics +import time +from collections.abc import AsyncIterator, Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import cast + +import numpy as np +from PIL import Image +from starlette.responses import StreamingResponse + +from vllm_omni.entrypoints.openai.protocol.images import ( + _FILE_RESPONSE_CHUNK_SIZE, + ImageData, + ImageGenerationResponse, + _iter_file_chunks, +) + + +@dataclass(frozen=True) +class BenchmarkVariant: + label: str + content: Callable[[bytes], Iterable[bytes] | AsyncIterator[memoryview]] + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("value must be at least 1") + return parsed + + +def _encode_png(*, width: int, height: int, seed: int) -> bytes: + """Encode incompressible noise, so the payload looks like a real photo.""" + rng = np.random.default_rng(seed) + pixels = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8) + buffer = io.BytesIO() + Image.fromarray(pixels, mode="RGB").save(buffer, format="png") + return buffer.getvalue() + + +async def _drain(response: StreamingResponse) -> tuple[bytes, int]: + """Run a response through ASGI, returning its body and non-empty frame count.""" + body = bytearray() + frames = 0 + + async def send(message: dict[str, object]) -> None: + nonlocal frames + if message["type"] == "http.response.body" and message["body"]: + frames += 1 + body.extend(cast(bytes, message["body"])) + + async def receive() -> dict[str, str]: + return {"type": "http.disconnect"} + + scope = {"type": "http", "method": "GET", "asgi": {"spec_version": "2.4"}} + await response(scope, receive, send) + return bytes(body), frames + + +async def _build_payload(*, images: int, width: int, height: int, seed: int) -> bytes: + png = _encode_png(width=width, height=height, seed=seed) + response = ImageGenerationResponse( + created=0, + data=[ImageData(b64_json=base64.b64encode(png).decode()) for _ in range(images)], + output_format="png", + size=f"{width}x{height}", + ) + payload, _ = await _drain(response.stream_response()) + return payload + + +async def _measure(variant: BenchmarkVariant, payload: bytes) -> dict[str, object]: + response = StreamingResponse(variant.content(payload), media_type="application/octet-stream") + cpu_start = time.process_time_ns() + wall_start = time.perf_counter_ns() + body, frames = await _drain(response) + wall_ms = (time.perf_counter_ns() - wall_start) / 1_000_000 + process_cpu_ms = (time.process_time_ns() - cpu_start) / 1_000_000 + return { + "label": variant.label, + "wall_ms": wall_ms, + "process_cpu_ms": process_cpu_ms, + "body_frames": frames, + "output_bytes": len(body), + "output_sha256": hashlib.sha256(body).hexdigest(), + } + + +def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]: + selected = [record for record in records if record["label"] == label] + wall_values = [cast(float, record["wall_ms"]) for record in selected] + cpu_values = [cast(float, record["process_cpu_ms"]) for record in selected] + return { + "runs": len(selected), + "body_frames": cast(int, selected[0]["body_frames"]), + "wall_ms": { + "median": statistics.median(wall_values), + "min": min(wall_values), + "max": max(wall_values), + }, + "process_cpu_ms": { + "median": statistics.median(cpu_values), + "min": min(cpu_values), + "max": max(cpu_values), + }, + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--images", type=_positive_int, default=1, help="2 or more streams a ZIP archive") + parser.add_argument("--width", type=_positive_int, default=1024) + parser.add_argument("--height", type=_positive_int, default=1024) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--rounds", type=_positive_int, default=3) + parser.add_argument("--seed", type=int, default=20260912) + parser.add_argument("--output-json", type=Path) + args = parser.parse_args() + if args.warmup < 0: + parser.error("--warmup must be non-negative") + return args + + +async def _run(args: argparse.Namespace) -> dict[str, object]: + payload = await _build_payload(images=args.images, width=args.width, height=args.height, seed=args.seed) + baseline = BenchmarkVariant("newline_split", io.BytesIO) + candidate = BenchmarkVariant("fixed_size_chunks", _iter_file_chunks) + + for _ in range(args.warmup): + for variant in (baseline, candidate): + await _measure(variant, payload) + + records: list[dict[str, object]] = [] + for round_index in range(args.rounds): + order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline) + for variant in order: + record = await _measure(variant, payload) + record["round"] = round_index + 1 + records.append(record) + + output_hashes = {str(record["output_sha256"]) for record in records} + if len(output_hashes) != 1: + raise RuntimeError(f"benchmark variants produced different outputs: {sorted(output_hashes)}") + + baseline_summary = _summarize(records, baseline.label) + candidate_summary = _summarize(records, candidate.label) + baseline_ms = cast(float, cast(dict[str, object], baseline_summary["wall_ms"])["median"]) + candidate_ms = cast(float, cast(dict[str, object], candidate_summary["wall_ms"])["median"]) + return { + "config": { + "images": args.images, + "width": args.width, + "height": args.height, + "warmup": args.warmup, + "rounds": args.rounds, + "seed": args.seed, + "chunk_size": _FILE_RESPONSE_CHUNK_SIZE, + "payload_bytes": len(payload), + "payload_newline_bytes": payload.count(b"\n"), + }, + "records": records, + "summary": { + baseline.label: baseline_summary, + candidate.label: candidate_summary, + "median_wall_speedup": baseline_ms / candidate_ms, + "body_frame_ratio": ( + cast(int, baseline_summary["body_frames"]) / cast(int, candidate_summary["body_frames"]) + ), + "output_sha256": output_hashes.pop(), + }, + } + + +def main() -> None: + args = _parse_args() + result = asyncio.run(_run(args)) + output_json = json.dumps(result, indent=2, sort_keys=True) + if args.output_json is not None: + args.output_json.write_text(output_json + "\n", encoding="utf-8") + print(output_json) + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/openai_api/test_image_protocol.py b/tests/entrypoints/openai_api/test_image_protocol.py new file mode 100644 index 000000000..b77038933 --- /dev/null +++ b/tests/entrypoints/openai_api/test_image_protocol.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import base64 +import io +import zipfile + +import pytest +from fastapi.responses import StreamingResponse + +from vllm_omni.entrypoints.openai.protocol.images import ( + _FILE_RESPONSE_CHUNK_SIZE, + ImageData, + ImageGenerationResponse, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_NEWLINE_HEAVY_PAYLOAD = (b"\x89PNG\n\x1a\n" * 64) + bytes(range(256)) * 8 + + +def _response(*payloads: bytes) -> ImageGenerationResponse: + return ImageGenerationResponse( + created=0, + data=[ImageData(b64_json=base64.b64encode(payload).decode()) for payload in payloads], + output_format="png", + size="1x1", + ) + + +async def _collect(response: StreamingResponse) -> list[bytes]: + return [bytes(chunk) async for chunk in response.body_iterator] + + +@pytest.mark.asyncio +async def test_single_file_response_chunks_ignore_newlines() -> None: + payload = _NEWLINE_HEAVY_PAYLOAD + assert b"\n" in payload + + response = _response(payload).stream_response() + chunks = await _collect(response) + + assert len(chunks) == 1 + assert b"".join(chunks) == payload + assert response.headers["content-length"] == str(len(payload)) + assert response.headers["content-type"] == "image/png" + + +@pytest.mark.asyncio +async def test_single_file_response_chunks_are_bounded() -> None: + payload = _NEWLINE_HEAVY_PAYLOAD * 200 + expected_chunks = -(-len(payload) // _FILE_RESPONSE_CHUNK_SIZE) + assert expected_chunks > 1 + + chunks = await _collect(_response(payload).stream_response()) + + assert len(chunks) == expected_chunks + assert max(len(chunk) for chunk in chunks) <= _FILE_RESPONSE_CHUNK_SIZE + assert b"".join(chunks) == payload + + +@pytest.mark.asyncio +async def test_zip_file_response_streams_whole_archive() -> None: + payloads = (_NEWLINE_HEAVY_PAYLOAD, b"second image\nwith a newline") + + response = _response(*payloads).stream_response() + chunks = await _collect(response) + + zip_bytes = b"".join(chunks) + assert len(chunks) == -(-len(zip_bytes) // _FILE_RESPONSE_CHUNK_SIZE) + assert response.headers["content-length"] == str(len(zip_bytes)) + assert response.headers["content-type"] == "application/zip" + with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive: + assert archive.read("image_0.png") == payloads[0] + assert archive.read("image_1.png") == payloads[1] diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index b7800c13e..067098f78 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -11,6 +11,7 @@ import base64 import io import uuid import zipfile +from collections.abc import AsyncIterator from enum import Enum from http import HTTPStatus from typing import Any, Literal @@ -25,6 +26,8 @@ from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers _INT64_MIN = -(2**63) _INT64_MAX = 2**63 - 1 +_FILE_RESPONSE_CHUNK_SIZE = 64 * 1024 + _IMAGE_FILE_METADATA = { "jpg": ("jpg", "image/jpeg"), "jpeg": ("jpeg", "image/jpeg"), @@ -33,6 +36,13 @@ _IMAGE_FILE_METADATA = { } +async def _iter_file_chunks(data: bytes) -> AsyncIterator[memoryview]: + """Yield fixed-size chunks of an in-memory payload.""" + view = memoryview(data) + for offset in range(0, len(view), _FILE_RESPONSE_CHUNK_SIZE): + yield view[offset : offset + _FILE_RESPONSE_CHUNK_SIZE] + + class ResponseFormat(str, Enum): """Image response format""" @@ -235,7 +245,7 @@ class ImageGenerationResponse(BaseModel): image_bytes = base64.b64decode(self.data[0].b64_json) filename = f"image_{uuid.uuid4().hex[:8]}.{extension}" return StreamingResponse( - io.BytesIO(image_bytes), + _iter_file_chunks(image_bytes), media_type=media_type, headers={ "Content-Disposition": f'attachment; filename="{filename}"', @@ -251,7 +261,7 @@ class ImageGenerationResponse(BaseModel): zip_bytes = zip_buffer.getvalue() filename = f"images_{uuid.uuid4().hex[:8]}.zip" return StreamingResponse( - io.BytesIO(zip_bytes), + _iter_file_chunks(zip_bytes), media_type="application/zip", headers={ "Content-Disposition": f'attachment; filename="{filename}"', From 65055774596d53362ef8c8340ed155ff12cd0d54 Mon Sep 17 00:00:00 2001 From: Lei Ke <1141466880@qq.com> Date: Mon, 14 Sep 2026 06:20:59 +0800 Subject: [PATCH 047/209] [Doc][NPU] Keep AllGather for MiniMax-H3 INT8 DLO (#7444) Signed-off-by: KrystalRay Co-authored-by: KrystalRay --- recipes/MiniMaxAI/MiniMax-H3-NPU.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md index b4563c1f7..c6ee0904b 100755 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md @@ -252,8 +252,9 @@ throughput guarantee. supports the H3 native `tile` mode only. - RainFusion block-sparse attention and INT8 quantization are validated for T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA. -- Online quantization cannot be combined with distributed layerwise offload - while AllGather is enabled; pass `--dlo-no-use-allgather` in that case. +- Online INT8 quantization can be combined with distributed layerwise offload + while AllGather is enabled. Keep AllGather enabled for SP>1; use + `--dlo-no-use-allgather` only when collective communication is unavailable. ## Additional resources From 8e738ce21b1e2c5389b665ea30a308f7eb528cdd Mon Sep 17 00:00:00 2001 From: Srinivas Krovvidi <194645829+Srinivasoo7@users.noreply.github.com> Date: Sun, 13 Sep 2026 20:26:44 -0500 Subject: [PATCH 048/209] [P0][RFC #3747] RL rollout serving: session management + world_model_env (#4770) Signed-off-by: srinivas_oo7 Co-authored-by: srinivas_oo7 Co-authored-by: Zhou Taichang --- .../openai_api/test_api_server_guards.py | 8 + .../openai_api/test_openpi_serving.py | 14 + tests/rl_rollout/__init__.py | 2 + tests/rl_rollout/test_rollout_routes.py | 167 ++++++++ tests/rl_rollout/test_rollout_session.py | 157 ++++++++ tests/rl_rollout/test_serving_rl_rollout.py | 371 ++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 87 ++++ .../entrypoints/openai/protocol/rollout.py | 111 ++++++ .../entrypoints/openai/rollout_session.py | 107 +++++ .../entrypoints/openai/serving_rl_rollout.py | 370 +++++++++++++++++ vllm_omni/entrypoints/openpi/serving.py | 28 +- 11 files changed, 1421 insertions(+), 1 deletion(-) create mode 100644 tests/rl_rollout/__init__.py create mode 100644 tests/rl_rollout/test_rollout_routes.py create mode 100644 tests/rl_rollout/test_rollout_session.py create mode 100644 tests/rl_rollout/test_serving_rl_rollout.py create mode 100644 vllm_omni/entrypoints/openai/protocol/rollout.py create mode 100644 vllm_omni/entrypoints/openai/rollout_session.py create mode 100644 vllm_omni/entrypoints/openai/serving_rl_rollout.py diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index ee141182c..21725183a 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -80,6 +80,11 @@ _EXPECTED_ROUTER_ROUTES = { ("WEBSOCKET", "/v1/realtime"), ("WEBSOCKET", "/v1/realtime/robot/openpi"), ("WEBSOCKET", "/v1/duplex"), + ("POST", "/v1/realtime/sessions"), + ("POST", "/v1/realtime/sessions/{session_id}/step"), + ("POST", "/v1/realtime/sessions/{session_id}/reset"), + ("POST", "/v1/realtime/sessions/{session_id}/close"), + ("GET", "/v1/realtime/sessions/{session_id}/status"), ("GET", "/health"), ("GET", "/v1/models"), ("POST", "/v1/images/generations"), @@ -120,6 +125,7 @@ _DIFFUSION_APP_STATE_KEYS = { "openai_streaming_speech", "openai_streaming_video", "openai_serving_realtime_robot", + "rl_rollout_serving", "enable_server_load_tracking", "server_load_metrics", } @@ -151,12 +157,14 @@ _MULTISTAGE_APP_STATE_KEYS = { "openai_serving_realtime", "openai_serving_video", "openai_serving_realtime_robot", + "rl_rollout_serving", "enable_server_load_tracking", "server_load_metrics", } _MULTISTAGE_MUST_BE_NONE = { "openai_serving_duplex", "openai_serving_realtime_robot", + "rl_rollout_serving", } _MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 82ee6e7dd..46ff0f822 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -440,3 +440,17 @@ def test_extract_actions_does_not_iterate_result_object(): actions = serving._extract_actions(IterableResult()) np.testing.assert_allclose(actions, np.array([[1.0, 2.0, 3.0]], dtype=np.float32)) + + +def test_drop_session_calls_dreamzero_close_hook(): + closed: list[str] = [] + pipeline = SimpleNamespace(close_ar_diffusion_session=lambda session_id: closed.append(session_id)) + engine_client = SimpleNamespace( + model_config={"policy_server_config": TEST_POLICY_SERVER_CONFIG}, + model_runner=SimpleNamespace(pipeline=pipeline), + ) + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client) + + serving.drop_session("rollout-1") + + assert closed == ["rollout-1"] diff --git a/tests/rl_rollout/__init__.py b/tests/rl_rollout/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/rl_rollout/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/rl_rollout/test_rollout_routes.py b/tests/rl_rollout/test_rollout_routes.py new file mode 100644 index 000000000..c6a9e2e56 --- /dev/null +++ b/tests/rl_rollout/test_rollout_routes.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Route-level tests for RL rollout serving endpoints.""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm_omni.entrypoints.openai.api_server import router +from vllm_omni.entrypoints.openai.protocol.rollout import ( + ErrorObject, + RolloutStepResponse, + SessionMetadata, +) +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSessionClosedError, + RolloutSessionNotFoundError, +) + + +class _FakeRolloutServing: + def __init__(self, close_error: Exception | None = None, step_error_code: str | None = None) -> None: + self.close_error = close_error + self.step_error_code = step_error_code + self.step_request = None + + async def step(self, session_id, body): + self.step_request = body + error = None + if self.step_error_code is not None: + error = ErrorObject( + code=self.step_error_code, + message=f"{self.step_error_code}: {session_id}", + step_id=body.step_id, + committed_step_id=-1, + ) + return RolloutStepResponse( + step_id=body.step_id, + next_observation=None, + model_metadata=SessionMetadata( + latency_ms=0.0, + steps_generated=0, + context_length=0, + committed_step_id=-1, + ), + error=error, + ) + + async def close_session(self, session_id): + if self.close_error is not None: + raise self.close_error + + +def _client(serving: _FakeRolloutServing | None = None) -> TestClient: + app = FastAPI() + app.include_router(router) + if serving is not None: + app.state.rl_rollout_serving = serving + return TestClient(app) + + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_step_route_rejects_base64_image_payload(): + serving = _FakeRolloutServing() + client = _client(serving) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={ + "step_id": 0, + "observation": {"images": {"front": "base64-frames"}}, + "action": {}, + }, + ) + + assert response.status_code == 422 + assert serving.step_request is None + + +def test_step_route_accepts_nested_image_payload(): + serving = _FakeRolloutServing() + client = _client(serving) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={ + "step_id": 0, + "observation": {"images": {"observation/exterior_image_1_left": [[[0, 0, 0]]]}}, + "action": {}, + }, + ) + + assert response.status_code == 200 + assert serving.step_request.observation.images == {"observation/exterior_image_1_left": [[[0, 0, 0]]]} + + +def test_step_route_requires_action(): + client = _client(_FakeRolloutServing()) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}}, + ) + + assert response.status_code == 422 + + +def test_step_route_maps_missing_session_to_404(): + client = _client(_FakeRolloutServing(step_error_code="session_not_found")) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}, "action": {}}, + ) + + assert response.status_code == 404 + + +def test_step_route_maps_closed_session_to_410(): + client = _client(_FakeRolloutServing(step_error_code="session_closed")) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}, "action": {}}, + ) + + assert response.status_code == 410 + + +def test_step_route_maps_invalid_request_to_400(): + client = _client(_FakeRolloutServing(step_error_code="invalid_request")) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}, "action": {}}, + ) + + assert response.status_code == 400 + + +def test_close_route_maps_missing_session_to_404(): + client = _client(_FakeRolloutServing(close_error=RolloutSessionNotFoundError("s1"))) + + response = client.post("/v1/realtime/sessions/s1/close") + + assert response.status_code == 404 + + +def test_close_route_is_idempotent_for_closed_session(): + client = _client(_FakeRolloutServing(close_error=RolloutSessionClosedError("s1"))) + + response = client.post("/v1/realtime/sessions/s1/close") + + assert response.status_code == 200 + assert response.json()["closed"] is True + + +def test_rollout_routes_return_501_when_serving_is_unavailable(): + client = _client() + + response = client.post("/v1/realtime/sessions/s1/close") + + assert response.status_code == 501 diff --git a/tests/rl_rollout/test_rollout_session.py b/tests/rl_rollout/test_rollout_session.py new file mode 100644 index 000000000..2b3252a1a --- /dev/null +++ b/tests/rl_rollout/test_rollout_session.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for RolloutSessionStore (RFC #3747, P0).""" + +import asyncio + +import pytest + +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSession, + RolloutSessionClosedError, + RolloutSessionNotFoundError, + RolloutSessionStepError, + RolloutSessionStore, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def store() -> RolloutSessionStore: + return RolloutSessionStore() + + +@pytest.mark.asyncio +async def test_create_returns_session(store): + s = await store.create("s1", model="dreamzero", mode="world_model_env") + assert isinstance(s, RolloutSession) + assert s.session_id == "s1" + assert s.committed_step_id == -1 + assert s.context_length == 0 + assert not s.closed + + +@pytest.mark.asyncio +async def test_get_existing(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + s = await store.get("s1") + assert s.session_id == "s1" + + +@pytest.mark.asyncio +async def test_get_missing_raises(store): + with pytest.raises(RolloutSessionNotFoundError): + await store.get("does-not-exist") + + +@pytest.mark.asyncio +async def test_get_closed_raises(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.close("s1") + with pytest.raises(RolloutSessionClosedError): + await store.get("s1") + + +@pytest.mark.asyncio +async def test_close_missing_raises(store): + with pytest.raises(RolloutSessionNotFoundError): + await store.close("does-not-exist") + + +@pytest.mark.asyncio +async def test_close_is_idempotent(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.close("s1") + await store.close("s1") + s = await store.get("s1", include_closed=True) + assert s.closed + + +@pytest.mark.asyncio +async def test_advance_updates_committed_step_id(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + s = await store.get("s1") + assert s.committed_step_id == 0 + assert s.context_length == 1 + + await store.advance("s1", step_id=1) + s = await store.get("s1") + assert s.committed_step_id == 1 + assert s.context_length == 2 + + +@pytest.mark.asyncio +async def test_advance_rejects_duplicate_step_id(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + with pytest.raises(RolloutSessionStepError): + await store.advance("s1", step_id=0) + + s = await store.get("s1") + assert s.committed_step_id == 0 + assert s.context_length == 1 + + +@pytest.mark.asyncio +async def test_advance_rejects_out_of_order_step_id(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + with pytest.raises(RolloutSessionStepError): + await store.advance("s1", step_id=2) + + s = await store.get("s1") + assert s.committed_step_id == -1 + assert s.context_length == 0 + + +@pytest.mark.asyncio +async def test_reset_clears_state(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + await store.advance("s1", step_id=1) + await store.advance("s1", step_id=2) + await store.advance("s1", step_id=3) + await store.reset("s1") + s = await store.get("s1") + assert s.committed_step_id == -1 + assert s.context_length == 0 + + +@pytest.mark.asyncio +async def test_reset_closed_raises(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.close("s1") + with pytest.raises(RolloutSessionClosedError): + await store.reset("s1") + + +@pytest.mark.asyncio +async def test_failed_step_does_not_advance(store): + """committed_step_id must not change when a step raises (RFC section 6.5 invariant).""" + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + + # Simulate a failed step: advance is NOT called + s = await store.get("s1") + assert s.committed_step_id == 0 # unchanged + + +@pytest.mark.asyncio +async def test_concurrent_steps_are_serialised(store): + """Per-session lock must prevent concurrent context mutations.""" + await store.create("s1", model="dreamzero", mode="world_model_env") + s = await store.get("s1") + + results: list[int] = [] + + async def slow_advance(step_id: int, delay: float) -> None: + await asyncio.sleep(delay) + async with s.lock: + await asyncio.sleep(0.01) + await store.advance("s1", step_id=step_id) + results.append(step_id) + + await asyncio.gather(slow_advance(0, 0.0), slow_advance(1, 0.01)) + # Both steps ran under the per-session lock and committed in order. + assert sorted(results) == [0, 1] diff --git a/tests/rl_rollout/test_serving_rl_rollout.py b/tests/rl_rollout/test_serving_rl_rollout.py new file mode 100644 index 000000000..5474ebc20 --- /dev/null +++ b/tests/rl_rollout/test_serving_rl_rollout.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for ServingRLRollout (RFC #3747, P0). + +The DreamZero engine is mocked. Tests verify session lifecycle, step_id +semantics, committed_step_id invariants, and error paths without a GPU. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import numpy as np +import pytest + +from vllm_omni.entrypoints.openai.protocol.rollout import ( + Action, + CreateSessionRequest, + Observation, + RolloutStepRequest, +) +from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_fake_result(video: Any = None, *, formatter_shaped: bool = True) -> MagicMock: + result = MagicMock() + if video is None: + video = np.zeros((1, 3, 64, 64), dtype=np.float32) + if formatter_shaped: + # Matches output_formatter: video is peeled into images; mm only has actions. + result.images = video + result.multimodal_output = {"actions": np.zeros((1, 8), dtype=np.float32)} + else: + result.images = None + result.multimodal_output = {"video": video} + return result + + +def _make_openpi(video: Any = None, raises: Exception | None = None, *, formatter_shaped: bool = True): + openpi = MagicMock() + openpi.engine_client = MagicMock() + openpi.model_name = None + + async def _generate(*_, **__): + if raises: + raise raises + yield _make_fake_result(video, formatter_shaped=formatter_shaped) + + openpi.engine_client.generate = _generate + openpi.build_request = MagicMock( + return_value=MagicMock( + prompt="", + prompts=[""], + request_id="test-req", + sampling_params=MagicMock(), + ) + ) + return openpi + + +@pytest.fixture +def serving() -> ServingRLRollout: + return ServingRLRollout(_make_openpi()) + + +# --------------------------------------------------------------------------- +# Session lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_session(serving): + resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + assert resp.session_id + assert resp.mode == "world_model_env" + + +@pytest.mark.asyncio +async def test_get_status_after_create(serving): + resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + status = await serving.get_status(resp.session_id) + assert status.committed_step_id == -1 + assert status.context_length == 0 + assert not status.closed + + +@pytest.mark.asyncio +async def test_close_session(serving): + resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + await serving.close_session(resp.session_id) + status = await serving.get_status(resp.session_id) + assert status.closed + assert status.committed_step_id == -1 + serving._openpi.drop_session.assert_any_call(resp.session_id) + serving._openpi.drop_session.assert_any_call(f"{resp.session_id}:stateless") + + +# --------------------------------------------------------------------------- +# Step - success path +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_successful_step_advances_committed_step_id(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest( + step_id=0, + observation=Observation(state=[0.1, 0.2]), + action=Action(joint_positions=[0.0] * 7), + ) + resp = await serving.step(sess.session_id, req) + assert resp.error is None + assert resp.model_metadata.committed_step_id == 0 + + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == 0 + assert status.context_length == 1 + + +@pytest.mark.asyncio +async def test_repeated_steps_advance_monotonically(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + for i in range(3): + req = RolloutStepRequest( + step_id=i, + observation=Observation(state=[float(i)]), + action=Action(joint_positions=[0.0] * 7), + ) + resp = await serving.step(sess.session_id, req) + assert resp.error is None + assert resp.model_metadata.committed_step_id == i + + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == 2 + assert status.context_length == 3 + + +@pytest.mark.asyncio +async def test_next_observation_present_on_success(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest( + step_id=0, + observation=Observation(), + action=Action(), + ) + resp = await serving.step(sess.session_id, req) + assert resp.next_observation is not None + assert "video_latent" in resp.next_observation + assert resp.next_observation["dtype"] == "float32" + + +@pytest.mark.asyncio +async def test_duplicate_step_does_not_rerun_or_advance(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + assert (await serving.step(sess.session_id, req)).error is None + + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "step_already_committed" + assert resp.model_metadata.committed_step_id == 0 + assert resp.model_metadata.context_length == 1 + serving._openpi.build_request.assert_called_once() + + +@pytest.mark.asyncio +async def test_out_of_order_step_does_not_rerun_or_advance(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=2, observation=Observation(), action=Action()) + + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "step_out_of_order" + assert resp.model_metadata.committed_step_id == -1 + assert resp.model_metadata.context_length == 0 + serving._openpi.build_request.assert_not_called() + + +@pytest.mark.asyncio +async def test_stateless_step_does_not_commit_context(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest( + step_id=7, + observation=Observation(), + action=Action(), + use_session_context=False, + ) + + resp = await serving.step(sess.session_id, req) + + assert resp.error is None + assert resp.model_metadata.committed_step_id == -1 + assert resp.model_metadata.context_length == 0 + serving._openpi.build_request.assert_called_once() + kwargs = serving._openpi.build_request.call_args.kwargs + assert kwargs["reset"] is True + assert kwargs["session_id"] == f"{sess.session_id}:stateless" + + +# --------------------------------------------------------------------------- +# Step - failure path: committed_step_id must NOT advance +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_failed_step_does_not_advance_committed_step_id(): + serving = ServingRLRollout(_make_openpi(raises=RuntimeError("engine dead"))) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "inference_error" + assert resp.model_metadata.committed_step_id == -1 # must not have advanced + + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == -1 + + +@pytest.mark.asyncio +async def test_output_encoding_failure_does_not_advance_committed_step_id(): + class BadVideo: + def detach(self): + return self + + def float(self): + return self + + def cpu(self): + return self + + def numpy(self): + raise TypeError("cannot encode") + + serving = ServingRLRollout(_make_openpi(video=BadVideo())) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "inference_error" + assert resp.model_metadata.committed_step_id == -1 + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == -1 + + +@pytest.mark.asyncio +async def test_failed_step_after_success_preserves_last_committed(): + good_serving = ServingRLRollout(_make_openpi()) + sess = await good_serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + # Step 0 succeeds + await good_serving.step(sess.session_id, RolloutStepRequest(step_id=0, observation=Observation(), action=Action())) + + # Swap engine to failing one - reuse the same store via internal reference + good_serving._openpi = _make_openpi(raises=RuntimeError("oops")) + resp = await good_serving.step( + sess.session_id, RolloutStepRequest(step_id=1, observation=Observation(), action=Action()) + ) + assert resp.error is not None + assert resp.model_metadata.committed_step_id == 0 # still at last good step + assert resp.model_metadata.context_length == 1 + + +# --------------------------------------------------------------------------- +# Step - unknown session +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_step_unknown_session_returns_error(serving): + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + resp = await serving.step("no-such-session", req) + assert resp.error is not None + assert resp.error.code == "session_not_found" + + +# --------------------------------------------------------------------------- +# Reset +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_reset_clears_committed_step_id(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + await serving.step(sess.session_id, RolloutStepRequest(step_id=0, observation=Observation(), action=Action())) + await serving.reset_session(sess.session_id) + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == -1 + assert status.context_length == 0 + serving._openpi.drop_session.assert_called_with(sess.session_id) + + +# --------------------------------------------------------------------------- +# action -> obs folding (module assumption) +# --------------------------------------------------------------------------- + + +def test_merge_action_into_obs_emits_dreamzero_state_keys(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + obs = Observation(state=[9.0] * 8, prompt="test") + action = Action(joint_positions=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], gripper_position=0.5) + result = _merge_action_into_obs(obs, action) + + assert result["prompt"] == "test" + np.testing.assert_array_almost_equal(result["observation/joint_position"], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + np.testing.assert_array_almost_equal(result["observation/gripper_position"], [0.5]) + + +def test_merge_action_none_splits_obs_state_for_dreamzero(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + obs = Observation(state=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 0.25]) + result = _merge_action_into_obs(obs, None) + np.testing.assert_array_almost_equal(result["observation/joint_position"], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + np.testing.assert_array_almost_equal(result["observation/gripper_position"], [0.25]) + + +def test_merge_no_state_no_action_omits_state_key(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + obs = Observation() + result = _merge_action_into_obs(obs, None) + assert "observation/joint_position" not in result + assert "observation/gripper_position" not in result + + +def test_merge_rejects_wrong_dreamzero_joint_dimension(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + with pytest.raises(ValueError, match="joint_positions must have 7 values"): + _merge_action_into_obs(Observation(), Action(joint_positions=[1.0, 2.0])) + + +@pytest.mark.asyncio +async def test_step_reads_formatter_images_not_multimodal_video(): + video = np.ones((1, 3, 8, 8), dtype=np.float32) + serving = ServingRLRollout(_make_openpi(video=video, formatter_shaped=True)) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + resp = await serving.step( + sess.session_id, + RolloutStepRequest(step_id=0, observation=Observation(), action=Action()), + ) + assert resp.error is None + assert resp.next_observation is not None + assert "video_latent" in resp.next_observation + assert resp.next_observation["shape"] == [1, 3, 8, 8] + + +@pytest.mark.asyncio +async def test_step_falls_back_to_multimodal_video(): + video = np.ones((2, 4, 4), dtype=np.float32) + serving = ServingRLRollout(_make_openpi(video=video, formatter_shaped=False)) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + resp = await serving.step( + sess.session_id, + RolloutStepRequest(step_id=0, observation=Observation(), action=Action()), + ) + assert resp.error is None + assert resp.next_observation is not None + assert resp.next_observation["shape"] == [2, 4, 4] diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f4f22f95f..8acd1d6fa 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -138,6 +138,10 @@ from vllm_omni.entrypoints.openai.protocol.images import ( ImageGenerationResponse, ResponseFormat, ) +from vllm_omni.entrypoints.openai.protocol.rollout import ( + CreateSessionRequest, + RolloutStepRequest, +) from vllm_omni.entrypoints.openai.protocol.videos import ( VideoDeleteResponse, VideoGenerationRequest, @@ -146,8 +150,14 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoResponse, ) from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSessionCapacityError, + RolloutSessionClosedError, + RolloutSessionNotFoundError, +) from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServingAudioGenerate from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler from vllm_omni.entrypoints.openai.serving_video import ( @@ -604,6 +614,10 @@ async def omni_init_app_state( engine_client=engine_client, model_name=model_name, ) + if state.openai_serving_realtime_robot is not None: + state.rl_rollout_serving = ServingRLRollout(state.openai_serving_realtime_robot) + else: + state.rl_rollout_serving = None state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) state.server_load_metrics = 0 @@ -949,6 +963,7 @@ async def omni_init_app_state( stage_configs=state.stage_configs, ) state.openai_serving_realtime_robot = None + state.rl_rollout_serving = None state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 @@ -1529,6 +1544,78 @@ async def duplex_websocket(websocket: WebSocket): await handler.handle_session(websocket) +# RL Rollout serving (RFC #3747, P0) + + +def _rl_rollout_serving(request: Request) -> ServingRLRollout: + serving = getattr(request.app.state, "rl_rollout_serving", None) + if serving is None: + raise HTTPException(status_code=501, detail="RL rollout serving not available for this model.") + return serving + + +@router.post("/v1/realtime/sessions") +async def create_rollout_session(body: CreateSessionRequest, request: Request): + serving = _rl_rollout_serving(request) + try: + return (await serving.create_session(body)).model_dump() + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except RolloutSessionCapacityError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + + +@router.post("/v1/realtime/sessions/{session_id}/step") +async def rollout_step(session_id: str, body: RolloutStepRequest, request: Request): + serving = _rl_rollout_serving(request) + response = await serving.step(session_id, body) + if response.error is not None: + status_code = None + if response.error.code == "session_not_found": + status_code = 404 + elif response.error.code == "session_closed": + status_code = 410 + elif response.error.code in {"invalid_request", "step_already_committed", "step_out_of_order"}: + status_code = 400 + if status_code is not None: + return JSONResponse(content=response.model_dump(), status_code=status_code) + return response.model_dump() + + +@router.post("/v1/realtime/sessions/{session_id}/reset") +async def reset_rollout_session(session_id: str, request: Request): + serving = _rl_rollout_serving(request) + try: + return (await serving.reset_session(session_id)).model_dump() + except RolloutSessionNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.") + except RolloutSessionClosedError: + raise HTTPException(status_code=410, detail=f"Session {session_id!r} is closed.") + + +@router.post("/v1/realtime/sessions/{session_id}/close") +async def close_rollout_session(session_id: str, request: Request): + serving = _rl_rollout_serving(request) + try: + await serving.close_session(session_id) + return {"session_id": session_id, "closed": True} + except RolloutSessionNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.") + except RolloutSessionClosedError: + return {"session_id": session_id, "closed": True} + + +@router.get("/v1/realtime/sessions/{session_id}/status") +async def rollout_session_status(session_id: str, request: Request): + serving = _rl_rollout_serving(request) + try: + return (await serving.get_status(session_id)).model_dump() + except RolloutSessionNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.") + except RolloutSessionClosedError: + raise HTTPException(status_code=410, detail=f"Session {session_id!r} is closed.") + + # Health and Model endpoints for diffusion mode diff --git a/vllm_omni/entrypoints/openai/protocol/rollout.py b/vllm_omni/entrypoints/openai/protocol/rollout.py new file mode 100644 index 000000000..819981d7b --- /dev/null +++ b/vllm_omni/entrypoints/openai/protocol/rollout.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pydantic schemas for RL rollout serving (RFC #3747). + +P0 scope: world_model_env mode only (observation + action -> next observation). + +P0 maps rollout observations/actions onto DreamZero's DROID robot_obs schema: +camera images retain their DreamZero observation/* keys, joint/gripper state is +passed through observation/joint_position and observation/gripper_position, and +the returned video payload is the DreamZero VAE latent unless a later endpoint +adds explicit latent decoding. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +class Observation(BaseModel): + """Single-step robot observation.""" + + images: dict[str, list] | None = Field( + default=None, + description="Named DreamZero camera images as nested numeric lists. Base64 strings are not accepted in P0.", + ) + state: list[float] | None = Field( + default=None, + description="Proprioceptive state vector (joint positions, velocities, etc.).", + ) + prompt: str = Field(default="", description="Optional language conditioning.") + extra: dict[str, Any] = Field( + default_factory=dict, + description="Pass-through fields forwarded verbatim to robot_obs.", + ) + + +class Action(BaseModel): + """Executed action for world_model_env conditioning.""" + + joint_positions: list[float] | None = Field( + default=None, + description="Executed joint positions; forwarded as observation/joint_position for DreamZero.", + ) + gripper_position: float | list[float] | None = Field( + default=None, + description="Executed gripper position; forwarded as observation/gripper_position for DreamZero.", + ) + extra: dict[str, Any] = Field(default_factory=dict) + + +class SessionMetadata(BaseModel): + latency_ms: float + steps_generated: int + context_length: int + committed_step_id: int = Field( + description="Highest step_id whose context has been atomically committed. " + "-1 means no step has been committed yet (fresh or reset session).", + ) + session_memory_bytes: int | None = None + uncertainty: float | None = None + + +class ErrorObject(BaseModel): + code: str + message: str + step_id: int | None = None + committed_step_id: int = -1 + + +class CreateSessionRequest(BaseModel): + model: str + mode: Literal["world_model_env"] = "world_model_env" + + +class CreateSessionResponse(BaseModel): + session_id: str + mode: str + created_at: float + + +class RolloutStepRequest(BaseModel): + step_id: int = Field(ge=0, description="Monotonically increasing per session.") + observation: Observation + action: Action = Field( + description="Required for world_model_env: executed action that produced this observation.", + ) + use_session_context: bool = True + + +class RolloutStepResponse(BaseModel): + step_id: int + next_observation: dict[str, Any] | None = Field( + default=None, + description="Predicted next observation metadata. P0 returns DreamZero video_latent, not RGB frames.", + ) + model_metadata: SessionMetadata + error: ErrorObject | None = None + + +class ResetSessionResponse(BaseModel): + session_id: str + committed_step_id: int + + +class SessionStatusResponse(BaseModel): + session_id: str + committed_step_id: int + context_length: int + closed: bool diff --git a/vllm_omni/entrypoints/openai/rollout_session.py b/vllm_omni/entrypoints/openai/rollout_session.py new file mode 100644 index 000000000..76848ba31 --- /dev/null +++ b/vllm_omni/entrypoints/openai/rollout_session.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""In-memory session store for RL rollout serving (RFC #3747, P0). + +Each RolloutSession holds a per-session asyncio.Lock so that concurrent step +requests for the same session are serialised. committed_step_id advances only +on successful step completion; failed or timed-out steps leave it unchanged, +preserving the atomic-context-commit guarantee from section 6.5 of the RFC. +""" + +from __future__ import annotations + +import asyncio +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Literal + + +class RolloutSessionNotFoundError(KeyError): + pass + + +class RolloutSessionClosedError(RuntimeError): + pass + + +class RolloutSessionStepError(ValueError): + pass + + +class RolloutSessionCapacityError(RuntimeError): + pass + + +@dataclass +class RolloutSession: + session_id: str + model: str + mode: Literal["world_model_env"] + created_at: float = field(default_factory=time.time) + committed_step_id: int = -1 + context_length: int = 0 + closed: bool = False + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +class RolloutSessionStore: + """Thread-safe (asyncio) store for active rollout sessions.""" + + def __init__(self, max_sessions: int = 64) -> None: + self._sessions: OrderedDict[str, RolloutSession] = OrderedDict() + self._max_sessions = max_sessions + self._store_lock = asyncio.Lock() + + async def create(self, session_id: str, model: str, mode: Literal["world_model_env"]) -> RolloutSession: + async with self._store_lock: + self._evict_closed_sessions_for_create() + if len(self._sessions) >= self._max_sessions: + raise RolloutSessionCapacityError(f"Maximum active rollout sessions reached: {self._max_sessions}.") + session = RolloutSession( + session_id=session_id, + model=model, + mode=mode, + ) + self._sessions[session_id] = session + self._sessions.move_to_end(session_id) + return session + + async def get(self, session_id: str, *, include_closed: bool = False) -> RolloutSession: + session = self._sessions.get(session_id) + if session is None: + raise RolloutSessionNotFoundError(session_id) + self._sessions.move_to_end(session_id) + if session.closed and not include_closed: + raise RolloutSessionClosedError(session_id) + return session + + async def reset(self, session_id: str) -> RolloutSession: + session = await self.get(session_id) + async with session.lock: + session.committed_step_id = -1 + session.context_length = 0 + return session + + async def close(self, session_id: str) -> None: + session = await self.get(session_id, include_closed=True) + async with session.lock: + session.closed = True + + async def advance(self, session_id: str, step_id: int) -> None: + """Commit a successfully completed step. Called only on success.""" + session = await self.get(session_id) + expected_step_id = session.committed_step_id + 1 + if step_id != expected_step_id: + raise RolloutSessionStepError(f"Expected step_id {expected_step_id}, got {step_id}.") + session.committed_step_id = step_id + session.context_length += 1 + + def _evict_closed_sessions_for_create(self) -> None: + while len(self._sessions) >= self._max_sessions: + for session_id, session in self._sessions.items(): + if session.closed: + del self._sessions[session_id] + break + else: + break diff --git a/vllm_omni/entrypoints/openai/serving_rl_rollout.py b/vllm_omni/entrypoints/openai/serving_rl_rollout.py new file mode 100644 index 000000000..38b53d9bd --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_rl_rollout.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Serving layer for RL rollout (RFC #3747, P0). + +P0 scope: world_model_env mode, single-session, DreamZero backbone. + +P0 maps client rollout input onto DreamZero's DROID robot_obs schema and returns +DreamZero video latents explicitly. RGB/frame decoding can be added as a later +contract without pretending the latent is already an image observation. +""" + +from __future__ import annotations + +import binascii +import time +from collections import abc +from typing import Any + +import numpy as np +from vllm.logger import init_logger +from vllm.utils import random_uuid + +from vllm_omni.entrypoints.openai.protocol.rollout import ( + Action, + CreateSessionRequest, + CreateSessionResponse, + ErrorObject, + Observation, + ResetSessionResponse, + RolloutStepRequest, + RolloutStepResponse, + SessionMetadata, + SessionStatusResponse, +) +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSession, + RolloutSessionClosedError, + RolloutSessionNotFoundError, + RolloutSessionStore, +) +from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI + +logger = init_logger(__name__) + + +def _merge_action_into_obs(obs: Observation, action: Action | None) -> dict[str, Any]: + """Build DreamZero-compatible robot_obs.""" + robot_obs: dict[str, Any] = {"prompt": obs.prompt} + + if obs.images: + robot_obs.update(obs.images) + + if obs.extra: + for key, value in obs.extra.items(): + if key.startswith("observation/") or key in {"seed", "embodiment_name"}: + robot_obs[key] = value + + joint_positions = action.joint_positions if action is not None else None + gripper_position = action.gripper_position if action is not None else None + + if joint_positions is None and obs.state: + joint_positions = obs.state[:7] + if len(obs.state) > 7 and gripper_position is None: + gripper_position = obs.state[7:8] + + if joint_positions is not None: + joint_arr = np.asarray(joint_positions, dtype=np.float64).flatten() + if joint_arr.size != 7: + raise ValueError(f"DreamZero DROID joint_positions must have 7 values, got {joint_arr.size}.") + robot_obs["observation/joint_position"] = joint_arr + + if gripper_position is not None: + gripper_arr = np.asarray(gripper_position, dtype=np.float64).flatten() + if gripper_arr.size != 1: + raise ValueError(f"DreamZero DROID gripper_position must have 1 value, got {gripper_arr.size}.") + robot_obs["observation/gripper_position"] = gripper_arr + + return robot_obs + + +def _encode_video_output(video: Any) -> dict[str, Any]: + """Encode DreamZero video latent tensor / ndarray to JSON metadata.""" + if video is None: + return {} + if hasattr(video, "detach") and callable(video.detach): + video = video.detach() + if hasattr(video, "float") and callable(video.float): + video = video.float() + if hasattr(video, "cpu") and callable(video.cpu): + video = video.cpu() + if hasattr(video, "numpy"): + video = video.numpy() + if isinstance(video, np.ndarray): + return { + "video_latent": binascii.b2a_base64(video.tobytes(), newline=False).decode(), + "shape": list(video.shape), + "dtype": str(video.dtype), + "encoding": "base64_raw_tensor", + } + return {"video_latent": str(video), "encoding": "string"} + + +class ServingRLRollout: + """HTTP serving layer for RL rollout sessions. + + Wraps ServingRealtimeRobotOpenPI for world_model_env mode. + Session state and committed_step_id tracking live in RolloutSessionStore. + """ + + def __init__(self, openpi_serving: ServingRealtimeRobotOpenPI) -> None: + self._openpi = openpi_serving + self._store = RolloutSessionStore() + + # ------------------------------------------------------------------ # + # Session lifecycle # + # ------------------------------------------------------------------ # + + async def create_session(self, req: CreateSessionRequest) -> CreateSessionResponse: + self._validate_model(req.model) + session_id = random_uuid() + session = await self._store.create( + session_id=session_id, + model=req.model, + mode=req.mode, + ) + logger.info("Created rollout session %s mode=%s", session_id, req.mode) + return CreateSessionResponse( + session_id=session_id, + mode=session.mode, + created_at=session.created_at, + ) + + async def reset_session(self, session_id: str) -> ResetSessionResponse: + session = await self._store.reset(session_id) + # Next infer still passes reset=True; drop GPU KV now if the engine + # session is reachable from this process. + self._drop_engine_session(session_id) + logger.info("Reset rollout session %s", session_id) + return ResetSessionResponse( + session_id=session_id, + committed_step_id=session.committed_step_id, + ) + + async def close_session(self, session_id: str) -> None: + await self._store.close(session_id) + self._drop_engine_session(session_id) + self._drop_engine_session(f"{session_id}:stateless") + logger.info("Closed rollout session %s", session_id) + + async def get_status(self, session_id: str) -> SessionStatusResponse: + session = await self._store.get(session_id, include_closed=True) + return SessionStatusResponse( + session_id=session_id, + committed_step_id=session.committed_step_id, + context_length=session.context_length, + closed=session.closed, + ) + + # ------------------------------------------------------------------ # + # Step # + # ------------------------------------------------------------------ # + + async def step( + self, + session_id: str, + req: RolloutStepRequest, + ) -> RolloutStepResponse: + try: + session = await self._store.get(session_id) + except RolloutSessionNotFoundError: + return self._error_response( + req.step_id, + -1, + 0, + "session_not_found", + f"Session {session_id!r} does not exist.", + ) + except RolloutSessionClosedError: + return self._error_response( + req.step_id, + -1, + 0, + "session_closed", + f"Session {session_id!r} is closed.", + ) + + async with session.lock: + return await self._run_step(session, req) + + async def _run_step( + self, + session: RolloutSession, + req: RolloutStepRequest, + ) -> RolloutStepResponse: + t0 = time.perf_counter() + committed = session.committed_step_id + + if req.use_session_context: + expected_step_id = committed + 1 + if req.step_id != expected_step_id: + if req.step_id <= committed: + code = "step_already_committed" + message = f"Step {req.step_id} is already committed; highest committed step is {committed}." + else: + code = "step_out_of_order" + message = f"Expected step_id {expected_step_id}, got {req.step_id}." + return self._error_response( + req.step_id, + committed, + session.context_length, + code, + message, + ) + + # reset=True on first call after session create/reset + reset = committed == -1 or not req.use_session_context + engine_session_id = session.session_id + if not req.use_session_context: + engine_session_id = f"{session.session_id}:stateless" + try: + robot_obs = _merge_action_into_obs(req.observation, req.action) + video_out = await self._infer_world_model( + obs=robot_obs, + session_id=engine_session_id, + reset=reset, + ) + next_observation = _encode_video_output(video_out) + except ValueError as exc: + logger.info("Invalid rollout step %d for session %s: %s", req.step_id, session.session_id, exc) + return self._error_response( + req.step_id, + committed, + session.context_length, + "invalid_request", + str(exc), + ) + except Exception as exc: + logger.exception("Step %d failed for session %s", req.step_id, session.session_id) + # Do NOT advance committed_step_id on failure (RFC section 6.5). + return self._error_response( + req.step_id, + committed, + session.context_length, + "inference_error", + str(exc), + ) + + if req.use_session_context: + # Only advance committed_step_id on successful contextual steps. + await self._store.advance(session.session_id, req.step_id) + committed = req.step_id + + latency_ms = (time.perf_counter() - t0) * 1000.0 + metadata = SessionMetadata( + latency_ms=round(latency_ms, 2), + steps_generated=1, + context_length=session.context_length, + committed_step_id=committed, + ) + return RolloutStepResponse( + step_id=req.step_id, + next_observation=next_observation, + model_metadata=metadata, + ) + + async def _infer_world_model( + self, + obs: dict[str, Any], + *, + session_id: str, + reset: bool, + ) -> Any: + """Call the engine and return the video (next-observation) output. + + Reuses ServingRealtimeRobotOpenPI.build_request() so request routing, + session_id threading, and OmniDiffusionSamplingParams construction are + identical to the policy_inference path. DreamZero's formatter peels + ``video`` into ``result.images`` and leaves only ``actions`` on + ``multimodal_output``; read images first and fall back to the raw key. + """ + request = self._openpi.build_request(obs, session_id=session_id, reset=reset) + result = None + async for output in self._openpi.engine_client.generate( + prompt=self._request_prompt(request), + request_id=request.request_id, + sampling_params_list=[request.sampling_params], + ): + result = output + + if result is None: + raise RuntimeError("World model request produced no output.") + return self._extract_video_output(result) + + @staticmethod + def _request_prompt(request: Any) -> Any: + prompt = getattr(request, "prompt", None) + if prompt is not None: + return prompt + return request.prompts[0] + + @staticmethod + def _extract_video_output(result: Any) -> Any: + """Prefer formatter-placed ``images``; fall back to multimodal video.""" + images = getattr(result, "images", None) + if isinstance(images, (list, tuple)): + if images: + return images[0] + elif images is not None and not isinstance(images, abc.Mapping): + # Formatter assigns the peeled video tensor directly to images. + return images + + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, abc.Mapping): + video = multimodal_output.get("video") + if video is not None: + return video + + raise RuntimeError( + "No video on result.images or multimodal_output['video']. " + "Confirm DreamZero formatter placement for world_model_env." + ) + + def _drop_engine_session(self, session_id: str) -> None: + drop = getattr(self._openpi, "drop_session", None) + if not callable(drop): + return + try: + drop(session_id) + except Exception: + logger.exception("Failed to drop engine session %s", session_id) + + # ------------------------------------------------------------------ # + # Helpers # + # ------------------------------------------------------------------ # + + def _validate_model(self, requested_model: str) -> None: + served_model = getattr(self._openpi, "model_name", None) + if served_model is None: + return + if requested_model == served_model: + return + if requested_model.lower() == "dreamzero" and "dreamzero" in served_model.lower(): + return + raise ValueError(f"Requested rollout model {requested_model!r} does not match served model {served_model!r}.") + + @staticmethod + def _error_response( + step_id: int, + committed_step_id: int, + context_length: int, + code: str, + message: str, + ) -> RolloutStepResponse: + metadata = SessionMetadata( + latency_ms=0.0, + steps_generated=0, + context_length=context_length, + committed_step_id=committed_step_id, + ) + return RolloutStepResponse( + step_id=step_id, + next_observation=None, + model_metadata=metadata, + error=ErrorObject( + code=code, + message=message, + step_id=step_id, + committed_step_id=committed_step_id, + ), + ) diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 80a3839ab..562ed70d4 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -119,10 +119,32 @@ class ServingRealtimeRobotOpenPI: def reset(self, obs: dict) -> None: """Compatibility hook; per-connection state lives in RobotRealtimeConnection.""" + def drop_session(self, session_id: str) -> None: + """Best-effort release of model-side session state for a closed rollout.""" + drop = getattr(self.engine_client, "drop_session", None) + if callable(drop): + drop(session_id) + return + pipeline = self._pipeline() + for name in ("close_ar_diffusion_session", "drop_session_state"): + close = getattr(pipeline, name, None) + if callable(close): + close(session_id) + return + + def _pipeline(self) -> Any: + engine = self.engine_client + for attr in ("model_runner", "runner", "diffusion_model_runner"): + runner = getattr(engine, attr, None) + pipeline = getattr(runner, "pipeline", None) if runner is not None else None + if pipeline is not None: + return pipeline + return getattr(engine, "pipeline", None) + async def infer(self, obs: dict, *, session_id: str, reset: bool) -> ActionOutput: """raw obs → engine → actions.""" # Build request, run inference through AsyncOmni - request = self._build_request(obs, session_id=session_id, reset=reset) + request = self.build_request(obs, session_id=session_id, reset=reset) result = None # OpenPI policy serving is one request -> one action reply. AsyncOmni # exposes an async iterator, so consume it to completion and use the @@ -141,6 +163,10 @@ class ServingRealtimeRobotOpenPI: def _next_request_id(self, session_id: str) -> str: return f"robot-{session_id}-{next(self._request_counter)}" + def build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: + """Build an engine request from raw robot obs.""" + return self._build_request(obs, session_id=session_id, reset=reset) + def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: """Build engine request from raw robot obs. From 767cc7977e04dc1f1ae7307e630429e9169af622 Mon Sep 17 00:00:00 2001 From: Oliver7th <97155135+Oliver7th@users.noreply.github.com> Date: Mon, 14 Sep 2026 09:46:29 +0800 Subject: [PATCH 049/209] [Bugfix][NPU] Materialize causal mask in FlashAttention dense path (#7324) Signed-off-by: Oliver7th --- tests/diffusion/attention/test_flash_attn.py | 120 +++++++++++++++++- .../attention/test_flash_attn_npu.py | 118 +++++++++++++++++ .../attention/backends/flash_attn.py | 59 ++++++++- 3 files changed, 291 insertions(+), 6 deletions(-) create mode 100644 tests/diffusion/attention/test_flash_attn_npu.py diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index be78ad690..183bbd94c 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -503,11 +503,12 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch): # - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device # sync, no MindIE-SD import); # - env dispatch in ``forward_fa_npu`` (which op a packed forward takes); +# - native torch_npu dispatch for right-down causal attention; # - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``. # -# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module -# into ``sys.modules`` (same pattern as ``test_paged_attention`` / -# ``benchmarks/conftest``). No production code is exercised for real kernels. +# They run on CPU without a real NPU by injecting fake ``mindiesd`` and +# ``torch_npu`` modules into ``sys.modules`` (same pattern as +# ``test_paged_attention`` / ``benchmarks/conftest``). No real kernel runs. def _npu_impl(causal: bool = False) -> FlashAttentionImpl: @@ -529,6 +530,15 @@ def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_var return fake +def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None): + """Install the minimal torch_npu surface used by causal NPU attention.""" + fake = SimpleNamespace( + npu_fusion_attention=npu_fusion_attention or Mock(return_value=(torch.zeros(1), None, None, None, 0, 0, 0)), + ) + monkeypatch.setitem(sys.modules, "torch_npu", fake) + return fake + + # --- Test group A: boundary resolution (_resolve_packed_seq_npu) ------------- @@ -685,7 +695,109 @@ def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch): assert out is fake_forward.return_value -# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu -- +# --- Test group C: native NPU causal attention dispatch ---------------------- + + +@pytest.mark.parametrize( + ("query_length", "key_length", "expected_inner_precise"), + [(4, 4, 0), (2, 4, 0), (4, 2, 2)], +) +def test_npu_causal_uses_native_right_down_mode( + monkeypatch, + query_length, + key_length, + expected_inner_precise, +): + expected_out = torch.randn(1, query_length, 2, 4) + fusion_attention = Mock(return_value=(expected_out, None, None, None, 0, 0, 0)) + fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention) + impl = FlashAttentionImpl(num_heads=8, head_size=4, softmax_scale=0.5, causal=True) + query = torch.randn(1, query_length, 2, 4) + key = torch.randn(1, key_length, 2, 4) + + out = impl.forward_fa_npu(query, key, key) + + fake_torch_npu.npu_fusion_attention.assert_called_once() + args, kwargs = fake_torch_npu.npu_fusion_attention.call_args + assert torch.equal(args[0], query) + assert torch.equal(args[1], key) + assert torch.equal(args[2], key) + assert all(tensor.is_contiguous() for tensor in args[:3]) + assert kwargs["head_num"] == query.shape[2] + assert kwargs["input_layout"] == "BSND" + assert kwargs["scale"] == 0.5 + assert kwargs["keep_prob"] == 1.0 + assert kwargs["sparse_mode"] == 3 + assert kwargs["inner_precise"] == expected_inner_precise + + block_mask = kwargs["atten_mask"] + assert block_mask.shape == (2048, 2048) + assert block_mask.dtype == torch.bool + assert block_mask.is_contiguous() + expected_corner = torch.tensor( + [ + [False, True, True, True], + [False, False, True, True], + [False, False, False, True], + [False, False, False, False], + ] + ) + assert torch.equal(block_mask[:4, :4], expected_corner) + assert out is expected_out + + +def test_npu_causal_composes_explicit_keep_mask(monkeypatch): + expected_out = torch.randn(1, 4, 2, 4) + fusion_attention = Mock(return_value=(expected_out, None, None, None, 0, 0, 0)) + fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention) + impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=True) + query = torch.randn(1, 4, 2, 4) + key_keep_mask = torch.tensor([[True, False, True, True]]) + + out = impl.forward_fa_npu( + query, + query, + query, + AttentionMetadata(attn_mask=key_keep_mask), + ) + + fake_torch_npu.npu_fusion_attention.assert_called_once() + kwargs = fake_torch_npu.npu_fusion_attention.call_args.kwargs + assert kwargs["sparse_mode"] == 1 + assert kwargs["inner_precise"] == 2 + expected_block_mask = torch.tensor( + [ + [ + [ + [False, True, True, True], + [False, True, True, True], + [False, True, False, True], + [False, True, False, False], + ] + ] + ] + ) + assert torch.equal(kwargs["atten_mask"], expected_block_mask) + assert kwargs["atten_mask"].is_contiguous() + assert out is expected_out + + +def test_npu_noncausal_without_explicit_mask_stays_unmasked(monkeypatch): + captured: dict = {} + + def fake_attention_forward(query, key, value, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward) + query = torch.randn(1, 4, 2, 4) + + _npu_impl(causal=False).forward_fa_npu(query, query, query) + + assert captured["attn_mask"] is None + + +# --- Test group D: laser input pre-scaling in _forward_prefix_kv_slice_npu -- def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch): diff --git a/tests/diffusion/attention/test_flash_attn_npu.py b/tests/diffusion/attention/test_flash_attn_npu.py new file mode 100644 index 000000000..c971e251b --- /dev/null +++ b/tests/diffusion/attention/test_flash_attn_npu.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu] + + +def _bottom_right_causal_reference( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + key_keep_mask: torch.Tensor | None = None, +) -> torch.Tensor: + scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * scale + query_positions = torch.arange(query.shape[1]).unsqueeze(1) + key_positions = torch.arange(key.shape[1]).unsqueeze(0) + keep_mask = key_positions <= query_positions + (key.shape[1] - query.shape[1]) + if key_keep_mask is not None: + assert query.shape[0] == 1 + keep_mask = keep_mask & key_keep_mask[0].bool().unsqueeze(0) + valid_rows = keep_mask.any(dim=-1) + + output = torch.zeros( + query.shape[0], + query.shape[1], + query.shape[2], + value.shape[-1], + dtype=torch.float32, + ) + valid_scores = scores[:, :, valid_rows].masked_fill(~keep_mask[valid_rows], float("-inf")) + probabilities = torch.softmax(valid_scores, dim=-1) + output[:, valid_rows] = torch.einsum("bhqk,bkhd->bqhd", probabilities, value.float()) + return output + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native causal attention requires Ascend NPU") +def test_npu_bottom_right_causal_fully_masked_rows_are_zero(): + torch.manual_seed(0) + batch_size, query_length, key_length = 1, 4, 2 + num_heads, head_size = 2, 64 + scale = head_size**-0.5 + + query_cpu = torch.randn(batch_size, query_length, num_heads, head_size).to(torch.bfloat16) + key_cpu = torch.randn(batch_size, key_length, num_heads, head_size).to(torch.bfloat16) + value_cpu = torch.randn(batch_size, key_length, num_heads, head_size).to(torch.bfloat16) + expected = _bottom_right_causal_reference(query_cpu, key_cpu, value_cpu, scale) + + device = torch.device(current_omni_platform.device_type) + query = query_cpu.to(device) + key = key_cpu.to(device) + value = value_cpu.to(device) + impl = FlashAttentionImpl( + num_heads=num_heads, + head_size=head_size, + softmax_scale=scale, + causal=True, + ) + + output = impl.forward_fa_npu(query, key, value) + output_cpu = output.float().cpu() + + assert torch.isfinite(output_cpu).all() + assert torch.equal(output_cpu[:, : query_length - key_length], torch.zeros_like(output_cpu[:, :2])) + torch.testing.assert_close( + output_cpu[:, query_length - key_length :], + expected[:, query_length - key_length :], + rtol=2e-2, + atol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native causal attention requires Ascend NPU") +def test_npu_bottom_right_causal_composes_explicit_keep_mask(): + torch.manual_seed(1) + batch_size, sequence_length = 1, 4 + num_heads, head_size = 2, 64 + scale = head_size**-0.5 + key_keep_mask_cpu = torch.tensor([[True, False, True, True]]) + + query_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16) + key_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16) + value_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16) + expected = _bottom_right_causal_reference( + query_cpu, + key_cpu, + value_cpu, + scale, + key_keep_mask_cpu, + ) + + device = torch.device(current_omni_platform.device_type) + query = query_cpu.to(device) + key = key_cpu.to(device) + value = value_cpu.to(device) + impl = FlashAttentionImpl( + num_heads=num_heads, + head_size=head_size, + softmax_scale=scale, + causal=True, + ) + + output = impl.forward_fa_npu( + query, + key, + value, + AttentionMetadata(attn_mask=key_keep_mask_cpu.to(device)), + ) + output_cpu = output.float().cpu() + + assert torch.isfinite(output_cpu).all() + torch.testing.assert_close(output_cpu, expected, rtol=2e-2, atol=2e-2) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index cef6a97f5..d1f09b51a 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os -from functools import partial +from functools import cache, partial import torch from vllm.logger import init_logger @@ -19,6 +19,15 @@ from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +@cache +def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor: + """Return the shared block mask used by NPU right-down causal attention.""" + return torch.triu( + torch.ones((2048, 2048), dtype=torch.bool, device=device), + diagonal=1, + ).contiguous() + + class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True supports_piecewise_spans: bool = True @@ -549,6 +558,53 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: + attention_mask = attn_metadata.attn_mask if attn_metadata else None + if self.causal: + import torch_npu + + if attention_mask is None: + # The compressed upper-triangular mask with sparse_mode=3 + # applies FlashAttention's bottom-right causal alignment when + # Sq != Skv without materializing the full attention matrix. + npu_attention_mask = _get_npu_compressed_causal_mask(query.device) + sparse_mode = 3 + # Explicitly enable invalid-row handling for Sq > Skv. Equal + # and shorter query sequences have no fully causal-masked rows. + inner_precise = 2 if query.shape[1] > key.shape[1] else 0 + else: + # A custom keep-mask cannot be combined with the compressed + # sparse_mode=3 mask. Materialize the composed block-mask and + # use allMask mode instead. The explicit mask follows the + # framework convention (True=keep), while npu_fusion_attention + # uses True=block. + explicit_keep_mask = _maybe_reshape_attn_mask( + query, + key, + attention_mask, + mask_mode="full_qk", + ).to(device=query.device, dtype=torch.bool) + query_positions = torch.arange(query.shape[1], device=query.device).unsqueeze(1) + key_positions = torch.arange(key.shape[1], device=query.device).unsqueeze(0) + causal_block_mask = key_positions > query_positions + (key.shape[1] - query.shape[1]) + npu_attention_mask = (causal_block_mask | ~explicit_keep_mask).contiguous() + sparse_mode = 1 + # An explicit mask can create fully masked rows regardless of + # the Sq/Skv relationship. + inner_precise = 2 + + return torch_npu.npu_fusion_attention( + query.contiguous(), + key.contiguous(), + value.contiguous(), + head_num=query.shape[2], + input_layout="BSND", + atten_mask=npu_attention_mask, + scale=float(self.softmax_scale), + keep_prob=1.0, + inner_precise=inner_precise, + sparse_mode=sparse_mode, + )[0] + from mindiesd import attention_forward # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the @@ -567,7 +623,6 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): out = self._forward_varlen_packed_npu(query, key, value, extra) if out is not None: return out - attention_mask = attn_metadata.attn_mask if attn_metadata else None if attention_mask is None and extra.get("npu_attn_varlen", False): # Models skip mask construction when this opt-in is set # (FlashAttentionBackend.supports_packed_mask_free). The packed From be9f3ed5c01abb53e084fe9316655bc700c2575d Mon Sep 17 00:00:00 2001 From: LiquidGunay <61792575+LiquidGunay@users.noreply.github.com> Date: Mon, 14 Sep 2026 08:02:21 +0530 Subject: [PATCH 050/209] [CI/Build] Align CUDA release image with vLLM 0.29 (#7445) Signed-off-by: LiquidGunay Co-authored-by: Hongsheng Liu --- docker/Dockerfile.cuda | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 3881ca2e5..072b12dfa 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.28.0 +ARG BASE_IMAGE=vllm/vllm-openai:v0.29.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app From db40ed3a3179d9125408dce9c89f1c97b5a5d03a Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:54:25 +0800 Subject: [PATCH 051/209] [Doc] Update vLLM-Omni WeChat QR code (#7490) Signed-off-by: David Chen <530634352@qq.com> --- docs/assets/WeChat.jpg | Bin 278267 -> 273418 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index 0e91feee5b53138f32934a9546169b36e4adf069..f55fa90f30101c93944efba8b11deb08f0dc50a5 100644 GIT binary patch delta 237746 zcmaI7XIv9q)IAzS1*xLaOH@!$1VlhOu{?l)fP(ZA6_E}i0s%rsP4fw>?auDJv8o2^;#1IT2Ob`V?LwS-b9I?kEC-)ZBQ~Lq^kqfgATNU5HhxY5miXh_NS)gy7hOL z7IGk(Dtlr|_MO}cSEI#ojoFzTvQH9e1Qu)oMZ=OXC0f`uc&gQiTy4OO zm{&gO0P~(my6(j83HRWN7W)xU`x5ELQojV++NkCeqd@qRCy&jTc`Y<%Iw=x(=#=R~ zZb26O9qD{&Mv=|r)HJ6cUz(O&G{zwI#+d#L9_@(y10 zOk3An2RrYpHC0mn0nHMpAO|#i?^KKP`KQF?NaCa?!HHB zv*^rI&}W)R+_2>X{khtjzB6h9ryvwmOXp91sV|qKzk=;Z{;!CE&Ug}-3@kRuK6^tf zh7zU=Gfr?DG2|zwAdyNm)f1qe!ULNIdO_kc$3x zf>9g){vLohL&SztB!w!BbXsPoHk0f6r8;-3y<6_L#&`HJ#~~**e(=7%AG^+WRs+(8 zPJ=%x-0;H?v`<0jNa}KAWHBrfSGcB}Hdf~9|EqIAe@v$Lu6Im#iST|LPxL`qvnW4h zfpiMOoH!y>S0Nk}71G86Dwy)fGhfTGuph4hCPDzZx#peT5L4~QdO|ZqT0fv@ zwR2J6kFRLjPvNYwPxoopPeGl?0ssqHmsmUnP1e0;Y@Yfxx4TPEG??~)*CYf`$ti|~8deU;jc-fZIv0gZh7qnYL$0(t8oZAx?!Pj0le(S= zxxHT+H~d0u3lkNak%`JsJ+FzrimqS`&f4XHMIl-gBoESo%rX=<7`& zR*K9Byt1PWU-V6)oLR_COH%*b=m+N+a`JmQMtnM;C+AZ5m)aV`frDlMf~TPQY?~|8 zMA{WASg?)j2Y>&Z;i1`cM{r;#=b%$tK&#uA55+;&Jq0nU-bLJq!L-e>f1FX!yYE}# z<5-}r*5>3_plK|4W@8pD(_*JsSa$Hu5B%c0?Qc?&Xr;OlJ;4N>f|}J4QZWC62-1Ja zUy4ffZf;C#c~>)^3tRaRoxS+lj3XbFCWoBJSTKoS(DhgG_bu^ymjSp)Gq}38MP2fo zMC++TXrref4uD=~cxPcnEyCAEf5siYt{L9?6yXq&Z&es7;ttvuatduNg0jL|J1nG9 z3Z1K3fUq#ZSx#O?OVaMoV0=kgMu`f$BC=IXO54z1kP39WBnw8Jon+-F zRpMUU>QEeHmf=k-2Dx{mWxRxp2sX+M?zrc&G*T}tRY9-dFNaHw|zUl>E^P3&ea#zq2fu+dklxbnGTBEDJa$g zOGERVf*$%U(k{`%j^KU@vYMk$DnADQ|2?hy#xK3J>`YBeDhZl;3gS}dUIn-jw_!!k z#$2y8l(~nKlEQ0y9!VBQ#?tl?(QYdl!<7ghSR9o6q)8H%VhGe$*S~BeH%r#l8on$_ zkvCKeS98))fH)(^F?e(g;xgqi(T`$DcDE&NcW_eEcx;AB`SH5d>o;PI6E2+zjQac+ zb&;$=&u0Iy>cX^sw)sR7b`w?V;1`!QW9_7@H*96onRjQCUgY_%f4au%3H=GJ&f9Ee zBsiYjhkb4V6!@!he(8xydbs<(Y16(Mc2V$xK;T-C2Q=*kVxTUu{saWmpK3*yyGOu% z7uu`%-Adjk`dR;4sD-fgliG!^y+WR z)wxD3d$=c3r}f=zd72F6EW77?KZbaCJl9B z5jPxv5fe6%ibJy|msy`~n(MHZf{x1Ldl4aFQT9uevVrFpwAvs?oR!kYD^v}H9X)jX z;5U8oNV3?w+Pv{ev)LK#;w zZ)`MMIXHdF%Lc%5q!2gspqVNe{qJH=)SOL+ZJq?=tqznF^2A=VtpIr(9wwly$UY}r z6cud4mE3jZq)0+VeNj967(GRce6$~c&lIg2u)$QUIE!DvWu^XTPDOS+s9_tdSrl%- zYG;x~4C8v-DYv7~BMyd+-vW0Bw|MclkBvsAlAbIBoQT#8R(I~XK2mOuSlrLj8FTg5 zn|IQ%o3@b9W%C9hlIDh9$g}Tbx*7LM)f4t_#n1O&ZVwQ>N0A;{JmF-yo}yX(yl^@? z&vE~3wLIhuto%jg%TU}{QeKzfT&o@U-)AAsS=ruaK4k$MF;H9{h7Ay}FhL8V;<#|$ zsR%(V1Z4(?4VU;Rbj)bKGYV;2i}FjCAV@$y)1-zf>!V*#Sl;x?2pCU&d=(h=Whw$s z@gmB$&&t+)pM}{!3Hf(#{C7uLeXvq1*TeS^i{T#{lgF`lKG=uS*GEcS8g4R*p3j1E zd<<}2xOo^ZH}2IQ0T3?vcHQ88tD3;Ofy_kBq9T)1lNAD`OG~=6B?shQ*)p|$^+k0e z$3+e-8_7y12r`dZ+2E#Jy?@b|Zs)5y-!pfMr#fD~A`SRqH`DWoRD}52kOSH=c}yky zoT^7SAL3_nD<97tBcvsBXsrp`Aw&sqo)aLuA4HR7N*tg7CrWN^lPqX$W_8bOxiaCA z;RbbiDHKHPXc8id1)qZ6L3thtr=T2tp06{n+Sd%-&5cofr1}((4DHcjG&7HPM2vK`{e%BKAuA{g3o)TYw4xH=6A9e~dAm1b=?(AS&kKZ8+ z$jQT-_rd_z*e4x-!!JB!yl(j!n@olamdS0L+;#|%x+Ry#?xfX$P$S3ZP?X5Eg;sD` z>&HEh{-kAA>mGfHo7UtY-48y;%?D!tUJ6}sCvf07!W8TYJZ$fh&Y0jk-!yp~uH?zb z<4!?5YsY%ot+Z^i@3y^+bh>Zvc%laPkH^U$6u5O)UvfGiID3Fqo{9?f(8eBu(`;MwZ2AkVY&6k# zd;f}5k2Mf(eTBZh2QvKYzuc5bjc;H=^H9!s33~L%&&J{VcuxwarFZ1pW1<{~4hFT& z)>;flj1qFE>>KGjglVtK3rZH({DkFnlJ3;-|zi{90zpi+7!|Czj@ zt|$o8pQANl!H2Eb>%>?j<<`N8AJbyjDdrO0xlYtI~GxaN;dF*EFwgV~myc;gx5x8waaa|@pi2{0ejwFxVN)Rw%V z?a~Rrp6+)zEq-j#wO#FmxClL%GDyaKoI*)mvQAeT)BGyzS;0CAu+padq@z@?b64g8 zYjsQQjeF@f!eb8y*=1v^)i!$I5E7*M0l(Tzg97&r`^dKjW#LF(z7D0on@+vxb|~vA zjT5PhG_Bml-&dK@;0)`2%6mcnjA@>`MS^&(HJi413X~6+&-610ENyD3Iu4x@br`*} z*_VAA^AEPj;g+y2EJ0FFk#4Qs_csSjbj1n$Ywvg3hja z)9%(_rAzleoXPNgV=n3<1q17!N5lfnx`1O9_3J6fj?J=zw03ZuWo~0Jb+cmVD#vo? z6(xq4%02njk?32$_N!J>fRk!-{j5J%4r<9592(hlKaRz+KpitbWj*;b_L;Mzgm~@i;n)--Z6fQjd zVlR0iIVUzO?2w_$q^2x4FGsPowu|n9>Sj!>V!qcE|MK?BZ`c3|WHn;s0g^b{5n2G7 zRs86$RI(-MdiNulTYyVT6tfI3_4#x)y@qSqjx6{%zC*7_Y6Wqp@uE?*;@r}4M!2k1 zi0j}zy^`*PI}pK9cfzH6v6$tj==RAbG94=y!yk}c=lerw{IvG+p2vK%ku`h#n=E#2O9z*DdaG}_RXbCmMC^xWO1{I>vDQUj*DX2m^)b;7x>bT9I zrW(5ZzArxu*iZR(}o`S32$gzHEn?5G_F2B^Z$P0oBpi> zi~^)qE59qMww7Gh@<<+ax9=Z;~A-=-%>D1FACY93IwtQdUz z^P@cenge;)j1fF(3O$?=qlg`Mw9HK^z-`Gz^S$7M1A1I>;Zu591Cq1{ zmxW#Qf8X`FhdM;O=vt{ISet%^wk5yIEoZYiF-#wqdA|FnIGf`co(tI03iS|-?$L7P z_}=E>eVI$?M!+|?Pw|;k&~<^VJwE|VT}wvF^9Hc##vUpJ!3i{5QFL+0Fdqv@)SlXd z1&(TJ#6$@WdTUlDwV|X!fp{#e7qCwh!=ve;>cGMNB;P^W@w(lv@I?{mz1yqfD>3K< zU*3OQR;ovLBn^P`z>nu`nDmqZOd|=Am;ad98zA?!j?z4JqZZ0XX4D`fe?Evst$Ei~ zrASDqRkVI|`MYFB<&Y=X+#e=rg}A?sVezti7F2d-7Oe=(H_S?nHE_^V#R};{=FfGD z^!Ng;!nN)sZUr{iI|Vh*$3I@f3u(N6S=7a7C}BI4r)iiE6jVEBJhr5Dfmt{4#>i$1 z?zT?M1eI1PLrrKB%(Qc7T=fOtd(G%abKFBtZZD{QDo$4?bu8i2K5U4dg66!O4XN2Q zky#sadX;1mz&vO1mJA*%1$H|ePav(lD8>!@3VzmQo*NLQQ6xAe9KZ|Z=BiAmwPsHQ);@sRXt`-~o)Bc>frtdzH z^7lHhXgFXQ#X7YDE&=2f&=qrSW|hcFprr#KwW7W~c|xu93|0aYMp;mk_yyzM{&>bV z4ne7{FB?9pry%^Eu4)jnQNSaXP`(h6tgq z*cPT^W-I70R*;Cs^q;eWd3N|oWh_XtFT73E6&RTjj4DZZ)7173%|z#{@frMZz;^+d z!r-DyOYl&xb(19istMPvegE*+3Gf2NxAxbYzL;)QQ@C_3N8}*m8)kJw+o+0@_Qr}A z^2H^rML#qjdR;@p`L}yKoH@O#?30pT(TXHqj^epirmT zQ&6ZkoKy0Qdq@)>n_Kr<82bS5s8g}ug5u{MF{>Oy=SUGf?Q{OLphMw|?P0BUK**h! zloa`13qn$;IR!0jG!91&mT(C3o~SfupoH(z)9f2WKTQobWk!C1W7Eh|LiSdQiVY`Y zz{IDJiD-HhMM}&Y3!w-4>AV*jhlTP|e3pHk@epn*P`L#m*8y3O_EV7XWgplS{hR#1 zv#jQ5KcS0zGJS^2vBxAZaPYLu3K}Q*_l}sb4 z>AAUpu&kc8+cOVJ`HE!*dJ<^fW?8fN=a*0wnuqa!H?{D?_GmN=gQ_i^YEcrF9KN9I zS|4VwM3aGrpox0u2z5C+rE@a~eK466*`N*B-uk&A$HQN|qjl%#KKKU=~@#q~mIob{2tPkMIG2Y=fQ7SivWK(0jO7tQeX z=uYprlW#K;vN)Qz*|l*2H&!vIxj;9*Gq$xI^YvPa7}Cdcwv7mn*7=l0l3hD>@sOtk zuGd}dl2mfQN!Tv zql@lb%8-_vp~P+qgUO-UFoia(;hIO6uRpLToznC$_6hSR_GJ-)^6IFS5(k*{K}%Ep z*o4J^MfTV7`lrlWZ#Xnn1@XLG%Ls1r8SYqHTiar19?tdTpo`M4G({~ZW8<@=bp4l*gTA!!+yFIW1w-7j)>-t8?|e!fQjo-a(B>7X)FayB z6~fSO!B<*QPxAXFE%lc5KrT$I-`T|7qV8ST7yOgYz?i%XiA1y9BHl{FIroB}H*hO2 zWphYqsiSJcQCHg)?U({&g_dP`&2 zG7E20NqEqTvlnXG*VpDoTtvsK<)f>q7JhQu0QqfFbTmdcsMB^4+K!q8Q*5-KrhxU}o1Htt!f=AMl>Vq~6&A*@KvS=>ff> zNk5+3!zMHx9^_`kuEO_GeY?e~`I^Pvf9yBSxkV;N@@{PxLfGLV^YsGXOLHAwd{8KYEy~0$bMx|Tw0|5RLpS|=eQl4jPcp~%_fVldzr8%< zUvv3q{8}J=Thlbi>~!LP>&@+|@U^&6n^oE0#PpyNomMdgFG0?JJ9?_t9$-$~ZoNc# zPJY@XI3nk)C#Esj+)zJ;kG5Dakqo@N!lwwIPmeSN;AZrD@G&R5wu@-#itRHp~uwU9s$qQ?ES8Df4P%_f_AHxe==Ipx4eS#~8;eb=)pO|=^>l3!F$-mdf_;qphYt-+Hx(#cXE`cq zoqio!ltv5E-q$zi>%e@Rf)3y|7gAJfR$A1jf>_%)5TP-vxgMY=j7*WV^b)l8c^~3r zqnK^?f)}64BR6E1??;B^ExV20$Hy1F0oNjoe!`rUL;Cv`lupF(l8YM4E>>e!$rldXB!()ZuXJm=o%4wPH}k1jNZtjmaju$MR?CA~DLRb{L@^PiZuWl%-A#G|y|kuWf{HyT5YVrY5j)H_`;J=nurylj z^INOH#zRD#z+}x`jixdAd;)gU$1kmqE2ZcMROATDaK3%)2PvY%u_STsL=6Wdn%#ta zkdOEDz$#v7EY}{Y)Z@prDGQgl4bH797bjhA(Vr>P>AwZryZfsFisxEtgYdgbi-x+! zkF(`T5}yWR?${gNXj|=JrT;rsBs++U|~pqIQ5^3G<|*w_G@2-kdoxzq1HK~QzuMI9R^VEHJr=KNIx|9w;dqA4q38kZ6bMmQ zE8p4@^K)(Wa@%xt;ivuAQ5i`=YV8Oh)vmG-ul6S)z7kn65C2bm@gor^k#fjehh-3w zUz0r?Z#IQqDM@SZ(Y<_L`rOU%yOfX9+G_q-krT+H0FQay&f!J7Zx;otp8cd#5_-XF zw4nZ?DLc}{sp-XIn9V~^YnuwFqv2E7)w<6c@xk&XC(4uUKPhU1d3AMDu5u8!biE(K zuNmxX3^kcn&j)m98?pf0qfNZq%Lp$64u#Mnv8_B)jX6E@^|{|Ovsv?VlLtMmp0Zl* zUYD9+39W47XS6W^B{su!deIVYNxMizeO!6@8-v~+`%T}V>HJ`>CUIDNJ5Sf7(82K39M?Ji=7cw{G(XDefuS#TC@kxz>WU#S^fqDbP+bhM9a@!CGdZhzUwY z(MoEi$aaG-2NM1zP|6^U5ZDaS*$=SPiNQZ{EiNOt-#Zi$j4};`-y!*icKCD_)_cUq zW1NnA=7SHvxa{==7?>sM*}mAtM@Zh9c6;;h@7W2F+QkowC}himdf_N6Ybz%-Q3}!C zDCGoaswFA_?Q?6cB8y?S44xj&2{7c1kF5!Cau!jH0TKd+EtCV2efOXFV$b? zoDz{Yql=Dy-@p5_z>ZHcM*q>@br|Y7wE$wL3kjbE5294B)$s&V%s(z@HypjKiY__$ zswrZ~;5IcyLitfY0rMFT)S1ZkLZxMeS}0ieQCc35xbrTZHw$#>Yf=|x_UQM;?V%z; zoQ3ctmweh3(lR|q)l2oRh<>S%0^3!3nE@*xX3(G#ir0dqm?0e++*~db%B=E^P#fgE zcur9K$TE0775n!T#I1U1yFdtBVz{;3n9LnV)>3tEtG96$A9 zaskE^v5s*!8Qw91NWE|f4vT{l~Dse0JvW0g#Nl`qSBJrQtWv7!# zRzg-W{DZxRNaR}Ok~>&>Y|~2s}K5W#ZuZYbME^W)GMLSNTLdurEl+^%`Z2B z45seuMh3@u*R{R}ywP}@G$P~xM=XcM%&J*eGT}`?mm-mf{XGR2(r+5j=QgC?K#G7` zz{H+6gv{hR5K_ZPhcz2eIs$K2Yr3&<3gVP+wAD@)Fqaz_Rd8R73P?BYe-giwmbJ41 zwS!AgE7Vz?Kd@N#>Kkjnmg)R+XFlkj)+~yT`u!K#m1e|jW1YyIf$X8qC6UJHu&cQnl#DA3~4OTI@#D6*)1u*>I1q4kE zK+!(!w^L+P8ldc(c;~%ieu6q3c<22efdgPgFFJ=(F*MDJ=p$F*`dCSyRk$sEAMj9r z>~nm{;LIjrLkW|vYPdMvl&7Z4wI;8I+-Q`Prd7eX6sJq!1 z6}xT`JVTUlt9ES?_ZT_<_gZ?&F0vClfK3;`-Yu4gR*(fuE9<5QbJ~B-IeUH%aj8AW z5OYu2R04+zJ$4Ui(DMrTNj7cc|5D9*zRjquEKJQCN3-%{Ww17WU*jXL6KttfI0x*i zTB-9Rdh05-)uSJv&%K|$K+UjAGOVs1L3`%4q@04v$GpDRUuL;2(Nf^{UQ5+9-&y_A zsQiQ^Tm?537E*@j7n*!>XES7u;R9V}=kew#O6N%)CmuF~BQm&Y!#QN#mT+AC57Axk z*`?CFNW~*c?S^R=hY_9V@-L9`xy`Yv4BL=y|CEohndjaJ&L;{7 z(@2b8_254AUOuxEhA#dGy@SQM(vkD^0D0Mw)M;_9V_W=x{hznhjZZ;-a1>w^|jf*xZBM%&VIjLV0`01g=89~D$Nl= zkP~P$oL%`aS+fWd09gHwgCtfB1?#?kk6TKLIZ`)8$Zqfz@$?ul9rZMjWF&Pd+arhY?|>ajn}% zxx%p-mm9yljyQ%}ldjDCSK{vFem>QyPcoK8BvIv$zOU*53WFx!sYW(<5s~=sas4Ju zmBhoaV$SOT=jGqs^R;%EQa(V&Mw@ilh8CiovvRs*<XyWtj;8}mUn8`j%k3NC!H~?)(FN=!>a3uLAG{Tdj+>vlZ&2iqz zqcP#~Dnlt(RdD}(5Wf$&+J*7|Mdc@9Q~L)H-~R$T_*sq+^8q&Q>0t?Va>h>ow(eYJ10p*kYOCm9qGs_XsNU8P8;rwI9sLJxcIjTStV)lbypbZ)o`& z;`{X$;8JxdVrdF~p`^roYF&k$F||HKO>x6aYu@0?lO;GLSAwDb`w7B}yh7AEXdr=G zq$8VzNZgeuSjx+?p{*}}fbpixyFrBanlr%%>5Wv0i6&`!y*F0B0V|_F#xKgV25gQ8 z!|`nYzDEAAkBl+5;u1Q(qLU}6lObLHb1uNeuQn-&x6W+t^~LrzZ^GMC(3O4n|3;RH zGhe(X{naY&Mg3RywGh2BRGJ~MURKI|jEjf~tZ|2Xq1u3O2=O@w5Rq}E39OOnWOL)2 zoWkSP8H?)iWrP6pvsahY-F^^k{Ow8D4s3h{Kum%a5tLpsdzrUFR4+03>}{LWbm@t% z$<67em-75gefeey%E>E=s>CiMj%UcI7mdNV!~@fqhYDGO3|3kl(m{wT)SqxpNC?IG z=QtVpuqHt)v11=l?d-{N)5PSzP9{RhVxi5F0FEQ5#LIxod9n|+U~svxSI#skZu*J6 zo2D_B)mOfd4#`iCI@h3PFgp{x(es{KEBI7Bwme~Z&!b}Lj)y@KAhM&nanslP=3vxO z((cW-{gOnU%^$=moKL9aUIqtQ_Q}x4smC1x4%ftgZ`}l9l(w`X(e4V(NEWVB&{y^h zI-%VdOYG)@?bwP@2lXBd6h`+jtsoO~b9pXG&*~)Gt9qDytG(KPt(JUoeroA8BxJ8E zp@X@?d|IvDgBwjWfCQPO-=WPQaj0NZvlXhuChV(yDZBG(BVaK^k@(dM6j+&?O|)F} z&9Qq}18m1b)tmsLDaQm#Z^Z3xf7={uul$c2D?RKz+-mx>(|e^+s^X8OnwCq6;E_-zG1cEGvP<8+&cxe z0%5`enE)|zh6jwMgKCEE$){p z`?jPpeO8WSSD}!8pY(Pzp>xOHc(j!3@rC}Gxlb^LYj2Sy0on^df4HzHHoLmX zfS;VqDa+**YMl0xGbTfEA0<_NwT>&{_K;ti8BBTT59B{I?jQJ1J}=!k&NnZ(Pv_k? zJm?3*Y6_?Sfid(b6LEHO8o&7;o;U+d5tb(>;5iuKs$=rH27<6U-S4g;ZM4B3p>t5v zy+Q~G;$Xl9L*v|m$Og{>SZMo;E?bwuR=3GUOrnjvEvA1?q|z^A>S}KMi?R*iZM_qs z2!4m2;LQ>J5Qc_;T)L?c!bAb%wl6z)`W{#$sRdRA{a^`BGIntGDlZ9<%*0Z) zbN!PiGSw7kK)GUnFi*m^8%(w0N+S!LCbN^@DDd&oXwfKVIQd+&AlC4I{2%}0@gp11 zm}lD%ZOWwWgDEgoTnE>i-vQ%~Eg)9xsrB)nZ;A@czj2!!aURLC^X?y2Q+oUSuvHW_^oat7 z?Y?gB8^WHiGUl!}>w`Y6ZiN0l49GtQ7(IfP2Ma|NM;iG4>z)~F zsyN72iHl(I@nBamPFhOoRu`|8Y}ALRsx7ssJbtZ@@c86-PHz?c%Hy zHG(2g;&eG6^wHqQ8ma4Vwy%m>RZQ@{I{GW%|uCw-wv(y7?>vH{WItx5@@ zTM{Xe9*7UKE>A7QR9OkHm@@#26XoR=ft9FEoNDeIxKHb!V*e?~GRGGjNeLjH%!d*L zI3g+?_%Wu{Rn--ZweKVk9K^h%MMV_<-s}=dW#8}lZdA1IcyiuxPf3RNaNp@$?sRF% z+x}9y+mnZFU%ReP^FOKw-|kbU2}uk!+%w;|JHE08y-0@8m~|Q3IVORm<|Y=RUS3kx zoKuJZ;qe+zUt^dEF|w2C={aG6z-_1(`B`9z)_~O5v^Ji1!v+&t`|LN2Ii>=o!`8VD>C$YvyeQfKEn*IfD#1e=O#RHnwfYGey6HqoA}&J6(LL-jG7g!z83@3{Y*XU z(TwUsPsY<&v`MyjAh%qEeDfXM`e50ZSml9XK!-`cdw1v8gL5uYI)QvqRRNx65TY?S zCOMRF^TcYhZ{E}K24ZhEcpATxbHKHajeOloP+D^m6BrZz7`(yzd1336w z&aqLRlHA+%UQlHH%Jni~zq?qTh};(Mcn~e{&E@l|v&H0)z(?vonO}c4r$Kzep_E+< zFgqPDr_{A4lq@D!x4U0*;tw7)sS~psiA%8lSuupz-8xQZJX;N>UEFo&?~=+OiCG6d z%#aJ>eQOM0tU?q^mF&qHu(ZMMG8L;mcd{6Ua@D?L?zeQvYhp}o?l`#`S3*~9Ei~SS z2pJjp=E%Kzi>i~eb5z2RpNLAWgG-da`e9%HT#uYFsw3bRu&KX-JHc}+^oVNf1K$1SuL4hl;q?E@g0#X zf+^*nZ!LJNKI|;oSxr31u&?_QBCxtmVCk=%fPb3??ePf$AM@cxxvTsJpHj4{t5Xa| z#A+j++_qi%Rb=n#uK;co5>DMB4;`H=Y{2iiyWoMe;9pU+EYD$MFPBM_UN)w7%& z=z&gk+JOaS7uY2f4x7g5_jrzU7QW2;K3JePsiV_h4!HxT$IiKMa)NDXGZhHE6~1uJ zWV-sY_04c&#!}21EydE|{~ECkN&&t&Fc@b4GQ!KqJ^<#Lu6KXYmQVMSY9mJRMa^zZ zEvnvq%ax+y>gkg|XExaiy5RV>#^9^%uo6|GlXi`U&A4Dc;9fyBz$9NNa4*YB?7kwR z^B(eri@_(e{{s5r-{pzmcWJ;gmHo-NTfh^e1sBBi{Fjv)*hEQmiWVmQKKXp1bJK%z zr0J+sWdOt=L0*!Zc-DG|YiDSn8gf4b#Gb8{YWmv@F+6*tY&4-#Ab`#NoQUVsVFCINEorejPPc|5}0nF+*5kuVEO^nEv zL!R_I!^M&dW0j3MP~cYI+QPe6B0oME{$0%YzQ{RL$8m`^aGV%PVr(^$&UTcsbdudR zC;VjJ5kGg&dR?ZUZ=vxqbSlMyEc}s#A~SBPGt`8;-Q6x0bB@eNiT-nuN!Xr2HMZYk zxp(7Dn*2U27n`h=5muxha$j09&;I79wP4;RU41<_Ny4*xS@S*jqyurYN4;80-Mpuv9ce(VmEaIF! z82F%U`VW1vf#KIGO@$nYGIMxf!P<0##5JD7rNO? zJbMblv1iEnVIw{U=xdIE4SgFHqkg+KxLR@GPzf(m9lB*3N9(}m1d-dfHOJIy8*Vz! z&jlYUYsRxkHkf?)b=`XEPu|igNPTnx0uZ%1$YT3R1Y^9!hR_880*medTp;^FVf>Vw zF12uIIx>2K;c(pSkUeXu6ZMSFfc57P@?P}(5+kIU?pA(~|*!RMLCwVk35*?TKF{zC;Y zy;4l@9mGjEx#;8?{5BK_w~$i?!qQccY^&aeHu}qU)sEv&d0g}-&d}Vxi+#2!!f~*} zFMeLQT9OZgbkJW1WQw(r{BSewu_s!rQ|T}0bId-yyCMjA8vw>hmDuCJ46}OYr=4Oq zeC`Cj_LO<<{p*48<2_N3Hs4|4K?Z2+?#E7omo9fB@{vHg?dAV~=h{1;+C6H78h342 zJ(#7|ggZDS+I{y@?rnOHcXfxi-7C}{FK}$SE7SESO;fSURkK^e~XG{N=jnr z?}+|=ZTl3c1J6) ziDB*VtBti&Pk=|>rn;4~Q$jRs1$3dmA3NKKIe177mP1Dh3s#Pl!nMSLDH?WeUJ^G) zc$~UPY#TLk_1f3qd9l_3m=^8Q(p&M`Be_Fw?f(5@x~wq;#8S?^^8EZ&;z>)>g_?5g z=C*va(}rnK$GS>F)%fr0;(zDAHNWzR5@vf@Hjo%CxP7aHr)!YW=a!f&NPcHnWt$;gFhqCn zM~*8$YpS>lrA!pFIq!F7e^PhF^#&y1&QC0qNE*!Yc-U_=n3g4Jh_AxFPiXhd{*M4# z+WL*|L^}d#59ofCi%$;G_%%Q9fp8ODfDTHgXpYgiX`OXrC|s|}iN;9Dw(GSBe-oK5 ze!2DHpQj+Nz-vfPs3%-|$kL3M)^cYHb_#<1^tr}Kb;hB|x1p`ko(rhWoMHz*_7sl- zrzC}{>(`?`y_ezA3`VSk%eDCP2RiCtn+%~q-ucxR@3#IF#Eta_l3_-=ri*&R39wX4v zFQ_$-3g%zVbm!xo9PeFhbqR}f-4Qm{1Yv9>h*21&ypV!T$Gkf!<~E(Qlz%nd-&XdY z#icstWu;a6Y$6c!1)mE0jJ*>9D3jAjeAUeZ9AIVqE8!a5EV(1|jENto5jX+1PgJL<;pj zZgKh)%xep`=b6pUd>HOHy^95gi;hx z77{jAmbqzCdCTRJ^+4Wnd1cIic?#=adywDD|BJ0J4~Odg{~p=*JzI-W2!=byQ* zuAFo3bD#J6dcWT9d%pJd_t6te$S(SNcvYLip%_p3qzk?wwh#V6aS}bWrEs61`hd=R zBQ-m9pEv!`D|#~T-ngIj`c`)Lh^)h|7k(#klsWo~ zeGC8dv?-Q9we$6TcSvlWr&Y7l?Osh0^sW`F$SYmXIhUu&ZL;u=z zwQMBYwrMddNT2Z>)_uwTOsobRS7Qb0xfvHx=SpnF={?7|J0nFz&8>auv5Y{NzW^>U zZyPOYnWjFJv}vrYq~FS5KPhpZDmhhA5qcvDt-BLOH}{C56vq|bnk;?&Ev*xJid!h9 zYMUZhnNG~t$X-fs)TXfByklGaaDeI!sK{%O zp>=rWK813X<(VCY&PTtwV1sZh`U;c;TQ%N(Cl z+SKef7|ZNZ{4vt`*j=3 zn>=@p1&+106pP!IJ}cI&=uOSOT69U^NMi1m$?$lNp;^zEUs=nlOV?aH9Y1M7EhO`> zC491flxgi7&wST&Ho{Qn#~IPBpci)4zH{=r!=jW^pa?k#)%C`05CzzSO+{mGih9c( z%PP%No;&V4{1JNmu;290B@sirVKW>R-8c)lMA*c^lnugD6IohTTt1x}KM~+hn&}I< zN%qRNlS`l1So7D4=|RlR5HR5nXto8k{2FYq7qxAqf;mvznEWYbN>)=P1>yb)ck1rZ z@$BM!W*~e(I-_ZkELfcFfiO<-x%#5cKv@p4xIAI#**tJDmvb_ z?HJuO_EM52ns;IR)(=$89t(|$XOVppdLjkEx+|%ZJcONWo4F6i-SIlP1t6oTB_*!S2YfCdAMlVsOjn2?4EJ>4i8r>sn`Tj;9LKnWr>-N979(97fPk`mSig`SF6$8JPJGy>KKYigF zBz|yvk4UOXz_GFWN<`2%bGaiu&|)L2{Urw%#N{u$NjLyk)9#+MR8RF&ZLLhtdNh7F ze>rf)wcJ!r`dH%VKf}ImKSE(sN`9mJ&>`H7->V55C960QtnK>Dxx&viV@V$S7l-_9 zuN*N|nv37|+3GO#&)@4^M!n_FFNgVY&53d=awXkuoSg#*$lZ2ARZk_uXFg(Dliu1h zp^GPHLi@vG19ya&$d&+dxG}9O@5Vmj5Q0>iZMPgmBb=$V_Y3Eze7%u$?@vHzSW&F3LLT@p!jic!QUHID*=cW8h zO4Y-Dml}9$6rbtUJ)kzWF|oWRx2olzT9Bz>;2=daC794ayE|=-3Lv_atgrFZ{c zu58TAjM4SIiM*#uzuJv})&94&Lmh8FL-QakR|QQi-L`23=P%5dUMf0=3W4x?+U-d0 zIlQp91^a6}ih^fKK|dY7w3@D4NN?bo2g`Zc;iz?EJktgVZYlbo?f3%GZSbRBj^E;@ zmi|dfGx5eV9?-11?wZfYeR-~SrHUs``k7Lf?RY0Dc3Yd2%-;zZH5aNr*E|Sn)Ve*l zJ`|LM&v-EZpL;|9?dJcr*EZLF4aZl-yT%+Vwe|mSrV z$UOV`QWn9Rv0ag)%9}PB^6lwd#_b16fv&#+Nwd`SF9I*hOx6Ic5ftlL=F;h#lUYaCipP+eSWE&8YDfYy#WQ%}{4<{F*+;Eq?ogb}4N;Ex3JWbehv`0bqk zec7qM2xj8HFLm_{ZTxKgPMs#)KfgV+OktL|P8R>x`3amNxhgcO1{B5}X1rX*zxZjF z{nH{NZd#k~_9A6EiJ!T&fa1oQYx8^OIWmV)$54uI%8j%P$(NFEi}gxu{4uk-mUjOA zbIZQzXsOGld#fuym~lb4U=V5spl*U*|3OAE96M_?IMGYhtDSwQac;OGf9zEiJ@;os zeNan@+sza)wY_JA>MzcQ5PDb^cC3cugHv?qRv*d>6OX4=WwTEnKghf=Y89K!7`Tgi z_;X|ZtNY4QHv+mgYUems#?&u#UX+iXE{~3WVWr#s@k6Q95y-sY(CYFj8_qP#cEV2! zwg4}!ny~{ScrRGelrGaXI4MeX^lnon`soesXj5^ep&Hqh&qXLd5A%Eu+a!ETj)o8$ zKdWLF-u1C@-CCVGiHr@Znap< znP1+Nm@@QUN9J@3u#lcd(jQJ3{GD!1f3vhs#X&sh%#XJi? z*)e5WVnmj<3&G^4VcdS7`m29_I}Jj2a?(FIYi9+hU&5YHTv{LlqhFi6;s7oLJb*|8 zU$Ua{jX54Tp)~yBXpI~biAtzlB8w%W&`ou@-~H0gmrLOnKYi{F>Jxry&AVz_UG4w} zoo{FxIlOZFwx8I*$?mRFDeISuA3mJ2K55x~m8Ov#jE+|l9cG@Pv|_Brp&CV3@gTR7 zg{=cYzG33O=H8*sZS!V2 z?OiET5%5$NP3>~m8+vY`6U8-s4vKSC z8*9;~Sl#`pU?hBgtpm5`sPo4|$I$X%Zr9zRZdIYAc;BrAKmeq>bs@nVkn4j{pjz*m z&nl$`NMkLmMs+U?s+=$x4bo~dPOj3u=3J@4mv3=>Kj%crtE|X>I3?3({EtCZpoo=B zrd`Fu_t%z(E5gqL+o{SH(_Y3%qA7#Rpt~q}1#Q${Ths8ng{-33xMMrrd`ZtU@DqY} z^x0)Igp}QfzX+PGD7+ufG~#4jK}ixRqTn8L^4%5btQ=M;>p8pOT=<}VU(V_uDAdzn z@u&ktvsi-<3^GrKG9d496BVdK{4at}hN1nWm^ROV7`=?*+#imU&S5hFsUqK4rT){! z3OWt*2+qa+{jVYfG=D;{CNEP8%~pK0a{&2-OapT{Isy}Kj$5)nvPaoqLC$I>o;lw? zZiN$@t6yx&R)n(o1;IzUOCvI+(`Olt72ezOuzt6R4Ai!BxSY~E(v!ch<`RD-#bf;y zdCwzf&+DD$OP@HDkuI0bq`*|FlfYV9ft`BgDg(sRoju@~)dh)zT=w=%82zhPlZwm3 zY!@H;x~+7&+p=~1>R@HvTP6N>V+>#hvjuX%D5IOOAQi0u)(O_ZC(2Zv-0^~LvL_+- z=#KdmR^xa0po$`KscvNww;BZYc3kS;z5UPrrm`o9^%ekX{VGT~gYBheJt=3RiuyEC zJrU-4zpv)7!<^b9S`y!@J*u@hWOLVdJz&m`sf2d(pgNUeV)0!5y29#u-EH}VCF+=W zkA;?>&|%4^zLIEwTte>z@bHd?tCSnx01| z4Kpe5 zhhk$n_oeM$tU3rH0j3T#TY#?t8Cb^(ZzoQZ_NiP=U*^8zr;@#1wg0?SDC9El>OB1M zlh`Wv<+kg03OUy%|4TducdxVI_VLH zrEWoN@okW=s#^1I*qA}upf}!*DGjj&w0G2evOSW!1{JVVB>npz&?lcd5@I&q%P)vU z-#X>jh)59~Yt`t)*Niiz3%X~NDo1PHvk#HIn~GXO?`0hYs_wPAH}y=xn;reV&mu+h zUyA;kiT^Xs<~JaOyV6li$-jtpeX^1&G^PrdqZr8y>SAW##Pm}iy9)tr`$sQBLyu)U zo>uVpsg0w-8w$cqln3!^0F~Iq?Qwx11m9%*90>9gzIW5T-kCA2V@^y&-2a*LX3AGj z+{NLWX{fK^T9TK$goW~^&KElx#?Ru-=Oe_QEMO9V0lBr+VRdDv2a+=m)$Mtit)AuSneRYiWbLSJ>5F*g zGRszusW>?$e)kvD;~?E|1tksMv~YVXomqVnw5u%fr|*3hqv@9@!eOdV06o6)F}c8+ zcPWdcozx}@kl8&GaC18jQW-Y3q2!u=Pnx47btMzSz&f-KN- zW9PIlo_~X^1pCiZQ8A-PYZ}}n_(06v!Nlfv*gtkxNP<$|IIq)P4S=gU0$)Xv24S`0 z_Es&G^+XV=ZQV86AVh~Vv>JAv6uzzdG+|F2JB*5j_cXqjt&<; zo@6Z^AbFL)5|Mz$%>HqU-;7+z6a7f$@{~}=AZtzh`qWU<+wUvU}4Wq95 zwi?+rj4C%>4mfzbq29xOw64(}qYSZr*^PYW-6)_HQ)6Q3^^~kw<_${DuT5=HY{8|d zYr_>O*U_mS{Odg*B*g?xTdg)s_mFh_YK2eEh6&(--5*qVeUr0M8d(wD#Jn+x9QlEvx)G z{mn_?A3x9iXYU%Dro4aqbsuNNEZ;Ys6pnieajTrN)9wP@{8wLy-LV*v4}o;q`tZ=s zEr`@gt~*>_j_thCNH2zB!ox}jN){!j2ln;GGTm3F3R4vjlQVuB0!y`btIS^_zYNL& zg#+;eU%38`42+)BO@18iS>}L^a!&6I-98&wC}Vh%d5=o@+^4SkTu*xM+)WWmN_pco z&JExHWI6>di-yr3cH8b2YyJ4K>yyL@6|+ODAFu*%`PDXdL2b8L>BQN07!JsfJgM4+ z6k^N6b!ai)8RRf<28q2BN`fa(AEBu29#ogVD(9{hjfmyGfanT8puto+N{hed5%3rz zE@yu|D#YCsxit+YYCT|#enH(zh7YJP{a|GNqLMsqHJvJ1<}J_7j7tzWZioDzBIsHILv`LTpJ^Pb5+vWu<4IKyAyt&3L#D*-SXdOCn#v z35>iEW61z&xtsIIX6_YXrH9u#G9`R?Y$mJ!$*t>_8kDD8%V?kVIL8a*mfMpB8j*Yf zbBz)Z9XvRxL#e2pg+v#!OQ*f{W!m^#Kifh-e|7Y~S9`&+_p9Zk!JO{KT6^+WG~*xe zFba?+6dMU&?M`^j)cuST?k%2#^?tk3R#x#!t%W8zd#pbS}?*&G4h7fQkJm38$z&n(CB5xY`s_~j;>T;kniYYzL7Muq8GszCe z4F{3*@j|y-#F4R^KgM^jDu&eZyhsSkU+7fMj^O4Ed0p7Tc&Y^Ppqe8bgnYy(N)G60?{@&=s4=6BgKfC(hr z0lHEY)pt+~o4VOhalA0YEnderK%_@vzfm>6abTq8?*X=`E@08lGe9gy5C-_+ zizzY3pfrGcC9u2OxS-juj$u!PL9;g{wWfxThk7mmi=VV;&#V{pX8*8mAV86zI@(_bXwpjq9nkifeF!)39n6S7 zu0auCcDEODoBLyUGI*lhU&OP2V_y?!nU_jrX$~(fCByscnQxP6NY)A`U>P?OfJ2>q z`xmhuhk`W*SrJTdP0{^_2;3ux^claN-kqro&1%Djl5FcnngzRO1*A9^9Wx&8fUg4* zjb|?GK{7)bDEqdhBH|(Uo5jpY^_;-4ut4tI+%msI!ylECPFyy3c%CeP_@3TgCK`>D zZ&9GncaW#@c<2E!;oh&FJBpWhITp$%68<5E8Gm)iy0fa2>Iysr_Z?QcvBS;ax}z%CaB6X? zJlkBrz%w3I6rTPU!Se_=@exF&Y!kp1a1CgreJ9hNd5NC`c$WTzx<~x#apbgROY!!! zE&?|)jw%ngrx9W{V}@&ylAu9oM6ATuDc4-~wM)BnzYgbn>4{tL;13-TXk4BC>s zKD{sZb5Z#su_BxNETv=UJiaYufQ0zWyMRi!fZU-9JGKxDh=4bCj@|F~7s1Z1$zCMR zE{py#(}cdx6#)mNFAWeM5}vr8sJ#c8up<6N4BMukLY5u}$uG_eB7H@!$J_ozRQ6u@ zi}>uXNvr->(eGnp+5FFUu4AuK%~~8E=gGBcpy}N4heyYP>+e(_^yJ9XyR0H+4yi?( zo;W1Bif1Z5`it0K6tJ@=n8@Oxw(+sjP_tRSK!7BHqYK+N(xQJ|9nPoCcTavQrTmr< zHM*KGP_HkK3rzp{HpaAf&CN7Lu8M=j8-5ag^onUmQ%|NsB8#?7JyBG_{k6h|+Y?)z zph0~HHJ4^=>wyOvtHrx0et8mS(7AO})yLMyeF#A&i?J8+ID}9a}zdMoj=8%M_ z{<)zGQYrY=);?hYK-yI|7JD6mIzXIl1xOAEVP*^Qb!LBPE6_+iMv%Sd$#J3x-};Rv zg;9T9&R%AHit-icNA~F>vUQL(YIYF$=g8m_3zPc;zK+h#v)+t_UUaU9Wh%Ex-KC-+| zczSazy?|_7uMnr>gw;6}!F>DXg`9|+>(QTRNN=vzQj>+i^+PoCwFx2vMSGZ;Lm@`v zq(^lH?MFSvby+5gcDBuTcoyWto;~|!e*Ub)k^B3>z&v2QqV(Iu?CqHnaNpdV$ja4g z<66chRdu3S!YxW+Hmmx&Uow~PX%yjxta*Ep^}I~6Vnzt7Hk`A>p>I}H1Xk$2au%4S z>vUvH&99oJQf~%D>Dprz=aTmR^P-%HHzZ^n>RtWHUb-Y(cniEqkxw26ePjzsi_5h< z4(*;$ZKxtvr9=h1Z7+{-5Kp^yQ(x+s`5pVY?P$XM99wh`e5M!UT905`glfA7yWDj%1Aw!>4;DsgzIi z+2>SujCA%0RUX}wH*cI*{O2v`)c_Jy4RR*@ob*ong7Z4l3pYGNCQNK=;8ya-AWh@KieC>5{#4$w}wM0C+}`XPAyCh#w{=V&ue&b zsAQVm)402V>^f-{nr&X`O%*;ov)!>cv)!|U=`lG6OQFKMd?$37tWKa| zkD1r$0UbF6Nnbp)n|X{VPM|#a6YNm-vnk%UnEOMJkJ7Es>K8h3FWl2J_u$i@8-uW> zOfL|6VM5nh^nK4Mr3ddTHjEzb{?i5Krcpns_?BE1X}c+Ytj6toDr5Jjp%l?Y1u6hW zuL5S=gKSymL%M1^jyG+dM~coHr**t1%c!!(N@d>2sTkTq=T{Vy*t$dD^n<(O_B3|+ zYu{|E#kKmj;nlRo$&gJ;v5Xtwqy)}z5Rpe+Z99P*Y2$?A$Wm?E-rJ&G=-nEOqMn>L z6X;TxKWU&~bGT)6>agesad*LpV-K)+rUb;ugMk(0@&OX`rkm#3y`|#>Lk@IfrQP&< zaZq*Gn7PA*O8v-aCg4-SIM%9mP@iVUU zYGNdo_U(#IJw$M8@FLAhh$r!m;pfO;L^WZXhX0E&tA&*DG&DLHFK?m*_hR5sn!43# za?f9cyw)tOW3c$WTU(O0{)Zb}&v)&HjBwsnLV-{C4&p)(>`L4L0C@76@c0rZaC)FT zJxH2-KDe#kNKH^Q!67uhXWRI5Y6no2YRTB3uE~I>!A=h^oEte5Vr94xdM3ova8Rgo z)W!0ONeq!IH~`iurCP)pbA?z1&3Q)qg?k6bgwQE@~QhF>g?E^?TYRMvwd5npF6Rw5`eP-=4V|Hy?TSfQtSlfi%%d0zfEoAe!0Q zL@`5-Ra6^*Wa`pw^QnW7l(Q3BB**t*1>j>x1?=t0JAFESB}3)pde+`cI+-`gcd=XG z>;qqR3#b4N`C&fsGUZW!5eLqW5tmIzppA;22viXuwnqr`0s2hb1u*#Yj&DGieu$|6O55Cg z@IxSC6RfHht2x3GUMCk(TbjpMG%v&DpQfCZV;K8OjaLSuCjvAKmL4VfG6dNwMAkDR zk*+ET7uaU+wzxov< zGVsSeLJyF7RU%OKnm{{YZQ}AYU6271*+^e3z8w;(F(N<`*c`jjxm`6?2y~7oyYbIkEbuYAJ{7BC%szs)V zBW?HHk{!2k5vc6siDZTlxH1>l<0{ZbTLVD7Alyl6I$RV!nK4?PS^g>`n^JJ@!{e3@ zK~U9k*-t!%cZ1e@S5~N~NzXaE70W^}DwSTcl&1pc=U?DSy-_&hJZYs{meqDG^>jJM zdcEb((_y!n121k8EAWcgaVyX@mFU*xI689A(BzSuNf8vZw_}PV}MR=t6oL8aPMFWuiSgYPI?EP z4>enn$3v^jd9`p9-tIj=v~shusxH~)bDdMC1XrTWF;9WFh;19Nv!|iicf{XQR3uwC zSetF&62LU~tAK{u7ErJ~04(1JLtVEf=UED?AY_exjAzE4+H)HNdMIWI&}Zj?V^BwU zI=aSEW$`|rx86qI?_+`KPd5l7Yp{x5){dn&vI-xk$|}im42Tn^4rC60Dh~8ECQsnq zs>H8P9kpu1HqPcw4BocO1m8#&HQSt>%#==;(@tp745NJ+Xs5#%GZzQ{8{-Pr`umE{ z)kz<*s9*CmIB9723sj`BSv`1oKXQ48me4&(jKnF2JHnnXeQhht4D+6p``&x^QQk`8 zed(cNdiPqcXx@ksXjt~`AyNQuU=~l@&;zBxt+8ab^gOv9%-Z#)r89O|Z(D5WzFh!s zN-S8)JG7%Q>36)bEb|MpdV=JY!^H`rrr)Kv~}$$ zVaf)>tDt%)Q1+(y^R`bqHSccw&lwQjpsue?Supmo@Tl1gvT%rYTqv1)6D{LqB3*A- zm#&n1;|+w3r2cZ3ZApA>`%fD0V1!PlhxIvTG?3m%0NAInMVzf(hVuZ9H}OgKj}{<@ zGsXLIUq_2^#fa~(b68bSmD_cKc1E*-+B(4IPXNAlGTfC_?>QW9HTS-vO`55TzfJZb zM#;C8ihqBYIT2%H4-GifRq!3G64}q$2l+zKWD*M2B(CHZEKHi@(1^48ov??m-$Bi~ z>Y=jZ&Hh65q2mH?=0A~4zc%*lGA~^9WPx-#GLU(b^#q#DhN^RuBf{nAzbT3zK6R`) zj(O6x(x!bwqfdDkb~4`{tySfEY|X0?q#qA;DWN!^Uq#_A0NqM70_naNKv5xM+%8~f z&%@2YkjMkeqVkz{QIr&rjN@k&w1B{wUnU~l<|RB}x^rd{bXR$8Rn9nty@~8=Q{*!(*b2v)~0 zuK@1Kc!yz6YlSxw^j*rAG>*xuH~v=VlM9kiSyA9Iab6E245DEzUgn-w-yOpxB{hK7 ztMhAzL$CV|KQF`PZGSEyZ>OEbc{%jHpI#T>@_ukcE&BX|OsyNhl9-57v1iLM&! z1GwoCN`PreFL0<8F8bAbiOiFd#OE|HUHn+KzdL1p5A<+kAQ`A=q%OM`faJv4e5mqL z0+?WnsS4;^<6B(+z__`1QrEV(YZvO-jFY)@-yV;uB1~fcltV1~13m((#7DFF5=1Aj zG8V^l*OmcoWl+76vc`F#X}%jNepotV=E1YOLVAtCeHowon$zO#KCqt=J)t<92wX|i z?0`&E;Rkwn{Yc%12WK9NeeF`IPa1hC-cuRD^#JkpDkA?%;V!m|BB|Zp10V2?Z&r=02J{;!~2Rm0ApM< zqBc1R|1K#%OsV*#rQn}QYRgsZnso%_|BT`P8f?90wA$9k{75K~n~Rzul+V{Xni@*HCSQRP6uzpDO&f50R+S44|DwHY0mNW}OhVS8I;Jm>Ah8zRBDl2E zu|fZ@1tGN%26K?;VU~Qa9{D&SxlJ=jO(K{+_6@4#Y9#KMhuH` zI_rWy7qsPm;6-BjhtMNwoQN};k-UKg5a}SB!yfk$2hJ3F2+_!<<4n=nyI&UWIhJ9( zHShfS?eEeseLVV~>b#<)BS(0tVCujbHlGz8eEX~P?A|%v<>4|c!IjpUyM6NMUX$@Jb)WkoIeUhzr>KZ~6=n%2(80!U7m99QM%ey>B180oQ z@rCf^(|L&b3mp%Nnr3#NJUDpc=CB-}M26g_uQ8S>qHSPO8;a&{*NJ&Vt~P7Sr!SEV zd6RHz*phuKKk|x@lUhuJ-J)6EvY&(^^gQku1VJ?tC33{Qs2iAx0!Aw<$3)MJ&Iz=3 zd1bELqT5z)Jr{L_F8WAYJro=qMAJmL;gL z)5D|iVSbHQLb_k6R$r%cr%Ijs-c$UEB?JXk+_!KbN0HarACW>#o0r(Nc8H(JS+qFj zq${lIFsx;FmlvIBUwFk-D3P>BY!rU5V-At_fRj-0^?AgA@^Ba0hWyx@ms2fz`tta= z;C;_}mR<6nPWI!hJNIG)iavewH}Kgeu!V24>P-%r)Kilu2)(Y7bTcqqf8D@!&@OMH zm#K1BEgEyN*^^?E6Buo>{w(#>e%?%uoipi`_CXMDJBk*lu>$5qbYb4HON*z%oY$Sb{Ou z*p8nwCsuENd{bXt^~nA$#{u34mWQI_WcR9W=x%EQY2pQ2A;~<=0w1~B1WAVyjUltZ z@^JpvVb#W_l=FLYEFwR@^poG`EI!}2H{@ARo#oeDdLmg3A{FDSp-$qY&YVS$0w&IN z7Y&oF75q^Wo%snV?v*)NcrR(VU#`Gms6zt5Rl$$2=KX@9@7`*Mw)dJGVS;#-eUNyT zEsmSSa|4bAdZM=-r$yJ@dej2xdQw+ZgoDtZTP~N_7j?fl#&gZYIq^~ZjYjhm%JR$+ z5Gj{JZh(A4tR^no0FtW>4AJT-`_L1JFoC}TFoc@SA0W6veYBtg^ ziz$&Xi-|S%3xnu}13h#&%GE{KZ|~=9(o6Lj@0O0AVgFUr0CWpCG_31s&ms^#@Vl6o zV&`mgUUtqsRM)aOzC~;Pru#CS8DkpUT$$8t@g)L>}pk9pJy_5M5{DE6{0OxgV>5$MJ2ZwT{83AC$j)5f6{Kf*nVr|iqDf0Y;KTop>%d7@+4Y*hJF2u!mEX&C7vl6HA*7ke66gNzP8 zIy;T(8mq_fuIlEBI^@{Y#?NRbWM*^S;F(Z-e!2Yn6_=|ou8V&}0Kh$x^cO;qc#Kbi zr}navjWy^fOZZ7SwrZ&>Po;WDsblS8(&UNV4>X(8l|$~SMctN>{eJXh=>jOtA(LzT z1lhugtn6Tg*7IYEzOL_R1796vSesuui<{3c7UY%){76PrJ^k?v8IPJR!O)zcI27$F zd1Bs_rdI^ysFN*815og0v&(s-QuD^9Cg1lf{@3I68*Wx2hCb@t7LVD-!Nsl20{9)UPq8>c|oa2QliY}OQULtqXOItBzOw~clEUNQAx-vFq1(nmhRyrE$;WFOR5+;P^&+E3(Xe;0G8upLIiw~iia&1ZocnSQX9{|CZ8?f|l6Wg%;sS-x*?3r` z8Q;#99j-B0RrxgV#BZe_-crVxlIsw$k~hNftHxHe>o4L7AZTcT?`cc7OVCY3b+Y$- zE}vWk$1b)$Gvh9EE7d{XMJQS-3#D=2fvQj|6SOSGDlM<5oT*N)%b9}|gq3C;$zkGF z0ViL?-XkBAOWkKZWnU^SK(V}X>Nv;Jwg1Cevk3B2T&YR{LKTJABSpEy zCyD>KobLF|<}(21bvruiQx)w}4pscnZJhl)ZLBE6q?s94=&|zLPiKABZ5v1EMZreZ z40VhWdvHb+UMPen35=`=fAi>yChd>vxgpByib;lMcTaOHVZAti0gi$bqbHeh8*Hcx z9ChK#nj3o(u0M!%>@F~>pS<|`La*%Xn>xE>gC0-UTCf?Ux5P>eQ)hfxIEb!rt!6#k zj2^$UV=6Z{rRM0;2A4~|p?v*zk*~+reJ%n$_hYNq9@2&4{ zMv$?^#V;c#EEO%U^mi(|9kg&eg?u++mjm?Eow*m_LZvh^$ToaIPHxwty8}j<(DRyo zi$LK`JY2BUCzSAGYSLkeQ9WoBG8tF3+l@3UfPgjGB0%__T}fAg6d>Kv|EJ#4kD47O zB+qcczy@(G(qH4L$?~v8=l#cI@wE!ox>4PP2L^enuf2Wts3n2<2scuWW~f#ImWz+( zN6=DTg%bn$mJJN>s%YIT(4kv8BhM;f%fW*#wO=-%xHj(5TBGsxm6b!AI}Y=QbDLJq zItU(}rEw#|r2%Gj9qK>e#O_tp1_u~66tIN>Humfaft6%f%#I4MNI#K{gb?AtOY%yx z+nYdt8J;ou8Re}Vjq>`uAu8i*`IdiqLq~82vv57D5nk_AgLBAYrluv;<1nplnuDi}hs6&HY8F%h5I2y$jIB zw8tq$6mxecMSEp!jyPF=riGQ|-(tnv`HQLPC%nEeqxxmzSJLal6(0(FLyYh4UU!`O z$>}YN_%3DwwmK#ds8Sd7*`gN@J7b~0h_t0`zE=?_KK3r$Af9IHhng#fB3F%Xlwc_> z8X~Q3>UC=L*U>5p;(>Av-`)hf#AXopn!0`7<-3Z64GJjJgUP`5rXa=CYpT=M14PK) z*^PXUviE;k-~T;*ggePjE$Q@po$p%HPu%fv^&$4j@V#`4NMo}AM!8`3q|K*f2zf6C zo2^e;Za=EbP*xh@_#+Kg1=3)<;kCaA&%cPLt~|dpLaa2q^qLmpp6YZ4-B$i4dPn|q zR8mgtRy6M8_&SmHfQs%)hV=#60_AV*mu&Urdvc01-;HXi1__xQ98&iia_&HLux&y# zjNB9FG}gv9oOVDsNIH+>WlyvT*Nz!;(iJ)#Q**_OhAXrspz$l!w-C#jlc&{e<$CU1 zOSyCIXT-@*L_8Tbj-`H#Y2G=*l%rkL=dn7|X~@kcgovq?9e5Zh7076JW5}p^o|rs` zh=d#f$qv>WY{!4qrmW0VQi*%KwqL)itgNK1CWK0_mX^FXROu?;sU39KcQ7h1C<5Yu zfDekonShue3gj+jlf_Srt%SSrvD(^K?divdYSNs%42>^JD4HBN(UlaxH#z^-?|qZ^ zS`pAG49x-)Yljga&SoKZ)%Z?1lo2F5cH5^>j>Xn^yi%s0sr`X)0H?N&DwAG`vEWzs!)HBG5C(x$)eE089;(__@6Zm zQ04Ycit;RjNZ0%;4wvWWJ=7!_=8onx{V4yT<9OQTA%BvCj7gHkeRf(4xSxt%?O%ke zBq(}1;-!5~A-|MEN5Oi6H>+Szw!qTE67L;wcB!#$p=%4R(bD+4omq<5UD-o{5#~p0 zLJ12U#OEm3aE*PX6~&4-#Qe?!Fn@p&fKC(3q4CFnD4|vQ2Uz7O3$Szhv&UI+hRN)V zuKZ$}J>toAMun&E z@_?u8zl;Z8G_2ptI1Y&39>lmAqk!0^jA|^Z(=n!UPck*p>X{{XO`0+Z|8rl)DaAg{ zRvg}LE`X_M1@_i*h$%|Di=L?OFidlSOS>U+LBj6V8vHunI}3-Kb@#4zRX3?heL8zS zROoyo0tN|QC0Lh7I566giQ$%T2~j#&a8S%EPjvW@aGAA_{%#YS?hn;B5>Ed@_xaCp z&AsWy)k0lh|0;^4b}8U7k#>n3PZr!DK4vR5(O6keYI3H^ zP1Q$M1D<}dIPUXUF1LU7=@ z>VETl)3}d`EZTFM=~%=S?pXif_}{i+8qn zR|?6DPmC-xBusj^MEpfan4MU|q(i%jv(GTnkbA&aBZ|z5Wx%{vUGu%{j1iS-N3>As z#l)g-?`5Uk@f=m1K3PxB?W4OkgO})MrL_2gkG=&~_6(aq4ehq{uV_(ty~tmNt@^=n z&chjBwXpj;YvkagBap}y*{}0~SD;tK$X)tR;#jBUhrL$_>qdlkf=7k3YTH0nfOMI7tJ%c|$qP$b!O;|Wu9ki)S}DU{ z{>z-KjQHLwJ(IV0o6qy!m2luk{acs$1nAJHSfUV92GUqwp!J@mHPY*9)zH+X71YHD z%Yw-}3+@#O84dp|-{Q=OyZ*FafbGD%1vkHBccJPv;x-2-^r^84y)n1*oE{=w*C+cI zUSs-|yG+^C@gT;B@U(!($m;qh#M?W`dvHT0sHE|p`!4G~CJHrlr2MT? z+^PcsIvn7Bv)u3;X>{MDAam!RrwGw;pk;P!{zn@BlI|HOj2K*cX*W@ltsb|j<)Vg} zS+Kn>>3j3;F*%#3X(C_FU;A$Unt1xHrsl5@%@eCyKjW!09-@0H-?^uvl_O6mZ^vgw zpo9RxLZk(YcCT?Uv2;ua>M>DpCHw$YR0LX*k>{%4?$;}SD?YWUNs2CoHI^#gFP=MbEME*&y9c1gAQgYY{ zhe&*FD0Vxl$}_>_FgX8K@%H9NjSk4|_i`}n%w?x1?BOu(S1a1llwwg)@qbG-d#G1sVS8F1J&)P_A+ZDe}$oSfR+xAFNkKwxF_>+bllKaha| zOJD}j1LpLZ`|zT{=t{_E0Yl66X~IvoOif`ytZk8Z6im&@JBROb9?qPd2Ib$Zzs zb*Of%hZXZ9Tp|IJN~WFChBcMm0KKBZLvL^&@QT-4X!9|LL*Dq6AJ$cI3*hyMuIuA! zr_~Lg_3!FS2xgqD@rXGH3DA?&o@`osk7<8@A_`}#NU!{d} zYL8kXq6}_bPD8zb-~-7kP9eZV3_z@=;ATP;#9b?R(t8C}{qn+DefhyS)zfN`) zUDsEHbSf_?*;+qBm^=Yb*?~ybo8Z(hK%EeO$^J;(MQNs6L~l&rIw92L1_HRpM7tkS z=BATl?^drXCXHSH;(1ebd^wpZ5z6LYDg^tei*iGWFimRL54W5-!XAO>dIhpz;&DvT zK|;1^`-i}YN{bYc_s{VERH8nCYyib*X#m*DW)DpnPO?(XMo`2x9#Mo)^apu*_c1El*}gszVkR)$ue8l{|WW$$WFvokX8 zjCgI`cH@*(r8W2cXen#npJb-YJX_F!x(4ryABWdf48tm{r<50z0y^8rHsw_=r(n5S|@n1bx~PtjztvI>8b{1usPYa4*R|KmbsMa zE2;f#&um&wgvl(J?FS%F>kH9f1FC^Xbi{NvhpD#EekN8-z}lKSB!;ljhlY{Z{H{o5 z()7uXxE~P7is*GxlYs_TJPfy^|1aL&1RBcs{~I46`<`7!)*_TF6&WqIBvJO5%1%s@ zEz{g0`&J}Hs1Om8y=;>`TXsU3F_ehR$Z}_v&;RoMKELxn|L6BS=l?tBIlsp_<2ZB7 zJ@6>=2`AZDJ34lBW4GbbG z!X~?svR&2XzSHmOD~EZ{Z2h)KQ`16w1*``?=bf)$4Z2 z&4C-;R`1_YhM!E#uA9Bg3Ms!`{$XF_ZQvF@RNF2#rqpH4HKv^L*SCE4F30YJUe=(0 zz$2Z7@3y>EqO1{_=UfId9#}tTK@V=zWk6hXXlGZ3ul?hR0NyD1N6P|OhGd|uN(dq- z(z96*9*0x9?S@x2HN|F^eIk5HmU_bXe-&H08T99Rmfr=A-gZMT>@AiZN8{^stx3Fr zEoG*C=k-ge&UL%?EmRQ@ECOCe){Ha}}b;%coB!&u=FW#A;9fLCgSkst1ZnOzz8+ zc2i1P2OfOPVo6H2q@nZO2vdByuIexOki2wG-_43C=j7l{+&a`M-y~Fy1uk}!CtYL+ zYgzy;=OHJx^Lde7|OftKF5&(k{ z8boBRqR{t!pc1=cUBfhKNu{EqwIOLIQ%t$f$pAct=WoLP)zYi-x}0BGqgYxM3^)DE zN7S|XbU}BrbA+0H@YFs*6NOZz4O{3ELbfgPu>I#9Fly0j%YTTe^b284FG^@U7~JFjzB4$4_5~I`%{V2%yX)9Yi1uSfy5)1^Ic%|3SRJ%jhr7ejV*r5f1rR)ikYFR}aRxPd)dN zl<{I+f59I_5}-d!V*`20y+FU4&sQ|$NF6pS1wTg+5{SzstgR5~HK9@q+{jtNjWj8& zGV?2j|7uvGez?DmfZAs>4fWp&@00p3L=9JWh`zhc3JIUs)8Ys29RQ9*b%?N%C_YUG zBN5hUPE|>0x~X2yCUvs+@Xhk_X}ik}wkF{+J%gY2uzyrf_uDv`h&cG_-!HW!Mgfl9 zn^3d%5iNl=pzgzYCe5oW&zm>MdAxFTQ&mK74Lw|r(28~`+Q;%Y8VEe7+DC^0GYUm{ z;6|K^JMjD98B=Q48^$B$X9cESubbbE&f1|*-f+n|?1nzBZG=1wwO&V& z@*0W(%xDWZ8Zll4aE2A7d}?A<=kmSqo7+M$56rmzZTK%zKS$Vq_oAOJpC8;&X95qT zl@F|$44znekd-3Nl$)R~GQ@x+a02;>o&N`sY7ewot09UsnI?hjgz7@U!At}uTg*TQ z?+72W)~ovB4EpkwJ{>iPu$}tEu{fwui|ZZ!Kf5L9e{oO$H*X2rT_$U=At3mTy1SeY zFyT`iOwqhbGRRcgyE!+G8*k=T&ve;cm*A>RI3X=?Vhw$H*ob`}*ou=-s~zfwI$(KF zq6xtTbB0-X$1uDFc0*`?il_@*ujdAO`n@rQ6wJ{ezC@?uZo!GFC3I2p}gr=BB0 zhNBS8Gfx@p7~reO?o7`1#Jm66TU@B`78rG-Jo1b5qiBU|9_EO&PnZ;l;t8@x>N>e+ ztm04-O?o9hE1hrs>E^|jURudAis>IAC7aV9zg8xvhNNcF<^VE%A8wFH@nYYmE86{` zUx7P)80pc+rot{;VQMaf%>`(fjIH&=v}DpmXbQ~kUG&rS!RR_ycGz!j1}*S2q`3bv zelQTyMffZ}rjPV~k@Q?afS--!g9v;tVJeHsN2Hp5w1U+H=-R=99pwT)?o z)`*_#+pjfgSr47`n8-Q&$IRc|OXco6){tAr$povPFSv2?qteb4#%6X^*W{D+GZLx$ zsvhM}5RQNy0E>#r0)8JU3#!fB?nrSXn?#fQoCo>8c}MjQUU)KX?BiMg(EP*mWC6uw zFdutUJ`te?#?1gg`@bK7$?ZkE)bKg+`o-`>dija_~9*N@2~L>-(DU7^whVJNL5pQT9+ zL&k{BL}~oM5UpNL-3(G+(f#rd!Yti@R~hFr)B9=POL{C=K4BoeB3N@;7}mO;@v6m8 zqo^cI**QR^;DpqsvM*np`u z@qSw%%LDRKMli^s+z5FDnu0G)7J`$A15pcLd?BR~YT{P^0{`8Z689z8bD~8YPXFe= z@$Fn!-*u9-chtSNR$?oX`w3xh3$+(Yu`TF6TDUO+!X1QKEl|Ez3HcPGG)}AzK$&?G zSz*YMvYl8{*HBHpu!Sl8b#k_(qVik59CPKX#|z9eC6}tM($akg|EA4>84pQYB(OO3 z@Q=REd}$qmV3^z4_wrZ9=z({b!A&;}ulp3e;S>6MW#!NwndrP)#(}FpCqf%i$`>Z; zMqMF?II%xX=SQ5^5Qbd}YVQ&JC+D-PjgY%bt&bwzC$a!R15=lEt43Z^*HQ`}bZW*M zn92mzu-xab8)9E4*|8-@jtH7*obC+6fLTUNJ`eaT&GIWi6Tw#DAmxh7!SZbmqQzuYuD%+s|2_d z^`G8hECc5#8^9v-S8{<^2*O@88Myr_yfj~qc7*NSwlr58Pe61>_KhG+OT7iQtyn$j z$^G{~O(N{Sg2nJkqd)cPz|)XLpac<|iK`(xoGx{$5kJx%$our0Ov!#{4h1>t80&>m zZNF1u%Bi0xrbl%3pIuLIW7~Pij_|-_ez^X+dOZti ziGoL2KdyLjLxrhECw>5W7BQ$59Oi#*!)3h4Y#8R_mZ8@0YX1n(3{2J`&yn&KkE?O! zc!!!Hp>4DDIU_Lk@gIc6=KkErC$X}&O-EZLl%h=T+bgDDitbYKp6 z46cBT#+dg&OLs&x|I62v%R^U{71QdB^qrAM6cXv73_N(sA<(!Ldg=W*=n0^kR8Oc| zH+{N;-WTl?!LP5hpz8Lhdh|{pOIj14XB-D2s%3{1MUje#6pMinqGdBmeA3eS(Z0ZC zfRv?Vg4t$gI7dQ6IubjQITy}mJc`ah58;!sq4R6N&#KgYNmjQwpV5Whr;X&S{oa-P z(g7F}-(OsQ@7pnZp7=xMf$ODF-9F2H69-=id?dm;jbKh&NIuOCdkt;|6%m|r<)nc~ zVeQkpaqYkAn_Qbb+diIw`w}bMg^Yg9pnSS<2K8G%7;l*R;Ka5=^?UOvmp|e}Ch009 z4Da#-$jw)5jcbJaXrhai{rq-7;%7Fm-vR5R*i;M1EHiEn_+PPk`ziT}n**U1LwD}! zJ5i5yD&rx49~t|TIp^}OnuPD$`|uUE?EQ0oc0~BviV{s4_S)LZcw(SJX>233R8dyz zH1JMGWg3cRa!0?J?Zx>hD$c9;u3ozJihzVXt)}o~U7#&FoR(d68z~$qpS z@B=TuDL;XRdkTn!0UGDSogJraFRloGfkUHPS6dS80 z^7-7A&F3aMf9VOvTrY7%U%jaCL1;a`Pw||Xs`Mh_yE*InFoANepLJ%CJrd= z*pMx^_AI@V{%;#Zi^eAGQo^==kjcwlhbvzxD;_=Cb4k?dY9h1X;0x{@lN{Nk!s)4w z&%@K_9=;6vBlv5&j?ZNKQ?1Nf-mD+sG%9>MFuB*`us8!^8T$+ z9^huEE4HrB21o1z!I!1kW5FB za(Td@Z*w@V%>P64@TiThU*YsS!^>vB_C0|5)ZE;iPtQD2+>_JtiSe{4bY2sQyMO zD)`J=1s=P$f?S0#b>YKZIU$T!d^sUrCvGmjz%MQ`xktg)y+TqhV5ZTnfW}XI4GFb< znrLz?XfG{;Za{8~oX-u}(-M?P5w!qsbndTI60b1JjsndMY86MhBOl~bECs}|XD9>h zP}|25tV~1sE2Q`{V5YE-Zx{d>)a$qEOPmB26#~zGGi@)k6Gbt>B+VYDx=)IS%!*T2 zw~7LbeDG}9Mg1yPBMRSmwR$8lPxCBTVgkZC!?IS(fTLOr;1@_-udE>qTmzm(rm!Y_ zdLtniy)p=&!S!qrHIa>faI>^y`OTG;OEyqzbw0%m;>ArGN=zCag9jb!$D1^v;^n(p zQNaT2gRf7dT{-LFbNpGwM6f$ch5dq%%wKfGkBH4ML_`Au9G)M^esn?$Ty z157t6Yh937m=|d+b|XzFEVOkx@8xUaS)HEr5we%R*r({!0?y?iRTU9I{M2gIHmcym~JdjaEg=k2M5BMh~x#CqbjFx3_Lp{8$jTZ zAZlKUa;pyoN2)$vzP^+%AJCj(sarnLdh@9)(y?Zg+%mF@rz3H;A00*XMIf~J;H_t6yd zK@MMGqj$?_!8N>rRanTwe#MMW(TeM*>+56UFwf^JtbXz#AQZY*U_}yqkqK3A^(a-u zhT{2w`U=IdEUPkBaYwNRZpHWnJP`<1bB>=onf~hKBR!`>1zzRJOAnbWb;zP&IksUa1m^SX7QgWe{cr$oZ?Iv_0^m zC1tP^pA1PGsNsFLP}6pG7E~MmGj?@a+BfXP>Z?|4ev&X<%g>=z8cFrNf znx=H|H#Vr6tMD;rDk6dNE1)BFg6w`8RBIo7vRF9jun1heOCffbmg+9e%yi4>Zv6x> zD~<6KqmPIXz}XsU2x80VL;@-A9rG*dokk_#Nx7NIFz)c2*D zr7O*f=Z>5f!~>~{z6g`upC?PBnH^cqj_+P#-th?Ee65$F|ylft4Y;Ofd5P{6I@0Q)kS8Q$g;B zpZ<~*znbj(KW8L)Pn?zypI9Ej&SZV>;y9DYUR z8;sY~EC`aY@fWn@G?I?&Wv@>2aY@CQX60G|8_89vF21-8>;qmh8xA{&(#r}X4o{VkNbe4 z*mfw@#U#t|59UasQi{!Kg(F`J1v~WQkgxK#?3W{kQ>-O`=%0|DKt{KGZTc^qwM@jFCVE$w4k%&Z)>IB|1OY{^xUQvP82 z^&1fKaUtVWXzDd)uBKBm{i`u-bNeb+_(B160D_OBy3y4TVTGNXmOzwJq^fkOHKxTM z33F|=ToAkcrNCz}=;67D^FNg5CHwfpUI>UXig#N9>et#ZQu*!xaa-JCa|?IcThqgtCFn;LHI!VFW-v>B4cO#pEW;QpVH1BkC#E;+@?~SI;flqpG_{V& zy|>wj%iCP?-*H>Bcexq+@O{WHIZXZpO8)8_nr)ccTtGf0qr*vbF?lT|IdofU?3k(V zxd2+?;cNArCM^5?Hc$Rr1=E4XitgW(PQc02S!%ozkzG2S1bN4ejt_@O;pa`Num zpEmHbzI2exi0(7o_stjLZefSXiq>wgmZvLX?$8W9(|=2c&q)qm_drSfz8i{Q031z} zkq>1lpJIWYVjdUw)w{3##*8ih+3KSieA&^h4RT8=9$0@;BvS0{zlPN+81lI|HcD!{ zXjtxD)Rqi9K6w38xWZZ@t(TAu!6&4*o{bSOny_kY8YQlcZ8?M>uo&0c z;L@Pi7yS0yfJDRe?OWJU_75cNME5Tvq$AY&0;*t`4H|GKS>c$eF6b9;w zV-QWb1XLWl4Y+oPtR6HSE25tpeVs`^RMuMR=;}ns4E+8!K0uvkbKNWD{Ba4n)xcOI zKkEGaE)NtxL|SvvRbmbdZJa{d4UrOVKbfezR2;J^BzWhFG2);M;)vP~FCK7H9755h zL7ho}hf!ZX9#6jNoGFB$Uw(@_qd=q03y@=c5($vpdMus4B z67`PmO9=pREn6SF?c+rz@_uu<8YW_&dhq(VT)Mf#wAr50*dCQ3dob=ns*-HK#g0YD=@O-{%rRuOO2?`NP!DJ}c>dBT)=IRtsH?v4Y=lyc$*{yL?ZopH~LssvDK?4CP40eCnhuDk!>{#IC?)?J< zpQ!}3YU5bLXzlqT3PTyzR55~BkryND-WK$}-y)zgx7&BF{?PpXgRD#VI3&dqGqnvo zIKY`K5@t{8Zb$ROai&?HE5~ZU4$e4#b28-qZ7#}ss8vC3L%F#j5ik~l7d{IjRS_zK zCs%*%+D-UAX}_{gMn~H>ofP)-QSL1@f!cw*8(O7JCsr3dTV8?m6%VKok%ovfcKbo5 zoy@WF2Y+FrK1-xNu?Z#98Lgd|z_-~VV8Z>-Ob`+FJTH^>5*<&V*dbvB@lN8@+SV=u z`l<>|@n@)0>YI|~MZ2*(&U+GfE*AaK9#w5DiEu8`P)sv8*1+fyy}Tu|6$EEiU|nIQ zxj{iVr|oX>eXM!O6jZRB&Hpl5_u1{^>8E)ycIucx$Hdp2n9m zSt2*D%+413-iHB*2M`{y{8tHBlupPxZRa+cKKjy#CHS%zD_IyWEz3@BM zJLzWqD@Ls9k$WN$9EH2yRrOwo4EIPm!i$HjpUwTq+hq!?WO@U(}!hEqvNF*n=8v4$gHuXnV15mpV0@ zubogB7PB72>vQZ_>k&wzO*@Ui7R&J+|NV1~ff_m2^ds}&XmGjg`uxT5nurqo+X|9@ z)%2i9T}* zC59TS_goqgm&1Wrw8(kFLTk zZ{i)xg)+`Bsb}iAWG26e(tW|v=pd|TLuugivOoyT%1c;xDQ(n}NY{6x%<;7TPYuFYYU;;P@yN+PG&Pf3*%el`s2ZAXJ#)x5w)=v z>}pC0Wu9Nm>Q7?jf+A7P*&r8jB0o7_(%>k~kghpF_nANC+uc4ZL0K(x1r+iNQOEt~ zT?;e5ZXj-(BiyrLR;GLp*pZ(&svcALE|3YSax@G&VlIVJ{g?!8O@_gg{kn<~a%OP_ z>>j(}E;yN8gYp_*2Kp&{{rX;f<%ytKB>5a(^Ea+yvG-QuW;(k^!mV=23-&S0Ar$5I zF8!ho-EFdohpZVtnI*Lt``fSwd#7CZmvyaqB>N-J1WAcE`k-~x0m*M$qe`KNSKQG8yt^=S|MP(9iMtcpMLA zT||R3aS@>VFGHUBG}RX1pN=+{Lzu^iU9!E)8X5}K_31u^X0H+~8XNtuSobLAMcy6# zmCPt0EX-m-5O?D@$hsAdT;hN_(Exq*6aQs!&c0*6Tl?L(j+FE3w6hHn`Uy2_0<*48 zP4Ty8uMT(F7=Kc@oz8dHbOmS&_9@1*-IIY9Q>38U(fj+t$uV_$m%c1Jd35C{ku_%p z|MtEzvF@?$Va;Y9`L&9M^)53n+fq1~?HK^&$v1$2+&4hBHe3%{&7s}G9%Dor@_r55 zCKo68)Ai+j|Tmu23W zHEG*)eF@vY;JW-#=#lF)G=QQ}+==vpL2?%3u{pFcT@?|1v4ktUDL(pmf*H5FoW>99=h=UA{*%Drdx3d6 zg<;zG16InOn0#h#uoFG>g~0E|uKc=P{~y8Zx50XjVf|iew(E?Vl*kAJ84>6h+sRJGNI_tYKu@xJ|zgV z5=s#vcQuo~`H_vG)gR?^jjyYkGR_bE-sC}@2yC_qdsJ|*SoM55>!JN+fBcxdA)v{r zo^}@cqHqrMtxaaa=0?aV%!(iQRkb`WWiU^V8Sl@b!On}B=Xb# zK!e`WbC(VvPE05w`-2z6UGFypBA^gM9CAfxj2F6A3Fj~rAx;@-)nCoqxJAKTku{&O0hYw2j7@`$Nm00d%ph3z?j;#M}z18 zAmdP!JJ5;;RlhsiH0~x&5%LWQ%?7MRw=7HGyX&7P^g_5_n-?Z4`lW028;R1lz;b;e ztmk0T^0!2*fFAJ3#iERm#Ln`u(23yz-bVwxy|O2SqSf9=vI(`a1wFE{XQdtb2hmeO z*98Sk=L=x5%*t9`2&gftQ8xz>xK;=ej88xc>y+Gq1b$K}kd_pA=LUaYn8@4o52S-T z&k>f8-am*sFf4(Ob^+^6&4JpAfCmkch_U`&o^@v0eg0mTo020DAD*nAIIzsa!h1CH z>uE296aX=w%G%19thEBPz*cR7T}WQ-z_Bcu)Y) zuks=P^IKRJW5m7x+Jzyx|E??l*DehG$7B3&?|E_zMY*^J8=&Vo+94^Q2F+DIw?eE5 zCHMP~GeOW{iR#lZk)7Jh^Ve>77LqlA1w>PL*wuGfq=OA;T_Uw$1~A-G>7b z*m{+I{Z03r!)EmjKF6eD7>|XO$-O1W%xr_W#--N1~qN;V2HWDMK7$gGhwn9CLjV6?k52_=wH!?e80;K`m5`AuSf4k( zJ9=pb?8l5T;pIMor#_5oxQE!LI$TW`ZLzG{bJ}a>escb=A1%POj9Tmh3ZPXE<(<-7 zUf-&>;O@nyvNsZyBn5mWRDD|3y-%!1ATWHA8vrTx7P{bw#5%iXYS2fb$Ob5N56lZ@ zNsY%=URNSB%7xsBJSO(O`;2?*Zk=AtW<13&s4)4jF7W;O_fH$u zMPWH)$2HBsKI0?5i+t$q(vHA%{$uC@G{uPsAHV@2W53GCZ9&2qbLy2@KB{TkD++vf zXZcvLmr_k1W`BphP=#cIg}0-N*?|m+1F5L=jDGSxs2JCcrkus>We%V~r?L>Puc0`; z*57xsd+6)euzZJT!J4eeqiyM@ssuyM<()fiTV-~}yiW#_B~l*wBL6UeX9VTY{dPcL z8)eI!A&4DDk;bi3Ts5s^{4B|7F;&ai0x`jkj!KQgZ*SDS^S+*HV$)%~_3WHSlD1#v zrBbL{s)QpQ4YfW*3AW7WO||HQG2676-amKLa%so9Rt%4&zMt3G)}ss7Gx^#<-Qog3 z>2q|X?F(ZPyIKynjG~Q4FI(0`Sm;|gO6fiRRAFI|JS-<6hY+yw`-^S?6B6Fc4g>|2 z`r3<&^N;i$uGd&otL*VqTu2XB+{90=LGUruRur`yumNdyq6O?bvJ1l0Tqr5Cl@SF! zD$o4Yzc)?v>?7%(KW^?~enxK(tGgsI;&K`oy0-@oHduLN#6S= zkj(Q<&uO%%F?*W1?3KOPm#%y1hL22DlvZQnLL^@xRE?xdx|D@+%oAg32cdYRq_)3b z+sRYwuEa0i(^n@Z>!VP)Zsv+NN+jE&z9?!0epcIF{x?#P3Wgs|uABmnXsZc`3L!AA ziYoT{C7W@&>K+TG;#Y+{Z82FHXUcRU` z<*ZYlQl75$X`ce920e7IJdQwcD7VNlUPGkb0m!)O_NcVTBu#)t)fapkISS!^x-4|7q3DR+QRo$uP)s z7n>Zo^{|HeoBd%;m>su~f2o<$m9(FpktmS#eGZ5x8QUwxG?3Knm?1NJtVIraB zCrr7@(K{|Doy{wBcpTIHFk0r4agVFyJP{ssF`wA78C|ak&ycV|^;R^se3wxY)s7Ny zrz)4U@C;#9{q`;YX*OsN=YkM=8h{pw^T-5(HxgC?`+<}U>$879Nw>zs zuCw#jV*z^-@JiBI5vvHX111->E}HN1evD-1lh=5i;F-4U*Gcz&+)J|ugrW0&{^j34 z#Yl5hw#)jR`H-K2JJ3W)YbA&+Ph*K*W7^GP>%XU65}8*jKmJaMWs`x9aIgtR3~`XW zItF1am5xbBFgO9LKy7TV!%3vC5-&tw9$tL1&dJrQSdl9Fl%EUXCAuL_7e2=1JO2Uf zImDcD66hS1!CH;dUG$>czk1h6oqF&jX{*>IzG>iq|0kgt?)N+oSe*!Jcvj&4 zKH>&;tYQdr${B4q6-wTVV0Dp8VaXHszn!~XNKHO=2s+(r!zp(+oBkZ+vE*=y&~PXK zT5(no;09)j_X~kX5D7#`l=CUa0Hq2|$(+I@c8G)(I8SWqO>BNU*ZtjjAY`~(OfGr- z(c3^1gkt(>l?t_Ar2rvL*azSf;G;fDi>u`{Pq>;~9ND7k(;`|&?Q9*_vD9}Q{od#e z;h$%yougJUm);+6a2GZ)>%>h?Qh}o5;J^*b9*3KOTE7hsb0ip8mH9ql3Q^g+<2D>* z{U0tp5j-JuT3AE$ipSH3TK)a~7Dp9XH$bldU2=c94`JViewr-SXK2%kW9Hl2?ApAv zB1_D@P{cy_zSPc3^{jy$i1fXI=ppE)V0jal6ywX-!NC`J(jtklDl(rxHAAk8o=;;+fJ z|LDNjMK;qx{xOOQn|*^2z=><1|k@?Unv( zFcZ<5Qla$lX}9q;7Lv$cM^G>^qMLyf=Nc)EF607HgLf{2B|b|Jg{W~W44}56sjuIV z2!g({(qpyv&bJ#kG*(^9XwK~~Qb|BdXWXS1J&tWvIIc{1M1aj|1^zOd{y`81nA$9i zV5pS?cp$N21d6E@H~e|ch8juDPcgVfIU5n=$Lj3R>c|^eeUF_d!5fF?|c?(6|zvgDcYNUklOy2tP4{$Aztm%n z$m%l+aq51|lhp3SS^%Y~?iBaG_r_QJ*!9In?~dEHLGr3GS5ic+g=|G_uOL_$oM5+( zfae)*Bz_s)`b>d3QAn)aDAv}4GdY2A? z!e^w9bWvq1*4ePED;aHu{6mM8hc_~_`s%9Ee+VhAvtJE(^o6(4<6Y^8 z=@$qAzaj8&pEQBF@&{7kgM~r>^-ySDzAsHTCU7y_IhK)=AF%qNE9~IATmoy87rEKQ(ye1w zm7<$gx?DdfWipv{%&}9kO4O*7=a@(gB4{M?RNH^Fopi0iYTIKcP0~RoxcfuyzS(3 zm;HWLU;W7HZ0~HS!|*Mwkxwh1Zx4J_jJaee3}C4RodiyU2r%ZoT~=X28zuNt91Nz| z-paC)naD6{%^X`SQTTaRCK17;GzrY{pk(7_xraCSQ2orop_P}Tt4&>BnOCge`n=!P zzv4boR4ud?#>TcMaRHUF^~Z^wVPn5OQ=PGKkpfLhRor1)obqN8@U zlHBHq2`KR#>&ZJQ;x-?SSi0@+Y|cK-zStM9fm|s{_EJXah@F*A5JpXb(gvJC0J-Y+ ztsBr#=kCbH&sKC%Fv+l)@dC%goP|WNtQ3MJB{Yt%M)`worc^1vc@20xUC>q@LF!GD zhiQ*hKm9!a0;w9)qoVVX2@+hCF!5DI$Q#4iPG!AFK_xUH!%I4S-Yo6lGPZc$+<5#q zq97;2-{Y`^yI=^3;?25qXV=Qj6QN!5(u8@d88DK@k$QG+uzg0~#k6z!aftnv?&ZM+ z#`|g4)(}d8P@2w&*cDz|i>jzzuIuy?sQm@Om$;)#OwEVSE~Dw5gnB7}FD5Ygc)(%w zqWdx);t0;bZ8a9<6j;Ww1Exyg8@d>Hx)WxVde(p(PLrv5j>AJgUw7(>p7 zRVKfVOwX&Hf2VYjg&F(vg#h_tJCpMl&B51xXr`N{hYfJlP;?mUBMqE6*s8)BRCD@D zP~TAmzxf_Q1qeI0o zpc0I;069yA#nl4<)vStXC+r76K~ypb%r4I@{P6`yMslr-sts62W1J8Wu)8!b=*>ZVKLo3z>Edb~gm;51wnH7E@c( zI5J}^Pu|ac!e71#b9Ar4;*{L!sJ*xEsjUkz6%dR{)FN%Jo+$$S{#4+n`xs+^(%kzO zXzI&!bD%xYj6_Tm8Wd5wEvQ}(Ha{4gDz$W5cUrbNEcm&)8>tK#yzS*)STB~5qE$t5 zGi?Z@1|r2|aM}qCt0Frw`@%x!1X2wm=I}_rV_|9M#|;h@xcG^^=KA_*UpcQ=Z!V|m z@OzHPz7LX)gfb^^aB5*qARR{-<%YUg+ac^B3N@KsbdBQcv#E%y=yD|g`!Y;oZ3mCXi z?9<~0YGC06h6sNCnyq&4!g4wSKI;r2p;r4EWPBbHepdLC3YUFE zdo)W`WC=`_ulfIWvNh{Jq3^bE}NUGJ`?%YDQ^_qzvjcoVl< zV*)oGi0=fmn~GO6F_fe^jjXxmQ#CFc#7pDBr(6eDyBHOh41A;ltfgLUo={4Xjojrb zxV-COq(HS<@xf)M42fG9cwAc#yu<$(1fg~` zxIe$L{l1?ge!v*Wi}&ZB83BQC2<@GGMGnaOSTMJ$~G!`I0AcGh+@0VW%P~DSr zS*z@%eC`>|gcj|yMV$pRzFOf!-ZrjZEFylK-)pd%z;n<#(&3E(u%PEl@ZUB6XGJo( zilj^DfQyDvL&1#WEg!1Oc3RCn-R>?~mND`>sJUe{tG9MsI5=D(I7y|?gZoPI!&?F} zpyryg+oRhz2+(a{Un3a>O)lUnm#Etv0-W@N)T|C97yX={zm;VR1(fm1)ZeF0@Px*a z&PpEXm6E^J&il5=vE>>;oWAQZDl#YXMRwf`+{{)4Tw65w)}`A!E4w>Z6Vm?0W3$5X z>{Lf32*{cppGG3+Y5;A>14NkPBsPGNiXy#`Jp`Ll@-I*-(h6a=vp!I8*%!#$vh=D2 zo4xD}dar`9qwSe>bGD3;7BKA_d@hkWeo`e(J84kZnSN^00EG#aUY5d2sZ4n^MG7^G zSsysq;l&pn4d}v(FSS~+f5Ym*I&(;VqB%zWcl4{NrAYx+v4n+Kwvnh512qkf^4BE{p>tO%-3G39!xZZGzv;w`oCv8y5H!QRE&{P z#T#Ej0M#QMm7QhOjP?F)gE_P%5CvgwTZ!HE@D>Yl+j!tl8r^2IV z&IH)Bw$I;UrDI+1uf|N!8oJH+-f18w&hR%fI+94SMMmxv$7~+jm?=khHQ3#csEoHL zQ=55TUU49pZLeMG8KV;{nbM`eY){so6k};(fdqh)9~JMDb&|Ntku+bFo7=RVm6b10 z3tx+p=ufv)@Tq69zH+zJv%`r?>FST zOI|jcKU_DQ%g;RW+o1O2%b?={>4V7qVU+C@R3buDxILa8N83Uz#3daFDM^vgOT2!!yD5%undxip4GqZkD zS3-{&iv_(q#=AOr{Dr_=#plodChO7>VeaqWzYmnLHYv(8QDM0P#*jgg98S#j7Q~FX z7e%G@$bfXGwi93C$&;Dx*D(@u(L>SGP0-vnU07rb4O`FqAM41Ih6O0yCX+SJF%6Y1 zbwhLJ&qaP?wjF7rHyX8mf5-P(6DuO()g~SbsR( zy1?3b4g_~ILFPhIbYvEPa*A5m>TdwcJS zUNbo@x&ZCDQ&t{29~!-~;$NYusS#Xud@hQ!KmB~yowm|cbHgY?0RdJ8^g!KAE->P2 zMz2Fw3a$ls@f4I`lU72J-^m}3=KS?+Vn2sH8-ME>ky9f1BA5{d?S0twmIS zB2T2i*F?32L0DCVE*3{3(1k6|QRw-?zYR{8{&H1$2#U-++;YQbg=(USd|3oH z6be(ivzF+30>k%B`cN&N*C#|=%JOR|$5Xi{5no)5O~1ypvE6v^@{Vc;>!v)ENS|TE zgNe?d+bjz-gcj5wW2G$d9T+yh*+UsC4aS2>+B0g-JhUzkhqn{X5g&R74joTES=xxW zUtgtmH;XRH!{iyHs4%(o$juz>1i_q6*8|jx&0LM3=Vu~hJM@X&*O>RutqcxH>8q*I zzdtRzev1o%^sz=i+=1w_IS@^Cme|Uh#|3Zt$c(6Jum5y%OclJy^Ul&ZhxN^LA?X$u znPA%@jCG_CyP?>-$mAyETSqv$T~@Ig)U+(e{E>I`ZoTiV`hS>v^JpsDw{Lull5rzL z#!Zw=xic2BsZ7a`GRLkcQ%Ev}eHk*AgrXEGgoMl)cjhU{%r@<9%#^*$aIuH`ce>x_ zdDi=fW{ridZpxk>&&GBr?r?n6WsgxhAh|G9uCJbUwh&c({jpoEktZHx+N0A@5rS9oX zrWSm1tyRn1^UwY~*Bkz!F<$@4_ML?LXWI{}u4e7GM9(^Hg9|SrnHmD}sVE|ODUxal zF&gK#%ucWw8*asz8((0Ym@;L5r|X$mEz3DICXT=!hMEQ+>P_#*1mEF4)wZIxdW}aK z5J%gmWPM97Cn5eU!3R<|pBst9B_-7y^<%PLYL+u>y2gL!41%+g_Jclv`Rvk^IDx!2 zu+b*_9BFBHbdRM{~t2e zsB~;m_LN@mZ9o3V!^wEpF%6A3epbwgL?0bgB9wL#$)E@g!o4JeW>ke)OWg)b;A+c5 z^kXgMf^~}3tB-a?Z50$7rG=wTV4X`2fF6?sy%XB>f>>}QT_k{~HseUQ)n52HfhYfj zlYGuqw*G|)(M>lF5>qO;mu*&dzSE4IQBFqkorOj`?OCy;E@=rAaL4`-E~oDKWyN4o zR85Ga{aRrtj`etIMviTBf>fEh-bG9HZ)P`)O1_zK0e|}wDocG_GGjXtXfXGo#M{kB z>NLjARl|K^`hDpEo7UP#_aEI;XE-KFNa)T3OI`z1J;84&f$>PiFf_^1EZ#0JLW={; znj}WmJ2-i4Pk1V}rCN>ndT>2?WjE9W`DmPtML5{LsH!4-`nV8;oLHiBhS1VMsG5bY zFwG@(V5;HBIWrzt*K8|@#ceylzt*8iTCJ?h$M4V`bB5bGO#uhxMU?WHe-`oWUM1Pn zk({*W+JA*u{RmO?%ca_oE^92{$H8o8V4~3=aiC(+i|?0I^q%!or7;V-ttB*76aV%O z${uPm1?o@WNaFxnoudN76=GUHqUcp`bnnK1%DF?+az)k_Y}LJ~XXaC-xb+2%H9OyW znJaz!sH6S1Bo-pf(N3(|0rNaTeHif_Kl|%uuI(kKXk|$z6Vv{w&u5$8mmm3YP>LkC z0Nj3p{kAQ@`!H{ZfrWy^lKw{YXi)xu8;7?L%!%Q~)aJIs%>>*^sS*xK#u;`y7LPNu z5IpwAFAvQvMc?F(x{!9q(u6X4kD1Jd`3y3F+ozHO0K3)9?D0LjN!5hv7Pg6?Vk@Pa z@kXi+GTHifZz#)K@G;HDPA1m+tbpL>AXwZ8(=>0m1v;V*or8!>KWYH}>;)DCqD|nA zApl7n00zIA6F3M0`EqoG6KwX;t^8C$A?)}^zhc09>#&c=c^dz;Te7R=t6NfMUB~io z>&oB$iazfOYU^}o@D*FM6ClZvLAGcvRerG-PK>Vlp?A;en2dAcwV_x2_U~?kh61y= zVGa!84v3NjD1hzbS|Xi2AI7k`)kgGk=80~`dd@uXUoXKJmUBOe z?bi6qsmbc<8*WTwCG8t_;%|hQHV?g1HmeL34@S5=D6)OT%(OgVxn6DKYpI&>c*vKv zcahMY@Qdc)jgFE5S0YT4Bl!q9x#fJlG#@EF#KO6oxA8buE1FNsDKZ zO2|%g!>>x-&ov!5EXF~`(UB~`jM=GnqxEtgAE;L3ZyFpEz-%Bd*yGaRFO>|vG=_A% zZOQu}NsbphcUdU`63((~40{Ebu6+RMV6Fq0qj9dk;Hz5Ms>#}f&}u5K4O1X&f4X%ujdNmFAR_Mr2?ZL(K}i-$|Wo8lNmWL}ScPzC3bW}4{dJdXr^_Vh4$x?IFK7kIu!Tr@k; z9YN*D@n(Cy8602*%q$-J3q2j7l%*+49ov1$nZruE zzn`XpY1@G4%EEJ(T0j_hl7T}si3DL#e9l6Q= zk-347TQph^;dhLO6jFojU*c*zTbS-?zkG`%q~fx@VWlINYO83k=^Bnko<5?vqp=l3 z5qF-%F3NbSk&61`e_fyZU{Hj82DN}_!Wq4%j(#))i-Th6M572*5e}t@Qw2RGp%x*a zK8%8ow5a`=uIEoOaLQXIa}4y#J+L%aS3dodL5g6`?4yUh1kccsDzxVzl5?Wgmnunp zAV{!VuD(B^_tC9ZEMxne-j(<`;-8Yaw4QtFim!y!4tE!Hih+=T3G^u<6L2n|L?_U} z(kKF}p@jpD+d14v(>-Gz--SHye*aQhsMX~@vZGws#p~|ueCyLy-;+0zc@gwLh_Ij_ zy9#h%5#u4g9x@ZBoGOyZ zjT`?n*-O~V+12~@iz}Qg&Iz$1bMXjT4K1zc#2^U95@Y;#xQ*CJu#?GC!+DDmVK$p^ z3#52tJ4b*e#$W_e2)U4Bx56**m!r+h`856=t}lU0314sf-E6YR;WA(nDg~MR22~8i z88Q|mG%a~b^KXP{Z3FiF=<4Xd`)443mPQo>0lN(LZ$yox6`;*t2hs&43qZG-1jzPr z&}32}^wWQsW@&!NED=J_8;syq+aVhNFqM6Dr+2PzmwT&Zj|?6Ngh3fq{(6ZaN&8u>4+pleBU2_v3^7&0ZZIQjeKF<|Y@-+)g3Awqft>=a4j(V=+?uH7tq;Bs; zwfeYeJP~>@nr_E(lr8AZbi18lN9Q-E%VjB*4-rV{5jeuYj(1I~zXC;N9KfT{b1|a} zSX1!l0fLrkC>;8XTi5g;lS^BzR@G&FTH;+AN7*L;9&5$-X)ugCT+DBmAK=* z;&Sloe>S{-S>UHdzwrea&dV;my5{`bP}BtBW(}r~F^anvpnYlgTW@Celr(lQ@LlK; zdec%{7}O$+!Pa%i_H{rJnH&(U0_}q03Y2UO_UmJ$>I|3^wTwCP30MuDbWWAYM*ph) zEG4E98Noo6uwDV~TCf@zg!d3`Eo&uf@2V5##thW-boxl@_8mZ;rKQVJr+G9RI@}_u zUqj%hm{01=oe|s=m0L@MC4oF?9syJnz1ujT4s*E229BCDv6JXrwo2Pk!;F7kPW9^i z-TpbJj74x-<@h9t>mcKz6yt7cT7DOUBt(C_&w)s~>$0{>Ufu2>yClxj-m3urqmh_+ z3WXV>q6>uW-|tlObnXcF2>yDX;>#ZO3%#?(?j$dPcmv!sc$xn1oE8xGh>R0M`Pud4dL=UP7LRz#}S+zl9B=-EKXp zqSn@az8S;xSNlwFAN*)(I~3G?2&6nZ!!nNg`7N>^2djF}&3yoae`Dng?HtDXr(Y0i zO_w4WD7lTGrPNE4Al%oYdyPj(8LwX!16|?eQ<(;COEpPd?arx2QWEzb*Rk^HFkYn# z3lotvKC@^Lx><}v#!o zK>O)!%cPWV>A%m~vfS!w2w=K>FiW_Tc@98Z0_Ya6DPViXEkR2xW8GuD57vsLTP%V2 z{}yN3--uj+%AIQJ5$S!Qqb%(M^fCGoI2^T(%%i#t)id8qr@tc7RPgxSfU4i^+dKDP zLZ@<*vZKvQ%2_|rZEjiLN7q43z7XNj%po)(8{!U8CLd{=;UmJQXY1^*WSe|+)Ka!v zRqJY)jp@z|eBqps@qOOw0@qgp7X$V6=ynviyLf@mIZN+DZ=VG5)lV}_iUyi0u zSBP&1IOC4!4emufk3fY7!>RWScOhO{!{?5-nr71PmNE4}N|Fs$-#$%yQ;#NOajyRW zocn;d)80M+S1+&N`OczOs@yP?QTenu=R<}th8b%?iUHJB5C5UJO4fk7Xp;s;KBh3`l&&2B zmKvoBm6JB;EWMPg%US%Yim~{m)&!3hpE_5of#t|ck!Cf1kh9y2^MIgX!mIOdPEHze zv?I`#`{ zn*!1dpl`DWqBlAs%%I>yzhX(b1dS__rRnLv(h!9E2)e!d45Y>;Wd2w3Y zbL&FT+@%JJDDk(Su-waZne!$`7^4kt2p;8r6TyLH1D$pb1zILBx1;VOB>FgBEP-7; zML#`stMT>m;rJNw3v!{A>V}aEPNhn^n$~X-yGS|*$S$4_2cqa-ACA(Wbvf==_`$wE zZ=m;MwzXDJdw-2?f=i6)A_x!|sN4`O413NU5QJ+fljA4mT;ew;(4XD@jMa15|76ps z?d}t8Xl-;(7ME?j$@*5Cw@T1wpU;xymeo=+(X(AlW^Qv3n7va+TKv=2w>{-yY9i(S zV{+!S_U+cYf7<(Fc8wAj8LIgt3@s6Kd%2xpIuiJ+q{3}boR%g4#!FrD2LCeSX`rQ-6Qx0=hXV}fXCcenYNW|b#r0a^X z6tX!Hmf2`V?oYdPFcQ5A2}a{8B}aN>RVyph51m%6C^eDy`SPqojB~-y`KT1X6@so= zxMRF+4Q9R+=*rL-9Ba#FA856VH|*t!b*9-Q_O8p zj@u>DGd3UY0l(Vx6f4+=s<=~^NI7!rZ^V2zC72E!1_8vO-9SeA;e7^szQEd^wsc%8 zQWPM3DWa;ug_t$?0ZpaeqM^(_r})!=3A@+8S!r(lk6b~~I>;Xya!<)6>)haw605%8 zF|qZ2gmp%XYqlI8a*OQkzVkJ%-c4$`v3+JSr3tFAZrO+vT5-dQl1C%`)=e+E@ei6p zBSpn|x38p1NZub~Ea>267{6kZsRgGI`=Zh>K8%T@>b?)W-|QAuAE1%DX?3qY*EWJV zDNNqdh9E*TTJuN-pwM<09DTj3L>!&B3PZggrp*b?YAv+DPCo`FBC-s?2+JZarJD}^ zJVG|AlsOy-zGju7VM zmIz;$dYBeROjWH)1+kJ52XPA zFN49aW0k;n$>IojbQr4Kdj0X+5dBKpTY0Ri`!q17xW=^iZXfVNb`I^!3y|Yw&tJI# zPfuu_;4?1sI<9L#O9$m_*-YFl8S7 z5(F%$+p|#{8nfyFRnJsE5)CrPk!)sQ^$P^E^;Z<@KC-b5yx>AIYt2-swjM50Jyc!u z*f#yRVd~=XFEz0pvBTUE3#P+ou4c@f{fFlRY_p48Up!sxDP{XhKh@6j&S-xSo-;!< zs$rx!JY0?3nU<_n{!}Ty)g?R~L2IKu2PbtDD-4YleSnswEN3i{tU!#SNm~Hmf=?mK zmTUDsSy?^XaQS^R;lQz9w`_>l*prPqq=xs>3YZvdAf^5Jz+>8=&+VuT?&GQw!W~?r zw;0r!2t;DegPhlnO!4jZ3B#CskcCTepO)5?HIroAm?nne9|=p^jphbyMY9w$O88gPX{_}4S*_56tFB;Qx0E;kuCTa#jod($dw;(V#`M0 zT58Hvn@H)=+p;XIWHZ_5CVZPc^HHc|2|Sn1@tRisspLZx_Rxxe9Q%_9*Gign&vV9v zp4y)xV%g#%Vi955T<;q}v>xk*_CoH7?^V;0n~8MeY`R1a?F;&61NDGRi^7ztZ3K%hv8Ra2deU*e>?`6ce#X!CUA%7Q zld$+X;N0hDVuh9xfKV+;NWm|`-W(b#sz>KHFDYxZ|nWkekbiv|}wI3QUE* zD$ZF`zt~?fsoy%uNcL{dlpoe&nf2iC&>K=zDY7@tIQ3~m;r3}0Pg#+|Y(?3xPa@=Y zQs_uNxODaoEd_E;Kv9A%>0M-pr?~~IY7q;MG;)Xd@zyCSJ#P%VwqJFperNv_Jlc$A zhX|ut6BjzrTos`x$O<--eze)Ezb;-oqPW$dVehGS4C|I5@g?Z!$*yJ@>AZ%0$<3x7 z6iI*MF{(Okt%QW8^J(!r%8~_|UKXpY79II+Y@FaI{AKNya+6eXY2*>x=#y(o$)`D< zvl37}ZIC63I=qnjO6Mu>aUO5$@m#t(;t0a z)Ayeg$}rtC++iAB?xXSnI`S%7p_V;(1~*(iI}>XsR5P_DX%@K&9*qe(9gI~Qfca{F_ zNcFf$e_i0WcC9&V8(W>=*EX#1=q9)_nXf=iLQtbB)O5bymozgg&_olOX%`R^9(4Yp zxG7Z|J+RKV7&Nlr7+^SlR@VC74gY6Hk0b(azvz^00)V0;3-|Hx0F*voi=auSZHJSa zfVImW0oCl(z(_g|@b@1Gp?%`o|myoU0R5 zLo~_L0yyRgAtY5kw~9VIlS<*PqmR;uDX3SY^q0F|+OcM zY(HAC{*<4m4ty`rYv!uZ&5XwGQ&9s3ONal2QCJ+tzH`cvtG8dm;l&}>XzxNiW|2JZ z_-Lr6*d2`8js}dDR}4;~(_7@f8EGNTh4ImNp{7DGZP=H9CYhmy0D);ZIa0rx6xtD2 z@v}mhgXPz|x;j52t2;F=;nS)Y65a*yOZ$i;ly_4oR!~#POcnAG&J}mm7PcqisM6D- z^ug3nwBul9Fq=o{c|qZCR@~RGOT~VL)JCiv1#J$ruPmHT><*)oLaV~vww&3owL9_c{7VMm8F)n1fu zLZnGHS*fSm`F?N5kGSR4@^r?%IpiaBz$cma3P==QfnV{&x%aU^O&6g?9$=;dVWY&E zM~x^r%anY3J1^fo)%up-=ep>5yEf7K-)EDr9@Q~qeOu1sALacKM_Af`CHgx0CZO#m z(neciBwZw?!%E!|6{*uWKxn8guN!onHp<~i@i`{hK6^JL@1ER`$Y%^xa*OiBw=(@+ zk)&tc-S@f;&U1iDD%k%f8*S~t34kJryo_@RwU}$O+pJAXP~3D-cC7X2_){5-DO=$= z>T-ASqJlSz>^M_<%`zsV3Prj*v)&LK%C10CH?OkIxu$GfvaEew=Z%+V^+kf;879I5 zRf+>RhDX4vN(09WGPy%Ez{Z|gFJ0h@pJ_FFguXFMg~Bd#dHe?aN(RS2DV=2yZ5;Af zd1>YP^;UTGN#@^em>;0*c65i#<}lJIi2X4wSUooB46+VE&pF8^^% z980s{!FnKGz3Y+rVi!b(Tm#ri)TZv1Hx;#vqNZW0(Ik+;QhB^*LFL7m(9ufsm|aN7 zi-9UP;?im#PQalqPZyUiN3^WROQqYqdOsyod;E%E)khiWQ9g`i`vZ)lTl;h6Og>3XeQA_1m0pk4fDWzTt&uY4;#oW5A0! zL!P}y92&}F2pJ=}d~bfAPfM+(r_DRamL4$NnJiqp>+EHDm-`u#mth$qxmLGw__kz` zY(yC^d>#g3Tqvm3j`L%B?(?fx( zsy8A2##4cIsF*#)N1KIp#vaCns~$ZWx9syBMGUrsoUhq5$$y(a@&X}}jQgRA;DZ@~ zk8KLk8C3y24iKrbPpse{X^RaX8!6dPna@bYuLRn4LAR0y&)m^V)ydhM4)xgC!+~sq zu!4c5tO0CO1{sCr2Qna>uRLZ$?THkOt;0{WTl~n%id=j3G1r_|1L&#O&@0<;iQg=e zOpLn$C94-uEhW-Tykx3E=N=xuVg z_)_)wr)6NyXYn)*S*%F9)WN`m_?bu#|4GIGTfX8z%MGLq+Qy6};&i?-FsqCj0 zCv!wUJpulLr36&`-h(>y&bp`idMR1*8p*7spUm7$ml`cm?cs1T1G*7T4{9@(ey3~3 z^^aVfsUYdL>zj1s|lgZ>N=46mC-){5d1(#%MkSl*$EGsT4RhVs*s zQ&Y#sAh&tL=@S&)^NfeAEYIAS&s#8{UNEBEgZi0=aU>`Q`t5+gKB4!cDv(bVA;P!5 zn1`EyY@I26vB-x%CKx+9ApE+Jccerib!z_m;*$%!KI_uI-ijWC&@X6LLiP{VG3uoe z906;#XD^z0f*4|1)fc7PWY#1SU5@M+*W z=M5$%KGl{AX2k7M?^*jo3eEI8XrYO&K)Vzvc;7^cJ7}ojEvAFF_rTAd$ z)m-dD@aS$F$)qn9MY@UvDpf5I`^fj5HY^8ON?W>mU@7UQ=WVB|nre^0`CHsl$CmD= zNf?>r6oUrJyc!ORBaF;Eq00}V2}xO%zwP12HgO@y4^x=HP^;zAY*PqtIr+ixENwzi zQonod#UY#~U`qJ!et&fvR7qHXqjrXKDH2eV8vRV5`fRXGpmiHYr0CS$n}Y#G^XaU@ znvX9g`^SCr`gUig@+eVfx6IHSxv*+G;;m&+I0XorhTGP+HjB=+wf+`7o~bWYUM^+~ z@fb;M0rHjoq5D{krTvCr6^^wcL)mIbUIoW%57N1Yoc(m}f>>gxkI4P}f6_S}q`r=a zKzkPHF7N|HZ-JF`&#b2L%2gBXU-p51a-zYP?L9X2Wy|%Yz8D=i!w|h?)g*`@V}QLh z)q+e+2Zg_|eG}|MDh%1r0>Q(>9RMo{Ar-tA4U4zfKuYKCQOCBQ(s|2=9|&6150L!z zF)V=REs7gz>cw!W_6dS-a>bsuS;1Gm%O8jR7rR@GBWC@R&uLGN*hx7|W#+XV+^H>t-D&PubYydCr6#G>}0RlY)BC?MLb%Z@D!`J5AikQbA3RxpR>N$pwr4jl)-7eU2tW%h3Ur3@dEZhZ+7yyi^gOUTrM-CrFw@qO&H>T8H1MP4yJEC>vf^Sco4h zJ>7OMO{u-v(t<&*1~Q$TcR?`z-j1gf{&%mJ`0p|u*y{iMlmEU@_G{*Ux^oTx^U~d# zp>n&>IpX2J5pS`RV$5(9xunG>M3a=*5_)H%mCai%&)YAw(7|4#_ihqfkd)<<%lQr& z;mv%k0nP~M(q@|eA%-afcpX?m0O;s}0>T!&~sm}e=RRW=AT7wjWz*PbgKC8UlNkXj+gR@@&XAp5bE1iu} zaLW5$B8~Y{cz^M)4_OO;{FP;XI|;gA@9`2&e-%e{^6igFOaK82Gx`;KiYh2PE0Fr= z=|zP>zp z#Y#WA#i2cCgi9Y;5%m03l;ip#{@O{h#v7&2S`mlEBL~9J0bno=0)*GwiL{RpA;(b! zMw5fv_;K7?Xfo%(o@Vgl1f@Q+njd-Pr8Nyat)4fZ7?vGo_ey5mNqUjyA@>JOazw#r zvbSGDv@GD}avaaVPNc%NZ2%`eP1?5E{F+*8nfgwxQ1XIEuaDf6uEbSJ@4F+Qh9pJ) z4Z4Ky2%^4i5Ew*?dRY7_XGML_C2C!@Zpv|m%|BUGiA5!N!gc;tSgZS>3Ga<6rGb+o z>GNp9eAaqOXe9ep7GM4cU=eexMU z{==7OO2=ba;Hc~sd|#Vkh(pfuj+^RYLMU-cmMzp&9P8?6sE~Y(T}fT?04MSs65wS- z-ho6b9WJ39!1kb-UK8b(=>lZM_P}CQ;e`8!PK~eiq6ahu?T#LMC3M2g4kF1U>uFl3 zF$A1>K)(^9Ow2H7&pcK-pz>I_!R*>l9*C@sZJv~0@1B=^FFRk(ql2*C0E}-6BwT9} zC=!&N`h@5Iq=I7~bANcJBw*Y6_k*5-^#@re1YaLbm{}WKw`$Tho_qHe83wL0>HXo- z$CWd59z5!AgoO1v2soh`I-4(9a(0I*47j6&x8uoPN$*!NTy!2-}8VoKAu#txYfd(95AQ3NVQcY@X0FprK7i@KTxQt#<{5s&tX7Kd3R>Vt{UY?Dha(CmF(4 zgv6ug&~Z>1h?EUL3?o=Ih^wfNrYPb~20A7FoUX>UDO8{R+6(Fb54Jz+W;DqMq;oEy zX-c|`gQZRS60}A2ZlWI78<*I_b{$vho?Nb#`BGi;qIlrML8&DE$w$`Rhwgn7kJiy0r?@7sMHD9&i`p079H{uH<_8lVeVB&UvW)goykZk>Bpu1MOv z`iEZ~?LeQ|kW_doe%s^0kX_*1eYHCMQxftFnju4dz`fM-{)h4l))idE-yza9TpWgU zIcpr`-CXLT_f>Cs2Kyt4Q?-j zO4fP+aFJCp`=UL2)!?=&g$3-$$hB}%jLqkPO154(UGlm9p<-fqrZBfc*{3^Jj4UHnaIIg z$;~6XFYQ$ZueaPy6u+)4D;W8LiGRHM{?ffYQ2A2?dlv(&F9g_Jgi=VrKnOb5U1+4> zM3Tq7)KqhoL6v(e{NNT`mF{0v62xrRXV8Q46i zsR##zhn@mrb2J{^k;Ot+T2;fUQIOHOWUtr7l?e?|Q&}l~F7De2q7_%#q`bGGcjC9+ z4JS3761@oZF3;j_12!ELgGU-fx}q58crwvw$eA>qMOs`nt{EDcO2`rLU3qx-^Vb0< z*QB}Xj2FJirDuVhhYyN$1r47a-L5X9RW?ewX)}_ya4j>eu#02;k2RHG(Icjfmd3Wz zNY#bf>SX=o^HGTs`Z@@w6F6vXKR?&Tjw%3UF0DbT@7yXB>&Q1+S?7Q?iW`^H1Auft}rbYHk%N$gCuV)HpUoqjIDF_G$X>H6z~u_teLt zt@rgrJhuIlUds62(P3q}DZPOC*E70TNJrixov8GnU_%tiTIr^5GesL=2;xlhE_f-t zR_@S}Khq=6c0%nH&spNFnoiqmo5yfR!IXYJ406|lO>hG6?oBn?!nA@4Fh{ZnwI&K& ztkrLwI{q}v-@91e@%Nb-sV84=+q`*}+njnYry5k6{?$>i(%#Vz!UKbJ_BPn}eGd@> zaTydYr=9i)@He?B{-vrk)crTn%i(m*t-pGB!(-dWc{f#IMjh0UI#s9(yhb$K1b#{n zXiEp0%^iB@G8>Bd>fw@a=dtR>k!I z%qhXMQ8LhAmcN8K`qy1=m0*r@pX?SbFV@8y+emsVDN_||1dZ7bSdLDj2@7VhqW3x* zMG!lcM3sWkv&I<2vbOMmg7=0ZanMthtq9bkTEouNM(ZU=G9tv8l3Pm&$nF_f0u78% z3~B-;Q5U$vLCzc+&KV534dj!HayU*`RYu%ScAIh4bh>!(RU(4ZV;KW$PJFE$i3E!t zIILPYC1+mQzqP@!>|)xGddNO_LzW{s$V*??DI$+KvJDaR0eDa5^YlqO{YsYr_4deu z3mzJmw$ro)`dig6ePT(V)?EwzbcnOvs_D#)dj?w$6mEzL)zS*lPPnL>>lCk%O+ax4R(*38 z>RtTGV01igZXSvOq@8{1MT8kK*S+E*vZErO+r!XF-znFfX_0|<)kE*C4!#SB9zSjl ze745t6%HqGh`!Mej?_zlCC!DqprDU#cF|cLB~!MTp}~hd?5Cs2e-#$htpsTqq_ysui5Cmh?U_#i*V% z`7rV>|F*TjA;}}(zqQgv_B~OHZsFH%@4r-k(My(by3uuv`&qXE z6W9I2Z(07clKr2l#>w|(AwF|p!?GUvaRj#tWjsxIzieLFwy-{F)qE20IU6ghBg~oq zVUz+`fjj>*RG={)Y=1Sje_DiAj3(sZj)G&PSDZr{jn}#$zZa5rxIn7^nAG{pa)05c zZEr7W+Bh>l)AmcNP+@opH3Huju?hA>u3i_|P9Lp0I~F!f4@xizJ+WIuy-JEqH8+Z>meROX;gK^_P#xNG%t=+hy%sa zivQh5|DSvk`oDjCWgArwE}Tt#yf^tU>v4!tNno}`G=rz;bV_f?nY$PYlB&@QyS(4@ej!c)AmDnc#Qmi4Xt#1$K& zuC&O3K!9YALzDax*Xllo=#3V?rzLD+ju`mlPhLuIqiS3ac$n{I^Z3<^BGFVM<%aM8 ziI)qn3uw7%CgzMsAz%98Y~TeE018$$vwk{&l-CR$T%#X?%UegH=PvQ9jr1|c;z!PY zNEH;rAKun7|5PUZlJx~oSoehG--yeNQrJE^Bis?{={LK<1RJzjtGV8)PP6G>-^dAP z>AL22B|akP>zpd}C9qmhG5ww49S#6T(BFtFv`e^>ZV2%Q#|9b|EL84UHWPb2_$Fz? zVEC>S&xE6G;QjB@vr`@VMm&FYhc+E*^ajQcI)_bQd=vaUjdp2}pq~i57-CpQUYM|V z6XOncosYD%cZoJye6?s0CmH!IDth70DRR_xuLl$h90SN-ce7wJ?IpCD7Z0T->WBnT zAw50HUjRs#0AHo=4nVuCI2Z7-R7~|~OCNl4XU&^96N}<_Fs_Auj|Y^F$s4-CK1ix$ zwLjUQvCBU{A^T!ryvhw3eN*mg&;;SGq!~g@FJ}PFGI+&tqoKbMn~^2}uf19_mPY8^ z`j%osN+doRyxcaQL{+%;ZOs%il4?sfS~^jSi$O8hOOkxezEqGBqe9LS?d`jm%3HFR zKaX`NhlhPDm5sjZlk#0I8w~t|CwW7#W+4$;bD>xQq8n@+X|xr?l|wRFOV=;_*i!J} zHLE*UXaD3xd-BC^v(gsPri@BeKd#AD%Td&zrX`4#t`y=+8jT8pEJ=B-(@g!Ut8Hbk z9XuV+8L2Ensp^4}Q4T3V=50o@ED|;ed89}>YZHL+U#8{a);_|(U`%%`=OJcjz8HuN zHE(C_oju7&83(aFV;9cL48PL^eNJ2}o!yAXaM{h#xxLx(k-qHwi$UH}FGh}f4nOot zwdMKCo>hGTibbSTK};W;55{9qd#7=%(2R!Je$xB1k1BfgZ^TJWxbklVCS8D@dAD%) z2mn*9(hu8%+vqV2GQop#q`vrib&!r#!YajngaPfep<{CknwM;-FS_U1*WqIsrJbWbPC4`p6Tq%FBvzBfR2U; zK$JW5@B5?3=y^EuD~{_yxbFdsEj2B*T*B(9jq_3Qh|cb>k)88m&?6!G4;&#P>vs%# ze2LyB4AEGAGUCz1Xga8L_lv<`VIJjzs79!`$GBfhnWRV@`n;U(l|qEreyaU(5O&b$ z@+`VLC!G5t_2hGgWQSNwK)G*ppg5xXF9U|tfZr<~2dcz`F5Ip`R+}Q_WHR`3s@idP zK+Xo@S@`4zRWuH?z1V0U!KOq##DD>9>cvLVYg(c$F7k?r8k;@A?A6`JQlRv0d2CDU}rw!OoRQ%C)aC&;?N9a7pPP9z!niGY`sCan9wUi_;UP7# zW4Y~wtWH|m+wTm#&}PEi%y=TPq|J;=t!OqerjeO^;$vs2>aT*g`049Pic?pmVtd2K zTn?msh02**Fj1KjCYTl!CzXTPYp3Emuw(qg^X@M`oBi?bn_-d>%Dn+9=FCfIs(=%n z?JCKP&Qjw;L;{Ho*bUikb|f%dzfK^=%kwSe)ayI!KaNAIMn1ery9vj7gk)vzvKaFU zU(CvwA@d2WwE*5Wb)p$%$!Nb&^2hAossOZX+tG6Ck(40%H^7jPL>QMErXHE=J+}~{ zNikCxSJzUW=sn`>WyE9nhNJV6=%?-rN1^|)P!KHx3x)nMQ6P~_7{~1Dc6?3T+X9vJ zEeI26en>rxE&|$|E-@*Fr1r?{KwqN-{Kc|?we0e0E2q?(BJLkq&Y7MqX~y9L_Vc6{ z#Q=wqy@oEEP9$4=7w$WT;$Pt(;W8`!Th8?%6!#?ui$8TXvQ z?Gs@Eu8>Z2Ike74x#fyMsw9q6Okv65z$vIqrPtf6syuBy&GhH3^0c)Pa&69&Y_Gw+`a?qw8jy;o{w*w)|IUq@P9HV_;s^nPLh5& zZ-4lS_6E21)(yR9@q$+TgE`#|f=Z#LDdg4ahBDZUH1PfRcgu7Z0;%o!{ib8jEQcT2 zS~2Yt)H;F+&p#wWuV6_!iwWh_WXL6QZt=U1wD(&oMZ!j9!gq!@GQ=)EGSsz6bGFdb zyvN7MsK!e}VV**yhdb14dtx>1c@KuBy7GWw5l0q5N1}gDQ?L9DA)+p$sEL?d0B%FV zaxUQzp%}-YiaZ)K0YLOf5F`$`{bbzH*hj1ZsFi#SRYZ}ZMejsaI%ly_jY(=9b5LBpv{%7L z2h4R(SWgG=*A5;IJKia zQ*f{{LCMGO!^%jhqupiR{B%Js4Yi{OlyXG4Af)QUFQ{=dx=8?7p?JItJz39T3jrzO zNDm=uuJKoJIT%0KplPgZ?XqD8s;Q!&L%~!^W?ok(;hJ%xRK~ZxpWw&D;Wv0poaefv zlEZbkgHQUq1<-$0zxbeO>07lh_U0rXb?66!W70e9I7*;i;l1s0x=Wi`oRHV+Jm@P> zze{aA(Gbn%nsh$L^|836r8&ZUpMkgc=UXH&&FBSq9AN)rX4B|=zzV4TiLe(YA#2kH zD+3X&f&F%P^tk?Ce-iqC@umOqon`gjlxtLT_!Y4*L7+8@M=N-iL+ockW`_m&;M7z_ zn~lb=r;rhAmEZ;Tr-zyIi&^(ELLUYzRui1et1Bw1(^pnj^rDa6JpcRh+P#k-0p@0( zC@2h#MgOn(>fIWU=e2@Ethkj3OG~dE56BOkFr;_EFUS@&)r>pSW-$D@JdV|7yenF+ zhHvpXiM4W~;CnwN=K&%ux!{2J9fq&i0bE&M%gh8D(Q_#Rgpy_k4T^f;bcD0ko-STVSQwms0TBmu)W2PtwLdks|F*Uw@Yzl zwNH!=PV&4nI_9UPq*Q8wWa5z{U6`$gs9dXbmfN`@SIEbUHr5^T1NDE;T&s-_Wktnm zh$reL-M%sMc*)M_)L884#M4uxEalT(=<*qo`IO_T4lS(lEXHQ)y3LFYReQ8p;g`bv z=^m!Iyz5G}N`yigLzkmwlC*X0GhuL$S)XBc?!3MX&Kt+)l4nW{-Vnm zkgGfFMB|&&tf!?0x8gz+d1no$Z6qtAC4?{6Wh3#FJ_ASZ&>;>6$kc03TTqYq3WPqsOdhP(;h$VV6R)C zTRh=t!29WO0N0O=1@2CU7AQaYhx0%w4Ndcur28F#PwXtU-nsev||-XF@`D~lClgno)*dxa+r>i?-SDU zG1JhpcfrS!Yc;C4!;nH6eN6x;%sBdTLpmOs+m46NakSM{ShUkmRWTUVs6#pgk|)$1KYLuT~CK-OH7YbMN(%k#G<;l4;yk@ zRZqo{EE~Mqbl8-NU<gf$7QI%|R!NL>mcigx!L@>|KwwD>O@V(>({gl zG~ows+yhWd0GO~O!L0sQ>=0Cri~zz^C1KN&um>4pGkEMkN{d%$>fj z(f|%b1hcF_X`LP~(1D|hxS^@VC1V;V8RxU$<2Wj}2ceJ3V_-l&5Tc&gD%4m#8*>FJ z35;G8&M+ce=dH4eKACM}8GcvxfX%|=n?5tt>%bJ>(!WCpx`eyf5pqd8v=Ob%PgYCr zPOEjyn-%s5ED4d0m3(?HnldodGL>u`IqXMpvGLm8`{xd;)Jia%SpO6gFra!8-RSY zsuP@`=|(IiDIV~-L?<()sT@bQy%HNe%ceX;Pj@I}hztk6KW1qaU#UtutGw{*#uF$r zhmJO?4t%v!vjCYy{i=8Ao$sk+C`xJC9u~*T1A2fOYXwOwReu^WO5db5P zj{AVsBo=KO#@E`Y+S$>O7lVSJfV9k`V(ZF((yo!rysjlrxDTDtM;W+d00|5Ih|6HT zizHMdTU2O@{!w&$es7gPE&4-;o7WSctb)EXOFIIJGZZJ1MFh-+KMJlxI8!u#${mb{ zuf;#Gw@NG6ex@{ve{FGU=vs;?I>q!<85o@hoOg#DD+9?$IvTYKF#;ifzu*AP$6mFd z(Qhh4{2=Cveu3Zy+*jLOMavCVHI4aSC+E%xkG)(p=b-d8rnpZEyWGGX3_BzJ`xbC? z&*{&>Y4!x9VtbnpdWwI`iEciQv!1bT3@Y48K-r9pw_EJ?`=&2jk#gq1s@UYCE?=PN zO_%dwN}OeALr|+Gi^sj ziWf8wiMR}i%s|Ig0W3Whazce-Yur{d0t?{J)752?P#+8MY?^P$9BJvN`t08s zEbo!7|H(t zP`Y9snNgcX6mu_smwneuu&kpY(?T=L4QWl&0%frh<|*14)+5oBRm*T*?h`pD;+S=R zV{|WiQuIPG4I4lpc8qSwxYk?ucm}qA*jH%|xT0r2)Fs4RnsE@@()bI? zJe>aE_@9MZF_4~8$V9-LGQFRg3-3=#Vw$^WOea1spch>j&al4qiQtV>Ui*GiW^MXY zgjK9_Sm!g?7OZ_<=D_zH@R%2rgzArDpG5v_bCpYtUheD5kgxGxE@f*+^cAxWf;g}Y z0lsK;MKg1rDZ);R2JIM;sG+J4S4CLnts>E}BPaagLDeQ@`4R2Lv z`wJ2)%DE)^5feY)v|bjKLAn^y4SG-#Gr8;wmVOAx`$f?VSx>g+#c-80FX|{p<)M`uN}i~cKi}@UbR**VW8cTH!3s?V-kmqTaS$wCCf;bo zsjQCYrDZf0HyML5Bu| zeAWa+VC{ZidDi0GA@6tuH~!Tv67AjNSedtN0s=&E~G1P@Z0O50M&O1-iiB z?}(V;zE`9z#Mju$ZWcA(j?jOvcLrElX^pIIdJMzBW6l5Gefrj#ILRx%WU;tJX}a?9 zpr)`D#{&AORMGZ~;4)LQ3CW zu01UhfRi3PN>eV$b8)c~zCQ8Y+SVHpBNi2<~m>b13FHbXio}JpD*O%iks> zrEAF2-Jd@<=>6^O;$FRiV)GuO>~rTX9DW&F82idCzEWsQHxRov_WrtiPPu58*^ENW z>V_f{Su(^yF>y-a6={251GE}q$PZZQ0q60DNn=I6UfP$!a|zWvKMMz%TJVEoX> zlrhOC#VZnmuhL}Xh6xe%t{&e!(5``!>8?&G9ly~NKjfI=1I2S55arrX_}=N+YcuVI zouxCg-5{d&--)Ps4|4B*W3qskOz0^nVjOnJpM$fe^)8e)nSFy!gT?o8n6$rkaBM{9=H z7aCGc$Et&MCYRX`pyY1~JiT{d=D;B-Q%P5Xz56S^2IT|~(tz}c#mE||viw8DEv2;* zl74cl_)#lE_R!b*h&`cJ?~lIX8sPhC@TI2+eQ#HBv1&2$l3|E1JW1G@$yt5LaHM}} zGL()YV`-dA))Oj$lEg^IcnF%Z5%F6EC^lI<8}J|_MX$@g4s^Z8_bk_Bp*A6d zdrvZV`(E2vMn&o&u>c4ZTIt%2FrOVYm8MZcHzPt3&BsS|_}vz?%Bsz;T>0`n=(?p* zHO^AYl4QF-`fsuAQO?*xIWcak5R6h`*n!%558eX+$Qu3tQvXZbv;jjrNO-w6LZ;?U zvl<~{hTa}3p~I6A?9h4gPi>~+c9`L}UP!7q``W#MUI9o=S@SsJ#F7{~r*4Nr9|oTV zR1K7Cgz2I#fP*UNcy}j69k*^l`#fgsPdO%kyWIsC)u3B8Os+nRdDa zd}R(07Ac5nR>HM+H;Yke<;W<1bG6RXFShgRA|HqWnvd_to`hPBCK4kp`}8sI-z)rW z0uJB}fmjMvL2|8tP36dZ+XjMpo{DN8;-#yYR%Fg$cF&u?9V#>n@R8g!^;h(1u;k?n zk1%u;_0!f|GL8a=qy|&n9a+61bsjmbLf#tj__Ob8x72)$w|B{+^NkkM{qbLV9>?$Z zmv289#tSeNC6^s-82rpTV+0TEWDbeD<9djMM6e&!oz4G+dA6|AMZ(xis&n?KT3@+> zn(8YsE$w=XOVz^u_j$wuz$o}Uz61oK5qLkMcHB?(NiYG=VC^#NSrF8HGwWc=_T<>X z6PIO)ykcFd3V%bXZacIO3`Y-O(o57EF^Cm$my)j^rm^6f7(5BC5xysgRm!;Mb4kJz zx~CT2Q?}d;vH%INlV`fxEl0 zTDQ+w(ywxHvGNYW$yaiNDY60!!~vBP!2YPuyh(K+=}tuI03H?A!l?4A(Q3|-x5Miq z=*J!G&L(+A2ab1J@|67$CuCp&0Qw5de~e*@1CdKARj4IoAH9S;6b0l+f_z0%-%jGN zmvb(dj<)U}&}dEgrI?>5;|uii5K9LSfC)H=I0O>E4lsEhgc6c_)TB*XDXKCzz7luM ztX}bz+*0zbOZ*>zYM`nwwg_mBcQf|_5s7(@S8jUl zsrK@hdSOX;*%NC5cO~u*JD?UHCW;yQ(~hDm=dAfs2xO5Y^3f1I=7sURFYCiUY(4FN z555MXC#OzWa3orrOy0dp)7g91cn!mp+G&9r(|Z|;A-y@3Aefox7{-aLEC_#`dsn(V zeIjMvGHx~50~*f%SVj7CR_Nn`=nSK4-{B1&XFCK&i$LCooKb*79f4_Wd!qu&f)w`# zn3Hvj)&s?>RskJrZ3sL*gfa*et<0RT)Zt8<~BwaZZN^_BK{*e=hQ5_8fg9Z7^ z`%K!doRZ1pPTkE{-^yusQOwSk3xS&A6#f^r?hmzE=$#DdMF}n@tRQlnV3x0U6mLEOz=V z!@?N2f;!N+oOH8cAp08`cYk8;Y}VdkSfnyOE%g5i+@vwncGO^y8LVQ!nX;&-4uIR& zgFrh0wC`QZK%IjI>~^>T#-kA9juq7mb*1>kGEsho{eJn&0S_BI{)p}~!GAL4-*tJk z{7TpvZg~^*KCC^Iq2kJ>3H`mY<0)Go{DbgRiJh;?NYLnkHF1o#yR)||jk%IYC9X>s zhPQs?`Y0U^<4EOueaSiHQ@7$L0AS?)3RqaMbPY{fG+mbw#5_S=>f}%2r8~7jdzfcj z$nTDf{8k_J_IX=gA4`_7O3TJM1{_Q2xdAg+Qor_XFYxU%?NsJ7Pl^yIg1mc) z*F~h1&8Wzy$}jKPg)I`j<)kQ_4}bq#e)|uaP3)%)rX&cQIipx_nt>r|y{dzy)B+J) zycuPq6^urDsjqnQL;UjtZ%qqzzrTb$g4O@rGh)1bqQlsKH#)g=Z1=EE4kz6h<_}<8 z!>M&XbS?K`oWt8T=A3Vg-PBTH8u|r#s2P#k`Z~F<(_1@ceRZEnq|aa1J}#mOnCQ}i$~rdakI&Nj zz%Y<#C`2iY1C&fdv*XK}9A7$xgdJ zwX&b$6MOb3J|(bB{}u8RKpUVZn0tS9P4kLw_-qm9^))^M@*Upk? zA4kSJ6uq9v9*r)wfcp3axg{qS<_I7%djwC4dp8F8I2AB4qs1jlDcV8hWG%Aj<;X+M zll}DpyQK<}Ux>pcpkvvNrEBKG0)6xWFbOgPQEHo&9q9qaaqXEdo#`2Ge`XXF>J=Or zwGgUoAR3jw8@ku=2p|UKfFo-In^Hok;@*KZa3Ua}sge^BpeIa~iv(p$LC+yp?mOj| zfD!6KN3XY4Wz#XSJ;+0m`Uhj-=%T-%$rrW+KJNi}WC1eiKb<}R{!N~b*Wo1PAXe0H zer6*&O7Dj9XPzCr$59bIuli{*Ww7Hng44ByI(PEI!w$H~_qE!F3_`760q$Ww z9x5kOIn)s1PDFCu#IaX~hZBspVU|Qm`2<(=G0lQRI1J8Ea~rvzRKY@i?2@wzP($U zUXwgK7*W%I_uX9wUkR}mtddP=v&r4q}F+*NTl7rv-BCO@n=qXb0fGa$n%xh(J5=_RJzk=E<t7yvqHbyWC#+t6`K$us44{$z%`ya9^SwK8 z6)cau-5s4 z%;iJqBtJ!Oaq~L9xVAzmO_5t{EXm8*PU$G{gc4Q>6H{V0-_9_?V>3O)hDm};9TPTt z8)dlI$jMV^OXH)&qTOLUXq(?Gu!O*ztd#74gN#?#o)t99i#HkX_w%Ma z-<85HjC>)e7Q0*sq_Zr*Xs_GGn(rGr;cDfVH9TjD$I04zoi2r_y+`UlM%X^i=y&;= zmBH(cAm z`4wZo8CeYekB?ExIP|>~qglS>9^D4eIC*a9hvC%m@TbZUCm!PsgQc1@8)$KEgK_ebs^ijpTymC?l_eT+z9d}OsN!AO><%aRXc=8Ww0me;h_*J~EX*{b$CTj>D9zQRnG$Es=#3+QhKdpA z8e_BiOFRAC%|pAR`EF0oaxuo901?-Lk{+QeXNn&< zc|KU0sj{PntEAcSFu1F+a9JD;l;JW2J2jss2fsd`+Qx={^QgWd1%Y7y!xsz7LZr=q z5Y;&>$*3-nnwCNIZy{TNUXJfC^)-Pw#xkQN2U{3tkiTbnt7}!em47Ynz7}UHnjo?k zxp?+as@CQw=#`Mzwbuw%$p1&qT>&g#EmN3%!Dr2d8cq&qw}^qY#+nbhPg}<}qdzoH zc%wg#t}Zdoom>;VrY&+@wRuCSLinjL=)oXDRxV?^E!Kb)aX^$PGPc01b z02J+=YYqK3TXAS-#F8agy%ftS<*8ArKK5zE?Lh6Z(^I~LF$xCK*TknZE$vVI*$SJ6R_t2x6PcGJ`DIQORWphDRRJERik?zd`6gZbGaB-;fp?8U|zBaEEYtWy~p++HOXc-gfMT)Mj(x>HK0Uc$_`XHNwsgl{|9p0~7ej!p_H{{I; z!x3OgZ&$JMs_^CK&xnZA6IZ&25+H zk4^P~=s4|gv|87Ov`$agT3^1T%6dK?V6`~S!B$1NLliT4kPOXtpoD)9$p5F*~2zh8~gk1BH2uZr1j*4uu zBW2Z8j8@I`&<4EZS*l{Td;6BpH&j%l6!UI3!UEsuVq{R5+Nq5Dgsb7O?c0%+J`j=h zaK1zQ>u0(3szx5kzzo6qw)bIuxpIo1-oL6uK*q2t)VdxpZQQ8@WX=d2RUe(&ME-b* zq(Sdj?oH1*8fXdj3Z83I?M5q;^<{Lve~f~8K@X4nPVHJF!Hg~^Y&BvuhT*J!m&*kBhK&l z*vO-2w1pr(R|VVmE}DuDuFqS99RsBEX|WBe-%PnB1y7Z^{x9Yyg}Ox$de=yhP)lcFjg*A~f=!?AEccT0HU-lUvbavZ@%RuNxSFnMY zQg;DgHlgNZBPFR7;`hRP_>vIm+q3Yyce=49BX4@YwG}4S8=U@@cuHpv$i^z^4gk-d zBhd`NgZn6nafP_>j=95)W>T6LD<1diejU79nEN-i1(Bl9F|pyslon=a;MN&5b=I?{ z%Yg6EXrmnzeAZN|%J$FAL}cnhdxGlj*R3(v^3oC(Ytm0B@4m^bQWUpsCo(}46F#Q_ zc+NzUxZ|3+eryLs$~K-GU9Xyb=n&WF&fBiz+jPFP#jh&)h0O*?%kie33o86S8#PL2R;NM9V5;^)!a|vZd>#D*p#ahBe&-Y zMwxi1%Qt=@e&Q>Z0-Xhr_{|xZZ-iRZf!d>ENMY#+kY_RG!EP5BoGVznIs2|Jd5yDa zc4+IgJ+uk`CSBJ5kYzlf1SiSYmn8Etu>A}%MX67XKrlXJgO{L z-4N4dwSb$vm$&9}->WpSXfdZeqlF*}e<4&GhcYFI0Q0k~cZ7Pftav8QfnWVo;=YcQ zQ|JhjVfBE3Tf)LXWRkaRy#`R52M~y~vPaFzjr49xI9dHP*K5s3YgMPJcz;Ka1Xq8g z)}^?F8SJ^7sZ+6zI5|FTYRyi`>i1c%{f3<%1I-a1l~rx(9jNt%8D#?As|3L!yMRaL z-ADNYrC_QIGni;&%lXBnfaf%exj@@rP?Rg+N+!IRS0Xgbr$NxwOWeP}vEe$;6rlR& zbPmO%_PwRRuaRNOTPF7w>6FsM;=C|kcvXRdp*;gkobJ4JM@*9)MgUSaAzh3<6#(g*xBM0hl(a^4@ zn^dr!A1Q(pZF;Ctjh{`L3keb()9-36QZJ6&@#ha*FMWH@AkRbm_bpAe!;tdQeQS7B zfeLzyAa>>s?Pb`)lR)^~X<_K`gAkz+eDJ%OA0pz}sqcG((c%%GdIbtoG!roQhqbJ4 zm|kxohI-C|@2y8se18L3JM4n%Bg%X`0l7dOic_gt9M=$n9*&=U02<$Uhor3HdIlsG zZy>m@0eWD=H-wi(NQkb7>~h^kRpZ$xZPS0_mJT$To_Zo1;`ORMenD$UtZWW5rAZ;AON=KZjGKp5*7s7ZJO4anGBMgOe$n!$>bcMEt>`J_| zgOV(m2ndU-f9(7ELCgBxf>^}oH-sEi44}^%EWb%*w+(NDFzo8VXWNg@Hh5&%r+-xX zINO0BNt+DIUcGf{ht2H=z{X@SrS1T6Q$me$6eX!Ehrgsr)Svtfr*5{UVi~j;lxOs$ zTeabk)zO4VvBev|!C=7#U0aB#8NWlA=&Td8Am9TGe+~(X9?!Jei2QA=Ki*LFpus%h zWlzE=l-zub?1G!w^E&rdBRV1LS61sB-?-jkT%#EU4$of}AFnUH-g&dYgJ!*^JMg=s zRPgAngfngd4*;In6wQVm$pASmyeA1bf!7SIBVLLEwcFx-OK`5kym-Tacgx6|oK+LY z&-!s(x?`*wLUwj_3gwRwJ5=9KaoxwUgT z?Sf5*?De0^9|*mu3hFgqhbD7c#0hQCHdij$M+BF-cBg9KH(`+g>MMLQgag;!A&sOX7y41iWzpcpyo}k5c8P`I16gl69TRI5D@Pg7OTXpYVazpU6BRozh~bw-EQre zDK(K>i0*62*sn+aUzJls2t$y-5dTP&NRXR9x&aHB>0M8paj1zi9n}cF?V#QN!vAWWe>-MGu1^Rf8J;w%0oTQ#4H z{klqdcbQCwh2#T2nO#H6mTuHUST zK_|ot<`o^|=14p}QE#(T`mqB)#x)6q6UnS(3G&I;ufEc=tjGMdwv_lU7;nC5_?U9c zMf(1WzrN9=0kVlpot@G=pYx{^ntl3*GNgFa+5SWbGk;}Bhq4VoTV)`; z#sLpjBntDA7`IkXJa{om^L$rEvryvE-!rpUEQ5F!>QmI8J3T&p%70(aQlw9l zFh!_k`{X+JZt2xDEf<-CPkdW7CdRR0`0H*|_SQ_Jj|J)p z1uG9O%}+lU^*oSon5?QTRjjpU|CH@l4pVr7BpRcqyK_K80GV9twt3e_vd>=n`OOjY zk+r@UAK7Ec;@$^Ew1{r0E#C<-!B>o%KL6H^Q%;VaAJbIm;G-gr(d!uS`$Yu1;{<3PY?%zJ>hPp z_uRHYInmik>6O1Edem?!zK1TnYz%#zFxhq2H*<$3Pz_7;iQAbE=R{5D1iO(~20om0 zpSFnx*WBVeGmxj3|J~2eZhSSFmErlicI^q*MK+EdX)ewVI{yMH6hM=KdbCLMaSAQ> zRWCV7!c$sqSfCB5a%DuNY#jt`5u<4-nnqn$UzVmA%{XdW!J!Z_Du;57ak0g1MlJ@4`J+%G&%tam!SV+?-H?^0x0BW;vvTng z#H3wx=rUC^w&~1!LxISbu334{-u;q_LCYUDO5Bin5*1@n?_bw!b?buFsj;R(0wn_@c2f>^T{BYCR@S$^^{^Ri0S)p{OfvRA0z4x}Ay!m~ z(hvyod~buci%p9?1Fzp4jfmrFg8Vq95e+p=fq~!3Wc8snQn6<$X9Y=ILgsB)jC34(5fy z?|BEBoL@8+o7LBZ?N#-bdFTztfu02^fcEo0g2xfsM|u=V1%THsQ|6gEcRTQwYUd@F*{g6pJcgw_H1BQ$kVmC4^?l}g`#+Eibr)dA#hE6_fWqd zV49t}QNJV>CDvVnfXFoorg$a>xz%?##Csb{4{;;X_bn;ZB4ESk$ z{lq)9e=j=KG?e$qr3iL8_RPOg@Zm?S2+(eAI136ei66nZ>Qy!jQxZ7nqZ&KhcaGZ_ zID2aw1{&y%wRLnKhF{3Q%3&0ed5-Qx1w$PNuT3tGr17l| z)AB$K+Ov#GC1_7rSxur02j+nZJrq;_b1~D9TsMq|~>htPv_-tP`)TTO86(?_*_ z3)z5LlXEsDSp|euVE};rVgE9C7LZM+z@3CGIXXz4_=*Uyd~*XU97D{)@L-6U;FhU0 zDNwT7R^HeyPspyr-^s2GADf|&xbL6-KCt5E%Ojo9k_54ewCC~={ZI}wpFsn@v3w3Z z>}=Uji3I=vq%x$K-t_!ZT%#pLq^#h!-oZxCnWA91Yjx$7rVi;R+QoPe$SE3xs+vET zFUK+^eld>khz5wS!H00%tzh&c&9;RyHG6wIZ#rc%y~{bq1r4E}ad1n9c1yXlubo#| zf@VOymB7#iYkJy$0BJ_=r7DB-;<%qg{DTexwDmlqc+-%5tY#vO1=`wB{V)~;xT z03IB=efluCv|HB?fsoMHac8I(!B~#{MVRuTl>;)U+dxf~30j_=_E5Zz+&i7*v%$A3 zjI7Ujl}RR>6e(J=kD;yZ2y3-+Y##y9Hw4UIYX@O769VxBZKn-mQc}c_FY7I%E4cx| zDA|OiHB*;B#wA=BQUHIkc|_38=?hdXFfh<;l4w>vLVCijb?8E)-BaHuYCy@-5DLtx z>}EbevO|-EoBts4zD$sjtF07In5VxG@RkciwS{QZjp5CeOLqOVe)xC7VqwJFI{>L9 z>)GgBJMF+O_`&k7Hkkm@aZlt7iiBO|1QN+YP;9f}aFYO$5fpv@hg;^cBOY&44{IF^ z+aKr>{C=PUeXG1%=E**>Q0{jb$$Qy_AR>W$6g7#3a!P`FzPSyGhHI+sEp;1 zOyk}6`vpzQqM-k_!uqIS8;UX26Fy%S0})o0z}?NwCe%Wc0HcYd_`rM;#9*@yEEm(i z05@UwrL;5cZ|od-m^^6N_|%lM!}rtscI|6#$ERmZrze@&z`8x!;SvZt^Amt6Rgwr=gj<5AK*KT`g%XwDsWw5qo)F7H^)n0-rm9k z7@8Gc1>hBf>=}OllhYvt9QRhV-d?Cx!GSjU*_D-}V}^>+Q?i?}@OwSA0zBig@RzeX zHT@8WuWsK#*GrFk3ek6u$M^sF7n{wqK@o!adxV~VwynlE?~YbA|hE)a$x3LA`^}P z;O*9800!v9M2mOFsOWH)di;4fJ2PkQ^Vajv$@FCH7atOjXkYfrd3sl27-(w{=m%?8 zzktKONpFBp)BT8zP!IDm7@)8HI97ekU!=mCF~IXQSz=%7%|R_tOH_dOfJo};3e)c& zL>)JryX}H!!%mB?ZGcCd0(14mym`T@e-HJcbe^dP%|+n(@Gv*%yb#0Z^S%w)M2?DlLxlDb!C)o?r*}Qk+ z$;a$mk9sRwN^xSM@Itcf-Y-4xm(OspUmpWS53C|QF%Q@#rXZRC9oQEh%}$Frai2)0 zZB6oA|J(Nb`DGxxEc+mk;2dTZ^Gu)>E&t8VLBh@yv7JfC&Y|mJU}I2XY=>K62#kDu z$@C9-gBJ$wCkk@%nHD0~gcEc4n0JOZbWF*f)gMNbO1E@vE+&2_!HbT;1I2fXPPG-f zNjy(8vJ2I#-C-_kJ=qMdz)U}5F2VoFfjUfwSCu`xyVA~*)RRAE$wVB0mh3sZ(u^toL zv@rrW*C9iSO8Y$f&ke;Cb5(~KyH;`0^zUT1gZFoDC}ek}#v-;JDiEDOK7@hWmVN(0 z@DHPyX#Wk^pzIBF3rjyH+KQqd?T+#vxI*ks+M1Vepx^!uXE)Hr0>36dAIMsW#4ZSy1w^pe#i&Pr3h@ zPR|dL%Md6&s%o2^N8?ldWP%XXFN*<#mj9=_`q!9pV`sTneHj6GG zrfbv4#Mfha*Mlxnh!YYQD?YaT%C#tW==Cl}>D@BwcFr1&aB)GPKn`Ju`D+HPegn#B z?54mVJS|&$KtO?lZz6#gL(DK`xY~XV8I(pRbr3`px1g1{1o;U^c-%U{#v{eSaP&~3 zX5H1^4wpPz<~aN9N|jYkpg*w$Hnj}u(`c4}rY=+iK5oYXYUTc;w?;zkN&D)()q= zt2RnKl<+AjJ?oqH@hLW=6QJK9zOu1F*O-8>lu*=G%AjS;JHYju0I&!fll@)r@iC3( zP5fkjv1-qH@44yr4jXIQ2etBBP_j|U<1;TfCD|T1wlX-O*4>~oh61?N4#tKUU{m%Q zBgCpGg>+;lzAH2<9B|n;NM(DPbF@zDR-Gb3tm_N{b^}Yhi2Z9`+dahKn}U8;3^t?6 zmV;n^&mCRHl3>P#p8@mL$^;)$Sxm;Xea{smctkV6y0xbJ$n;MN)EnN)GrZ<%rPbr8 z=fQfRCjhqnfJ(D2T@NLIn_ID@?#`@k_*s)logjlai$6(8&a=ol^w1sl?+s<|ULj|E z{ySqvpwF5O;|9RKhBotRCNv+z6`4|0!tWG>UyKaAm2wtcAys~cv+3w&-^|&QJ;_jP zw?sYLzULtzWQl$S-GsK4p8)iaOc}H(jCg3n6t#dau^8u|$us~+eH9BUt}yE9Zd6TD z?b;Ya%a$J3{xC9%%1gglOU_LPO z&*9GP>|Eb_cu{hpKu3aH1`>){AJ+OK_~89I$>=WkO?_0-Jka<8sRO*2-~jT5nzH5h z&DX&72?6UO0VBX))#{g?SLv5YVpS|BN5jaa%cDNdghty8*v=UJ_^Cmld|fZQuEO$Z zgu-8N+rsCd#pNAT)8sqw3-ZGVxyyfl(LnAL=7ZtqRqExqZ$8mH-7dyc+~IwPEgk)*1Un@I3w&XvLAMBr_7)u$Ty)f> ztd4Dw?<4u+8XnaL?tIhg=#_*X7p6fG)U;k1?o0#|pl-qkpl^!V(m?3Y7brkJ$i=2q zvpxVY6hwT2VgpCLKvP_JqBDoL_uX=c)b8|$2@#9wqG7Tt{QY9PsY~DarVyJaQn8xT zP-{M++&_Qx9>{Sd$h|tX$Qq0N*DGrF+%;N1v-Jw;48AHD^G$VkPUQwG4$JV>2u`+> zW~1zhqm;GUVve;^{_%k!!_h6~VCBG~M=GQLEyKY&1r-U=eAO@UCo6g2z6_5F_lnnIin{&75OQi3I`8cIlH8vAXIoDL05_wU*y5V^;qr z;hU+|8-qM9OU+$eXOaDYX9J=o0hq^g3Rrq}e+$^)6>qB0fAz+F&PAsIo_KSOH_^TnZIs}r@x{`lhZdzO}&rh5ej^@{EdqG zy*k9;W2q7Kn)1H()VgX!-Xa$`xuWf{lAcI}72b17#ay3B7-oaUxEk2AGPac}24^^Q~1 z^BUe(ABjyH@e)@*(J9=#We3jU%a(pg&P%HE_W|o$7hqBadvL6P{tvkCH-A7|vD1W# z_e_w}hjl}(@3M7R^#pIEApRO1OBHH6?L|n`X&k^>P|d!cqXl?~p9>irhbSiEvgfWu zwI3T&DZjMdBkAZ{pm0{!m$r@#LT;ZDVf%K?shm3UYluTo8U-Dt7q zao^6TRD-j2XlM~#^G@tZA{PxIl=|WZoC&ppyomUC0!6XHOf{0SqliNX&XRHtFc1Dx zs!}U)sTw#xOj3FS)tRNoo=iT%@5!GRoqw-10ky%AjDTdS83n+v)nTYT0c%#^$ISxw zVXB__YyN*8zp`;?G8g7C%!3Cp6k`=(xIg4DwKTU7i$5Q`{HLQ~?xmojwauRs`tddH zKhwhw`S2xaE~R?buGc<}XNUfBmyLlEl$sB0IsXot!i-_I-gY$x0jcHyTC~|z`V2X$ zHPCBp^4#X^ZQsVXIK?*uVJW9WO}Oh|Rt z^XIA(OnanTxmjTLDcSUZ>;qqW8_;}(2zK8i6}v(qJru2iqu|~X zUZqBDY%I|ozKp-nJcwJa0%N)dR)QwJgUU-mmOOlJg#m$+lnu8%0I2DtNvIpets!vl zSGy;<5NMU2pVLTPt+oEu?Z&FVUmOQ~s!$LyPr2Pt*n$v}@f`OO-=3WI5-GrUJ9il~b zOKkE0n4PX{mD(iJ#0i665@uBu#&!W2ndr3d9$d$Tt*oH<8XF>EItbof4JM;dNptd} z&EiZ;dXQY0hLVBVfFdr_p~v4g_sPbIL8UcES!Y+5J&(EhByNjLzCxt~8Zl`8?|Gk9 zkQTYvfa|y&P+hjRNdN$lObo?cvm<(_2EAPirxzs4Ht z_Y3Nf7k^fq`>@vT=(k1^FWx`Bl+Wz>2NCT4pPcbKyzx8yl9WF%5_s2CC%P!GlZJqs zHsuJURSmH~VBGqUKupTOQb2$)m35Xviy&t~{J8GchpD$8DrPKSeEz`svSj(SJ4R_n zt!ekUIc^kj^}%`|#uozG9uRB)P#6}3mpkD|Ft|Eg=3qN@;A7C#xeP#ATBE@@fjt1C zHi-BXV6YBDPW?%{mK|j+_S17aGNHh$% z8Bc8?Y^V;P$-W0k4;j-AnMCpNZ>&_BV?#jVWH)Nee4A zS{C{3W5QlwW=_>mET+e*GbpPQmCJ9x-s|_SoX#^a>iRrrVm?#TJe$BeIMbq3vXJ0c zqrs;(!5)hU#)D^YAI+f;%XR;6Ld3HAqHch4fXchTT5z(n5}JZ30k-2BNaQ4)-0B5E zRmBUYE*0sgan6anDbLH70q2&T5f0&`T*x3&P&5Y0!x6ns}486)O7eHQ0-@dJgLR@Lvx zD_dKi?vahlOBE9PU(nr(Rr3{8WpNreQ!MN4nW<`v>iuqFNIPa<6U`1ny17)!&kQUv} z;Q*|`J{=n%9%+*y1zQR@l2NoDbI;X;Ct(}+@i`w6e#wl-{ zr9(kL%*KUer|jB5AViG4N;ImgO^^wHVAr>PPxJ;q;giLj{Qn{B18emvc^=%RLm{1Zy|&j zQjuNub+U&nWh;s?3R!0=wKM!dPbCp zd8iw;-#Z2Vab`uh8>o)ncqEHwn9^4Rq@#D!m1xMc9y!N%~n#f{uG>Ys#(mRtZ z<@FOV)DpN<^i-!n4#DT*OAF`@E`t`G_1#pZ>`rY?K;p)}2Y;>Dff;8Hfcn#!-4Jyz z0jOObDFT<$-|yf!Nl?fB*WllL+yIK3E>EmHWI9g}4uPbQ9;0ZdLGS7Gvg(-*9hZ|Q zviPksJ~o?+rGKr6(le7%8yNYw9nR-%1(P+9JLfP-sI@oXqWFAL+FEHp3w}TrY^!P} zbH6l=&v7@rm&TTMJNA}p$A%a`6ub4-VrM>FBR!i?Jd2+-6#@R&1sDZ-B~f+1zpeyYeD+E80+A6LmX zmr01eUF~$(M&4)R?mqjeKQ`zQ6kSRZxD+(FgMz1az`+#D4jc~UPO*-B_33Rh-CzVx4N_YVqFg_89L}Y19sUmWoVyPi50z*j z3=K!(MwzOBSfBV5oNCC!g>vh<5nFyIosu36A?7bHdetk%c1W8i+@{!%4T3 z2SHO{%JPG{q4*s;_36;m#Q-rg1Fvm8#nhMl>N;XfPiVfvKnePI2?oC%DY*n(-pZc@ z=7BLr67DE&3L;xVNN_hgo5|lc<0|#INX1+0=R142r;@#?PBEedu}0nwj_$H{Dh)C$ z33k7M^**o#Gr*8b-34|gUo(goMDux<=%#iH52-~g-!)l*mDiKPLZctF%tDcOG3_)+Y9E^V9Pof zQ=6l5Uaq12xbgnjJbQcgWR={4Fd6QJLnEs}|3zKx@;$RZ$9+cffk1IH0?ig8O2iBK zNYUMja*g*=L0tVfZq5J3T}ytqavWfS!b9~VJ(h{kKLbW0s3ttO1 zYNLdNdbg1j@rmzNG?NEdh461Y@D6rvJTP&-Y|Fw zI@$vX4pgiN3oF07=FXjZ8-XhH`R`DPFat%P{G%anHBnqYFd$}?akD4M2dcHB>6ZS@ zI6S8KO#Yys$kV zOhyB#6zZ-b)K>Acrmp*`t#76j0x@q^y<5p*$`Qa zBzt(Sr&BaYBR%MN^Z~lf0*`NZp72P@|dr-pD$ zJ7OhWbcn`AcZAs~C=#I?Es2$OC1aw4AU`R|z5PKgrwkkG%Vh3vI z)ra@QvOEu}{HacOzDvl?+`kI?xH8w5<_WRwd@w}9h+-}6$ZZudY0Ekx(Dc*EDcfIP zb<3*Hlr3`(ys}S_;Cr56#r|C~dOeX6aG^bFkv9%`urPb}Tn1>=146WZ^ECUu3qeAR zr%;Dk!v3=(mmA`&k9BA(U%5?8XE(Gd766l=QrFoiIwX{lS%t^4m%kDAoQQr`y5{Xp_+@NpCOfVBFP0IHvFgFvu_z5zDjD{~et@&^OR zio1|p{S@{dv7%`D{Q#}Bvd=!*j~hn6zEgRfb^C3N-%adzl(SS}TSyv3gSp=mbO*}A zkIJyVS73R2ib1IH*_MaqKdQ(RKi132qQzfiD}0eL9#>E0(4+#UIpC18euCEsAgqjP zKRyA*;|z8(#mGF})Ag00mB6g=)4W;ste1_O#?hTrk(Aq|X<_C)2=4~>l0A3!K-)P? z4kO?lv%xjsY&zQJYB7nL7@=>zuJW+r&C1HiAf9vQ#wX+bUc6f4k~X-|As3Livva$R z4aWi^xBALnH9_C0=G^Tb%a~x1>BR0#J8-uTbk_Rx&6?A@-SnO&}sD*WwcrFK;9kWFHF?@YO(X9`5U>kW=RM4b!X3kXYL|jF7BJDw1Cx3;U&3lM%S#jaAJ_AK{r$EvJ>E&;=(v}> zeSk|r*W8(W7lA_jFV=fy4A6}*Ueo<(Hc*GWo>~Q#=G;UPt}9%bnikZa(i@R2Z#wtb zjNy`gNn6Vt0l8288v9?Y_HKrrr-1PkuxS9A0!6yoJR#-*$a_-um6yz}mKJ@0KB<&% zoIiQShPj`!;=GjIVF`(_J=`< z`<07+&eyrY`=M^P+*8fo2XhH>3Yky0JYP`urQn%j3#0~y;ARXOww{G31B?=&>ANA& zQgg^a6E}h9VJc$f-{Fo>@Q;Bm7l`kyX^*_AS`g&NM)0@`t>P4jDWLS8DFc^!MsKO$PF`)vo1BrwjDi3jU zh#5sGJJh85?(6`TxJI}f(7@@EM3{oXc=}yJ0cM(uxgSHx2kAB(FD#9WY#*Z-cRti) z7#oELWz64Tjy1P?UGX;_t{TpNk@j0j0g`7yg#kPD|C(|HqZho+3BoN@EK2+uJVdlS zgl({=Xh*jy9mp;5_I@)wb;a>wfNbXnVLMr>_S0zZQ|vN9mEU~mr@{A5LDw;HgqqNX zUNuQ9t)qD=A`lWeSM>9NvrbQ0hFycYlYZBE=Giz&H%Q7S?883Nfm5QDUw!@uBsIl6 zKw^~W-KVL6D|YgvH4(*!#(FIDNgZz>ZETU>Z$=xVBcXroCj$gH5xx-dIiS}w5C!WxM)2m* z?!iGqE#$m6loy>VU0t?119iD$RaMP(b0;rpcSboFvPGG(E*z3!A)EH%8%%pKF$SO@ zw(dkK>gGzk;?R;o-LRS3DCD=07P%|reC#`D#^Pk`1wZ73@~@1K2b$&q8TkmK7fSN; z?8Nh_1nKHyEGIBO+T()iKW$-9VY;J5gq~hu1B|c zAAzR-Y?^6H#xQXP>quP@khx`3Zy|9jRn@3VT1fhdS!Fv2DW^cfQf{e_hnLCug$Tts zZg*LVU@bdQ2iAjh+Z;hkNbKrgM|_=kX2$J4jx52!BN;(5s`pB&Wo~n@L!eN?uELx} ziqxsWO3~fP99!DtQViXA3UxO zVYw*e3@}znzX3m*-R2Tq(io`!Ex3-7Och`}$LNXAs~GoelA|VFG`jCSIB9wRK8v%{ zPqL|u)r-uG6S6g>f~*zaW&FO{k)!|Z2#%@I#k^@(;onr?3WPgAU%z6(35n#%wUbQ~ z{OKyMU0o%&k9LYH?I$If^9~&`X-Y9Zp;CKniBsv%Ea4lF4 zo2Jx@I#{y$-PrScbFB?r3PE_DjHE`Aq4*FKL^19yE%lmye)&v8cI$S=uYJ0bGEt`P zDWalJ0jk^$U5zK-FIrEbXw35_;gf``@WEM; zFO(}*VZNiwDgexq;QP@h&q$wj@f!z+<=;*bX@CdTL;hQW*E8v&xipWEOYkeo2&th9 z62$g7QRE-52eMIwpt|}ucehiWTti=TaJ@|BUnv*AR(wkK>cUZi7_g$V1SDZOI^-I- z4%*gWMw%w=B5+^>aUDnz1GW%9xQHU%C&fP1Dn8$`&K7doo6+Pc8~N?q)bx=?zbEHQ z7B$cPc+ZuAa9=_I@{nDG2>9Uq!J@{2!_S@s-Gl(Aw&vD9h_Ikl0?&%;F^VicYO~Vs zz~G~sdLowUfG#jGXTV0e(eG81iSG`m`uxl0mZ-rxD5@D{dwn-gQcd_|Ah~QOB{SKo zLQ=Bi|N3*1pgGa=s#Qc$E$z^vgUGP$XMzUXt-ZH-e*wL80Qz@cVm2Ap;H47aWDqsa zYxnW8fsR$jqPUZm(c2GHW}_LePBb>Y)dmAt!WIe1j(}93K>W4|pR(%(^(TCss}>vpQItCJ#760$)l~sbwj6KOkRW?c2`g!~ zwLY=n=^ityDXS?|Doqc%$L4k-LTzOW0w#-mQ}7l4)&iLqnYoXg{_7X^R+Bs-Yv57k0e#zL*+ z_hX_eFFo$vE-ndE%9KM*Ix>L}61T9=ppMHpWc%djInF;@JW2b-+N)DwcifFIK8U!5 zo_==x8xa{M8}MFCu0rD@sWcBER7Gg`+Ny;j^FXTpa3O1Mi-*i zi=8w$QZ=roY2<&(wlS4!`GvZc_Q9CF8#9?RZEY>g{fYov!R{Gi%cy|-li5O4;Bru{ zpI>_*pIk9zQr$|8 zs1oEKqxfnG|EV9%`PBGP>S9lczT@o&*cV<{4RB(8qySg7l6GCKlQ+_po$UYWVZqN) z*VSe`6h~}oj+@_bGryiRt?L_e{;~0Iy@Z*pDxov}P#hfv9ENh{nDF#gLUv^r{!wtY zb1+}skG9uYCDUImIj*-!r&oTr|5{-fARZ~hU$&1o30(It)o-B=A9~=2vXZ}YaaU8e z38P0P>=xS)kjosH91W{R(@&u2{DgTz-#-vy9~MPJ+rvY?ulwiYdAtCi*1P=S4NQ9F z*~|N`r6}!9IrSA59}nsJgBU~R1nwf*vATqk;z5tWw4I`?{WaUN` zx$3sv2GpZ*cxfz8iTof8RiOq;kI+SKU|JJJZOA3#lQq1e7sjU>+zBcozvzm^xxpnp!k z`Zm$lboYk2hXtJBAfy!)lytnCpU-GoA7DyeH>3bbjD;F(KYwc-61cKIlD0(KxB6%m z9VriRSYSXz6p|ch+Kqu=1!%qR6hdmdwF#CpR{n zRvQv!<=E<*ADi-*7Vm6pNZ_AdLk>Z(X1!J#lj{xnY)K6(m2Q!2SO(yp-d0iO!eg;Rpeqjga8#-n{1BA*nY>vU>-)A`ebaOcIeolO^$pMWNhCEiO=QzC7(R``FPommsRoH6oZ0e&n{;PI- zs?6SVYlK6-u!xOn$5mmEQ-}j~j33OA^L9j$M;(WxyQqkW8HSh+Vohlf ztmMaw?^tB$lu>q)&^BGFnr3DGP{Wa~4|DOfK)QztGfgL6240{SJo7nr=uyIxe1T#1 z^E6PO<@p1A*(ALHup3s47l1>cfdqxO$``Iu|DMg>NPmNM6_0&%O|d<=KU7ex*bnD}kha{tkE_Z=FobWT`2n!TocW=#k_7bmYhD2Mrb`R~rO9bSyJk*~CKWxRi@%=pbqw5|hNOz-%~btv zH28W=I}|o0Es}0@jD}Q-t4=x#U7JZ81Q)*y*vNa)Js8j%xeW4&bCD=src{3cyqEIF zX=X&&(bd%jJ=@)Fp?qU{Loi7yuB7AHy0Y1~Y8(hwBP`s2=Mz2#%ck3xEI;^ogWd9v zsLm9~8+6=t^CA!{A$+avkDK2X5{;rv)VZIFA!rL=NQqP%Wpa#^7Y&*%fILSwSOF++u_3D*K#+fq4FmLW}N|17~ZW|Bj4e9%jo;(B}KYj zPr@-%+QJQDzXs(pulxF3|J~5XU2P##=4qI(Q@rymEetBBXPSy%JX%t!apE1qJq_Ig zS{pxr>$V@YHaz3x^=WHDUV?>iX#pVEB*D2&Kt9mlgE$04Otu4EUb(mps6&!KEMy0i zwlpp{m{ivXl-RuJM1?At{_U&q{kKZbkM@alqD;mw`zHg)Mdx+dYLKrHtYf%99R+-V zmY8g+qIUP~2F1U^#MHXoz>8&z6O3@%yF13MA2>#*BA5q600Shy6R9%1T_Qx2jE<@s z5dnjm(fmGQnX!d=p*piYi;WFQ90d*8qGI_4X8+|CD{?I8_LHX6N`MM;q+S9aUvl+Mr& z^4Vhh#(x&IyRiR!Ko*-I0F^QK(+^YZ66y}AR@4Gy{qpt$FsQNP!q_=Kt`94ZV#MPL zSY)_UQk7m-$}-whUNIgsHNlnE4%>Q%5S0O?un|M#NN^&!Gjr2An>~MDwo3vHa;eKO zb;B0mC?$2Tw6unQt^W1NyRbc0=7}ncPH))rCncwlj!*{=M1+acrwxsx{?pm$Zp<1` zuKSgSk`6mb>~$NvLr?yJ9`Qp!o@lZ$0Y0V-Zm#Xf=eu(xlcM>?RyC5Jh+R&XBcr@D zqdNO9hsvV&Iuxv-^N#YhyF*4crQ!9@y|Ku(V{{*Sic?7AVmN^MAm`v%;J+~pH4v1KUZV;-Z3qVw1_z@+0 z)wq!&Pd_p_8j8OVXU#jkq3T~y`^aVFTk>$8`r4i&ECVh6y+n=)j}>m+<3&8{R@7qc=hJ-dWFK0nbV|fyhvFAo38_7L=uY zK`yWFq>kTEO6sW#mV&=6j(&K1N;e|?=B?|TH-TF&#&`{OM`po|b~nT&w7eE8WOVB` zFEz43a8C;j$kZZeUZgA1Okjc;>@?3po*KV7OAX6$t8xRZ79O}rVuUOk*>M*l+V*^WU zc$O32Ux>{x-81G>|H)p<&W(8(H=0t5Lh9YFC@X7bz=sp~m?^RkwCVhdTA55vH=1sw zts7|1(C+d?cs?p|5&&q5ZJz(IQlddu3k)eH%#qKdKBH|wqz~Ui31Gult)h)a*HtCT zp9hciDfMX%Pud;NkNuLu;d$h$y1%;Gd;{I2ly(I)zebQ*fLuQSX;bCZKM$F_a)~d& zahUzQY4c32TT&$lT!^=G-1rZ6zxoW3Q`&U#E}(d7rwGJqyp!a-G?RIPWQb-< zjf;t&pjq~%LQE(%?bc|Ya;k+yJk`PWoaeRMN#(zvCjW6HAb?Qznte8>6t+7&8zVgw%G zijYcQj2EJsh?2hDQu0%E$pPy3Xfg3z<;xlNSSFlY~0lOe))LW z4Ovu%+(1UhUj0CX_UwxOd(ijTk6UQi#q`)tqrsvYrH3^zq|H!J}y-w-01I_I0P9#!_M5F337*jCjx=_ji2KQ zM2XG8I7Jm`O|`<+OFmrwNH$7c#Kirpx=<{q#lbxuBdIXN1vhsE4C1Vu0$<2 zAn85}S#4R3*)H7Y^08l6y25avIbC{}V09X-X=sjYg9Zkp21ps|b3?ZsvNp6m4I&S~ zrk<)PcX)HP-e?7z^3?j+OkZ!f`}EjVML%d@k2I)IiGz_sg8=3t--Cx$#!)*u?a4r- zzCb^aXaK-lE9e80U1z45o?a}pOW=B2;s6S^MlBh-4u)+|gLFh9al!yoA@#TOg2P~e z&PJ+T8^5ml3G3E==_gjaIb4r+%qgDKDez70*^ny~<*B)SK^I-AF1A0R4Yk_xqf{lX zFWhC0PJtGU6W0SW&e@N$9cHtBw*+qcbn?y_yY-n}47VjHOKRJ|fV5>InSOALS_+=` z5WBQ67u$-WsbS9}f!aJ+cTpwebDgg_m|k6K_T^=&l3kv+1?#c5=3+w1LK3AMv6fY! zQyB%zEUm550SPxRDT?mGHm`AGXgspu9%N8{g>3$SHE@Rt-|Iy^zQ} zE6SFkd2RT6q?Mt3z|n@bDI$uF1bHDTR1Gg9kYf4k<1)(%0Zx$rzxCEOz$o~DTqA*O zRr#gmB4sHo$O0({Yj;yOSx1M&G9L|)e07VRbY)#;T(#aG4%dd1l#uftgg3mfhB(t& zpJYi>)WI{Kc-QXvQ@%|XFQ5teXxBJfka@cdZg$ls*I147-Mw^4$REM-8Nv1tvj(<; zbrqJU?Va{LPCpFae}R=|^0S+cyVrVqWUpR3BJ?ur8#N(S@L-e7f=o8GM1h4;#N-^L z-GDmO{*26Afx5+@J~g1$JKw>DFS~e>n%;QXY759*?&Xa4{vJ+}eHinU^;VSo9;^B* zJ`_Tg$T~O%)|@&T;hr|Cz$KAxQ%> zPZPT+N=(iU&|h$Y@rtlkvJ0fTv~BYpVu!w0U?DqZCl&wqG>8q2fH(nl;nWT!zC%S% zzhE}RW~%aVD@E|tNrt>dy-0QfrC7JHQgkEndngVxXWks zbo=*oOE$BWiZaf(66$ZC0DAH@rQdx_9#8~(DiLzzA4Fw&@IMGkyNSP&|Cy=BhPw@3 z@*apVq0esf8AZf$UTOt-q9G}0%fY`kRHt*&BGK~FqegG{ArBFl1gtTZnQg?TO;DZ zyW$li_N2a7I3`_=f{(%-p`REHBCySQJ7zxKQabedlj_ON*TWV!3I~7HlLb`~EQflJ zpNM1bSE7M>R5&LzAOSpeJWbFKwC~D;P`N~DbvbZQvBKX2ikqGNCgyN6y;H`lQ1)3K z0>N+o^YWA^%@pcy{Xb37`MhB`(BwRgtOg}FRE$nv^o$dwzVWZ-*t8}yt$-_D{Izj7 z`MtZ#ve|sqp4Crnrc^Gl%t(|60amE&#lv>9Y?OUXMc9&u_Zr{UPW{;(Yvb+rqmcJe zL)HDgKk9Fru_h3wD1DjbrhCMp9ux+&uBvtJ)+e_+6!q>?uDIx2N*yaNe^uU!3#Ar; z#Bs)L<0sCKGW-p#)Xf%9094q09aX(ar^2}aGz8S+!o)n9{BVw*Bma&-69eJtpA@B? zo_yX7T`KUK!XE>&;yPfm`wdjVQ~<{iCR;=JeZZ(|^l73NvgQ=YQWHu#qJY=%KQ(Ra z{$?LRm=4i=pboK(*5GpUJ7hb@aa9>Y?yKgp+W6(N4g2@q{hTVKW4(RpR&jfu?5#Ai zWO5+?K^z;weVyA{BHz(hbLYAbkeuVrMt&}2Hz2rGm;Ot<=U>qy=B zW}vR(qSkebTlGz;TrufI=X#Cz-Mn04Hd6;3BM?9RgV?!f_XwB~_HyVacRCSCcE;>j z%=RqIJvNy(LVp87+3HC8)VarW}S zZx?g1>(%w9XZYgk?fmCyM`%}>eY@2re^gWoY(Z*>2uv2o-qU^B9yPuFU`IpsUcW@j z+hz&={kEd45uoloZ$cnklxG^*61MtLPx5)=Kr63=@n!odes~M+fg3UX*;$h+K(o_jBvgb#ME)_0ale zF(bGB(cKh(Kk~*kC~`cvJgg&FZjCg%maVB>+ZT}w-9&v6S;%W5Z@W_1W9xAxlAn4lPHy1thTTr85KI}ca^l5f(Mkb25H21|6P>w7V&`}>lj7Xte4iv#T zk7_*nQ)J@PG-MNB=7(!#OFtmzu>FVKJqQ`-nJ%Vf{DL+k*>R5fT)u=*jiKjhcVcJ~ zpi@INDhyS}I}J!k*kD~+2<&P4$>2Da{T!tyGkPQM_eGj#mc-;`owd-^?*5~ti7ijG zCKVosTbq$laI^qb9<<)E?+$JQ1;Z4yzqO&1FJ*OMsw?@k8Wp$uob3Ab9%pTw3&r2J+U5$7Z^1Vh#MuKmg53Rn- zbQn<%xue37!kU{LU$N0Wc}E+`j2VZa;$@vvZDaoK&X`{wzXi<--nNYU$(mp8Vx_Zz zI-@Ss@5Cq|tns~omRASpGWon~YE+P|h)9M!kTYO92A(|EDRr2#FO3v>L(5N7&-its zfuxCzv@ped@wf(bT&6_Af7)>Hy5+1U-5N-fAV@Z;B5N`kPkF$`oxe<|8$c-aLw@L) z-Qjlvn?%)D6mb%+Nh)^!n|{Wj3A+!?;Lj3pxqE?ju8$K3kX~g5d?H>rb+EWtK}u$PrhotT7UHc zcaE8LTzb{fEtfC4aMwm zRh$SKc?@x%MDFv|nRxrf746io-_!kM-`o%OY*U3-Vaf-?dhAY5H*T-|w>N17<-_|0 zH9xXmNX}k(Hz(LC{R%tM*`UqmQkVFvif>53-`p+m7^U&MzsY5rT=Qcect6Zh4#=+o z@A52f7arCahoGIK?O$Jn>CrlI{D29H2RE$s*EP*IpRNaDdfJ5>DQ|?SjxJ^QgYVtA zYc89c8|!(`LpCcGg?Gf$C3k32Kc^FWngKqc*Xg7Z>QScWr)c!rkPXlFeJmn zT&Crzf{mJT*{|c)xg}0M8#fTjU$%g+`U8~=ku z7rRmJxM5fj(r0$_tg0bfj=%l)RCBpnlm`ZuMy++xBtz`#C?j1^1XF@M6|a0fmCME4 zqp3M#!d=|CS1uhraj3%XZrEoJ(XNBI=P$xIL9S3wm&u8uxDoq0<~|Shi52w$qnVQ& zIVG&(S@Id(w*>i$h+25Lz4vVNuyXrstfX+ezp9xj4SXUGpd>`0<;QRhIF9CMFh3>p z!u!2qr=LR|f;e5~AS4RiRLX(AT)8yXdYZ@(znV$AVsucmz3_UlqKri#Aomls+1D@n zS>5Cy-;bx3)k>NEhI=BRec-VeP>|hiISGG{wd+(jU*HPD1{G0?Ps_3hXC7k_nV#_$4&1 zU)wP{MvZ=?{7cpN)>V9ed-hII()U@@?R2fL=fun0A@%EftFAer(g_sSuh5^MJmAuF z0#R$%{0saBc22q*#q^yeEFImuDTvi@9@o<(Dd`*ZPR?Df${c8_x5&NpX&R%_r!d6* zN#SOf>e?IbJ;>ff1muJ#1IHBB-tGjnb|XSq0+C6WSHMxi>n znmiA6i+W65_jwVIW1Qbv60&z_Oz!7!A83V~(IrK9ZA}Fs;WL$bLfSFz9}`;2QME_g zs~Rj`8E?D)@(7f@({wLtxO-;jAp{QdFx){(PC_?edO{!eOfMFYc%0!vs~QpQyBPOy z@vNh9vXVl4LRS?K2nCjHZK?=Rwg!OAmxWqHhLZFgD6_~IG0bBOX=484$ssP9szuH- zN%yeUIp1eHzDG-iIJvHIKYGqE2fy?Vd=(o4i~(CD1soe4;|VPcRz~Aak!SNxp58=X z)OnRxP^6r8QNn~*M&!Zw7nU)8zs(ows-ntGgIbU*^c92b}+P9QydRLH5yb6KX^Rw1! z4@m3me%L~d?Q1`sVU(=Q&U_f2(=yDdo)Xn`RTK6O7%0xp$lLgStmNwvWlO_-zB?{p z$*sU^DeOg&Z2^n_1Mt?aSkI2Xa@i+cPJ9RZZ!(tHQ+9dyeMI7?(}O1lB@~8*4C#c? zt-(3&Roe#{IroE#R&+vs2KWZ~DH&_l7I&nc4LXvZ9(Lpr>p>1MAO+;gc!6%LB?D&8 z9IBPU9*RheYO8!`;ao+{3twybT&di9#gxqq{sGy9KIu1?2XMEdSXt01pe+W7{Qo8M zb5KVIvQDjA#e&H*$jGRR)pICSEu!4$>Wb)@I&ms?YhLT|)ZMZk!yVTkyPM7|3494) zI?FN(4Mbl6u-Cy)-vi#{HzPXaCOOsepesy8#aVsrB7L*>7|=g|I{)E~Cnia=XJVWy zjBEU5HkR2JA$$UX*#Awj64I=L2rc)nE)oj6>B0roe?)Jz=}oK6W?(^~OvZr_sx=&dD;zPg$$8Pv6a7 zO1cE$|3L@VHu!a!%)S5 zWWbUxE-;URVF z6J>OxqI2?pwMvEB(?Wv~DZFfjc``xQqtMS->b`0Nt;nuMzeX|b9)Xq7J#8AUQ<`8z z8#(lkj%~cJURLyldfG|#LnPNYu9g$}=6Y#82na3AR2aj^+s}{Rx|uR8#2@70wzj7- z0oOBwga9Qj(W$B8q@<57rhcYNawJu?AbaM)e&KWBM>&Gr zvE$)ao3$yM)}!V)NRXgVSWkTIQ(>3@dqgcf^E25a+IAqJ)2C49KkCj(v(#|=VRo(Dw?iMnKndC>A7#2_c* zK9C@5s6%tNDVg;(Bl!;@GEaO0(`uwW{e5Dq-0^2f&6^eVm)@K zV0F+b>G1u=UHi0tnP29WAkzqINH4I*odlSaxi~#uzX~5dtf><@Xz|U%m$4eY`r`uA zFR!fWOWNAGbg@Q6Jm*G)g2WOVNc}_@)S`BSnpR3t55&8V?w<2bt6gv0V@gxyD6Ol{ z?$?%htgBPj(VTUw^dxhM#)LQ14q<*RHurhI;F;XTA92ETEV(aMFM!gEJ_NfNw_D7# ziCB`mys!93p+BTBiq`Kw8*g;^e6hzEarH^p_IWUnbq8z_5$a&k(+BU6klCs89jDa{ zLkkd^qnFC(Blm8_wn0mENcs$wdCUIJk*^gJx9typ_1pbkjM$dt-_yQG`eAU9GK3iW zKqRG2+so|91C}5>826sZ3!j@8k6$h4?bQ0E(9)h3xE)}jG@>{3C3r|CvG)2^t$hfg z6*T$TPvjrLF^DXS>TQ<=+whcm919f^F?2pWX|ZjdC%}rlQEQS_6u+vd*Yt@i7fcKd zvRi2(rA#S>v$!9;?F(<_^KMS}SvQQ&|5hoABI&Jv6Nq;H>wZ-vPiCx^78ONe4(zuR zHJ7)!p7%fiJ)%t)m8C@Vh#9tEnB05r zcAMq&yFM@j2?PE=QVL{lk|nS;y^g>bs`0J-?dke9H`~8xp*p1y@zKWG#>SAdBp-p; z8-{A0k;L^~%^9q3t8IM3({m%srFu2Y_qQ993-~4bh#0Ggri^6jho>hz#L=g za+K*ESW_~+BAimw3K^bTI_s7qX1Wt`u%pVvPV}gk7*O&NsxUa}V<2RjY5cZ!bU2x(ZmiJ6<0<2Xn*v`8 z08V{SY;*DRY6AX26jI1LSVQ^FiY>qS5zdkZ4u{xf z#1JG0%AY~oy=WnTDFOEG_+pZJsWY*QKIC+Fp_BA*!1r|kbyy}63}&_N>v zfU6A>z|3q=!r=j)?KPCAi*~12QB0+tG>Ml7EopjY@_NL&wd47BGNmf*fB2!o!@b&y ztb2K_AZ3sz+Shi4ndAtZat?VAk4NIZF}s+jC)(w4J$Md1i%M7XH@KLGrAA_tFO`G~+O!*IBZ|J8w zJsjaEn>u)|am`#3lj6`x)L^wTlR5SW32>B`D`TeTI(cDvOdlkS^|i_|r-(%1kgLtF zGebG{S?E3s-1#17^Oa+u`Hd*6yB+y=`D$>U?4IM`M`jL5;7>ZOug>|nGqzw^A2?sqhiTiBh7X-$;61Fh_l? zW5M{!RB03}1p2F&9XGM}7BGa^uqEr9|b(-0^yi``TJR+|6%fXmyCE+JAg^?dp05D?e|MHNdBUdlwq^)UB=yB(9KLy^l1Jf(K!l8y_MhksrBWJKl+ zua?U%y-|~16P&Q7jm=;tLka*(`2Ioo0m!%0IR+?ccyEeee|JD)ugL}}r1di{jxJKP zF27R!x(~zd&p`&qC+4ZQ=rhI3cBlFeb)%|_=%QIKm?MPRhR3a#1x#W9cIF(pocn<^o7jegH@@LV-Ta{@#<#G&fGr&8;eBV9+ZEvIN*`BN zcwPQ6T{`sMn0K~T!Mw+W>dE&x?ETs(-{c|AyFzjMW@=A(uqAKGf;mGVxL_Rz`+n{Z z>T@)eK(6dTi_l&1q`k1*e2~8U`M0A7k_xR7(+jd448|^=z>bE`+H*C9|EQ0$Wu$_M zT1r0%uo4)fv(SYG;F>R5mQmBxy3qO3b(D3o<7nMOEPcKRZ7PwX)PLZIfQei7Id|nI z-X{a^__ky-t$*$jhD?F%gDwwbt2TI(hYWlWhVy5;Oe-Ziru9llb5u*&Ackt@pLe{U z^BktF+bpL0?wUia%mN}b{=Z!aHiXHwKd)_?Tk7gPMXi-fa@w;rB?)nXgH4mC>cmB5 zp3GaT@?|Ag4M9f<eES}nm>uMGm!`fam{Ab z#xF^OZ=mFV>bkT4{X@qpu-~(R817F3k&Dqqu(rKd=)Po6=7D8KKHX=o3b|JKqB2%f z0s4cFslT#QbPsqVGx<0P>5D-nlo<1 zRT&N~+4O$DGn&Mqe`-UCJa!qJsCdQ8aeuK9YFb0o&`Y{CAzW_kTn`&$4~kI=`d13H-{I)AWL9Yw39*!Kyn)~DSr_0iyq z{Ih8sUNS0I?_O~qUPX%JlAl&*uh3ai%(x1sAt%E}@pN&_JdjzLqTV52lGYmv|0+p4 zzEWBI{MgOt^EwwRa{Z26(OPc(c_FDsfZLvpkax z$9_;xxsI+ylPxiyAyA?sUqXKo?}Rw$@*2gu>r18a~G=Pt*X-o8>5uU9S0_FCr?-ypNws~ z-F451buYqT2$-BcAVvAFbHMNuMW*K6zRL|0M}DhR9BE9DeD&h?Du8J*_xf(jGr7~j z^P?DrI92Tax*?-habPV!`If#snN#l8E@EX2Q6sNxJNNdOJNxem$Sn=yT^|1b;f~KC zVhKdP2v~iAX}Aix^SI5UV>K<-*?u%8U;Vh{WWA`BReUy8SwqiZuNOtS)Wqy-+3f)^ z5ZUnmL@fUkurMP4d@_bWi~-wo4h0lehZs`;uDMu4QM6e~q@x!GTkdatySaWbv9;Q9 zr*ra+a*I7_Hw7*BRY$Dzp<J0;^N?lv(jpD%rBfO@bCS3+_Ry#aV?kGX#EErLP zk4JKb@J~O0OM8U1Dqp`^o4ueh==Ts4{6nOJ^Ap$IxS>ey75m73WF-@k8oged!-b!16dQuZQK%9ezRvBVU~7LtUFIYr1GLa7W{%QDGU$k^8; zOZLp5jVv>zaWc#IIedTD_1w>W-`Dd;>gCm(bDrmM9G~_5dA}bPLC<8a2vnBs40`eD zfR9*UZ!uGe-s=zWV#cbl@t@Y?)}*+vl;Tt%^QC6jl4n(JSn@>{n|p>1bB{5c>UEqQv?Q(o{C z;a3JR2XHrZ+#am3LVnc_imyC?TP%<_nk^)I13YIrz57QG4z;+gCYzqZ_@4 zD|MzMRxc{?Xj$`~iZS~r&86BZXfB^tha!Jphcz|3Ay6ooSmh!I#-g|YUkd{j@iOc# zO5FD@^GlKQ{f481!J&S#&u;d|D%$H26j2{OxFRNnnp*{?_o{6vVqila+Oz|b8f0WN zyJ0Lm8V*R{bm=GP;54+|W?UuXVf|m=5Z0Yfn3t97>5(%?)Hp1C;H^)^?edm%2aE#kuiq?s&9(Gbi~)>frz5IJb{!(u;zfw$ z(m4cN_=16M82PIyjMbydZruI$=#Uq)kdrj6*P$d-X6ilIbag^Ht5q#ay%kx03H>{r znh6-Mm{mr8RIAIMjZ_F>fo*P*@1Z6iok2bwX+uyGE-P&q=_sW#;+&UW`=MzN?)ox%y^FDf)C>s%wT&FU~hah1fiBlcZ*6dIc1oP8)hV3Mtn{8^|6OTnt#* z?I6j!N3#DR%E$=BJ{&yEYNOKx$E9C%6ys!_->(O&ott;urNgXVyN(#@3GKe+NI@JR z)=UG%YY7$)27KZH7-j|Hpg|OsdDOj?6|B=@;5|KO`LPOh*M?t6j}hngz&E?u!{=t& z>qjoh!D2AT^B{gjCcyms`~viP$u$5OP7pF0cxo+jIxU2wVCP8P}sw_9$`` zhGIcXBG!-?!V!y`(Kz`f8W4vv?{Sf)Yd~t}$0VDX2Q5nnuV=Wvig_)OdFS)&j9E}G zphoE;Qch~lh_t(@j@6dHT#JLJY6sr=^LqxdZD4yCVsl1gSZ!KP&=mbHuA1^px#U6@Ayzp8NFUD*R0Y% z`<|)3x)rJwGXS-<;wh(xrqK)L4ADOeT%DS`7<%I@c^JJQO7*Qr-{r-4(+^B4#uhx2 z5q+q0ZIk1Oe4gKSFmPQdpyCgx5Cd250p3CJwER8_3h0leIW$>$jn(;KY9=SkWhAd> zs`wu4Q{kX?TSH)7co-V`_gbXFLeIR-I98<-sWsj0r1-JN`^6>Y;FxH1^7yJK0v!X+ zAIOBRKaZewgD1jaonVjo@QN!`f#hbCZd~3(sz=O6WvA0exFE!8cTf~4^REMl95pl& zY@>b5YIV))Q2^X2RD7h#%uNGz*O*_cf}i`ugJre9U;+y`+h`Q_E;9?|4+gm%-~w{W z=~gMew74}Lo(NY@qo-d?A`-4V)Uiuut%{4hOOScZZnU@!d^NBc;4DW21?>@hEvW;b zNVou+goB|y){G(pAAN!fP!F*h6(ozgvG2_h(og)0nlhXgm2?P_R}y!ekcNdG*m%ex|JZkON>uKyX6~udLG33iT8^v2G{R3I@ z$gwYgtOBXHW0!~?kOXcH_&zIA0v_UR=B~d0mpi2P)7p7nVr!fHtR@eC*k2vB_WIF< zoBBuk*bfjvPZAD7uEv9c_d0bau*Wm+vPOUt?F5v#V7$+YT24De1y*avcQmVnko{}~ zebT+9a5*p2H0_>`TzgXTduv7zaZoAR(twK#nGMyYF{DR&H|6!{W#v|Pp0k-5q?Ni?m zNjjj(YYy!5C~ojN5yTsSmw<3?1WiV!>Cnd37V0F>wCK6;Q&P0X8y_9RjisfMcqV)sRhH$b`Z(O^8@8;us9X@2spt~QbD_AI-<`u*gnn`S4g zXo&~e2cQ7dz3rS;eYyl-ObRNzJT&z)D*NfspJ& z08{rUv)WSx5`zbTxi^5?LRm+5h;NRC3Cst!{T*hYpmJ}^g-RPtb8V7sc39C{pkDFO zK>4(cuhf?!AO7#|!}ngRZQ*l>|C8yg>{t6igq3Ypi-8~F1aTTT$U1P>tDo68+0R%{ z4{f(e=W`=7T_>RTW8)b*o-J3Un|YTj9P@7$KdV~YP~H(|=*M&q!G<-odg>5pkvULs zb@N1a*qOG9xB7%WT{Cw6X<)0x^n!O%KP>5Gz?md=>GDAYtMVY>`zDK<4ju}H8^KF~ zIHLf=Z;ZrV7}F7u3A?x;S1u5*6n)~rjMdZg*$a@V#oA&9cbkWFh#9AoQ+)r7D^#?@ zlGwrnoOS~EbZN^`aQlw9cDB!*|4UrIr;dGF_8-m9OGXpqQOuUA_zh%MUmPLV3Gkl{WT@kvbn<d2fN&qu4t7$Wu zod9#8Sxc0FF(*d@4}Z%YJr~oVr`=Q9`Qfa|8!KsxdC^$j6~{%TOw4ttjfkQI{a9tA zdw^YUYGgB&=se+qJ_R$nhIo_y#c^cVM!k5PsV2rqW1-0%^996H=E%rk!Sb z_3R{sxCf9e>KRle+mc*yS=4yqqAC^A_yg+#GMP5*F)t*Ft~(wzh0FNP|Y8sAqYAHO{v)Nr)CeYpLe zGWri&U&kT0A{l|FC~`T5a=Uc(C>VF$%T7Jr3Q{XT77I(Xqxc!8BIgY>lt%nRED1%` z@3QAy`jqpuPM`hhXqK;fplj=~PAVHybr{n>u>$NS7Zku6cawqt%)1bz{oG`zL-!|a zQ%_txeOxy0SqL}stEbBYcL9;A)&357t*UP!{rKhVsJ&a}ES?Cu99%=&0B_V4YY)4S z179~JGEVt>2A7(KYG%m}Of8E~l*mr*cQ$RfcHw5z^IIqTkZANEnE!@mB6i$xFw}OO zr3_)+hjsi;L^xB5>niSuhu(U;I@tAPMcl6Mz;L$3G}*%T%3G7BoI7MT1}JuQfV}Sp zYFi&lz{8+$`)guFf*}PW3l2Oy^{z>esu7Ei{Wyo3d6gPXc{Z6%a+$st2m$zEmww+#z}2dqF>=8T9+?)2P{$m^eu; z87ZbJ!4Tf0j<%v+53Nj9hLO5*9Nd5$ELrxLoJXL4y_$h&g>47#(bS4T86T-srz}%I z@!08)VuOBWVqF=v>nw|WfoUn0I21tOUzQVL*5Ndkwu@7~KaGUAl65aC*NtiL4zD$3 z<*2p_WyvR}p4v^U1cJ_B%H3;TfFVY*rRt@Tn^2%53OZWJZZ-ve>D>+b;%8n&Rvk$A zB6U>p!DPT}QI(mjC8%!fl>}bo5Az;YpY8;gn^JRDrV1DV)B`a`0DFzSN{o){=~dy5 z-#Q9@AI>d63VW|Hnq`7cB4bFD^LTJ~YD^+-&*ah_{oJgTyko;fE)kAWK28^CDtGqg zKT}mZdLlX_Y&X%82!;*$mi0hsofx;>nH#xs*Lt(V(8bVi{9g9lLRg5a_4fp&?(wkk zFv&;4m#ih_xh^w-hu{DkDEcozv>jnqVKwLffe;+Hk5Wi*C`H3EP8TiBsid_v9>&-E zJUCOY?eyJVxSEGKa+44D9(#lNj3vh!nUk#>W+{dVv=Ie$o=T^O`+F`}3(wd$G(1z2 z&r*CYYfo!BOxBK5Gyq-l7-!}Rv#LclEP95;+Yh7s55!iO2bGeheM9YZAKhE6L|%Ye z(htkt`R0Bt`vCvV;%g@S`@CqDfL9y&rvHd1Gq(j%eyCS*el5L+Wi|0OET+RbFv-X; zyl}ZOGPSg5*Uzo{|4`k7#vVNO`Ti99FY#dEb1c;IyUA>za3aUuktx z5Ub84PLz9E45|k6U@asq9>vYt2eYG_8}xD{=AL z?`Fj20W<}*ilOLzT(Gj{>cI#w!izd0MA1OAYH4~frPEK{^F=z)7b{e|dXGb>D;#1R z52ZPQl9LY;4-C_s50a5~68Qw8TuX}&;WMO}!hL0z9p_G!0=Anv@f@v_1AN1&V+lqL zihQ*Xr8&syxy;%sbLiZ#kpNwdA^A4VderC`AaV%PYmb;ah^*8!{$I8_CV$jte?*dn;54U#+jwl?Cu!f5B7-Cu5;NmHV> zf9hTQkgB64c0>5ygOB4qKX({)G2m<4fO{c>?E)e=at4z7HY9>IVU2y(6%hVXSCWvw zWLHjFj?;<@wO%Q!?tG)Dt<9EWDyhzih_p1tY5pHZj_`>%RSqZ#C?HUH%DlHm^d8i9 z%d`UWCO+!Zn=>!eg7b&=s#*JK9S&AlKX*gP;I^v86>s(>Uo~nfoh9H7v-?Nr++?W3 z4MCy-{TZ(u4`!X3kJ3jZ|7}A|j6;;OnmZ87HWI96ld8MEH_XZg5Fs&KUbsZQzx_N! zK?E}${p!z|VZ)aY_lH-l56!+AwB-wYTV(BJ0(I+lFc*fOK09%-x_Pwf8vJBIx;AlH zJa$g$>+3785h%z z&1alaj>_E>*3bCX!MPiLP7U_6e2_=>g2a@>{JmmU!Wv)-9QTs}x zYYmYr@lT|VHu{|#lpSW+kM z5Tv1z9sfefKa}R#UpMf?NyvkLS8~0j3VRGji; zslH>!4HwBm2*XyufrCMw5W2rOmFe!1UGM9RPq(?GDKMg*7$w_b72oAqwT~O-MzN zGl>+JwEH03prHM=JJVsg30KPcgJyLHt1mpx7USzHgHbuE9$se2)oLBr1hsaXMy54n zaPG?91#NGDd-%DtA&J)oCPS?#@zJak*yE!t?d9CBh0fKKXWM+)1RD1x>%GhB4fuCy z8s9tb%dS2-xpf--1I0K4jLFV$L7^aJoh3cQ^NDNe2SeI(XW_xkoC%DoZ+xEe{dw?# zr)91cXj>gE9c|x+n$z+Qh?R^(!lnUGa%zPZx9+C$O=6Aysw9~-O(N^Z87neSpB2ZO zkqZMyLrcD>Uq8_=nVy~=b@1oUpZ#X`>G8g6We5Tp)rpfCGm>49W-6jzH;?@FKfGXN zIHT4p#glYw(LKXd=?1%E`L2ozHV#2EKRSibZZkc@;s-Uu=&)tLevqUzj)FV~xfbP$ zW!KY#7WFB{vEj_-;TJvMD<&Ls&#;-Oq3lf+zX4p5DVW%27?y=SW$mTX4LeGJTDyn~ zym-DTOC#SW_n`Y@OZ5s-VId)>skFVXEai9k(QOH!q^U$?NFZr)a2NfG4{i_aLE$n3 z8r!ISGxb@ezmC9G79$WS-zxbidcM|n^6YusUm?aa87`yG`28G$RpS{lmJ|$Yj}0~O z90i$B3mnNnS3~xT?M<6c9`Sdwco6Q|Iv8%BVv{VN9;0VIiB6zLl}2O+5zM4Pf_UK^4~zoV2!D? zOxdv@w`owpt8kn(lMU@s$O`+iPqqRGoF8O|bqR9>@T1jv4H!I*hs;{zffvgY7|KJR z90$l5ljQ|GuFrX$q`IT0?$M{%E%&KDITo_-D|#j8fA9)!#rd!oh+E!Fv`zWnj#TNu z?V_mMuPANS%k@ll#yu@|$^Va)H$@uOf0oI(n~Pxa{)o16nmc4yn4zcov%~L@jG(!@ z+6JH(x=Ml83cyXZ-NNzDPC3rGx%{N~HH7}|V7JJ&eJ;E$q+Ya3h_}Xc-}P3K>fu|h z$oDOv0MP+a&OZ-a7)l7#Bzim->5<5BF7B*>6Cp&{T}#~_t)qVRK2&YV^Ov6(2RD8H zq5IdIAq27!O}Y$g3{!i(7|>omBrUR&s=argHn{KHq$BcrY)izz28?-wKdV zKb^9>pcHc5>x!|bUG^2VKXjI&4-DMQnWqcu-%6q1rie%;xYirs(-*+}y%KSiv((2y zKR2%+ggo#qZhS;M*C+MEyxl34C>!=VNc1za8)`cNq(4vs5V-_Xb0CL$dL9spH(Gxs zct9SZas%(>{)$PtSdLu_PgFAN�(M-ovZfQVQ0(n)x`S^6bLzKeQlNAT&Hw&D!Ws z7`tVLWs$p;;n5iN&J4)XD8rk5$Fjv7Gcqjm1%zrYgJTWpuC1e0hK82*>>su=Bqw5D_HSK!NaTL7GKTX94x+`Nw<;L;*Hi4_^Qu1 zB?7v4NnV{8hB#QESOhi&Kp*UO6ytV)|5I8`V5+?;2a`jp$1t7)x_mI46#_v_ja*^{ z30Bf%2)fJ%GK2*jsb=rzpL@7cYS~<*Jz+C?89&B!{r?ERxnd~x>#!buUXvkavlyG!>3^{^Pl=+L;G~z1Z9y_Q;&}Jk_*J_& z&ceAuwRgAl`OQFlRE&WrmMEZKvQxRa*3bbIOzH^*YH6P~h@7l`%&v_udWI>;wiY11 zm&SQt*lkYZeRld)f8nlu;l)?LEP>gNfZ5zX;0wcvD>Gn-V-E!dg5>ug{!IJC0p7|5 z)RuW5P`^|$-VgSkXBWhm#r#n@0lJsInlTdT(H zO!ZGG%5GHJt__;$c_HZ$R+=TO&XvB9UO4bowTA5Z24w-*9a|yT*o$QZhVr(T7VzH+ zC)T^Lcwm@)!1+JP&<8)88T*Hg3>PCRFPe;F)$VLMixW#8P8<)+^^g!+fzlVV+)tk7 zjgdN{xKAG0^&GbZ?&mP-8)UmaUC!!g9Aaho=&m8D95(HkZ4JB-L~-$AiEXRVfd|`< z1EvKClDMo-H4a%djT;Fp@Hl05UUdC95K=Y!@j;F2u?xG7{a&1Y8GA~K|E48^b?6sh z#^ePn@~B}ML$mp7!#q??z)hgYGti{NsB?2HF~Ej*X;pMB4jRXZdP}c9a5t*|qIc^( ze|Rt9JsA8Y?aYc1Xp&T~#1?t%ez>v}6i39yFq0#G*|~MkSMKTv)SK7u8lDm>QcbpE zOxkyI*$R9o{QpnS|JU%m+V;+?Dg@8oLINt(fisX|XPyo44C-!a(IkMBOlp2AFk$g`T0sU+G9O3GXg-PRi;S;wchLTa1m zu@PF*Hm>?qc$Uq|>D-5eh7VU_vdlosD#{4S5(olD=L4+!55oO25Il3$z^P391=TJH z!3m`=e+q<&TOMStoB-deYxkzI z9;J;g2<*O##F@oT4CR8&jCMCtRCgwc^m;iojX$RJoW5l$LqM7!ZG z3mh|YdGPxN0CuBgX6Si;O`G>hm~uKVX)$v6^{Kyvgy z)LlnDC<^TaTa0{WMH;}`ThGcX=QI|uTsjbE=O7T-gz zKtU6P_mB-B@`!W}Gw(s}V<}BxCB|zdq{Ik`0U<6{`$Bl`h@8*nqm?>7iP?HnGTJPl zU#1fHL!WUPqVEAp+;G@MKw^_70A>#n%wu3RJfUw~$X%#k6Ar+%eP%|^U*=k=qiew^ znl2NM)u8mW9uHo^2~#3VCs*V%9CFV|Kt4+8U*36Tn1_0_s{cUnvXMTqVOnJ&C<{1I zG&mD=q}Ss)^fPdkAz>PS3O=QigZB>Pat+J0NVVqrVEWKR{JWi20k0%`^8Jvn$--<~ zLLh1bxAmxj^$|m})xk3lo!v&W_{rc`Lt*8tu0d&_A!&+x3a$+h4bZ(EQiL|m86FEN z=^nshXdoAfiCB?6pEp*;hulU3h&9(O6)zFPR=#?vNgrXqBWPH1>!!db;!XmsuQspG z5c66*@RZ8#049?dd<(xKPI;8MEgh2iRemOl8{3()-}&gny^HcucLCJl;WlfUNKS=* zC#+93vHI#`0KB3|aN>tJXdLghWJtdrulVsE?x2dZK1z;QuHw3UB&I!hdCZ%6tLD?+ z-qOxxrovVf!YNpq($vwyOL>3)E|7jzD?L^0bH}?YdwhCdXHdSnrGY)W(vL*MPru%O z5P|DkM+H*wwEcd>?ZCYEw`#CHzJ`gWYeZrr#DlUubZTe$UDikyrUq2jK)*2IgA8i4Tgz>lm0=R z997#|BdneVOekQj6X9Usb3)*wMKGp&!nq{AJnTAU&igJc_w?nZ5|ik8J^kJ#RkQp! z88I^s53WB0|C=R&{xv0rJMlpe%g$vnjjrmP#>%ZtWZor!GDUPC7<^}Nfv;&}=G-0v z?Ja@D=Im}slHJ%}>7Fkk6{Ow=8DA{ZvYfIHj;ThysxJG%@tn_!wbzP%{3n=We*U9p zD7wDk&Kn)=HvN>1-mkG^5BvMC_rI2p7LZRw&|X2T{eT>w^!pJYmUdzFmnMJ%x*CF+ zHg6+MtAMa&RjAa>4U3{q#9%|bN_t;^8Jb`9ow=C$dD8QO*}X|Mp7 z1pu^$y6`i|PB}el+A3cmmYb@dYED}#o~I8&V4BMR!^psQQc-AzdSZ=*%rj5@3h2Ja zLj>w6#0G#_e_s59nAkvWAN8YuC$ys`Cs=%q%y<&*IR%>wt(xaEt~J**`}HpnC(cB*4$z|c4r)z-Oc@mr)-Vol5Lqash<~po$ogu zz*EM@Srax9@qR9T34g6ERh}q({G4&7p#go9?JCE~451*VK6)XIA-l8rRAKO8PElq- zGZ&a|9!!0@Kussq)Bc0G9l@a?vJ2M1s;24z0eO1o$_p%dg}H(SA+A3j zm=;VOsSD^ygBy4P0HuM-)C2F8zn6x46}g??zF(>Q<ZBV8OfD~B9{RV^>m&T*C2(>)1?U=bst%) zQ=>oEx1`#?tgg+TpX1S1e(7UWq~aSWFFw87LUIr{iT=LMkUjsEa75d;jx+VT}V;gcZm6>YvwuhBp z0vlQ^koK_g^ItXQg=J&{PQE(TZy-L;Vs-VeBsp5}gQZXxZ%}ul!I_jx2vN4LPz@om z&W4?VqnxjPv5ej}t}OXge{tGv>hg*AT@^{k#lpTqCX$YF_XSOHRp9WUiDbZ5wZlt5 zvBf9OX~SOgu9^#)F|7Szsi5|deK+~|9KN{k`o!ji**-6IeKa@gn&y7jrFXJ-jP0Kk z+86>4N*=>c(>sy1@iAAa4KX^z_&kxMi-|ujP5{=35?`yUyw=~3xGFL69S{a<40lO$pT0yA1?QBxP@tG-BGmAjL$SK;N2Yq9JG zyNH5n+qb|geSsn;;3ZMyqPZh*FqM5r)qEkvrbArVTc0JH`5oS3?$P>-;JLaLbEB?Llj^Dh>dc#a#)c->_i~1XR52a&j)uC3Ebce+R%J-Y)1-{?GJDZVo5AJ$dhgbO{r4NS-FSH z;TXc|AJ$&5kq^EBhHND(3>*S$346w=Pyamw&t@t%jMalp7YNBX#X6vMGXFlfv$1$% zATZFU#IY|#XGH&u{YguYrIRwNVzknBB&;7u8-k}iG`ce>5?#4>O>Pbl2L1ATaY+&C z>ZB-0cM1vXKBfI&XIZ=_vc!)|{DU}P%k^VWY$LFT2rJpF>e98Ge9?J&@P4X@#mL69 zNWw)|1*=M9^=ASbk7Lf?*G4c zk0FpKo;HkP{T>qSY#owm2>txK#r}W&_U8s06Sec27=uvT3Qq0wVQ{tqWzEL15DgP+ z2^%tKK$;Jiql&;g~v_^KeF-j zuODvjS^u?w3`dL)+u{|Iw^s-XvFTC95x~+&0!u>ymgWa7{HZ0dx9_T&{Y#|C**lV-;8A~7Yl;yGdt(Y7HwZ|c#?b-)asq0LBDUL@?S)!*f*fkw+To?ROk{| zx6e9FzXtU$jbg7J6i_tl4L*SHe(x-SVubdxB)Hmh~>_ z>7gZWc)8%Zs-tXE>d{OLXd1S87`c>(je?34-NMaJDzBj2^%Pab?RBc(3>5!1!_d!f zt8#Uow=-SFwoAq$=Dx-7B~Dwem7)Ix8~pHJ`fk<)P`VsSZghJ}ei|Vw94T^sc)||Q zZ`PT;^`81p@>5Uml5szBq_akxp3jTVMg2>Fe?w)}dJ&{p2hm+9DXb>tV-8iZjU;4r zI&xvg)kN0MW~BIBb;PeH0ZMOjrJL>tYDwI&U3#Zvqwza=Nc0La1G&2y_zvlX_1M3a zA{U6me^b;P>YCoPRKJagczP_u{#wQ{!8mgVMBD{#Mfr(UR8EQVE;VRtly)7MCtK7F zqy!d6d7hTckSGR9QWd$$w|rX^8mYUVl$AW$g`>{8%v^4+u`pRVW6Kt5J0um54xI3gxnfrTL=*aMepT5FdeOKmGqlYWm*LtI7v`sudc8G{} zkCvz;IV*pR_2p}m z&nHrf!--2y*<3+sSa)^w+^wDX{f4bcz5jH5nrZGfbLLkisAB49yU zZSa+)!kwzBWgHYywpOvfU14bqkzkr_3p0**H82jf zC&b!)!fmWspb+R^F)K7gMf#nJAPvXtNy)wa(Kj6r0>!NFwN$1`OF zQc&w(2!J4+uoTSVR|L}|HYrHJi=ApLN6AYn3yZ(}c68-%!RzY#gF!lXwx2)$&?R`Q z>uRETYv2C!Y(@~8N;ruej%P@Le3UZ@{zeaAfhwWn zE!O-soM-oawlBDP6UpJ}#LvI)-WDBv4wMOEk|Dnn9<+WqNgxh)a5<0TlTPHFb<5Yo zwP6L~<5ZV(l$ReKeE8U=e^ezS0^I_&W#acTEQ26=0K<;@#=G$zB?W)4_|vz~j0_F& zZ5R5U!g_&Ez3y@5(^@8FD$+Tq93sWeCbmTZ&>{vi9EY)XFP5eFD_T^?0&s)jLaC=k z8?-D%V@}woXWf$#^k!cLX&Bfoc63K))+~m%7H}+JZ9W@?&y$4U=m8Vri0*>|x=}Uc zI&02$)5Fv=U4QL@fL%-*p$Q`YvVqmpdTb_hXn7QYcIbY~!OYfr31f-)u7>T(``!~1 zBZHw4zV%;wE6Rm6su*n9KF18~Fd*lmoI}RWDZn5>;oJ*rgUXZHnC{?N&3#Mg->1BU z(hD1+e;>M@MQnVK_`#CQhAARZpu2!BOV}Al!th~(4kVRn-VYIor%lU95g%EXlTMW_ zUmWl!=BK;Ajio&kCZs=;L13Nnl$#B!`4s{T9j@e5{+k|kKmI}NN3UtFg4~a@5H6o5 z!Fm&Iq}tTX-AN_bbJUU3F9t(+&3j*Gtp)hIt?lcRw=iv?l&V%RTg4zo1KV5jEt^xA zPw}hw=!Y3+Mv0*A94EIA-`2>$DNsJprq{|#bju@OwfB1)OYa z<@!cqoKd5DL&L4`D2oX&z0_Asup5DSC6^;5h~)e^VYqT0mE=sO`v2&ljjrF&ZF=xo zv&5?V9UpHK!U4LZw{^Z>DBZuj^VW)#B14GfShZ5N7Lg+5{usxc;Y$MCsK_HLsV~Z9@S2gnq{1z zY{qPqTgB&CCF;HX?Aa83>`94_*xPfl_O^(xzpZgtC=3XZsWEFd`0jSYj5}0Q%A2>T z!TzC~wb`4}j(UBfXLuAnIIdk2=_P=YU`!)R98PM-129r>XKr>%nG>M6c9wfQZC0Rk z#${Q~v?%DDBJO@+bj9&hLm~W!6KYnC6Cs8X($aoO0|R{yJP~1vT2@BMCa+YJv)wHa z{FV?X$e#u-*c*73YLu8U0{mNu#l>)=NRBcz7Xwqroz-0DPL`Z5h^#s5w`Khro6?r! zp2LOsQHhiLziq%+8f>6K3!=_}I>_!tT>tUQb7JEXcCg3Q`Rax8kf_>RiLCth6>&GU zYn4{YY);xhM!}dNR|JG6QT)&h(VB%&hUi5a0QoXBtj}r+`jG_wmeo#%WG5?+7m5x` z#y=C~RF7BXW>=J4!jyi7o50uDGY(TFVo>h?AQGKvm4&`ichVLJ!ZU+~UuMsEq}f?i zez;|HhWQP>b<1FAgy;?MP2V9F&pOzvOUz=31q$2Xi(wk;d74>R}e3VC3{%_u<;GrIbO>KX#viT%faLojqH(~M%8U0D2UGhvVS?WCtkk%z9Nx*LC|)jv69B!13I zX&2iT9suO_vZj$6H;r}vK@jf{8NS3q46H{2BM6H_@IfL&#CM_#FRW5G2J{mmaX;Ee zmxEMt4pe2+Xm5%;`;G06N=Q4Fu_XqyDos$Cn9rdMNhYPC6LXjsY`lzMn*}}z+u=nj zIz{vEd!;Y!9{My*KM&!{@B(NCt2a#Nb8PlQQZC1D$g- z3L$o*X%L7i$Ty|-*i_dg_i23ic=Hv4X>*cIHxC;U3LusPbJD{Y#VZ7xC zT=h|pYW|?;0aFq=WBAV8 z#qiKW^BSt;p3-SSkDECn0=XU?g{oIAf};Tzd;q(LnQM&1eJ~bEGJzw=B>pf%t3;io zzL5%#hClYZB|dmc7s*2+8Iix<{)bX>?-rT{jBd0B?jra-a z=CL<)4Zhvtu^2Jq%C)q^AJLbTl+^ue4^rwJaBEkIv#}n;J4C3oYH#&mKR%hUSXqUgLv~ zK$Kvwiwwck9@;)CfpIwS{#}`I;$DW`ZM4HkrTZW;mLT($8x+CO{`)_ErkCODVy6*Y zU&a2jdjCavPK4UBP`@9Mkf0JW25qbtPJHFd$^XbE3T-e!1%V439$9~I)KWlbAcPx+ z1X{$fN+2>2@89eQf(f>tD?H4I?=xJ=o9r77n0h+ey5zJg18X$BzEV66-2X+fS)IIh7$Am;W`4;jtL^) zelu2~TmJa7W|2$j6K6?1M8&JC{m!MDA&u-6hIkTe5ARwm(&=emD@UaHgUPYk)zGtVe1&{t+ zkJHcIwn=FAS@1T}YAlz}@K;r-YkV7Cz@-K{AB7y8_t?V?dI2H=lYydyW7^k+uxG$# zy%Od|HyN(ua~^HS8P{=~!V*o(9_-Ij_D-PpSF~Bmc%}9r zkaZh7Q>Htmc$^J*(c?tQB~&69KLEn8MBWX|{q4puj02oihT$2!bFgD>vcrRh7HWGd ziaA=^ag<~A9B;q@A`hUJ`tRNM-_P(%g(0^Jte%@PXqBpUIpdPb+~DB$>AdQ|@@2tJ zH$L69j82yX3wrcv-g}TGIjj5DNz}(kKre`-_@fx2YCR;1-8!to6IQrziPjjO2ynbE z_&o!!&g5Or?Y!wcHLLs|Q|q#?jIH_gBxedP8vRX^Aqj-BJMp|ka^`g9#Dm+er_JKK z=1veNbnZ0W?_w6Zk4xR&6THvp@wCN(yM5c+P^}#yNY}_hrU^TQ{Q`8tHMh;O;SaMimAUHL8jBA^sv3ixId`0MPiad?NWD#J>rJKgcL% zSlq63zpxtCUM{fp`r!zwcF!IUKr~@iOdH_Y^1pa3)W#=7qv07EcIv$PKxS8u44N8& zndD+fOYC?TZit*r+oUBeIhxfMSwz>X;v%?qj$UMWSEWWSOgMluF!Rq&-@v!`Xgk1a z;HR`ZAn|7Dz%oalMsnEU3eAfmKF~yvY)z{R!b0dbJAwB|*fI8D;=Eay3hVwl=v$IY z1Xv2hopRzZVe7O?9JJy&4sE<3T9>h-I`TwE7)bZ=Mzg1-OLs<&Wxc3ruu$}p-_>=A z?HNLQ0)bvYt`338id$H9_!Rv@SWP=lj`?&2cW4rJr*6D7GJo|e$OpPGH?-$yzcAP9 z=bhfm{S_a1*q*;Sig}mCIDngieuFgMS)S*JXM8 zI{uYQH_P2L{6;49qMmz}eHUAsqVacP+|cCzt}3*a8g z5(e-fkz?8ZDuX8Hi=boP?DHgRIJvM5wc&{ab&vLb0oElnfZPktuz2*zK#d;VP82DP zWALbSAcbM61=MOFU(rCQkEs4`8-)ewM{!~gEx%bvp7ZVRy|p*%I?FI8Qi4qq0(Ckq zh7~zdqC08tFDc`ka-Lh~On#3XV^k~}hO#G z{2hMbL+=u=yTO+yQ!cTs4(_7Rj!1f^r`#T3m(?f)mbfkS`$Y@r>csqyGC&HzOrxPT zuNDawiS&9R(OA{VTOY|+fs^P6R!O>Tk`l@vFLU%nU)Z%oCHk4`VqVgU=w8pMXX1=i zfE-3D7EOjVT#7lx%w4au=9hlwF$xhhoA1=4(Rb#16P>QO8r#R2Syv2Kyx~-eTHC$Y z2DQ~8f8Qs}G_YEo@vO%I<&-y3s8^u%z{QZ2CWF$Ps>-Sbp8M}m-0^sDHbwm$GG@d4 zM=u}$R#mMeWOd?$ut8naOkYv0x+cQwF}E%@C5ctqDAg`JQs$xD^ zWOL}ip8F{WN7Nc1VIVs#!}^GuM%9OdjX0MYi@xiTx9ryPoUZPq?K1Vbj4@J+Q@0 z*-|L^dI+WOTZ*?r9+yMmrpRWD+<& z%u872tLN^=z-rtB15)46D`(-V(S1{OKa^%==uKlkDz5=FqIvDqf}I!4t|fb zMs3|QFn`??r0i6uF&^rOh$c7*KNdMMlkJ43JY|UfgE$!;2Z*8oL+*2>C0%3`P=-pDbnYtZoMeKl9)2)@1ePIi_6~ z0P9IlR7Z8>RY~~51z)eUq%KhUDX_Ibocf=%O4f`IqB2ZF zZ_)W>{`A&@bhT^T0rey1O6jMLn;sy>uJ{2xnjFwMlq44a741|5@9iNPPF~T8k{Ioo z4RE*}RbG?!7?DV9P4Y{&K)&3^y(f8OzfhO1fp01Rk?z_4yChnDN%O`&!=O|d_S`Ny&b_{-qeiA-}`r^$nw*{Bh#sg4C z@&fPVpjsC=keBhx*35jMYp~sm`OEbPGa80~u3_rY42e4=%?(pG- zgWW$Q7$O%+u0JCO?)ov#Q@>@>o^YFwH8Au&w5 z`s!7qv(pK%BX{>|56PXI_T1|q=b?2v{xVWJ!(|CgyboS%b{RBe^Nl2~Vy?I* z^y}q$F5q4Ie?HaUjR5?_yvOeA1b>|$qPy~htzi;c46f$}>INiD?(39P^Hdj)IQexW zS9rhKr37_Rb#%&K<^pAuQMpRir1Pt=p56MzHU{MJihuj#|74-*GZ-BB(4uG}-d*G$ z3*h{832r@U!|@#@guO{RI?b|H_>J1|ckMdx9MB?x{qTv10{ z##PR^+4Tb0k>y~)bwH~l^FfTgC72@aM^m{H*7JZo=V(Kufn0bv{Q&MQ zj+2>;5CM^*azgEr7(P~KJac%oJN^vkZ!0;xH5f~6g>?D7AV1$3xM zI~<^Cn@<>;8T`A1VLGAC2qcfV_FRanwrQm5RRrIu=(&C@{icZPK-?RNDW_3_BPsKKd3ob6*sZn*YT;rP%!?n2?Bmbj zE@^pGeDdD{Ho>&c{0SD%>T@2iEKS3{_OmwvJUS+u%Bp_L2qDW6I|F?};EBhe0ZDH( zq7FniOZ*=U<~g2XX7lL`&5`*SXhpmHq`6MNfEeCS{by4S)scNQMYip_^VQd5d^tjL zEMp5Z*X^Z0y*$&1T5`wVfrtbALj`^Krl_`D)1cfUROVInTWXS%CoUTRj*@ZDc=W%U z^F|oj3vlH`hA=as?$;em3RYcl|0S?i=h(p8JPBw*+#`$8XTmqvlXH-eh)%;$08K&h9|TioI~lCZvJ~Kx z{JRQ7fPWCvnuj3b?Bd6Gt3$P(ch*q+y2CmztO?-#`I!$tW-MxPKi)cmi4|jGmO%tn zyeIT2kHo!h)Qy67wWHE-d*F@@{^x#qNQrWlq?Dj!`!0`UCFqSP>cQZ1peMK|duuZV zBEBE}B-y)gaI#tLPDb*q`rf$Dt%p}>a`LTojhK$y-F;tOjfz~@o5MR=5W|OujE?{) zXtrRjzg|Uck<2v*isa6F#_HAN0)?&{2yFM^7S40`fC`QctlvQk15eF^;n4x#CA#VN z-ri1CYSM#l8RR8xUcu*?SqL;ro80fOw>qfTpUp`NI_2f8gxmn8XbiDUB68hzNe$@i z*uhmK4Ch+j^kP7K|KD3bXJv@TE;=0K<(9609O6y|jBkdFYosnP*eh82nD4B-74}Lo zB{tb@4pQHhX#5TYn?NbziqH-!=eqxf+weYB2ghvmT z>-DJ1*(S1KrcmXHawqGsFde<=mJdyh)@t`a?vZ=#>TAR8Y(tcDd99?{Zy?nP+FjDrk9SDS#I_h9QJNLq}`Dd9Uvnkb2a8jK1*`#cQ#94 zrZXSp(t0&~!~1&lqjmRZY4xh*v^C5^mMbCI8E3&v@jbN9Cak8pzDQTZo-H8YMU1r~ z&~I=DT4M|!HO;lnSHvG^-7z>_Sglu1 zb;eU`^>3)YP}cmQSWEB1bN<~zTsLgkmF{AOxEbQLTmFFanXz#|X?*BNT_2<=(l7O~ zp0B;;?ek?CAA#r@WIEuJQd)9@u=eGcV_ebbSI}$Kqy$D;PS$o0_=i-Tb==V2&}bGvaqYE~l7P^W)Ku7lVa84eiO&ZDjAXGu zU9o?Sy2(P`=z7Tdh~|O<8J~>bcMBDrkXU+oxa+>ez*!<&u^;0M@RpodLJ<}d(5)Hk zpDH69duaw1>0q(^Ja>|C$vQZk|E|)7o_);e3k{abC#9V2?{ctX9{-n9;r2Ji-l0g{ zZN~0RBHda&TYI#ozx-ppz1y|l_L1u+K|=AG4fEt8+tv(2)Re_T-BCZWC|Lsodi=j? z>;fHo8`vxT^XGfRZ$BQZ)pJd9(BbEMsZoE?E8VF@GI+Ad_TTLo$Q?uRhx**~$ry(4 zO;TJ@@z=<1WU=koW;pUs(4pm5M`@pKqSi1^?02eWAQPdS=f|*Pt*I63OtPTGP%HKF z-tG8t-^UfRWVhDc{T0i0)D6htaQdsdm4<)$Y4NGn@q!(k%lbTv1oTUwU^vs)6?wJI zSf=4XIrLG=sAVcl$|lc0Kxj|eR>P0WHTCRm90y8}MIcuK6hv2P;8*SP0RVqp%5t$N z+6N_lmn)XkaoN+=B3n&th7IS@(Z<80bYNI@7%9R8wG!gIwarp%ZzWD2H!pXn3;sks z{p08eKH*xsirEJRw#Tm8CH7%fO4oMxiXfo5B$3E87c?cvf+6~|R}v>Rh4vp!zkp8a z)@1}}RDN{x)BCouR}TQao5O;4QIAW+FP^X7oqFKlb+mmYyC%Vy`vYz-Y-2NqJyyDk zLya{(beJ7DHxVc@J^o&qUOymqd7tX~bA)&mCLd}6X}hc;-3Zr%ZkDT|)H!`HhW162 zI#A|EKmM-X;On2-(D{Q;mcRBYZQ0XEPm;a|bJB4HA`6_(R^W(gLwn6T(w<@BK>|Q!?DfpA`Y_SA zwMttnl`ZY0C(;-YwW z0V1~NZ#I&YHO2LDPdfWa@tae;$_YPh=+3yIB4cJ^x1m2IgoBh;B;Dnx=5hN!DclcA z{(7V^`_0mqEJ^l*YlP+NC)nvtbFUGBm3V^Ba@5RTJ!8E^Kfiv z>*?xtSBWb>10HCyKR$Y)2?Qk*pQ>g-SjV9Zd`-~9Xaw-z2eisFI@ z@xM%P0$uzkFZ|jrWZrItUGk8eBd{@@E9QkfPy(yB2l|TTVCZ!_t~qr2mWbbl9RD`^ zGYwBw)d?3p?`yT4w^VNf*(;pLO#??+X)FcBWPre&HbX9E$aO@5?-5RSOV2em=yqZR zHasm;WD`f8`(1e`6n+0W?`G3&KE_R;-#%f+}TI1kl6srbITe57WosajqC zutB_)@MMmcmwY@&49~tA$2SlqfVn)gNHJZE-U8vk;FDiw{4Mp9we{HNZO3Dx5$k~K z?1d1AE$Ck0;##12@E5aTvE+kwq~XP#8ZcP~-E?KL0U2WGTu$9H{=^jvx3l}+i zntoO3>GOloJfnFSFdudeyH43Y8dP*GFMSfGRm@10*@7lgWm_!vPgLU zVk|zOP~RPvr5@{6$DSSwDDNH*9l7b0P7RRaySBxH8{zR?;{Kx4{rVlrR6Z1dDsnN4 zC;N^rV62YPLVdilBp|s}+j{eZ&MqX=0|Oqs5Jxx{z8S#UJop2(?hVw~+N0sx?|ziZ zRX~FHSI+<0{+q+-z9g}FpUdWk6@o3x3bD(-vTgu^kGpSyVfpEtg)w**;m% z>-i6+9kIK-CtUi0j2uyJl8vVFpYePZ`H4&EMG0>wX5I z3!Z0J`rK?zzTjEb;=X@-w?C|3w=&8SxZ9bJ`LwvytD&AIV5~495CTT_o0~q#hf2L3 zt>r}EmchK_zj@<}3^9KeZxWCW;LD-hm?7ZBb?#Ir{_3c|Q||XSpgQpR%9a;)bl)`# zHf4KWXrFe27T>K+%lVzAbfOO06aQBTdkxemV?Zb1{@U;S9S?gj@W)$6@uoR5(`{Zu z0~KqT@f977;^z^Q+>2&nyKBZ zR~Z*zg15}LP6~&yLh}}pBTG2k>GtyJXy8`#vqAk7EZ5ey9N^l@saLCs{% z1cRWi=MlxO_go6SdmTgE7o1&x`&{-+)=iU%cI9tti9UxAU%@m1X=I3YgL&Ur86S&N z#=3*#7i%lt6n?`uFT3%Qso0|(EHMAtlZ!y$-XmZ9#D8b@B<9t^2n2Fi&Wp@)r6S+` znK)HAOKbeFhq&&U9Qz>D(y}8>% zSF2)(jk8ec!)0x=j{+B;@R`V0sh|gKdhQH0Z6@9p?CWeeVPufjF`yJ+zjS5a{Sz`V z$4}uVA;PA9Bh+JqMYjI*1mfBIwi_Z{(?=Hl*>4$m%Up}rfk?B!7!-^|AKMfMr2{MQ z(bd*y!6wWMytgWzSfEG>tW6*Rh~e#@AtV8X`rcpt>3nrme!11;fgP%hTu`>4^nPuPiBKFJ9=0&&TS(z;p9Xp0^HKD!x)W&rP$Z+vDaz@d^}7 zm`AKS;YKQVPd8`mr5w@(NVA}42e$)a0YH+F7w}wSm-kk3Or3=YIP}#!Yn7x7S`0c~A4uZN!5cZV`ig}5Ci^f$@uGho@eP$mV{(>TyVP2VhsMJipk&wn;q;hk;!{rXqk z&|U=Vceet=n^OEH_}<^V^iC_XhKdE{R<=i&VV4J)bn3CgoGGt9l)h)T`FxwY6I+%E z8t?^rP0;nlOc_)(=&vpV;1GmY1<Crwkv9s^k_6_41B`Vp2BzN80?XRevAOatu$hy=ZKB+^TpN1$jlcX9vcPl2)sCY z^sY3X1ytOi*UZ7~|gNJ%!W3?jXemtUsFK_btWv_y}3e ze>8bPT(HDxtyM2YryOe5cmamV0-zUp{9@~bkK*3;*Bt)D;&`E+3CMTx@6p8!6pZ=4 z$iM->e;OEYvIIfBDkdAKH$DOP^$#Km#c#Y13+U>-NPz*8uaUjKW|HRDDP4a4!P@nj z(|tFOG{s1!-*DM?IzG+$q(nBf$=ngD8*OFC{su4lAn+SK&YD3(LM$<0L>x!PB(Nmq zcE|SRX-5Rx7N5c6JBw(n8PQ88zi;qa_?2JfIJ7GJ@G$R7_P@M!55W6t1W9c{9ehw9 zbmukW+Jib!3C(SXE>Hccev}>7nRknx8e$!7nYO*#`G%uBaf$nnDh56XZ5X8j_Ga)v zx= z*j1!F-km``!xLi{DGvBqJv4j}03N|LXpa%J6f6X9Iw$nZwmc#i-rjU>z7zS;3tF;T z4nF85FP^Do&U4J}tqJ-RLeL9ZO6c!T;D|0v6}ZpsEW0A7cm;T$?R zCkV$GXw4o?5%-8Q7pXJJSI0D)UAcY2`28^rpU*9NKeF_3%NK|2)_!to#8hxV`gqE( z#o2IHy8hAtOY(0k7<(>3epg`$0(2T<<_}6imYbi#d-tMJjZRdo6=_jhDuDu%M85Mn zAO>m(*ba!LKHu`fUAtz#Occ__(d&jrVVwlpD*S}*g&U9q`gi-`iu2T*foR>80G*Kw z7jk}t9gv6eMGsj0kWCqBZfq`AXCo-FWlS3&)PU6xKZu81jUXsA0@VPtIMxNWCD(C| zJcFsF-ID=~Q<1t4W~0XYGsRN-1Y#`T_Sa-L#aMpDM^=^~ze)n~m{c6y$`9JzbO(}* zXe0$?lahkQ{POaHs>8w~IlFPUlJ~cq3LiLN;v_KhYlIlg*GI%~`uvp;3~{W9|6F8<-=rl{HP;IOq*;auRB6IbmpQOTaI!_) z*7;*;BInVubDGJ=oHQdzL;@K9KJgl$V-oaBAb!$s6^Cm%! zX4;hbUOVKw|BiqfNYG-#Sc(uq3+;;%#$JT0$iV`i{Y6%`Rty~SAQC<4P)WqIh_YPY z@+{TC%UA5}?Sawf81i)osA|k?oTywF2~pXQqJHo!GS9q+Gn@i5%R{r(NA2bAny{B3 z9KR1R_U@cChv6?&Jnbu0A9VFlSWBv721Bh1_W0(;-aGsFZ_QNBRA&pdDkie9P&oJ6 z=SK%mS0QAw*dTuvAE=^nWnqKhKlD(BB^*a-?JL@ar{H_7Z~W}zLhr_mUX8pGb|*x+ zmqz!9ITbN(VTh8?kr&QU|BPA(?~*H!ft!GQYCRj$;pXfFYjIYs^lq=)s8i7(mIL zDn5U8`wME{(rAtn8H(0lU(^pM@`}IL%Tc&16LXA#{}&AQw*fF1fTTAuKX+j8S_W*i zV{<|+pn1bjD0QUKo(#E&*w+Sf^<`(s{PbcZ#`@WJ{cwK6*NHP)?L*7YW3bJTyu`7< zk8b!B^nckbYui9skj@YFwBRq}*|4&KA)TkfND)EG9g{|N2lltMwdH;fOqn!qIQzD3 zXqGE-BnpX<0sC|p9?)p^;$|0tE@&KncejF{u-}5lWGTMBtK0IPbtLT!J;f~N_0l`f zI6>#Z*XmBuNAA*e!D*65u-$bG`hZn}Yw8#RaA65N@MborKa;E}=-yZ9yM6z#W4zNx z?IMHNRN=9h6U3R{;LbTL!G4~CRakuU6y1|n{$RnSxJx$Mz^9G0x3KV_O{>edM}cmM z92xm4mov6YBhl%nHvqKMoR#heB7v>_QOE&EsCmPp2d6pa{WAax4k*};&TPJ4r|@Ai z$Mdp%?CukwO4jM3mc-0wIMzgL(jqB(ZDh0fKwTJT;uWJm@)lSDfg3h^mdbBU z0+P}qL)MzC%zO%t*3FOGv@33YfySd`m9(1L-w(62=n~`~vIf&j+*K!ZY22nA`OLDq zUKD&(`t|lK5@N`3(mCMqFAm2+U#-_bB2#Al`cr>sw}pD!aRhIeFvEwPHr&%KS5L`H zSPv~cG|exof3`J!e0)~ylWgbVO0!e!(ZFHUkj)WrcF>&QY^{JOzTv%?Sv-qRlne&= zcTulv?>QG6icZo_D?me!5v5_lBluB<@7G&Z_9Ok>B8}9Y-5$to=NN@{X~=)G zYTMFy>Yz>#+xo5gulQm?6ic9sI$c1&HaCbDG18SJ6a-ppI#4AXW;oklL5)0of9_DB zE#9;DLOl|G%=|j_|Jr2ig(%m<*B74w4Q7fXCe?WFIA`bWz^|ne9f2Ieks-6?zl7&` z9^6(bXk8SiM&VK*!qm*zB3Oxm-6~C@7TN<>p7!?b(QJ|<{^ z#U6yd9OyO%3&FnfAH?dSMIAJ__rLxFD0Khb8Ic$QFA6Ak(5RaX@vG#`7-MCuCRINn zpws$^&l;)F;`Gedtl9Tx%<{=!4HKWGAKPQ8^eWWTtLK0nUyi~*2wlJj`9E<6s~iM* z0kDkyaD6X9NFW%Yuiz8#T%;LgstQBFLAy_{jFO0(G1?aS{O)aZVb#~Oq!UVU{8{I$ zwhdHEcHtiZ3qy+6_6Bzu)^VnWg^K7Q)ix6=f6US4e!5KlP;cf9e9xCePXF}5^U{X0 zQ~ca3&i5af`at_-bbuk^&k{6-KhU*t3&y+PJc#PrSH#&!+3eXkse$9c8b;~@PUMx=PRi2#6iM2ocMy_hs3(`Rdd)Ur@6@ucd+Gl?klIVbU@qFWS1sUgYaCXyYE9zI^*EC6#)fjf)^{G|aQQ(Fkt2g8LR-l! zMFv2=oRQ3`&hRC{3$C5aC#)mToSWawjRk9o;w9OW{nIL$DdNdIUQSO_MdQru?e`!M zhcUMyLOl>IM$NMK1Dzri2ScA4+=r0{4wrzJP!p!+c5b-OU}U54dN6&d9Qe_)%I6OsI78nF*O6 z2z(CX&rhlc@$J549jsF$Q(qM4Y$-p+O`AFebO+cOSHr2I#JX6W z)BQ^M9GNXKkQIMhI+u<34h<_plOc|zMN=i$!Q1870@6v6Z|q!p$idvO$ENLsXIsMi z69wG&uID?jdwL_iJdpAdiGp9W@x*gue3d&)hJw{MH>^OpDTa~(CDpJD!0a9S z>g@+-XT1|UZJ=u3aBADu!hwxvkFv}>ugS6>P{_%da>BKPy-ft;cbchCBQnic6im_t zFV#S=UC{l#3J`feu6eanRQ-2J?f@ z<=72yZ?9^5#RkuULUT5Gx)?D=nvwh9a#r@1hT6v2vJL%Ry!<~DgCDYsD;_jKAOHUX z8T3ChCa@-BZ=i{HSpc}l*uVn=MFv38+XlG^m-b1hJ#q2P7pM#I<}ArNmO84F+23y- z$|t7ECzj>(a*nb7AGs>t8%-PqQ#mvJ(CnEdbRQ}nFWebL)&XN0xxCM8WFVc4XGjfDf4T#grz=%heL%#jaV503lxwgG zjoT_eJNvPDT~)BUn!aH6_K5iBns|iFq3Z^V@P0gFuRgVhn6S}7`8H52X5>ngMZf9Z zHN9bE>DaUVbJ2a_v-7A7#YDn20-Ca&=lV9Ja01~OgMMlN++Q|XLdxJdgtPWR1lt|| zIeyS*bp^P0cJ0<}b)4Lk4t}?da?Pt4scZQM$hk+$Jib(fd3yEh-8bh!7wHvKt_5#g zQJ(bI!&8D(onL3pxOd( z?gdsYkbD45+B^AeyiT*W_Rt|xut9{I+tS6mw+8N)G>Y)fI&KAl8TC5)d1p5Z9jM8B zz>e}N47oI7J0rjB<#L0ECLe2pn$Xv z=%&ZL@XyDhykEs1%c5B8)?c6}eRfs-DfwXGi=N&%J_e*||Cz76N6wH? z1c8LGuWlBUANBho9x!VbU$c|~CFtKi^**i2UtinZFS$K|b~j%t$kSu1zs4uEV90Y} z?Nx1`&8PoGpxPl(sb*_QU42u5OH1ECytfpG!7|~~1Tgg1ez+P$A$>r#g>GL2!$O;{fhTrF|Cvo# z$GacHw?xqc5tv`Jx3AOmSwvKe#_#!(sT$3AMHORz*Ej3+5NhG@?2|Lvv5f5dw;(*-l)=$M;>Bv|8m`?8$oA7- z`1tvSa5brH-yY|Ol?wV6uVU7USGsq`XTC|YkU$@&TQQ=THUS4y4Np_j8>=JsuMzZE zl6#??RIV=`lb5?EgWP$=wFl4YA8yD+7^8fZrBNRMqnIJ7`8q?!kW)^OYO8*ooA&R| z^$}F>cJ~lbDL_~rR@x*$?C_fTE6R3qQ{et^{6rAnqZjYC^CWq)5a(vLzF5D-c_~#P zVz{R-o@cZF?X!37Ok4WlysZ89a%E5qrxO6Xap=_IdhfoBXpC6%C|5)^*bvE0i@V`l z8!t9%`U{obrgu$^iG^M^RjWGq_5|<~hyvS_i^&>5i|wez5}a8=-}eUJWt?Ib6X$Ou zcl+x*>IDaK#yt=+8gg6a>pC4>^5s!9B0HHy0I(~W0kB%k7l<&@p7;0qt}3Wq^V7Gb zYdQ4OJ4KIPH$Lg%DA9p@!HmRExK7xPYaM7ax}WkK#khS;&}uPN%J)@ ztpm8S#y!MQV9R(L06dh$>?JE8UfnZ(n(~6y^3%;)cKjW*9Ot8Oqc%Zd?0_BW#ubQg z8BJWnh~fx<$Bj?y7G@lS{YlJ7ta!tO^Q4if`djagQ0I_cJ@UPN>7^P&$B1l$+#7M? zjAh1Qu$ks?%BVz_qtra`NfycVJV=zIUa$`mMR zgAquAG&EfrogF3|o04g42ji}1$%lPrlb6wY&s!q6_fIQ}pPO=~cd;p)Au7Fuxyz6$ zAF7Ec?RA$bnmhM0ZOXIbfp}E1)7@8bN2Nf;$Fi=hdITK+&EBzKrht7d!?-x}BDh&Y z<@M~v_VT{+qDAB7Lx@wKO{6@yY))jyrtU^UBpRI-Mg98dj&Acn4Axn&ucOIM(N~f! z^8qn-D~!W&t(YP8fv$ozABTt~?y7KF6iLwEni7;kGVGH0X4lo~@6-3CcY6F3A9s(T zSklm}(1+9|MSV>?kp?=bf_Dx#!-kZuz~R;h4WroiqG(kQWyIsFN{Az~-8)4sh#w*h z%o21tCIP@@o`t;E>CPcvaPKR zQb}LnxTf;ORK2P#-qNjimkBo|gQXYmg*g|j#XK*;d!0^v^!gI=-K=y!xttvYHEu1V5Cg-0hY)y=Aj)Xzb^D^^kYAg z=QD1)Sa4>bokUrTmF&G`^TIEN9sMyd*US5U$(LyM%vE=9@c2#@Q!XyTMv^@k!56wn z4%TqKJw7R{+`Ji*`aazkT7&*F*Z%$^saDC($}um%q~wo^`9Az^BYP0^72)4Yni-N9 z(K_Qp5mangxZ+^goj`Yg$oJ`It517UpAYOA;!E6B2_`K#BS0kI_6EkC$NXtCunkEq z5a>wFRryJ@lqV0@`0mBk-TI?cxGj0!>PXnq&BsDqoRJk9;AsFsir?U?XddAPV5Xjm zZ`1Fgw*@y@)##)I8l186+Dd6z3q5&rxi@m<{8gEm;YrXHf|%g^?`5`uE|O~Nx-~I4 z+2zvoHM;Xl%jS+VHxLues;^xi-U_~v@^#msR7zyfy)FO6Y8ae9F_IwG1Q~&<97irZ zrQBWCx*!jR?IQ0-#F{1D&6;oL@3`E{b}+sAzRYz?V#eA9+HK}jBT>2oLw1Z2Mfur# zJ2<&hkqoJY+`L-pG~6W<>KU?Knj@`BeT9^+Ry+7o9&zgdShG8GjpBpYD0qpq_^FyW zkklDT9$pDuM8b|XdA6PT&Ju?kbTyhJ6v7X)T@=x_6z?K$O0J@j3{AlQp;cA* znYrZSO^xdH{_?OC@nN<~Ds2x}kkp&U`=W87tdEOLxN^EXYm#LQ7gQ^SEJn<^woSPo z9edi_V0pUIZlsm}ZrXrc^VcR&g~k21vEhGzGw<1ZHlcqH53`-Zxj_WrvPY^hW${V^ zbss#EAw>*C_W>bmiF_l|J0ULysjDsdcg`H|^oiLoVzi+yBAO*G&b_F)g^K{VJ`gVr zOolTU02JV7z_EKf8eMLnMqEt6P~3v~U_};BUpa0jKRovLqZ%VCzn8x4oV`ii8dgxP za?{~62OsZ|AWNYmN{+~1Xg+fVQo)n;7SjN1v0#Zqi)+9X>$nMOyiokl1bFA5;uN!b z2`qstorjmA410@^&6MIaEYH{_w=_NTZ0nBZSK8|8@#fm<(Y>Sx%c1f=@BWzqXARPi zx{MM3y~2`Z?4yJh@uqKNTjW8Q{BQKr{vpvyk4;t6++MnREURAbD3DI%{$+paU!6t` zJn;ugoHc`Q8I55{Vy%GDZWrn_^*7n49Tce(lnX1G@Qm%iSg;Ww>vKNG+A|**t3+T-TV1>Sdy7P-LTiknF%)+TKoT*O^6!?X)1zh{-n-7AIQ|bnHDiAKCtj(%Tip? z?3A~+hq%3!5*tY7D7S6aL7N}gL!x5A97Pj~@B8e`y{5v0-%Pu@LADZ=`@@PcX<-Hp zS9)^9EBeZ(E_=L0d!ymwR|j1FK@4aLfbx|vP9(w@$+5Vo+)}Nd@U@LK-kP-a)Y5GH z;it^xdh>o_dH4Z^gkK;w5P})1I}vECfLzBA`wRwhn1T_2mr122V_@*q9bj`K|^Yh7>jlPHkhS(k2XMms{FV;Jzv$XwP ztZ5zI#5pHt9lE)F)(CAc!m7jo- z!>MD6^8sPz>-5Xs~*AJu#*1lGnq>2;*~K@tweVmhz3ocZ!Z+xZA+G zS4pUzUlfN5begf$P{rDU;s@dZjfG-g5scKOr&-RWlEQhXhE*3V*`!gU#sBesR>9{4 zbxDd>@OsD6W^6Gp%OGN|&s}oL2u&K@C+D`33kX2J+)cl0USXCUdZZX&-0NSN7kU20 z+-s#$Hu~MXkRG<6yz|IiT#&mWJp}B@Q{71L^_a2TM<6D=owC8utnka*H~v<)HTTE& zt#cK_G2*^^RKFx%QvH4k3({d>7hEs;*Tt|5hLUbD%c2DNAz3X8mZXBg^QP7}AhnX? zno-AEkRV}Kjj1$d>5ue_$1}pNj8!43ICk^Ywg8fCT@uxjP`8#S-@ zCo-vtTNp#A@M&#-;4m(nEtL!0EpPwwn$Izk*^8hDCBvmO40yTmkSBhQeaSo{Nu9TEZTmk3;p<`jNRB}I7CF@KJC(%4tDo)dL%@+z~X`kI!zl~P^!CLLeC92 z*COK_)EGa0XKgs{L|*pbUAYq?lxHWOm};jk-!(y<{D3J`11Xj>uJ_Pu=*{@_v7_JK z7_|O_P@Gxej_!`Ahs0FYS@wNdC0peFn?g_W>tu`e`zM zpp+5YB+yp{4n)~3p-ap(94FRp#t2o5wa>DKyn9-$y>5rjbNz!j-{PFN`FvDBLwe)b zh2}&7p6`OUF!ikr$po4kv&;Cf1rZg$DBNF4?nnu5YpKVN*Z%77Sm^&9hUD26GYjMww0`!ZhvA|wzE$1}u$qN`prHGwoz3;`+} z7k6_LKIMGpM9@9;6I%;soIQNc?>_w^E+#kd)#p#t;GLE667_O3zIjdd^dj-`O`D3( z7B^+WajMfnHS?LVY)aa0O*a%2AB1^bxqpHq8Lw^t$u#XZ%>~dFuJa`gn^SMXdv%@ zTKWPaDO5vS#}QHBy>o(7>cK4r7X0FyWE4~Hk|X64W?p|igYL;V%*1!|5|%RJjg{)D z!EwRhPGrnJ4Dr|6Kh2jyM9j%-gmJ_mO);7gtt`P5W+K*?TmzkP>V#DU^lvwRZp|*w zOKMIr^o*40;3tQ&0Zgwo8PAFV zA&8mq1WX#M`LW8-9w=C`Vx0tn&SZ z|7XObrsej}MvYosHv?O5p{HZo`fwDNreG*g`s7uu94_sLvyL)XWve>RSaW>d3%w-!PG3!FT8Q8Vs})9>kpB zxSWh&arplix0z3XGoU*_+dHXw?Nv6hy88WwqPn}rG_(XL1ANsH;ua+c#L-<{bQ=VF8xxh?h>G`3ASn}A@bX1DYJ#&2H#IKZ*v(86i%dyQ^GwHaB6HHiAg9He*pQZizBQ~yeb+*fkynA`FPT>?$a;{0@{5i|I zhY(o}{V{=+oLnBqaraPSw~1@wO#fv@|qHT#V;xUO~-QuGO** z<6g^hy}NNG>Aj8vC@R0^I;h8(xq}tRuMrL`ZB*H}e&$ik*u|iTg!NNkDUf&EKRLcq zZdE;DNK=mI#jb~SWQOC&-^sM^8?X7Ee0!B;<$N(}QnUF+FT&~D+zR_Qz&%xyWIP?F zhT~8WR+IFgK|Km^BHjM+kJSd;x3I6iMn&Wzb;vAL_teZ6a)eK;N_}chDL*oQXAA5Yd zQ>~KoN=S7?WE`^~(B{D(jh6Fn4O-fgirJX!Fh5J^6f+eYPAUkr>elsXm%gML+~Aj+ zr{#8P{ktM`8M%ZR01D5{ED(h#A(-_mC99H>SC?AjHdA>pgUvBtt;z=0~uj>qluq#Y8rni=Zk3tUcuDN^_`az@=*j^<19X>U6p-@@D$TE3={}e ziy6YGCEagtfNIIg!v67Etlr=G8dmY0uddMB7ERgp`j&<8Jx{KDdC47e=`Lo#{&ocz zISR%)F*7gEkl+yDMf0YmTXo{uc(bdk^UvTr!P+`Ob)vfpS~>!;K__pN$@#zi)GAX1 z8sdO2@PBEF|M#9a1kU7tH^skdvXCHfJt)8$fcJUD8#t`FIpAg2==YdyXyPLbJ47sV zkw~auT%(9*q)hbO?CO6MLhDH3AJ$%L{t%>5mX6Wrw)>F#@m*q{)hWeqNH{K!26&`T zH34vl%k2QZ2oKo)bWrK6izTol;HggigP2V4JYVvAv5h*5cVHs>; z3ZtI-Tzf-5`>NGRXU<>ZoR;MSf9~dyYoLiV-UqD9}XM_b)d!O!qzjv31_E?o`uvSTXOuqZ<11)?@d9OmXEW z!!Y|oD_!XVvy^^f(=;u9biK@tch`@FEAoHL!F;j8Wk?xn!4nhEzj&$K%zE1f z;U?swx00bn8R(N8kTshsb6lIkg0&3xAGL0Mn*@DSd(%Gix?qui70upMc=%iWtLFQC z!Crg4D*SUF{7Jbk9rEx!cAlP+)tmok_yX7jRglMVO}kgkP;&l*v|vTEQDj`0)zuqi zsns=Y0u?g5Efw<-%-wZlNfzHIKo~YmAR-R|g8e_B3t-O-#l1mM98u9|id_oRZ@T648I;LNR|?vHu{@cP0jSsH9u|li(u^y5b%xmiy(+(LV@=TA{$lB zV$Oq`w>vsuXE|`R1VR#NcwdM(Wb?bX3+@gaW@9S7sO`m65LCb5 z9zq0xF!aZ#nqr-NceCJQT$mCrzkaujQuI8z?6GHk>ub4h0Ny@*;1C>(^(teSQJxjn`55FB z0Urn2YPi=OU)6m6%Ew-AvJqbL-}W_C=6z^k56F(*#gm8OXWs+;o(J$?9o(D*5>~f} z?-Ylg!u=?2c;fYlkwn(w$qCy2(yAuJ)=HBR>ChHvNdbcvu4e+eo5zKrcjvEI?8Lnj{ zc5CA%77l~<@@NVUHEwi-oOU|UK~3Y<+sHZ%uflg(Pa@){&#%v&W9vU+f=a^YCFigy zK)+hoH9*t}h7*Vu8_q<$G{Y2bF7bDlrNmbVHPU4H-bMt~)@m#}Uvb+?dtT&G6~rrl zIc)H6DT{v&{M;r$O`Nyj|A)3WkB9R8--bu_t+C5Gwo=(zY#~#rkVGO`r$V-*LPnNL zWZyz4g+wGXMbJ~JfBGGiH6%v_)6^!fen=YC%I{e7O-{m=79!%JLq&3T>Y zdp+LAQ9uH4Kxa|fP-SG3Cz*L5ME`zvA^J}7vXkj%r-mH zZO93qA@?j9S@#|ky0y3l>o_HW;%AycfMyHq;0k7_91X1$Ff2|rjQQbC6R)-14RCO|$K+yGY^ zMKd}>lgWE#{v!*}k)5GiNjyvORzj2-@0l zC2qU1pjAAwrjHx_#Sp7UxJ|)6V)tjJU873SL~_MP)@3egedS@Me<7H?N;awc7JkcP z?3iOnp@50O3@2@KNA8NSBb-N`MtFU(K-VV%tfV|NI}&yb`Z3`052`Q;_fC{%!4UX^ zaB~9d!4}TU-hg;}A<7M)b!cGsI?2M}L-bM6NB)(p_qu(nxmID{W7t!B zCc&RyfN~_3sB)Lf0Fd{rnd7G17is^XlvqDb(=~oT&I_6WurRd?;Y*^`Vgzb~46e}% zTK(|O`n5h*YAP4rJB(iCRIifY_&9r;_x90T?=3NEO0-XtDA`=YxMt0;t8na$8|Rbp zyq=_8>J}2ZfbmfC-;OxwiX(EwEJ6aSJoExc@2xN-St`Q@^XdXq4EKWRfb#r#J#0Jw z<(YQd?65|6%@MKs%}WL&^0M-g+^Icsig%;O{Av$VfaJOce59>?iVOZrcJii(ccORy zc1HJW^lw_$_LaHEBMUBnhjpheR_W+kUWtrS3ctXP3hHIZf=E^QDb*WkfBz4JT5qst zSXBOTKzM;oVdM6NRXOr^WK#V)>ku?xsJR+?DCewLo<&1dpeNT@*RoR z!CB|ZOTeu=>ed+hS=R>;etQt;NQZ#Dv9dhyAC!gd(3Xrd_yx9UW=t&*YR4?P z0Blvm?&8pLJ6P@lcVgD&32yae1pAFj>=6DII7T4%wWwivwDQsQ+8b@IT@bD?x1b0PMJf^7nwfEMS~?q_}5O!Pb~PFj1ezEHpeeA1>;4Jd0=lmNanP9#_P#zd zp~;+@Wj(B3%Q+;gGKsZ}GUE0AB;_5l?+mhk^Y+K9B&s(J^S^C;Jq&z!962|=1bpt* zBw;Ye;%vKagv-TIPhw$tJh(f3JgzfqY$Q>E7L#P)Pdyo=mZ-Kn&&T_gV{2@MPS!qK z)%KZ3s2SyTZJ^yK3Srr?q`w+}g`pe7EhbGIc~E4#Hi`QY$PkZ2xIUm6Vj85BDE}pF zhfAxWCRJu}L{DpI@T`xa$8t!<`(*!%#J$(fp1stzk4rud_n8PsC{sq5V!bgCtOiOs zJg8A{M<;<)@vFQ#PG(F1NU$|MMPWPe28+kQ<2(XUix*44H_vIf#jaBPGwb1WigjR? z^TA=Y6EQmVm#HoCgtMHkCYjVXG2K3iT~c+GKa*bp!+bjAqGVFn4@4VgFJpAB@{PIu zcvKYT^zlyPo^QHeQGnxKf()EHfSIiloq3FTJDj({n2Az8cUQq6YiT{T1=ieEW=5!^1!C zDf!md&iZ$rieaSaX)qxOtWJR#vDpg=f!u-GLRLGn2PTtq%)n8^@c)+O8(2yEV@>E= zbWgBsNIKH0n|rZ}_mGY1$uXe2jKa+8{lCRNcu)ksL|6U}s_c1HDam2{aHLi9cMup! z*tbS(^|7iIQXS6(K)(CXe5IO)Go5wzUrT)+jwaXrj~FOrCk8?T+L_ChK!Td)4J4@B zhzzkQ;IaVvmkmn_7lK;k0st*=O?$I;9W5ovz}BnbG&K6wKr-}Ie8aiB)%;5;CXu-) z8yPccR9NP}Iz4qDYM3LgW?G=3QQZM7wBC5JvfCahC6bzO)LU-74AJ5lUyiRI{$%b9TaCUCivP=(+j_e=b-&*CNvG`%fWknVz2#3{EtDJrzGDwz zeF;Ca?hn~ZQ3*A+8C>5s4asPJU!OK9A0%;mQQ#bGb#L~U1S;ht5w?>BS;GrDc>f2W z0YS52ih(PUu+$4Fi8w-eaVu7aahQ66Ftdbhsls{M7v;Pd?Ro9>mN3vt`&?J0Xpo9N z5P1aJWu2kK-h*ySf&EvO%t6>Zuzdm|Hw^(;2#u~sgeWHmoISG{veh{m~ zS1w#40gV`qnd{&-%ViHT^KNM#H-x$d_d(bu|!|f@u2WS1a*^;FVnyctf19jKXhUL#~bel{bAd z=Z^pN_|bN>;HI33*l}{=oh2mMn~FjBx&Fzre6s5}v`GOT`xdy^dfu2z*BEZ8iaJit z0lHw`m|df=!g%{ap?%NnsjbJGTS+w&@6T%{&5m~_G(Cz8`&eYk3z=jv5b&X<@~G#1 zld-kFOD@zkN^ct!)2sx?k@MiwFQ^$q#AyG84JR{;rifMKPsngzfr%V`Io!_rNK6tm zL>Ei=f?{BBan1^7`5KT~i3=yi`+>_(v0Yc&#vxOAdqBUU-y#MQbfwLKC(b>RB#B0Q z+y8#-uypFv)RcpsL4M#Z{M97qlW~|i)%Hd)hJ^*JG`)0TQ5=ydp4t3zr?+j=5nB2X zMh_dBrkbLY0s-OuuLwidEAyB4>jBvfme_zLW25RO;qHU4CWi(y-yTV-aa3?*|M9T2 z;YfkurQ4iq4)uhmyY}pYR$>?udD}+t_>B1CMVffq)B)PQSCyWI>@*RvY!RL0y*_#p zFOGBz9h3Sj-ovgGr8FaPjG0*Z83Sup(s|)m+oEBeQ@VqO`vDaypY64wcAPaqRY10YF@M$!*=@^w9VU<8Mj)QQv`eRk44le)3-Ge5 zGK#0F%_8|EeX8oCtzzRIq$}2l9lN3A`kW7PA~3+wXlaM2CM#n)!2d_Q4mm^Xib4eO zD_w{v9Cb&m^3`uyp}vZZ*OW66`|S@Py~2Ck_$Blm`675SB|dV_6}Mr=EI<&)(Jb$Q z9He$78k{*@07tfMdInz$W?z4{^wTKgD!1Hr1d!N?;1*@}F{Ehh)A5J^-qw0>@7fAM zO}78{oKt>Cw^EeEGai)0^$RGJEiEByS1#a=M*w^tzCP0es4#S1PHAcc7A#z@yQF!# z_9m_Ub%!Y`HADB4^t~7S|3hvGn}~tc+?On9!ZT=wsO`*`Tm@Qm<7=q(=uDD(o4qyU zdH786!m!S!7=8HrOCKBSB#zrhzBv7;X*|4&`|*cf4ySHokwIs6n)^coPYT2@j_cn@ zfNCa<{rj&uwck}#TsAg7>*{*jCL>lZ;!G%rdfAh}vgslN&KU4DJH&IyKQCfhaEEMW zx?`+3K{E*l2F^YVmnqi}KjhBs^}HJQ&`&Nj?KJG>rgZD9F{BQr!%6X2f^7$oa`kJ` z<6&^a1>`Hs1dqMXfcT1W)XKV#vPJwpp>Ea9qph#w1^0v>S-$^YmGvO$v)GELxWSkOP zuS}2m0WhfwdN&Lv->(^VZxi*Ix57fV5icYV`d{znJz(uKp0Q<|DkRoa(q%z~g!hgQ za3j0p=tJM2`rV~;>`^Jhfmvkb*1aI2I+J)j43ur!iXq6$cS}2OjP)NB|6Jw2HR3H` zs5z{WenUN)amRGR^zID{5wEeVNYz@CfImL!zj1o{-ox)q>^*wi9@XOcM527&nCv}? z>B|Tc0n70MM>QuVVXA+%Tf{>`hI?T&4i0_!C)?T?GqBin#zp(t8>esSg)0hUUlQaV zszaPlK61E3O%mdfKbGvFm(JG@ulFSAzq|gP&=i-SEn#`&Tq2j-RXhm8IV-&31h8Rp zkdFv(QxQJMXzn2xS1~@tj5*&Fv@UbXzl+}dWYL`PQI1{p5(}MLty1ZQOsmX7u!=H6 zum?;|isq56?QscWKIo2Y11snn9!l$Gc|R~2!}A-TDfC12VIuut&;(n-t@npF#i24^LpQ-*G3P|8SlE^fQ+K&ujmmTpIfK zng9QD?f;WYL;qn-{!i|E=h_I%Zuxa|Pwu|soXijZ$u2{m8eXE*P9A)<}n7Y8@~ zF5l|kuPD0=y5V#72naZK8z;rZWM|o9!1!xJgmh5`%x4I z4ageyLpL$eSiVJuP3aFAGeD#4?=lzuDTHzf&u7sx4Jx6ZxxTXhE zx)`b{f>2*D@ID_$A@~h*9z>TPh@_`jX{;xDERfvW0YwAi+z?}$aN9-LS%Z7Z0d}{@ zvj)N6yNZ=tkv-l#nh~7IA*yd638T)>KY7p0g@M5HKkPb|ZMW?+$TW6O02`NwKY`6>c|ahfTV(OYcaklKQJiCm(Qa=at{0DOUF}C4K)Du!)k*wfX8tO z7=2byOQDo*U4nNIKWvHcs4_!xgsh4kMW&4gR+?4uCvBaZ`<65=2j13FdA3;_Hz<<% z)8+ZT!M&2F-HdPBJfoWZm&F7EpGQ*Y$+X@S;TiXE0#Vp2yZ&L54#%@M3-?1Oa_|0d zbTqwPQ+snXV7J2g*ZOz4J57~0D~Q4nsTb%BSafQPDqeu!&s4@uaSu1KUYjvl)cl}w zVWfQTz=oau&@WFnZ@Y`mx+%}|?Wk8u1$f9QOBpCp#=D(p3|l*r0)ctbgQ~gg=^fea zVS?1K$ZLkFuUd{pPxU)IIi249u=CeFR@G5m8&o(7GZ+Br;HUv0MkUUybO1`ocW^!l z7j38kVmPe0N8P>KW@PhyzplKi${v=Nj`QOatZ925eOA$X*X{9)$EYxE4`>g=iJarr zYc%=whT|0mtX+w`_u+6o_>x@II&m5(o(ZrZJEK${Fz+ z979d!luA}l$%@kZ-wS;L{j{E2naNx8yr?;>Ysqv?iy|)L0W%-CYN^R$u{}Y#5=Z%W3_t4@hR5BgwprE zl`HC6+|J%5v)5M?>z3`Fq*;6ji-ua>;A&Ue(L7L#OBeU|@QWOAyc%t<=| zAGlB*v(!*=hXE;~lo6QTo_k~@fQ9TU`A2q7zk)c&9{(lnP6V%3{xbRa>o>9QmG~NK z8pQ%{U`K}_B`nnrf)#z2w8*{if%;A#2UGYk4c+C8p^YSHf!(>)^)TdkaQTm~?ojaY zc8}GQ!ai;FWHnymC!_h=s2ECC6fh)O``fdL zfK=tR$JW^#Nz4(`3iaYBD^2UOx<4n#f=;6siM0?9N`*p%bt7j+60tW@V8y&nUvU`i z{DImVx42Y0E){z2OVqo%%Kx-hgk4m8)%fz3yjgX$8{1nHD8K$2bV2_! z23!_rNZ4(g;#b;nm@u`FV}l_VOTA=yV%kb5dumHdLgrr(#Vt>`C)U|Xy!0s(2_L_7 zd#^~!XY7ASJc&1mq&^_8N9?Q|O_@P=pyTp{npi2An1+VLwigyyDmv40jVo zd}^P0Ha%fq-_%Jp)4$mgc(_e6X?BUFZpP1@wmd+|iYrqu$i)bHKl;iWlhlvcxL&=> zrTi8dBo{= z0z&}UM}Tsm{2XCOKBdNVadzpohaG}5)bn28v#un#24fV3hjI)!!B44Wquvd5 zam9|zm%*<|W^I~c&uqoTxzw#=N0`@=hL0+W*zUkOVlIGD+I1gnEhi%N755J+b)J9p7rVnq!rAzcaE>$D&?rk!5yVCq)7G(fCn8!-1) zC~cD>jFkmbSAI77aJd9wUFZJ=D1#>qzGh9-v0el1F@C6Wu8Tgm)U|QntqaOqix!`&` zsIG7M`uFvbw@nY0UzxZ(Rq}L{NhD<`rskfktb$0hU^rf!f@KJe%%ntwi8ECCLJY^~ zRr$u095dVDrLpDd#r_E=hciAR2j`7e2&tnzE>b8?RpYfb@UHMcg%)nSmsF$9Pf+f_p`Tmk>4mE@PhP0x-ZjO6`%9rZNcBzm$%KUS`UG;6yW ziE9Fyv$vk!)qseerIj_kbi3N6-6GF=3>`idUsQe2Aja}Ib&x~$Ww?RE4A5o*_Q{(l zLg*o~&STL4uPIUSV?i(MRbh8j?|yMIaOTC(U$0BU9;)kFi<%bST=I=&qY zi<7fDoH$z=ienHt=b|{ahP{zL96d*@Qh$CkGa>guSEC(Sqdy(ngNkLBh^LE`m#Lm8-H8B0U@n=mge@{Q)!g^3;@fbVBF7OS(zCkI&q! zk@mObFr8;dL7tkRF%be5h?50GTQ(q)0rjTEYc4ZrZjXCUi1)nV?V<5cqzi&MRTDPw`=g(cTp@N zwCl6(f9yx6`<8d{7uF}t76|`@4x4-yFqLTOwP!q~`R<^2FMBy39IVg`)ga%?Dx7fY z*_-jWL8)HI{)FQbif72i8czui&ST>2?W6IVd9XIrwiUI)5Y(ZQuH5w9=VN9(NUr7P z)fvm%-u?4l*Mqe>3J1ppFc^B`v(4MHT%^dcaO}%USgxKS-AfCm)0p*6WxHXKl89ZHIB1K4ZsB-v;c?^D`HbexzrbD6tA6@hO*;n=|EMx}Y(Bok;k_puv{_ zUSXEVNX4Do8Q(+XLX$CXr-G(^wsBdr+AGS4(3-}xM0fYBwdIBS=QG(1QNWoLjPN&t zRayo?JmBD_pfPw$4a&Wk{1%EY_n2|nxL1SbzTvT8+tll;&SE#@hmuJxZ%-<+vSLez z7?WUWQ|||~Fw4UBjVW{_`I&B#ZRj&AH#ZALhhc` z34J@iZ{CxC$T(O#<-)suaO*Mupjapk_1RNrs|J5^9~8NsQaY`tjlOaDGG{&?7inbD z8fCl2x-;pb#w;L!Qufxyba*W~29c%S9Pwd7sl~sut+Phz=rb%K3z@E$AI;op4d=b; z%9ppWMGVg|Gk^CT&?@Nn1eM=fa`kYJa6R13TOyzFc(sg6F7lH?I2KPNr2&<*tS*f) zBl;!`QO7Aok)#g}ZJ;!|rwJ=OpB?KMZTtgXVr zh-4l^bbuaU<{eZaiTUwO{J-)x154 z_cifw^*>m{;#-?>T#8NLPG3BPTByas1`xfvdbn(^c>ocezS5?5yrbH~s&zy#n*9+r z_g(i(J8zv1?_bc5k zP=g+w;>TVAeX?sFko&%J^7PEy&FVzH=Wjuah#U9{S<`K6uUw5Gb^qfk@2E~bN; z`r1{YM=1UM6dGg56Z0k6gyc5E?%+8qobn60b=o5*BO~K6DdNz{%03NPtMbCTX0TQO zYXHKDqU5hyqs9KGdUnb1TAuh~Jn8f;A-&Vi*lIuP7W`jP%1)%RGb-+m@T35H?iLaN zOKb;%5@{&4@&wg6iHwW6|MOQ=-0=D)v60wN{Vt7r*-Z7dpiJN532q+fjcKU$#!Id%PM(d~*?k>S?SumeKk+EqfACrmUH=4F?S)Ba@cDODVs9P*gUa67p?RZ9eZ#-RR)1s1 zOkc&I6Jr4lkl<$BF1?Xe9oz@!J=EnFK+IuU;3QqD_gLlSMdMlR87^9g(9=7Ki-Hl` z4L)qYr29C(LjNNo*{Rt6cSQ0V=nDn{|NILEd6*r$4XrZ8fRkoq)OVex9DP&Vxn`)= zgTK7CuILKMLsraPzwO;s^4pVg4<6+GEZQOYB6X|jk%>Ne37glaT)Z(6LSj$kSBJ)s zi72&!J3Q^X3lD%Cec^2S_Xg+~`U#d)G)=Yjjm)E(PP5E6^mgzv?dP`J13CCLPW3HO z8isr`?#!L9>Q?enNKtt|Iq3G7uWDEG{w%?dVBn?ztSWf0-7d^rfY0PFhvv7>F_ZzR zywt}$s?~Hv7?v3`PLHpr4QIJiGpEgc&TZN@e6Gu>jS3FB?$k?9l^5g^(M(Z1D0g?E z5i}JMMHA*@fPOHCm%yT{9sLKDskMsx*!$PYn-ar}u|Nv5L&L-XWe!*wwhtigj2L5} z!EX{$gzQ7KK8)1n3E^#Tj^xVH_Lgbie>SGwSD&h6^*OC&PA9X&vDUX%??sxT@<(WA zavlafe1A#0@x--bJlQr*@GDU9D(YZ}nkG zo!9{4Zo4aEWvl$9x0K_TU8K{}t4u9<1sWe*MgQSk(%F^>ImoPOI$pGZ6*XWd6JI7Dy@ zTwJwu!xn+K;sEzz50QXdTq?CkzJy`wsIn>K$umR(4hHRMFKd=<7TP6*YXxYv0|Z_vzNRCOHB9zT zuF{CRj$RbAGrKupq7dcuNlr{x;&`ue_B;f+qv3tNA2Fl6T}R+&EHLV8+BWNvVdO&$t=eJ*NM9KAMP;g@` zIdf~}M3r~2(hXFPT>QrD{_so%&A^ZPHpwiEbl}jc!T5+)Gl>1(s)JHdp`h*DR}_Q~ zQ>VJOO^+**&i(jJca{Mxhbx_c75ZZ7;)el(x2cV5y!EvVf4765e+RJ%c(-85?NHS( zrWCOa%MgP${=5ad%@l0HdnK%bkaSiHIvn9=W;aV+!DBxHHa3I3 zPih1>?_Mx9754W^u}An_zAWe;g77@ zTR?h;Ha!}HaIaA<=;}ioK1O8QHVhXmw(S_Dwz8~QD$iGcxIoPM#;oxYpXo=^b9dc6 ze-jWMkhpUyhkXO{qbn&^3@w@iexM`wVAX>ggCy#!D8lT0{dyJtwy%C%cJYtK#qU_2 zT`CL@1IVoz+j@NCt~D0u0JF$YiE)7W2r&Cs3VLG&#!pgLTJw14=f;q+-~jFGBOY;T zRwyjq^>cgb{=K6aCzKxGqVlNpURWuGDX9|=>fI>=*uQSL0_;xF_r(eT8HIV^wt}iF zXD@_f-&Vr%&{E^J5;6#cnr#Y@FNUHqWFlRQ+Bh7d5g=J~vgo&}jExnw(CTybsUJTi zT=ng4w*;!MG-Wk@W|6eo(U=axEVXUlW`I^x1U~GW{Rk)&jdjs2V*GC_#P6{-3of11 zHMGeovy;8KDyr)kWn&xYnj|waqOxj)pbf#QOUdtFV&9s=-4;Ge_FDIX=xfl(R&3tfY9%UYuAR7m)&YjVaj$KY)FiHqvkRo`l8!lS!TfI@i;TXUZRm+mawJSWi{?+X0Ne z(ZBN_4;a8~M?6M)?_q@u#eWarJT`?0?10ncjP>8gs zvz@^*By@JLpKCBkOx=WVOzj^diybJ@zcbtSfx-O>hSw|@!hh(|gJFB%1Vp%U4?fE% z2_mFCT;F!+aZ>ktKGvf&atL9bkyE+ z=B7Fzy8)~$J?s$Ua$ktX9d8=8<<6fG&l_40QE)-~^45ji-qsIas+`Y#?{wgC{b8eX z>+WlbnBQKysjLV1%8qc&pMg*r%q>ssDi25)Dohzss4M>+Yq73b&_RD9vK;G197WG; z#u~~@8`IF8*~QVy+xzk%vd@ukE9S{R@oDUt-tQ8cd{(P$4(<`y^XLv{G5;S3#pw*_ z6bV^ffJxwZZ3;y>VgoIFL3J?zBeIZrFWd9oEybvx(aTmzir1{Gqb5$M9#Q!l$85sR z!p%pa9~EV1=|={CR$bp{xdTC*pf2~oEQRU zTiy>&y3bHk@J#u+-d2R0n)DJ(@tu1`)o%0j;hFJIJaIbdwqZOs<{6eIX}nkrBm#ZA zd>MkCd5o$Gpsck39?DVhAUi_=?!Y)3u>i9*&cKyLxIryWhT!S6Txnu0`-9_M+3*gY z-aWyk<;3e4q2hG8#lg_=EAS*m2M&!p$wb>cXxqiULziGYJsLSfY!7wln9_hT)NN;) z_6MD@0RFST%r9xpzT(cuQl^t)jF zBV59Ir>iwm@;Y$*hJvtiORi$DZM9zVv#W~l1m0WC@0kt1mpH?gtS5{Eci!R;k@BIp z(~H|Qjn*(hW??f%qhorH{aAOPPtiiQR^6+h7Y^O^-WLWexDw40-`u((Z|1|2@oasJ zm=3}Ek;tXvl*(;!O<(3-5tfbO2cQK@ssQ*M2aM$Wn$fl@oNzi@YTijpDuu&yg+G#8 zBreB=nEE;YZ|v!N@)tytS(raY3{MtFppiFj~ume85U!|DN7l1Qh_+@snV0aCxwTW^9Zg(v`P<3%oxcSRD|?j8x`_ zu`^hxwu{}=;jR|Y!DJjQfQ?(Xy!9PC8{FPSbB+F}(dM)=?N(M^6&?A00;3LhAzTI9 z%6I_D$j*FU7P1TW4J`$Z#d*rHIWxaF`>g8n2|3p)<#0}|DdY>#<

4S=KAxx^P=pSz=yg)Z+#E0e~e!)Vr9)UppoeUJR3h z-!t4PL^1nuwa!=GWHVmxj_(QKRw^;s6oKmMN893ec&zItdw?#AFK|bRVcme)CA45Z z&>q*(%iB`w)aA8RSM}0vUtV#{)#jT^{IhAQkoWT{>+77lm2E5VF>&OQ4{Z%;t;22I ztf`EdXWa&|bp$vQEw*9&;Pr?h%K}ayOu;))JkwAqTUV3+=&37}0U83!V_)d6`xP`V zc2bNpa$I;UZ)qr*08j@gpm<`c^O)R386x#e5wzObd>)vJ_r%%HV2%^J&4L5#zrc}# zi+aq${D!n)>alHk@czohUj;yN6jug%clcmyH8$$`mnuuLOS_)lT&A%7z0~#I)G@=& zKUHETRpvt&Q63`oVE&!~P_@MpFYK0AzskmMfcU%#?!6Xju>w2nSu;1YR;zb6<2r1l z-02%TB~_nrewa{>ZiDFeA8H>x>nO#2Yi~nL2=&))5#7W*)wlR5O=kMPvp-+|$L!C4 zkN$vo?cz2+)Zz${GOLi`TuvM(gQpR;ppy^!i5BAC6G`v$>P8lyZL5VSQNN`_+ZK3dbnJI;+iB153^kvkwh}M3x1zy> zWV`EED2_|&{pysDsCH>cu9+3|0)By~<&2jxptEYe{%%~Ta(DFu6W z=R<>)|6Z`8&d0^vNhfT!^R^EH;y;URN8}SCHQX$jgLZ^%DduwaKgJ-Vo26cvE$1yf zmDsLppV#yi<}0~ur2;!u{)t5R!MiE7V@#Yx0=) z>Aj-`^po(lZ7KNN<|){L1o+^^O@VILXn;@QSgaVuCqMB0@AmybvFMv#_WHklg#XIW zfd|R@hllklXefw4tmG9~O`IX{s|BK7S>eK4do}MLY~r0kbC@YCWd9^oj?ah$@_Z}R` zPucPqZ!4#^_Oq`2Lb$+^a{PrcAXcZ~#PAgu2`^A=38Ewb{qUjUYNV%bpmvgWZkRS! zRvjI*FHYuXpiAtG^%a_Y)*N;mjqnXnF-yYRCiv1WSwXm?ECkyhViUTrXiIuZRucMH z`z|Blixs(oayLo-2G8enzocS2!oCCc1HdI`!H{SpVMkNw+0?vtv{=a1VMBROM6fx< zrtZ4}=)BO6OgOc+kb;yQN*u16yy+0x)p0_V<0LbRmXflLiMj6|m@7A2o5R!Z$#%H+ zwX?Ce)O)rizS&(r*hqfip(^&or_eNfJeY!uRR_oAMIL{M%)YAEI#K?>P}VhkZGLlRNBHM5F#2+uL-Or8=^8xef)=LF{a3Fd1<$-^H(pB%@)089w&GuV~z0pCRoDFsYi zEIy$4MjIag&;aCYI>hN)gr+Hn`4LMFkBV3&#W$|Y#((Y1_)6L|s4QDqRXu=pn?j@r z^hSaRE7%#-2U>QZwyJ;DL7vRz1FT>m44yLEYuKFP$&i@Ig!3x#zvr+^IN`23DLZog z5iT`y#4)!dwi~8q+Vp$wHy`C+gp$XYTfpe-MU?T;01m=@*>u#f=CwI_Y4^Cvpg;5Y zoSLuO9qf3fr{bm78oSu*-}DwmShhaoA@)&NtGy`fAH?jVfni!x@VMZsZdtO zmt9@ZY_l|sZY9>BThM6O!DJg7!w@GOmmD{F>(E$t_=lZ+>)ZB9&9W;|Ng`tRlxCFD zzXJOeoe#T%$sS<Kq^O%lfUdbRB#1Pw-*e2Ko0R8JmwZLZ$|m*kyMrGixC|mxsDOMr3tpR1 zQ%1P*x1CLWDL&VA;NT*!z+6eJhYK@KgMPeyo76jD=%qsIBxf_YhP4tPqQLJ>`zYp6 z$I`}n?;XE82GOSOsoggMU3!ilR!Z;Q7modaV`o@Vd)w3R8t7BhIGweF!}>J;z!gnN zqvOOLmwJ5%i9fYlwm(-GqDK+_3sen;2##Jp6sG1u6ejpaH+03NVw`J{PGmF!SrCzX zmq_?>R=V?z1jo<4Ct08%&Q%~eVhV^qFdSfec^mGTzkj*rwdWQ#QBCWa+O5wL@hxwR z#8QI&gwEycQR?eb`1Po}E94)P!#WG>Ky5?#9s>2om91FV5KVE9)kH_*YNn%Z1|*#t zXvEaP%p)P6H zs3c)E#Vk<&h?7L#0K3o&#ev+fT?1Ci>1cjm<|SEcC_!t^FIX@OIGnX8=IQNn(4i}@ zQL?H3(FcyTWZf6&A<3OtZt(92!)Yc-oD3n4>p>s()qJL!T&yDwFhmR9<>r0=2@LEV z&E@yQQU|eo;J#ntck=0BBI7QR1-vhP=MJ`z?^q6oLPO_^W(i;Z-Zyob`xJ&Mz7D4T z{i!GZBmeGw9)=dk$o%NPnJEU~&)M0TX$%KyW!DiWjVOkl$Ft?diV)~&*=*b}Uj1V8 z!r>n2yQtB#)p;jCP4VzjpOo#-- zA0Q~%Ml+P=x)zZefj74`mhUsasr>!(0{sU_=h}kvXaQdXuQlI894J`sohot^fs!J2 zXkaa&4-M+yZ~Kgo^k4DfRpGfYON0nzqX^1fj$4(ez4?zD~5p)cldoI*t&|rPwNR32sPSY@|Mi!4k>7h{4 zJSb$YNuE~FGA$5rwyH10j+}FM!{w+}bKaYieg|L3i5*P5&}Smic}$itU`sm=KH+sZ zoq~?WP;I71r~b4xbmkX0dgTeMtI5)q_8op!s#1Q6OWkK0bN`say~~$cI2Ul^{18fX zjmQucN4VdUrs0Emi|KC|Dzv1zf`XX$=C)^s7N0!z_c^DOD*tY^TFjqUHI`<`s+2@1 z2u)M|4seU4nE%Z;_xr+%fTlgY1;qK7$Xtc*niv7Eup>0vBn`y|8fUCJd0SyhN!Ei9 zA$3u)5IWqRZsPZTa>`C3&o@gNavf2jjbCY?ZnmH$yu^-ndk4AXPfa*{m1sMo;LxbT z`kRZB`hqKfI41njP-3V;lWEo^iyl}hckDV&_*Ub?*Z1~!St6-~u25?rLGx`T3Rc{R zX`9+dMLuX%$7CPZzZkvn(&{!Z+w~L5iP;E$2;?#&<3JIQC-D0@4r*xmxLkf0Tsl5p za=ym#^Douqt~Ar*-b*X6b^+xqkz@g@5?dP8q^=BVj$?&z6P{ybVrwEPmSA{(HOT5U zIqSa4)%Am*?XgTN?!fXfbiD6SGz91MxSJ+?ItyN(@&2{cxzy~!|6}5P*ZCCto_?YS z)Em>VTQF9ii_XD(>WprVDQ!@?VH~?CZ|E25JdiZeZ+ELxkA-< z@rGG~cdXLk9=2E`m4e+i4^R=u*eoGVx+gfAApeT6V}LyrQ5^O~01oUj+@JL^&{?zz z8Tih)Ag!r)x;%Dv?PZW+@|9VOEWXbt$82KciF>jCAv7oluKl^8KHOsiuBIOL$>$Em z>H1mtL6PT3t`gSc8na%9ex`Q&SJ$+mcg-5qIg^b*AhEd&ej9W-S~M9TtY&Lurj{EB zl8)bp_Ut3v8=ESpC_nNUQJ)mT>gf(m=@tlfNgh;`6%hK;UbU%aL{EhS3#Y$&2h6On zkN+3spkcDwRJf!~+v* z!9^eu-5M_K&jX#tP>tD_mhIeLgnKFiIe$)8o@&!ycf(67QSflqc04A@le0>ku# zyx3+Hd?2N$k~*9vGqg(qasvPT5)#t-H}`FEuEhfOjLv-beE8YPA#t0x@z?7%DN?qX z8i(}`8?!vlJ-L;QQ~*k^X;K$M0z<-bGHm(+aFQXnhryh;li`xOmKspGRx}N@-}~#6 zj#-f6>9X4rsewxqqN%X~%+Bb&5qkmt7-e}JRSb%VBDN4vE4U4{Ck2x=7=n#6MUCVf zoMm6YvrvU$uF9K5hkAu9AaAyO$GbXb6!R6IAAW%@L@%ABe!=|RpUVzzeFMbd(iQ0d z5!)B0;MwnthwS?U#!XxgbsOdd$qagsc3D>ov0i)7R`)#Pe;(ZI^{P)j(mnTaB)Y1E3uZ0B6)2pe+aF zQ}V{#J8})FS)KeY6*{-tLtmEeuP2us@WS?6S=`PR>?338VYz7XrwVR?># zh#!pOo{-&fVY1BV5Ml6DjG1b1kxJ~Ux4nNOp`^3^j@RXZzgC_)hkjBTW}~i2-Mxrh zc8pBsVw?lXkdYx8WG_p2unE|-ag3VEJ|e8wM!7(w_99$AXQWafw6|n*!kv;olDD1) zs6Fpg&K>it>)}?+so>yRR{1Oc0==k$aDj-E&1Wa>WGjVzCbT1)a=S2+P!+lD2?0@}`iG$E>ifVDLrspGoj2qK z{Tq&~>l*3+Hj`6pCh41fS!wBvg!|3SeT|ZgBw%B{geiF=0FV}O6}IX+Gpxn0KF}e1 zoY24FQC@VoBHx?NyU_IX@Mij#=zRzAdx6U6bEwU7Cv|@LKQp60A;&4^`W?`iretP1 zcKAww?)C306)vW>1#H=SCfp0;y6&TxJLRMffMKqKII%v{W8eNz z%6q$1a1n?+&8@+(#{Knp8??L9;9RQOms8%aB~{P9PAq-TZ6kMJvhkM%ke>6ve*D{f z03R9vCSy!#$2hUIO9c(_O^hyVnK!M8CtRr47VCdPxzt*cW_+E~s0b8oF3l>K|dH`Gp19KIS^QkV1m{0w6JdR4~Hep?JqYfs;(=npsa1FL8UwP+at z|IGG^nYO+X)y&^Cs1-4C*IRO~>VAjYFyVkp^08Y6$_dQWW~8Aa z97xfs9l!@We-D`Vd7{+sIa1}Mdj9MQ7b!@9TP7(C*yNTA5NUWN0#}P$N~KiBW)f>< zvH%Ia?6)m?37L6enBVvdG1pFZZdp8gRYK@M;RIs0yHq~i<>1YDY&1mrIt}s?a5mXb z4;g59mwAaHGQ$jcC;%!dYA#8 zy-=D2q!iy5*d3qfH1c%wd~Lx_XmDTdk~VqsMi+HI&4YDgnhx^sNaHZJd19@_!<7x7 z=x7e-7cPvvT%BocaLQ=T^DB?dC=~ch!SlHr$9CUSx|hr^zH26_}f`x>Gq$^2xX7VG)=C&cST7Q{uqhaJEZom-^lTOw?bly{mhqD{)ymnVZdD?h*dKpdJmN96*A5skW_m z{zyjPZ-nk%v(7qi7fVQ+cjRjSlEeG(@5$rrD3Fdo0?{ED(9|BI79(6YiGSU3@ABSa z7u`seuyXxkIvE?4%bxw)mml^TeGW?5LH*NTRWt`C8r?nClmbIb9$Nme1T!g`w(AGT zJ2}#KK+R-79iP>K$l-*>T}~xfi>J`vEn38ptF$TSm$~V-9`~7T@hE1+e^)YY{69+> z19+n@d*yZ+waQu<7qb!7G$UR|&&@j%?G>oKX7Y-+*`(bKQtT1-Kbx!wUpg>2DP!rB z&4*B(78*V?-J1Hl%j~GIH=WQFA+LKO!a{8{aJPvrKzq!u0<=dWMC6A^zkOka=`M@$ zO8QeOKuA3_l)*_m*yFEbxVBS4+!#>$?gacNi~*aW)!;f zIQ6_>-*xB_#JYMBB28$FUweXvEyENfj$bG)z5~DvqT9jS7r+S+3EPj)L-^aMF^k^| zeoM+0(^Ni|>;@V5F$bTs{_o#61e+Az44GbF7aiyM^7OCK9r!bm>JKdqmjU{4xgW>T zkp|da#AIZB18lW2=GzV2_!`2!LG7Z@8w-DUK=(%yYe)Ar*HqmL_V^V|aR~I-4Edbg zm&MZ&AH!o;^)4Jz#nN^wFI53%J#;$+9~~e9P?-mjG_(RMK^fcB==3s1=(x6vV86OoFCWfI-`rJ9JJ(I<)x(~{PD~xM1o;80%^)% z$Tq){G;`FK@@oM3@u{zBS|$yJudBaQtDls(Y2TPT8ihhKj6NXVHpbrYv`6%=>$Jf1cm> z_k54%`JY3|kkTEDpa6B{x*tE72YYr}kUt;5X&@%O z7ru+fUX5pnC^ETGG{McQ;GAhPB;jf5vh^s#Xc3h>JYYEh8i4>tNsL8N&)^avOvBpl{8w4DvBFM# zSClHe>~-HKu4ssv9H%D>&Q-d|2~eBhG66Y%j7NJX*N7guCca>Cl*vx zo0ZKE1xoMaR9v0JR;FII7Pvc+d~@%&xyoB-Obi{@Kz?%^*f%>El!@3)cCri|LunVA zcC32~5R>v_kiwK&&0=$)_sL?-sv$O@{P}AuM9lsw%M72@4nR%)ux0YT2hYsK{7h&J zVhFW$BSnpsBTjZo!$G9^lJzQ?!cg|!zGj_ok0vsE_u=IqP5*c%{UUl#Wwf;yR0Pz^ z2EeCR!A0svHT_tjfHg?buwB5bb}ElCrD4c}v{L`;p;^(d?N{6FH9R3YJ#}}3)Q2}7 zoU&*_ygi&iWZl-S8z8{?03%=xKtB7-s3D}A9n@SKEbT}$b@0zc>kh=Z^!cUzOR^p) zUD1pj!=A~P4UKyoPMHo;km*RCU2TzBS3Ijwic4nTIXDX9xBKA`a`z%g#rfWha63#kEk`bL-e)>Km*|}TBHp*+!wM>A zi;MruNm+LUhlUK{jdlm^*yh2b-*SPGWXt0<(-$er8T~oe+0!izZw|Yh zw-@;v_^@1~No4lT&hmex3U=iRqSSykA&zQ+NkL)P(u^e-r(m;}{fmm!)#pbZKt=dq zDbbAhXFpsYX*LvF4A|-I<&fGR!FZT!MA;(_JsE}Y!cj+LVcolsR}e*wnEW?ASB$*^w zfRJKdj{@6pW0-QEh?~*t)0L5eJwg{d*h5d2zMi@R*b6D4oSs3)o$dTG&96!VHSVk3 z3eZ)|1V60P&E$5Z6@vpq207OP0v;&$W2%m$HLi)YT_cOF1g_IDYTuk(CM^StQUNFG z#;wY3RX(C)(DyWoGcX?xQC|uac8VC+f{8d;P!r)dRLy(Rvh5SU!Z=ytFYt@nhi=|W zUYfKj37dWU?CG8lv-YC?N*ot~`(}FrwlM_fU&jXlE^7g}svpDe4M&n}5fR1bY2Jku zOSvtL3KiGFc|m8+(ksl*OKHC+`vq&k-ylx4%VzX~?N-DQfAG%$5qs(d+#?xY_}QWn zl57#nJQ$H>OU?iOJmlrCR*xA;dU;Le@1xa}Vn|cu)~r;6{+gW-)bJhyf+GA!Ah-x% z1Fsz^K_tRM+}_X{;~9cgn?;>F z&CoSM98$R1Mje_mA8!vSG^?0*c0D_47Np{N>{Pdbm7R%UbZ}O3I=^of2AOXbAVlU9lQ=f-X zU6+9Lou^u;*f^u1PvO}W6&@+g7=6b{h-yNB_2`?>Q}{5jGpbaLp1Onrjx5sJBc{xd zdYNi#%6$88|MMeO`|!Y0N8>0C?r#lmyxeML>1u z#(dzTmA=#^(O_3AD6-Ot+x2kxKCdUQR%Z80e{6OlO6`fL!{LaaXni7fGi@geA{H+k z0-8PHJuXmp2M4gMPr@B9{?gU2nZyO(^$5J@mA+uA+?AoaETM4q!0F@QiG{14#(f8vPxH%u_2;5 zApAYxfKo8RNPiT382b*#!q9P*6^W=Lx&c`CCyXTbU7d;4E=j$}5P5bx2_w2&+~tDX zCmG;uKMqC)rfthb$%8E_09C>PVfrJm*OS0Vd_i09GH*P^LF~qd(CQ6iXf5(u>SS1W z5~SZ>F3bBW=fZ-1uT=fh(va_WuYI+7xhKps1(z&K6AGeCgYMZ7Xe17U%E)hCMAQfH z>p@ybOhE}DeK7DL$G*DQDC6fg)Zvs){h0t``(o@2k|e9w>hk3bmYvJt>IDNf1;m#D zc^%Mk1UVbN862FJ_}Yap{ZypW_<>JkOOd6i%ddUaPt{*k@G5gq`tts;DPEklBnnt- zStPMVRsiWofDZxbAB_}(fOd^W50DfR_N#hU8XK4#rKIcPPbhw$(ci#vi>&*5>Zao@ z`%jiD~hAuA;LCFsHY&s@-vLJI0KO#k-W+nS+hLE?bMS2lylejbiant*x z0_!RgjDavYUFl{~gDC2S6P>CICu#x=y@YMgwl&AiY%g)GhTgqNiZFR>CjaO%TccJ6 zLyq-D+D~KL9*8&#MqBweAO9;2t^=FVvOpBo!>icq;x_Ru_2ssGPHt$CH59=8A&lGd zkd>Xiq_Zbm%>F8bIO;vfuaMN+Wk8HFd{_Kg60}3RaH8PI@}U)}TbFXXM`m1q-V3WZ zQmmkAHb12Ht-V@$fF$UzT$!wOGYrAZR;v421`3SQ5ZwojgewXIjFsSoGML`odVgl1 z9KN^J?P8v+Uo9;IX&BFpo-De~XaBj=Vc*%-Z<5=iO&E$phE+=E^}YS2XMFdZ{X5Lu zH%mVlK)_-cf*3NUD~Oxmw+%f;yOs9jp{HIYg=mW&o2xu2t}NBv7{)jM7vHDybMWul zF4Et_pYWQ4A-G5m@1<~9y`Ivk+xzSD&@~_QO3Sd})UjDe-m3At=b5J`?Uk<_yfBde zj`gnxl6q5?@|vHKjueokC;TFS!KO+A)qj{F5U_|@IZC&Nx5)jyYkXx7D65R)dJmg+ zB)*eXl4hs!mIi;-hH$qkSgM`k$^S4agSdkI4t6cJ7R`l&lk3cZXlWvd9<(DAAN;09sHze!e z7e3Mfe?iaSe;S${YX}eamlS`fkCc6nGXgeCz`%y-NixBbaM$@typVeZ) zZ=vS9Plx5td=zukx_0^ai7PLBgy~%7&|ydtM_kdE)gmkQ=0~~!fQK}av=!IP-I8^V zO395|&p#!7gl~`eW>S?!p2NE#$lw^yx5P^I@eBORc8N^iT{}jg-X;JRC9rfg>eM?( zsQp?Q!;+fTqvpcTJM0x)_-p1$UgX^S$GPc4K7oRw$$$}K^^8OCqwBSp(Qg2B3T`XM zF$id4fN%_Ng(G5Gq4+5A?tpT*3}0|KzK7;#_x@0pqb22L;0(l>|JXcB)F6(`GM$|K zH7XuJ;XuH}FiW7{fIhPVLPnnTHKt|V%rLCC#^fZ?TsLii{SH9ip~m;TUui!~PD^#j z)zxjik*a?9^Zw~M7uDRox=PmCvzx_I%Bi<+A>CBcb_$J;>b)aGN1dQGIMu#l@>6%1 zEiHi>Pt;Qco7=PQdQrnxr&GS1_;l+Y$T`h(t=9ICA+QSkJK&~mXIvM7@=W|NJrQ5e_3oD!@-HgPdN8igodr_HcX|F)Wz4ljP&4&S@xUmHKFc8=LUSLqs zj)Y#h_%)nMwDGP_H#cr&N zeg4OhBVr|?f2|ku--^BemB+jk3M}~I=F}J%u>-8AwQi~hw6FVy+8Fj9V@nx4N2C~? zF*#wx4tZSuw;d|U0*>lwc!w?ldKLGkdef4kghBShZ#{NLGDT#x>-Tr@)s3U-61wUI z)~%ItX_@^BuMy#Oxj14ipc3?eif!Ku033@+mF_OTB+*s!R38Z zGn-YA-{IP4`_OCh3R;HAw_Yh%GCmAyvM5BAf>#9i+AR6O5u{!~3NW;2Kf0v(kYY9@ z*Q4v!I|sxUW2jtSe;i2Xs>|OWZ%TO=WBLQTvMb#g1fT~eFfIK%`;j0TA`XM{$=Af4 zeky?>;7InArQU!x3+XC06}}f3HsdvM)i2xbry68bTT~|LW{v-ivVFRDq(6*Pc)Jm$ z2^Ilzuywj=NDbLZdJr#wKYy~&&dgNv_$B?6Xycv~Iy*E<447U=;SxlG<$;J!>V zcB|zn3=24c`eRYIqzvw)ch9)MN#|+%oLv2;F7|mR&ezCD`W8PLQYyO0$wG&1j-J^F0L{ za^T%fk7oN@yVIaCE?{>)0XmGdWGHC(QKt}`wP=GR>%}rc$ttqk3%_#p`JR<(%Lcwp zoWEI=H#&2YEki4ArZ;1X^L9FyqY3}7O;CBe4*_s2H#nT^uVGN+V7pI+;>u`e(yGAF z$%1;0?Mer~Z7@$JUf`lOMJpz;+->hkh1ZMXI#eJ=1Wv{>^RZqEM1=xD!1R-K>qz2X z9}prM^4ub7jUf`ZoLlTQ2kYcFkmRr2^m*u70 zSq`+SnRB_~W?Hdy+XmWWol>LV7tE>|o9$ezfape!*@qy7qIe>Zi*fOLpDm)eV4psT zz-sc6q)m`@?1h-4F%1dAM;tDYoa2aEaq@3p@)JvO4Zw)T<7cqZIASuArD;CIdR+Uy z>79V*am^jThzWdzLbcD^M;o{w6&PRSseR?SKVZ(0$+bp>c*TSPEIeA^y*$1(T=wtY zjwk@cDIM0FJ@wmi)=o#Ye}u0I0Wcy#yMy{A)NtT6nm>wmBe~Jn7Hv4rwFN{P-Bvp* zskCw0-SBJTY4!5+VXJs~wg)fxFMvZ50^vbW&IM$nFf&)yO#_q^@@@^^6Gz;wvHwHogbw8N6=3!VRmm#D=f*68Xo6zZTv@_tk zl~ElH9q3E)A~pKwa+A5^9iYb787{}33rr_RCm$rfVBMYj6EKMC$E{!Q+(kTkCjfS_ zCYCjb>Ny-6dY^yO2=1tI2N_stZcRk#EZulDAk0#uz3O_y-ScN;1NKY$olsy!HBZnm zpwP$M*G9G2{InC-45dAG=Ni2Rlx|URx8`I8d}XgY6;gXHp4s9Vd@W~@+#_K-UmNl! za2WNleYYNgWvBo&YK6-$Bs(Lx$}bFiyVnQ3-O=3U-c@_*>~*C+%_Y`A1Q>)$vyH}uvBF6mgUDr$w^-kt&3ULs?xLz!^unX;xNq+VGiGqR zIpGrwH3S8qaLS^TgMj(C7RebAPg+aTQL$3LqqM>Oq$A$N*CR9FvHjx3Z|~FlMtT21 zs3;>3Mmx}F$(WT1U1fP3{gq1e4NPTtq_e6SD4yLp=y9(_D)V%_tS5# znwvFu43xa%E+17m|60z=%WxOUF_QS5Ax{@Xi^CAv)hNIewxIT(p46UZo6Oym5vgx^ z+Z_D4IO7mcwUO!G4!|?~UvRg5O3;6BZ^pm`#$p(fRqziWVoZw}28~i)m({(aL7s7+ z=F2kVVlOx6`JF9SPzWw=(=)eMAnpnH%Tl+t8v`%<9n3rNJ!p*|$@#Ouc#F!KK zI*=O2N22!t|07EU4uf+Z)YFnjI7q-3i4*6qFh~DENEn*#9Lz0;u4%-nHT@~y$?qIg z=bd>I|Dlg?tf(>Y(=Dz(6PCXx^e5`LI9;RwGS*W=N9K_#$I4Rr2>xE*hXCD ze9p|0Wm)SVyOW$V9}2GDL}cm#x)pjElPeOOXL05rBo5txcp^kicJkLTG5TAnuiQ8Q z1|;=aSz43ZuxbGNlav5xk7b4SbD;xgwc)~@<0KFs`}Cy(e>-o*@e=nrk~fd$Ik&4PodNL$k3j^4I}Lpex}UrV&t zCwX3(rJDjt`N^MPh`pk2xurmKq%PFpZ_Yyh`55b}XO;@)5=Xs(*7H8D#XJNS%ks}2 zX%=V~_!0Y}QxGl&-V?x}`UdWj=5J55%a;00y8O)PJFTw`sGTQ2-*r4G`xZJbt=q}> zO}<BKW(*Ty#NKe_K7Rv`4O~zP`meGQg?- zbgDLsy~S08@DKf|6fd0lWyZUm$A zL|LjkZtoJ42khlLbd-$;aJntyXxH#^m!!6^hjX)MXUWg>SX%Y$#GP)LzSr4-63|Cs z79#S20Lp?0vB4Jj`Xw(3+8s3+!g*0Iy&IaG7XcjV%lwO*Czkv%y+Dncgz1;|C)m;I zx^a8p?u*UQu4bV%KXqFlj`WDVPI330&JNd-ez~Y>dk?`BY?mkTx{^5(n<)IJ#)3BT z1DxKbe(LKQ-23Kjd)}FS8U2t%%%_ze#LEL#b!%DY&ZOC39oU_(Uoe&CW54?&2I;{c zi{oKj@~MpoA|bm`{LA+uMViO82|Ug&ZmtJPtgJj9Yi7QRjX9wZUU&B8o=%Wx$0neG z&DsL=_OwNtk@XHUr*TQk=C>j(4iv*arBH6?()*QMGZ|-A`^D_fC_;c+!yim4ypxU^ zfbdU@_f4t;YhMfy-0Yd5_QUJc?IE6r zdxRFZ_aUjaQ1p!u;1`gfPR$f(1#zQz?HVZ@pT@!9ANfAY5GEmcBak+&FXnChiu}_3 zsMg2rFUcz?S3*}Rl>oKP%P^28SVDLUL$G%4O*;I!bUDwxuy~1k(aqf>F!JWeS^3Q0 z&hq!GS!Ui3zFR-|KuV&Dj)>56q-{}psY&Pl^Pr`)xa>nN)2>Z_^h@T>uK?R;?{Ga5 zVQY$-`OyPSE0O?>qo8gxc`XjfkG8Ki4B&=P%QDi9ovogz)>11E?->Ig;zjE6<{9`R z73!fay9>Y`)UB?W50;+~Gc!m-29mGNh&#uy1%sIU*Qwt1AV_6r2TXuR zL7&uUFP-#NKeVN2a#hXeZM+>W(^uqNpr3%^gJ-6L3F+yqJWS20t$Lank?_zM()+pp z7myyPi({ZlX6MxA9PTYV<0}jj!aJX?%&bznRl6j{6TPAp6A=5q-{EKwc)S6D83Uv& zY!_kvvVed0O?Dc)1$lR1BaZy`nb1&-v8iU6RnW9JRFIqshC+2VHzEj>HF|Iaq`~Bx zqHDnMVQ}W$Iu7Pi+2&si_Ai_ zR>aOg9nrBbif7VT3dY)v7@C%VKKR?BxdAVzj4<2LG4X^}@0e&gx7WgfP53J3Azb|$ z44fcCq-{4P;Kmmyym6p*W>iHu_H+v59Iy1f7n|@Zj?0?U$2pzN>vrmKJiDzJ9%&l! zV^a%8(vZE5_qGx`g}{)u$lG$^fEtTS)ZXm*i2-}gl5d%ZD39zfX5JFwi-|_2Omx9(WoH|PbbH6BalXgLx{E%d4v!?+K>Z-=`53feU}H>DT^zlh&dT>o!A|;x*Xs_1_{kFp*`jqv z@?Oxtojj5gS@Ej=*?|+EO;!J6MG%fNX6>_g9tKFn%?}(Z< zJimx*4fS>ni16@JNPonZUS55R#qbs{0%C2A!vz0}iRGD`a`^PwTB!UYIrN0`yWgDh zY!W;BfI)F+2bU6HVjybTg|YRHQ``OBo7?lN6~XLnRwDVVKc^;_RX?x@6x(PYkGrww zchSaG{yB^YM1+1qU$z|KLb6I7?2*skC+j<2mPHSvpB)wX;q)OnYD_M3{)Hh1z3P1f zt56WI3--V;L~}O-cIpB2vaRk>Co8O2#^9!fC#=i_8mzd4z6gesXDw7mPpt(U2#g>( zd!n?IO|sae_9Aw!7vYWp%T+E*z2n$b2OpWD=s|a>6=-2HC?XviA)T3A=vp#Ry7Tsx zeTLkJM0DTr)_Rq9b9=2m8ESwMl1wiZXj2%DYR`w0_)`q^iV1H0?BB~nE!E>Ng8L;d z9eC-(Jsy4A;fTxV|8QdoHXm)Y*}(<`|mUh}4uEw90)o`ih^T^%M<6_ZD*DjIUb8zk4u`y<7vNs`yZO`9L&{- zGC^2~{MI`kpXqS%Ci-8=o?-430U*2T|-7No92c^0u~4HAL#m*%G_ zTd$&DufMpqR|~tCNu8~{QdzE+oc^!d(PU06@J=vgnPZs!+V>+t#Ng)A-=p_=ZW`s< z+|k3k8Gpj_s`zUY%(Oh*a-DLh1h?+~DO03l(2&2^iiXiajuZuq8IzhbC3-jj$PL=z4 zrsT(zBRr`HX=Nx5(~#Chh*>xk=z+)aqFG2r&V>OXk;M`BM!SeTJtDT6N)s)am6o6Q zI7$%V>$&wWS;BMQy<}m7AOnbsMly1c(`(FkfOS5tsjllDvmZyq0w*myUAH+8s}wHD zHi4ljV{sSQ{m;;t;g!%NH{}4W6UgIS;hx1tYiD2y#rTW~Ewu{cn|hoW;J>Ur0oXnmR3GB}pliF@_< zt9jcWT#JKU<9#eeApP7&Ijfw@@&T0%i@yEua5MnST_J%cpVneHlfw^UJTj5x11|0V z($;LAs|6#4%Oz9m-%#S<2S>prk3;T-h(I3n@Pr|&;^TI6a+;7_^V_L*Gs7#F+`GN^ z=}jN{?UAbe_y(s6dLfAojz&Vx`&5rQ3cZ+12gkv`FUFjUtP+}UAtXr0yyYmdu(%kxcw#?5~a8lOT5*n2<3 z`$)@I6oEG*Rzsl2o}{xZtzo@W*vY*lky8>ZT~h+p;F6pE9(lVPy%r_h{_tUU@I;8^ z?=UbO-~)rkouG<6ZI<>A!Uo&S_%yY}0YW2C9f`Sv1{J*+Ow^|-S-vQR8F$?mH$~{S z<2MwX)|!$^*d53mUS7P-&n&&XlzaA@)_qEYW1vk?PAz&yYXeP>CLHtBOa5m#h&Yq> z_mm@L7AQoi0n>9%&mIUa&C0HTiWb|`-AOREwVfP5wy_{EdQ&H#qnBGBPW|C}G&q!e zP+9xD`q|1t)Y0T?_V%ZKBxl-w9qt9KeSlQZuz-un*4D!_zmXgZE=#!=hF^j&hDhW~`Ud{M#V-+fG}u--&Yk29#DoKqXEW zVtxQ>y?o6yZX&9;3Yi+EJEEGyS?ftt&#y*4Dt_baao59I+f?$y^F@N3qHm}j>xrQ8 zL4;D#+q7g^35G1)2;R5c>uyN1hzF_tbCDGFP&a?3isT6i$b$1hNk+D*x7P)(ujxvx zWksFA;9}yLgUI#UKqZyhGKPaGky8&C!9|R7)UaH#1|NV$bi7;+t6%uFo${#biVMN? zP@L8E3@*nrcAQr&~h2ie?>|{#Fr>gpzyAf;O7NWg1yw?cT73mU*S8?YvIh% zuH}K-CV0=^-!feNl)s$712$|B@dN1yLO%o~*&XpnNgyHu@JxllJ|%uKM;qGPTa_No zdR~0*{JQePtrG|^h2Dzw6kxr~0$nusf*1SsHM%u8{DE|cNKKiqK3_t>;nl)&!rbAG znq+|rpB!hYN2l%Caug8k2!BU$x&<`Q4*y~H&YS&%uzFRe@IRj$l&ZRe5)0UO$~#5W z-;!hnB7*u0c&82onBzu+ipEtnE{h~O^dmhC>Q!tseivI;fbD=@C#qaZ2wsf zbhOwa+X%<9(ZAb2jJ^yvlY$;}vUSodQVh+DB5b?wxm3CP=XI=vg@qM9Rnky>sQm3Q zt73&P8AXMW(-zEBYhn;nWCrA0(l7ao%d2hUkq6ohXt8+x@C3ACt)H<1b+w-JP5t=# zP*ZDDd-8h3*HiYna}u-Af5{=MWWk?l%Xrzwv=~pCMi0MeL=fIsCh_uv4YiPqiCq(( zArH+Is;k}TCIqy>tpe7 z)rv!z)sN1bu_lL@NL+UOWWnTSr9(;;ctAnI;R%gX${^az1~Qt^wAGz9liQ8W>+eR% z0C4e)114nCpbE#oG~$13)}|RX6#nR`N}$&XLt?Cn^n~ytBnIqM5t0Sz14-GCFft5Z z7)2ZwaP7;oE)8%$(pcY?BIV`OtMDLNKHc*Bb7~}Wp9+(IPSqFRJ4K&70fs6!{Xt0(c^pSv#r!oxKE|#1d4)kX~em)afw|D$#XT4Z?ehgg7WM6H-pZLf?YebPi-(B~RyHs3$m z_6}9MEH%16wVJ-+)ncFEp!V^hJ%_5*z{$EBFu>igAK%H3+6S+Jp`&DD{oNjCf_At7 z%I{uT;!2H%>RkD)hSR_I+ClbDh!5~WEUUtd=ivM1WPd+f6NbLLErxz6)%|KC`cNoT zv(0?kodsDS{1#5arEUyGdUxd7j-W04%CGeAX{b8ukuV*zH^eH`9eDOGs@=5uOv*u5jjr=84-UdmMzEO>Rhmp+y8F6ABPw zcqu?SCP-)u-^mID4#VD z+WI3XDgd_%1+D!(&YyzWJQ#r64k~yME3 zXLbMOvWz5c>jB!{&nOEppa%O(gq5-0jPu>F`dU&>zj+A&<=x;wKo!&S3~2DRetn*OrGl>Hvvt&|Swu zJ&4O`lD^oH<=efI?}drY2}iOo1wFm!$|HX;_Q*FRz|z=jfVqh-Sy(0Y0SS&eC>?1& zWvC&~i*y?)O_YRR-gKXP=8++xdqa2M0nYdGigL2O6zGa9)n9{AC`-cZJlueWcl57J zfhX<}1anRCJ3_-L1X#KRVDN;a3ZsUgM%ifS>hLty42~owi5VIZaU3(pZ1qtg3R8lX z^Sdfv1+TJ9Wb)0?ivS=SyYmG4`L%N|U}XvX*Vv~$KsDud6K0?;ti+FQNGig%KToY@ zpfQhbGI4D@kn~hQxjv=fh!>ym6XFX}d<`OO{Qu@SoDgJ4prX|~&=_DPeddn`smDD( z^1Ry3=w)_$T((!UeR^qF?!FJK9&*Q}KSuF6ZdifPR13ljbsmQBm)4)rLGE;rE&`GF$8xpNZoJo8gW-d~qJ%6c%B?+pJ}6yuaN9bj}uFho$ZQm#S5m57~} zz;wh~Z!A;7nVbV8S**L8`~5WciJGf_zWEtGwp;v`a{K)Fl^04xZ=AKtSAZJ~X9)fQ zUAh==@HUMYiqy8AdmW|qWbV7GI&M~!OZ#1IZ!g}Ujz)~MoB|^YoqI$5cJCddx*>Bc(9=M!#)_Pd7$#_ix`vjy0O8(+)Qo|c(YNb zki;?|PK^;jYf<}Zd|tsB8fqIgT9g2K(liAtj13!J?Ov6K#g7@(@&SL7E?dzd)!q$l#5q97TlM*-K8yTio&Mg3*zcd>Bb8iVapHp8 zv&UlW(9?ofQ{V4HrwLS$9xxpC@%P;FdwYUy?KF!OB_IdK=Hw)xbfkQZS(fQ0KlJ^tLuHojAI46Tc$@Mi{@=U0a@{m$jL zF*;^%2fy%~dU=5ZygmG>(6(XTsCPIi9clg5W!rJ!KjODHKGsY$)qqjP(tPOFz#;u~ z`|Rk057_o|){!#iU)}y-ttFLd3;Lp2baV8fw9S|u3_6r}Gzepz)W?(=*XHQ3Z2#UQ z8a8U;#g*}n` zYazS98!5En&1DuW%2{RxRjX;^`*pXXO6OH;dk;G-1Ay(cph zz4Q6K8Fn5w%%p-peRBHR$adkCh}X_401mK0#J@&S+k`Gm<9R4onlg9K>l-0Ky;e=c zd%vx=RY|p-rcZ#r>KOZV%_2uu$y6w^>W?gw?>hK4J$RRNzKt4A`cz2mzc4!6OW>)~ zXzvicknU<;9Rm?!s{o2*beq2qjpY)3o_IAI8PI<912!ByDX%i4Zs212f`G5 zA6I?07e0g7)1F4;>$p#XQVV#;@5~O>ATKmbseZ@4ZOJao;}9(5Nqr@J?MrOlPAu4; zQgGy|pr-iVwAp2Hd}^myKM5lmVM~^M9MtZ}+dEms_uS|Sc1m;3+Mo7*`iNpWR356J zPcb}TkEU;42C>Z<$ln?_qu9l}I0WsB(CgA|ed6d+P{Y@_w@|I=qU^DUF}icBv!}}6 zi(j>3RY2TxWXSxeLm`1QDe7FnP;qHLNhPVINfVpi(7ae>DQFO$ZO>kxDa1DcC!xlX z3@K3NW7C0w?_oEgp}Gk&5L&VI^9FB7*QBOzc%Dq*?Hp6$UH5CZv);~FdNdV+*-^z8 zQJ4RJ{gt-x7Rh)(-yYo#Z7Pd7&~UsYcSzH zLn-Z!L@&q%*a^1dG~i}p?}=`n0_eOU(9eb*b_5y|cB(eL=>jB^apx3@OSTMPa7T$yfYc&@LS?qM5V&f{I;$}#FzJM zsT=StZPu-mhU_DV`RiwYn=CAMUv|v8;&x?l%*LSU6XZPqCG&%=LRJ=P6x=EuutvSz z?Ez^cvPVN2?bP~teJ7;X;o{G*doq|?xkV=qm_MB}{23PG6v(L=p^)(o`2y3d$B@vV z+c_~u7SwkqnF(D5n$(oO617jysyaEIGP?da2Dl zh~AB%+GD007H!*GU%$yBhsOuHyWVwC^z|OEA?0aajsYqjpL2^<8eYVhLpSXiGvLWp zA>ssX#*BodYyma|lwdJ+_@SX9b*)F&SS61 z_-B75VbmU=TtQ6EkptWVl5xNlGRHQxw>HOKh7(F%7A$zGlfpS0#}5mKsYD3y zyqUsp&lc6-*h8SAKqnvfJ3|XT=``OOq z#+$%eQLoMcMCvDJP|$x((ZIdJbNWpQi+MWIKUK?tNm+)=)t17d>( z@kaNi#!*zWI^fhn1h*?<&dx=JdV`VxbR`vbJq1VoUJ@CZUb^UiX=$w8?$p6Ox;_4E zkCL!25mHo^9gbZ!?A<|0jyKELNn`_rj}!}1FRaJWY!uPhMkL8ubs$^ZN7n+C-+OQ$A!!4wh7s zj3dRyX!2?8e#!n!<_v0Gaku!lugaFG9F4RE`SW1SJvy*YF}KmDn_Jsk-!zX!$=UM0 z&>S3=uGD$-WH+ENS6}Jj+4I5HE?`r)lM7t)a2zpbK@EK9GPCDi{3CRDh2K_z8ZD|# z@o}@H!6AiMEAJRcOl0mX-{FH-9aMpa?p~7t`pkdOLrgFkM}$A#+txS`^{Msxf}8qt zotC!H!!NTR9yptFz2e|GihAWLt3bMR^ise!Pu&ndeDXPy^DV$#@2eom#(^$T*X+k5 zi#-H>_umtM#K5cBnjYE6;^z=%U|)PpW3c->&)xS3C>m0>r+=m^!E0BrV9m!mHLh4x zFI5Ell+XK3{#KiNKy1G6&@01saiuegQ7Zhs=5doWDV2jAo1&~Rn!$VeFMj*B+m-v17hGzPi%jx$`58-Ta#Fsd# z-xfuSJUlB;7JY0i)lGuLBCLB|CMMtd+z+nmE<1k`#p>-voSv+I@~H*&&E-YT$US^-E< z>8`FZkboglZ0^y4rH|2q>sIt~9iSVuy;%XlU(bI2O9dOI9T=l`@ZDempUU6kbrc{2 z1!}(#+x8@mSc(B42w}7ab)n9$7>)@@y69-JD-i3;2c0Fo;Xo z>PDFIa4xyzdsYD1R%R!s&q&lZ|K{W%w?lT0WB&QkPu@r+DR$7&md!!#>W%?XvMM}5 zc!Q#c59<`9P20tM-AM~ZRS-lcBp+L3_v=TZk26ogQ!pzBCb=y zAUc3h#dyL8W+s0x%@8iZDL(99FP?p$8(f=Tezro6gRjW)^Sc|-@5{>f;;vSafWVj! zlj}aY4py#PQ3)Wfu`Y#)xT{eMaQq2Pv#w*`@KuIkTO!}5OB4kTV0rNGWyab0C$ zL`Ox)YZsxGcNsB}!@+S`^q*T|%B=j%C&N;Lf!adsy`?jTi^BAo|-88RTJ2$B&oo2v_37n5wBTNWrC!XA&J)gNl;k<#2wH+iS{4V0Rxz9>p4 z83KYmA~JHLXlc?%zSY$ebsM&@rP^vIo=YBGr1a^BvgWB&tC#$x5OMV%gvb&WQ_r#D z6{K5MZXHzlD>oNgZ`9r3e&ujn&H+D`6L&>F9+e?-eZ{XxBrt6YOm07j$u#bH<} zd<9H*m;r8KKkR}ktRm@zFY`ZqV{r-!{E;}yr64GxB9yU{jSBG8Y)qMriDNwYU&F8~ z5MBm-x7GDuH*HNVSl3(l_yU1I z57k9P;~O`;{d3yTC192;pc_LrXULKeH>|nJ`N}gjlxAM3I>^QH?m8!1NQ|bIiL@R= z{9`R)cX=*1tFN`U#`W#~GwVN2Qz&z;FqDJma-)abqM~%bB0sFk{}|F>ScA)RW#?6` zj>_!EQ)Fa9!NF#$o8EqI*Wxc0hQ}vW%zyyuhe=dzqD1Mbh2P zND`wY$D4+RY78!pZCD2+1ynh^h}XPww$J`N9_PTehf^*_g+5pJt!ncI*sjw^|GnnklU(zj)po{5x#|G>NWXEd|!c>qjYyk%Y z$}>T*n8A+ z8T*+zjT6)EZEtnc4Px)#b@}b_*%~kZEk)Ufvpjy!*ZsB*lUu07DrD%m1bTYAQM(!p z!Ya0I+FB3)W28uJ(|tV`+L63Y+4_-n`jgI{Jx*l-x%qsbb4(Ir72_0T)lHD7EE*`h z?TTi_z5xP;NO9!#B-BGDC*Xu08fNWYw5)5kv9=$be>whTrv0eY(PKyjpO1&55%m|e zl4IdqP>O5Wsm1+VG**X5<{(-UMpa>`kr;fC@NyPffJP~ZEq9ARnzoEL)mYmn=;1v= z409+vw=)KV1Uk7`ti(Wp@-Jvz0GcsCG5vF=i4jb{2jMjsvWs<(b8}{L6R)+mf3DRi zN*apXXwDk4DiN0(97_3cjO&N^k( z*1iAR$HOH(v&1K5_IyvTim}Ge$KO0%~T;O5%QEEtL9|oz=SUbxfctf0#UE zI~z}`>xmx;<%dm!7~FxRH=!jqlen1+OA^^VIn~{(R$u4NEJ;fC?XV@P}YkWS@;JllKYLk-}Am@{{oC)sW`AV=!+wIW{EH zyINc0+4ev$Fg48e)--l^kd@8@>N06lk4%@4L~@|SLw z+}2ay6rXL<_=7}^STMv*K`U!UF~%4fG4sb{T+Q$(`OTdUEsOE*Ht9=_B7N(hG7jAC zf{^;hoOs)w*50m0aD3?jhp!xKo1i`?tvGD|%x>`~SXKq@Eo0o1HW$v02RpHN#EzdXiT%hYbqysJ03SV96MI zlBjS{KG|nq;+}E@0#Ub2*d$T;>cFg}3XCT~{L?KZbrD&3{ZHAYGb)uj6CXAO&4sQ# zeZ-#s#OpDuVY&(gd(wO!>DaQsWuVpiQJ-!x+2mSrBWOiK&R6|L?V;P5<_b6d4{h%q z)l}EE4+a4N>C$UZK$NOTQ;{G@6A>HIgeXNy1VpN|Ly_J=KtTxz2vUr6DH7=lB1kU^ zfrubALCO)5IGg8tf8Xz$^}aK+X4aZNcyK+eoO80zzRSMv>%J~}$$~ODIRqe&1N|6l z8aqh)!}~(Mq6NTNSeyqoC$57EZHFl{#js{rqB*+5?>gG})hDN7o9Rx~vg&*43QU>p zPrOs+jpgq@$5BK@K!ABoEz$(JH6x@8M-q^Xl&fT!c5vXLoPQT@G_=yJ{zy9RTU7Xo zY?q!ohbJ|U0g?$p_az9_7tg^*=4N8BBuf+muGSk+W!G&-ak{Bm_j50KzEbdhy5^m; zYz&?Ha`2B+h8dURIMW@qu9aqpBKnDB-3hUx2ycO>+VbF9JR6l!1dl=+V%9oM*{E+z znfQBAl2TWcMK}@)on#Fp931E)|KaQ+QE^&AK{{V4+ITU*5rG}waq;$EwKvP$XxlvT zDbCqqmUy>SoWk`0vnLK)&po4fk$T&+zYjm?^eT7p9UGF;j`=ztYZLw>J=B6I_1v5k zOzNG@iqSevx$u<&$;_IF|3#q5GBJvS-B*UnUn-mWoqjGWo4hcy_;h_|Bc>_a4I5S}`^sFlh1ok-oC}|qjTT4n8llfYQe>HK2$sfmb27sygT^t6 z@f&|z|NAP-SJ}oM9>7)7Ny_Uy=O`p88v`>CWQHtIYk(pqZ!@ctvwUkB<&b%1~z7Yi zc2Nv;Q5lWlu{|-0A&e$1fQEf7#7Xz7@n?ooj4Hc?l8^Fp!|vNxm>FO)_hVP>WQ1o4QsP z4uAWk{NYD%m-CnTZwR-F@Jn?l7C<+On&3=n>Tlyu9R2_*}GLl*^J*|mI&_quZXvg3EBSvM-D!UxhDY?*&;or zU>{`-RR~;aY>AX+=vwEMUlL&uJzSt6L45Dc*Gq99@X$0P*>_IbVp-m`4A)rqT04gw za!e)o=zWr4tE86-(?)iqJNjWIX!Q|UY3^X<+(DR3X40jNH3jAuFYcDt`SM>#x_Gs% zm_;}VvHb(o7|~Gc`x%l|@}y}uM7n7jjfXfXE@a&fESIJn`H)LZ*`Cs2TcaR_@K*U= z6~>%XJ;EoIHIBqckdr^n#6s+v7f2r9KC@^(=bcHW7z)aJ`)tySK<>CeqmTSne7PdG}fz-9D6+{jqvXzlZ#G=EHvBk9-+% z{MyWHEJ8nga6d#qp*DPhHa2~9%)^HTvt!wcI6?~bH_hP zczxodXfE$Ef=+~NwnAC%9&m3WN^XI7!x+do`_ z{oZ;Pi#Y?oo<$uBxYnC#lSb z!3a9-Yali-!wPeFuvuFzRMv%D0yy#?&7v6>v*VTnUK+&PoO$C`824QacxU91tBx_Z zG@!OhaXL`hW2W$v(XC)~vcBuJ=C!;fVRNfw8B!R-Rm%62+K z_5=6c32-9}$Y@$lg>;6n!{=s^kx+-3IpSUK?4z`)>na&_ChU)-{e#Z8 z>S~=rT?FVk6agKi+}P8hTpytD5MFSCa23~VF18) zQXyohbp^{1WckMb20N?ujdj`*O~_I z4Ie(Aw!}Tfdo$bJ&?Uy>S*S?NCm|^;@eTeB;HGlgPgj_sR*7agm$_A@t0d}lL;1|I z?@5py_nyya)n%K40!kahd*2_qmSZn>|Tu z8#5+kz0TF#ct%Rsp+VDMR~#RV@ZS4*MIl$lrf;6Mluk166nz%ix7}8n&iAr;UyB0 ztJTF6E0SraDZj1%9Fn2i{`IfyHzliDK>G91sG+T0h<|t5p>XA*m@g-pR{27JBwQ*4 zd_X?bwoaa3izopUX&wB#_p)BOR&`gg?&UmEwDOLDa?y7(ls~vJ5YI4h&SV47WZXYL za?#XRv|}?tywseKBmTW6QKJjjNy$;Q3a+Cf8TXX%WVQqJ{?$Y4j%)N zEV0ocC%$@%0WR8wPz4{%4A;`NZdziI9vJqpy72SNhwoAwFJ;-PgH{d>Kg<*F1aBMk zxmsONF34!$=RL$J=ST(}>)b(Fq~i9v`l0)7nrHl^lP4*?kcFZ+ChFGbH2xmF z#3vad8We+Uag)Di`;%1?ABCI`x;ScW)ReTmB3xIYFrWByl%ZSYvTw87!0Q#vAfB+u z8QvoIPo|a#7renta^x2+(lj!MZVY(bOWqMKYP7Weu+EbZQD-<|8M5)*0dWDvOgRht z;DT(((@^Rx#O9(txfA2FT(aW5s?{T@(D~zw&dZ}oMu2F2GZlOl4XD)&xQt$WCFciv z(o_eM6tRRATKzG+Rbg&4!%i)!5AX71tZ9mkQ&0P`E~$7+)uz?5ohI=g3}IQFBG$jB z;lDr=`-|99#Jq-PD=xVqXJ0&-bc%DIy(S+Zy2Iai|72`jdabR6RD#it`ijj|q4rFF z1p)$?C6Y1VD%j~#iUq74c(BZgwler3O+B&zp*j0ie9=<<=*oJjbVXuD2ahe{gRyc> zRZ-XXIVpLe-s%t`S2XE1&?LV(?5>F2j)BW~ZW|6lRl0Bbc%3cMBlRR?6F`>D!5GpA zs#+*}Z0ru6gh6(%vjZ)@A~I7a=7p0rDf5LCM(j(2qgqaJI~!S&%eQBH@K4}dhpXic z=K3BB*y}^?DsK97UpA&O#=SrF$O)d?s*SKGsY|A}qJ_^xZ+s6e`PVU=v9vgxt5wG|`zahsBo_!+Px zv%$xk%>kd*2{q6kKF@R&%4wyEkt)_+lq}6@T20;lrgBWsSiI_F;unxUIuP!Rb)iYt z;AzDbU;*&iJMa!>-;)kFl4X)qv4?-5T`nIhcr`uq&D70@HKD9)@%zmeOE+0@SZJP9m%{(5@h?ahnke!fSp zaX*VOlY;z)5$2!pq}$1pJ78nHO{N^~ox^u>(tWb`m|HzK*)?_LU1Ri@q_69$tk%b# z9@AD5{&xD+3iM{nrl&UpWoh1lC3$9n;tg#8P=-Q@;sT^w`}gcH=ihbQJnNGX;`o_t zwN&SK8x^V;~!QXdosRF;a|v7`#{u_ z4q^4L0Le={7HFUjOOlNJJvFtm+q>&Vt$+J4{QNFR-Lc5>a{DsysYzr^0{Jy~KV78tjEm>^s_~!-^ZRTa$Ab$RlH(-I8V)`$K)}Uq;^5=R zaGHc^-VzJpy?twDw)CsFlHW1XD+1^?6JwVW>84(Yh4@DK5!pE5?{PaZRV-#Tu9fcQDIj?IxK$E82RI?HE{!ywhBL3fo3r>pm}Zf2-*);~LvYX|wqjKH6fqp?M) zf+|=)6W}1*B39hEu#Yfe@gH5+sT%M(Gvsi#CKEMjxBjUgV0S-Xy3Jk`e_T4XjkYw9 z^;L?&GW~(PKvoxu-VS=VFJ-*?XLbfYk-QQ(i;K4v8CtIi47UGV+x)gsC{XU5wZqqg zi%TMFZ@ct0#N=k=e#_9 zk}+fyRRjyVQY0Lxf2PW9R@0b4(~j}OSD;)#J3x*=GWQWcUyy zlddyhrjsUuMZ!s4UFTLGT7NgGnZg`7E1zP(BM&Cd^X0n{iZ>5j@!U`c4@7_yT}=Wy zeb5sN#?i)@n9d)Wx$=(`4~w-wO0=3Vd#-56*2D5aiUUK$!{?M}H>`>M(B6SX{O$$5 zko!PZTo-%`u$tjhu;ex``0eelPNYV~zxJ;hoAarC+`>6zA!1suI**&2`~5Ba?lr_c zk0&Ymh+}vqXifvq3w>=|M8X!Jb?_08d!Vs@ZWU|;e~IH{g_cf9hj4*AD1r3{2;%Tr z7(kE|_Bj2`iNu2ic{BdX9+4^E`sco%5u7e6eDGoC-0F`*ukPH{@vKoJ=%3~JQ!zIK zYjNI)=(NgqhN`D?d|noPD=Y2`oXd(@!aP{$b-CvD?3PZBv1~QsA)RD)M^;NW9Y1s( zMbZa+zUr63Ewe*LM`xHQH-Pd5AMO2|oY7CxYF|{z142JRgdtW?vOo-*6|6er>jALziQw9D zuZ?D{lZE}eyR*Dn^M9csXSK`i{^Mll)NIs*F02koIAnP%`p9A!W`caU8z z6lJxov|})T>EPr(TlVohM8e=dHj?>Y)MVzX^|Sr340WMnR<{hRU|Cv)_r zn1-Jj65$f2tSy$@#^St(O)u>eO~lOYQ+4Vu=Rvb~!OJ~rzUV;_1NS<94z2}xEjn9( zmV3YK4-ZhpYN9bag4Nm>gb0<&aJx=3>2eD_x%jc2@gMp8X7z19wJW{CAFFwFL=7KY z#py5u%vH4cZ(vym^`TF56kyzyY9M?Y7vO%$0!dG>CgIb-sJkbfD2GjnT4AJxi*V4i z7TGtySh2E1VKi?N@+nFzPx!?@!SGN>>R#doICKB_5r!_#py0pyz$iI z@AEic0+cn3<5t}~y1ec8_4@F~7B9WIUsigTDV-N;5zn0vDW|`guXV2gB4B4^EkLj= z08HvMzjX%*WMDNPYmKFvl6$jNJa=pYPf8?d%;6uOg=hnKLO6CK8a1&%>lA^gp!~C3 z3Sa?Tc`|vF*E*4_dc25zTno9M!kwUa2!HSEP~b7t?$##QQtvCCipR zhK-}q8Lx6Mk6cT<+R}{Z#SAGW|6%?fY^KedNPI$CKcEeZZ&{5l87JGf|MN ztOH`l&!XY9jV|NW!`2E$wQ`u*1E;^FRqF}G zh9i0&=-!RG&=mw&ID>&ZbejkP@4HSZPK(d~^sIcy88pLZ@9);vhijvTprI73nnvo~ zbZ^W?;*W$E<6yx58IlAhKcQD;AAlYv{V0Aj>Jo_)R#c}Ht#5;bs&e5fm;#6jMcry) zgId=-mqGPF4glTQYDqhf-QSEV+_*8Msk1Y7i>v7kG|p70J! zy8S|c#{3(+`PM3fS(awcxxWZ;qtm|vFyC_I^y(sdf6r&$NxJEteD-81Z<_U6hB5MQ zbkhteu%D8Mu0~PAYiFpOSMuopJ(z4`fIc~}5@ym4g?D{hleE*B%<-$fQ9dC3lVgd? z80xgebbFYJD^*I{P?o?JuE8V$acxmS4>}K@#cd9zNhuoCdK6(J5Sk!v=aZwug4h8c zsep3}_ptOJ+H81hY^bQAA#C!@cH3K{$A{$@rK_~RF!6D&*n?$$Iz5!y!Hd{@}q7~K04!SJ-1+_uV^P%+GTeTTbb85FzeYgfV%4T_#+2FHDIw-OsQgz&|ldB>JPYb#!|2k zpMrfTjb0V8C)Ky(IZTr&ChU&`qN;YX(_QN7PAEr(NjfI!O7eYjM+nhN{oovtrI`R9 zzsC9J@=hFn<@Q(n=98;Tqm4)klIM+4h2@)~Jl011PjwMs;suz|;7w4QJl2V7L&;uk3OEFiydODXxnzo%`CVXN%T(N! z$75R_*rG*qSW`~a<;^s-BwX!MdBID{MOly8qEMMmG{#tzt?u9h%mOYNNeZ0rDjGzI;;eSH?3?~?x;RGzwu_%|a zJ`oGjb{2^y_fFHiP&W9bvB;Fg(QJ#nAUc)%>1`ijTd8(NN9Ydfqz_Qd6tT6rM6jfN z*_F=>*WVEK;Qv^q>@+XW_?oUVA=Hqmp&STYAStDmdv~C5KQx%p(NFzn_t`+L+<3x| zU?tk2PvmYZ*zHZwIS&L#hfIU5FY{-|yQ*c4lKCv_8~%9JiOB)tlkAh z2Uar|#HIDUG*u5O$7Bmw*ngiNLWFz_`7%Qa*up1isl!fsF;jal-zs4$a|R z*6sFA9!?_nABo>d0<-_jygO1!rzt5gC(QtOJA5ASnDOXZuxt7$FHC`7-Q!b&Hs4-l z(*5_mWfB;*jQMJ>3rOXUFRVn&BIBn8s#}J! z3n;=rpO4>*Wh-_1hPz^F~kHnJMnJ(e&>Qwp;|;@FL0{ zy-Vl62U?G@5S&dt|1wkqnrPpCKyq5iAld7}5HMBU8q!W=^La%0;2`5N^_N&O<5za9(J3?YPE81)sHbv9+7jiVX1`lIn&HF9V!e{=_zpu5{Eq>2;lnAlht^T=E@ z0*+fR$>aR21V7f@2AN-;Kr)_G0?re$xVI<*7F2m>9wSMYaZ$U`Q6W73J~cVPwhOBP zk*nW+F43r6zAX3dHe}EpF@O1koMyB$`HzM+ME?T~A4gHR_WnU&X8W%j0aB24)*1C! zl^ow1;mqWq$>7c(x7~~;nAQ*6S(V9P)AtP3L;@E~azU+N@-U<7;GYTpJ$ozW-2K}a z{9zO6_g!7=ReB%3B_?pErkI3018Q{Yd$&S9U=Km}*C#4l4#UfB0l^F5c5#AGxtqdO ztGd)5B^KEuPeSHX^+%>c1tbOMT0m9%KdrK+={9Jqg+HQ5QzP(BzqEuvg|Om9awJZ3 zsHEGU@!MoW3ZIwv8Y(#B+{tflM8PkRV2eDh4^7hB12BYR$yee05|^E?OZCQHn1X2x z#N#qL$nsRhx1YVNgH+}qslOJflW|W_W;=p*$N=P|tMH${2#19Z)b9OvxY?>uR2d3B zlKcS$%lJjnG~ogZYDSQe2T>^|!2L<$!`+(tG*@Ap^ag{>5v3!s;nm7gRTo9v;gPsl z1T>sY(*_Y?6y!TsaKIQ|f%7-CKb<)SD3oWh_5C2lnpFp>`z#+mB1GQb8*?@|a&UUw zW~#?+ee=@=kpv#x=dfWnpT%||p+g1@U$v1>(PZZ66%ZMKHdKXhS2n|>JcLkU&_70X0* zkbQ2S09BL%$vlo*do&I+YK-^tM8yyIcfm6W-^7Ff{)(fAIXHyeF1_89NC^nXe246j z@Cow|z-Z0t>SbE{W3o*Xt=kiYm#(%mc2;zjq#RDLdzWHJiLFQu z$lYV2jc7D-XMhq+nwjIsd~>Mko8{82suKVC#BQ;Xk>MX)%n|@IN;|?uO@tO1iR%gSanW4>^ zs!#hdjs7MZRwsFy4T^aZuEWNbcE!J2lKh*-`W1Y^&uD?#emZyxe7_=pT>A}b){UxY z6MPH5*B$=YqWW;doj1N!QPP6SZ`letadNOEI7%U;!SKO0e)Kd9!=viZ$*!uiZ_svWd)KB!uEmhBu;M~1 za@?RRi^ej7$GG_pwFUeL#{~0(Pn2vKFA8KCSX&gP=IP(L;4@i1&LqS12MgXNs5!u?PPbM?kk2pVXGB! zbVl|3adQUl=L=`f^g#?h)O@Iw5jD>NRrgclz*gJk8DJXQ&xLrwPtfnZd%*3RLryqY zraNvkn3c%Xv5IVM9DVORwhi6l#P|eyNjJ9ipDCz|(@zv?m#vQEuuoB(ORnz;k_6vE zG^A5YgKYck0dtGPkhRpcV}`vSd*YNW02X5*kj9a; z&(pCl1on;zR0JG@=(LgFx>q~*7$_&fV#QMUay~$Onp%z1anO)LUvniY{ZaAfVyXiV zXSSN8tY<4jxE#PW{hQ|pubJ>1+DJx)f>X;C3JSQDseo-)1Q|bJz0C;U0L0#X#-ywd z(>y~f6doz%_FyOP@L#pi8Q-3<%k4X#mp3Z!&p$X7${1u1s@p0#+p1$w?J{6)(hqX9 z-=2z<*;6f!J(e^*dZIi?3`bpO3p6uK$l`p4yI@;4Mz%n`ax|4oVwos^lU94d$3XvXxj2y`lC9|b}mG( zjc!l(TaFgKuguh&4;r^{OVzzwSI-cZV!*eEC0S!1Ltp?$X`;-iky^gRuHH`L8>&@F zHW;~>*7;iZA@ONN0`u>aj<*|+zMag;p?jz`h+8P8L7d5;MhBMj0VszX?)j z$Xt6ggH+?oHctsJvQYb8$T{1B_Xf2UufGUUYYNzeZ&$$H^dBmr-`oT{@-y=7#&dnB z4F={x${j`hIouLK1=@oz(TsRPeu&uUmMIV#8Ep2z%r-q;_0g-33dF<!A9v$$V*XSa_wesARpbgD)K9+9KFw4)b*=JrgEmJ}!)r{7Sk@{(OU2T9iB-Eg zsdxV5q>ZVpscr-s0_Afo(DLC72Zv;(B#H>Z#(KMXR9;g~#jgsQ^!K|_n>nXo_e)10 zo00$PCzZ?|OQBd@QFk(>XAo@K9OVC6*^vLeve{7UGPIG7jz(`^mIt0@9qJS+`o{w&eN3p@?w_t@b=akmM@n4QoLw)Ia{)!5YIv1o)rrx_H9q$2sn~Fv8 zLmfwvl{hM*8Po+LFXAaC0Ui_)aJSPOIC@cT&R#buGCAbht+g6;|}_jrq_f~ zma?9`pp%A<(=VgwY2YGy&zdm>Mfgq=yQNv&UGDah@5|RLd%E=6K5tmb%3maXUIE1j z3$|=pC`D+ND$^FiJ-GS)z+hLq?$D}T{m^Q8>h(!hpR1Ac0g3JU-myteT8geV(&(DrBjhUeFX9+K{QnehdjX^mA+df&TvL)WIz-AAW3`!G(fDgrcz(1({KcaOw;y?a zajH#HNEm1O>~muvn=%eEgH&i9QkwzTy%jk};H`Vl!XcrkchFP;jpaQldY>Fdf9wN_ zr&)<3|K%^QZqHZ74l@qTgJNpQn~!eZQ+2ROD42I7B=IIEf?R`}b9{^zY6}iAZ7 z=_>pA=2+&hr}tbOXfp2ph&!@j-?=P+}hmZhjy{!_f+h| zTLV%~557YJ2^<=G4>c2VH%%5U?f{H{2Dd87`tM48PNQ>d>aHE;6X9c?E< z-N2PAMth}k2nrcy0T0zk`v_usml7vDWlRcrU+!uv+$d0&**IObGs-C&QCS!kHKA6? zLGRgME|J7(3T*Q;i~`6#ZCGjch3@vd5qV8@sAzIYpPln&oDp$E_mf+Azxf4lBPxTduW&|1&#`Cxxs}hXe!Y@$Ve~uqv0G|QYC)Bv1dXlj zM43_^{F+nRJ5)@2<-19@3-$u{Nm@I8(^jRQBIA<`O8GDG_8y?c*9vlLnTT4OK*D;` zCSJs)QEa?wn7#Rib)S=*;@0mYFAkgC5}KU%tu#lZB`}eC(KUFILq4EtS=n(?&#-t24X2Asgpa|kK0a3 z5>iDmJB6ulAH=Zod^m_#(ewb>&-0L;5MHt)?L{$+AsWP%ugsrnFLjyZd*WMPikdIP zcs=56H7!gn2na*Rp-HyMuxwTgitq|qw;p$_cI;I%Y(`kU5`4Nf^W;Y;q+7biZ9a$a z?s~sxO%E%3Ptt3O42ttA)8=omAmt&w!26qqoZrtiKtXsu5lI&A-Vab0cZf4ImQ5_p zV(d4&9%x6R{_PK6|KkOqlf(bx;?V!ESO0&$IQ0MF+4skQm`Xx={rm>nWx_(vS#?X6 zsthhpS{h!^kZ<)lFfMjwz;s{2?JI7)pLU25%mE)_A)IM8AjNkJn32O>NTO@QRjdA2vbM0V|%S+ou^J>egijJ!`+9boT znLy0qc7}X89lb=7+#fs{q6ne(^sM}=Z^mDMb4cb_2Z&7>D$5l*J2$GE8tOPMUN<>p zc`4QKQ%#jT0C7m13wsA54DVH@xhx7-H;vtwyhGt>+WtP&q7r05UWnL|!$JoqZx=^< z<5^5SJ;lCvj?`?Y8~I81)kxbX31MCJS^RwL7AtwPw(H_(ra*xYH19pS%GO~y^^TcC z6BI+Td&KMvosC}M;@%o%U^cLVyr~KNK4%_&511aULz@N98 zN4#^4d&A>vFQ@4wjf+KrQPH!q8-Ec~=G(W2Ah+ZR-EQn5)!-|8dP5YS!rr1@=w94! zEB97QKI^D<#HHsB4h~L|!=Uv0QxUWy*$Szz3R<4Ig-8amWJ4nkciGaGvdbTy1;xz* z(>sMphHM|il>^3RKnCi8r-;N8m0)R)RmM4`H~&AVp*rrHm-w(1_9j%zFBs1IXT``T+unGP^Hr(GKYob`zgE-F$)O3 zlt7{tMQ}dRM9GwtoTv)10_8$|oM-Xm#5goQW9o5$dDiFp)D8KCecF^ori_s8{U=8y zV*xFw%M3l7ZU9#jo8qU~Z5g>-tw#E&MQw(2XCyO)Y*I&3pK=+5oO}T2S%9C7h(-~@ zv_R!5ym4vzAUV2yMsP7i9P}&Qu!h;a9ZK>6Hs9l$jDO0f=*quV{^$v6zHcf<4pIIh zo}tK$sMQOAf|Qj2RHIcrXy)ddWIcpGC*a=N=BoymydTmY!U(H7y|MAnqH!C_1Xf&f z&Qz+kxqaKn>)cV!24v3Z$JZ^x85y7cz=nVdquifPD(LnRdT|Na>)`$M>CR?#;G@iK z4-VI|D5(_81@J9cTe}WEu57LME*vkxGx*Q$evUEs1{hLp>xWw+DA>$LC zE(0rCEXrC?8e2T zxqOm&Fi3%xqmU~m%j3~Qg`f*Qm;+7EQ&J5iE?ye|NJPGEFj^M|5)XEG2t;*8C)jEB zTz{;%!dUY;;?%R67k2HI!#{`bX7kl$Lyoe??v2k9uM0i!Oq~uzUBnX_0D2CWj3>NA z*3pW{1+%DRP2{U@{)2-PntDZWU0{xPO+#&6_@I!2R9)TEnuiw-xf>xIyzg-=OBlB3 zlm5#|+Q9Y3@f}E>5LTC6wk%J@0-*r8VTeSrmJ7=@+8U#4>s(2{$RcCDYfNsO_@7S2q57ffg~ zh&myiWOh>jW2>d%w?g*Tu|L`HX>dPvm`akI%Z((mbIvU!7mhTeeCvJ5OTw!otL23c zu3I?z8gYlonng?GMrYFq^2adoK$Q2rlxz2r9t3t6{@+b+x%dNMWTmB(fa=PN3fjlr z?xb2FMNJzw>FKgLk#i0XHN%D3aw&IW!?Ws)@)F!0;*K8t^HN?dB;33Q+QVFQF%W1}pevruLve z97z%o!_Rq?K^u*DeH3BZ6jqPuM3am>;zLy6EV4{jE^@5d{>?AW*%#_pNW*|+(t6xT zwf1whi9*MzOTTOCzRH%CzC3zH;m8wae9q}Tk057zENnf)dmj4X0?r3SK!pg5+LQJj z&mSTS0U<{27JeAlExZ%gj8uOxGKDh=cnQu2#);xZWNY&*j zH4rX2R9XH`FK{AC)lesw@FW&sh zmCc@;o<>!d7&_LaO;6uDmvv-NPf(WGPQ;Z4YPzi~Kn)lv(ku-O&kR;C!p7#F8dbpx z`^OEn2fUA}EbT9uva$U3IriwvvvIkEgXXuDGen6JRIuro97ASMf|?)~VpQ^ow&F}~ zHUUB=%&>uOWkY4vSC9Fj9b+G^G)rmcXIA=B2=&7NRlG^Z5dmJtRobAr@ci?D;tpKc zcEJ7gm*4v_C9YYgTeMRO|vMN4A6`&>{f_ToJT*w!3df2I88tr=4 zG}K!};lP)rP28m7(YnYZ=IPF??;bM{AiDHHd~71;_6|Hd+vnPUwHj_GpFP*aGWe6p zgU8m8`)kAr*+D)Hi3b;$-r-h2rL&uyZbN;Ex}y|M(ktlrtEHy!V%P$t{YtA4T>Qv9zh zH|0|5a|iK^uVYx+D3WVS=`p&hzXQkbjN=XHj-<$yx{fNz+Zm1q1{yggROJdDypa0s zD%0-gjljL5c6uzpM)xPWWc%k8)byH56<*IaGaYJ47 zeHXSbW|E9ACX5`oGtx=+T|oA`KWTb+g=(3An1v64jaLS7<2^X1(mRlXTD^VMd(l9BmcYEu6*(Z=RfuhE7^R8Gtp#sfKLE4i=I-DICqt<>NUc5j@1Sx#6Ai%T(h(Z7y2oajlQ8c;%$~k-Kj?n-P?74{?=g$3zcO19;!!RU+A|8%4SrjG3*E zX?`?^;(-lpNqy55o-XAx#vxyCPt_b%kDb5g45x-62;^=)>ahVj^wWQW1=l`O+&qh< z6sJi~WanvbQ6t^I4*g1WnLC(yUU8$x{I^i^N4w0M&u+7^CzMxEq#W^A%1Guk&g0ZN z+6h1mFxG~OZeb^9*3X)9V5}U!G=wzj4NZKrx!tY)BNS^-$Q?v>aKBAvo3>Jue4R8O5;`N$;7?KbA)ohsgdhfmI>aW*K7N*vB zz2jOe@BEZcK4r0v?13n1>s`saax=UrQt3a@lgi+4O*>>y&VUGmV&+7MB3-TCDW=7p z*N&^X*ID$ZaPxG$lSayKJo{*<7N|4FoWTE|pT!tYmfd{T%#^5UcCS&oHaS|{(!0r;FYQc1vOK9%?#rNrh>q;b{P|wm`kxnn5mDdaui&f@ zI2TRhGoPXO;xQoj;HDzxio2AOkbGJqFup5$A+A};DNiY;!fsS?dOOe>U(-r=LjLhP z{ltdLL|G2HFWRU)=#U&-M6CnNK`8N-(;=f;T4P)2INI6T#&@G{?pOz(wrY4=Qy+Ii z_&701@WnO0YrGH6bX_r3T||)#K-KhwKsPEHPrAVwUfD|4?aF?EjcvYL z2v2H9i=Y;hp`h!ah6;W_7yKIFgj59;>X@BQC5p5X_?{z>&&{h$q_Thi1M**={~sSW z-!TIRm;b)<|M<)Q&4rIBB89M;6Zw!>*s`NJ7okw`*nI-orGKk`;(y-*3FqSJ^btVv=r=CU^2(fWERHl zzIrd>=SF>1<%WZh`4weB!s}K9^5Z}0-&8li&(f}0@2LWr1qA?)&aNn1yp;9!0=*E# zYpz6YpZkI(@et}0G=0i0?4Kl*F1{c0xLad}iOsJKTS&EM>nZBVN1Q(XNl3yA-M0=7 z5QHs|HJ@GYY%FYKYBvoJ2#J*bhAXyO;un4~|Gl>IpVCSHyOtOnOTnrEl~7+jA{$;h;0rgAKh3wFKgly=)CLpKD8e6TE?(2$;IKW z*J1sbkIG*_E3t_i>Y>{e(}PYzWMTX&D47DZ5`mM(;senv&DUfo+yF%H4H!&J$DK^z zXTh_bBz(*!RUggiA3n%5V=qnN)|E=P%Cwtf;rE}ou4Aig$!f@TiGaMlk2$710?zXu zM94vlX|jSWP6sA}z|Pu?DRZGIl5=r$iO2RJgF!`A5$h(AIqqL@=uk8jg;4 z40M9Zwn2Cbx1)rBS^6t5vOU1G4J~d3@-ve9q^il`pKV{hOL|JB8LO0(#FZo%$*#o7b=GP-nRB+g4I2${^hbq=EvwrKQM1B=P=v6~o>!mqYez>3fI`R?^U9VP>3zo5tBS*y zeC1Qlcdmds2oU|s)41M5lo!)R6*S6itkXD<;JJ5UER98=%wMDx%mis7^WccCT*bOz zTg(i%RMjuv%;T#9nUxnHmFvxY5?_oOKdnBSKH$8{LQaCjyemjN#99xw^iBe$_bV<&(s{ei$!9$J3wztD!_Y6cGpb>(u^n^?)s za1>U$`1xL0=F*RQRFrYdlH;LR#z^B4G-zTi2 z-AWzO1oS?wC>y8_g2P{X7N~|C!*|a-#&fAAvn_$lo+c?~tlJcLF2V(xU2ZD#_0~9c z#y+9o^^%YggYX`T#@+^=&6jc?fKkG$fM~bl&c200?YRdUa&C<(ZB-}+pB21|d6a)O z$`O64_hX*y528$h`BpFhY?#jh4Jj0jBMt^d6$TQueoHAF%nCI#5x(maTYWS8<`iyp z*6paD@bmO@V}!n)a1pEhEG?NNl$j|p`Ol5S|4*>6LT?T3>j>Yjc4M7)3LxkM? zyly`p*ISLjQN>P=yzYw%T^GvTy&0!((rWWE)htF@ki8Ii7$Alcj^+cdB^h~8Q<9i? zjJ4NBsmn?sx5I0$(eGvKq7G5U4rct^u*;4++4! zBlff@zHo6qjM9a5$p;i}bO>NvJdPFQ+gMEMm`cxQTZ~BRAu8`_A_ATtqRretf(e8K z?LFR$x)v`nn&=_;Hi{|yJKGP)eygBy!%zl)BJF+1_So)@S6x*+IJ6SfDq>S3kh!BE zoJkCnQsfnNj1gO2Y9xa?x4~-qE3iL7XVD&o46+-kv0mr4gvDFkqn^1aw@YtUnTT?A z|5Ur^>X@OfC0M$4CrSG{9mNIZxpXJi}KgMu{xJVy{P<{UPk^9yegrKpn7 z(LU=!{t?FL!>+puWv(Xw7@_Y*4}v>`JJs!(iM%L_$M>} zp4b$_z}qm+4h9K<_op`qVOC{6WT%_V1)9|kHJ|bGPO8h!x!rs3mP?jPphK<@#};Jv zNRo9!|FcerF7!X2>%Sihv@Tn-@s$RhoD;MeqD)e=CxL3;7J!UD0%Y~TjJ8uvV`JvT zn(!^IkMxhlT|+w%S3C#thiPh1s}jmtDHP4OwSI6J%iJ8SRyMm_SCea_uG833{XBtN z>PDxC6qH((()wgmQsWH$t_8@-{PBcFEeT-uAN164kgqRd7M9VY0RJw#4rmJ=kTy>ackeT7D2dIZ!kz; zPR{hO>(1_%bgdjWjFX$PW^~ZU{O3PpU(Q2MfeyUnA{IU;31o6RL4oRoSCIaI781_s zpdBXrMmu8cAH8(hSQ2rpjWLjaSbnKbUq`-u77Am3IO;NPI;>$daN_>OVCoorCtIQ8kyfzYO>^nb|E0MJoE zCgn{*9wQIaJ>%ovp2Oj)kX^7)-U_kADt{oMEW_Z+|9@%;DvksNcFnd>^Q_xV2G%j@-i zMM6Td&G^^KL|by}>aEW}4D?3A@r0DV!#wkL+aOJ0nTK^M8J`JnksS~d{D#dJV@f3K z_d2az#AO|8t@vD7CdruVP5V|_UgpAjC*e8HHLZ%D*r1&ne^v`@cM@YfnX^U_lB^r^U`K){**-99AOSp;-K zUlY7H5H!vpFnhpgw2cRI!P3aIIir8Ul{`I9iIurtQGYHWFCjA8+VD_>_r?YcX-zTM z*C^SY433{fG;*21(_%A=XQ%1bjjnIS)BP>8E*gtEdMvzB5PenZky7wF%;vMO{i|!9 zRg_+u6rO7A41cuE?1Z@WIlUNCsq3~Q%5}Qj7k~Yl`BwES)hn+pR`#h1^XZOpgIm(`9JM(}asis047C_S~kW!9m(T{q1c-u>p)Js86#uQ+|`M$(qUBODNVbSqSp-4!S%sC4u9ao_sI) z0nxG-e0s_~58ofQ0Cn7gh|8#0{glY)J=*7Uhe*nT3;rwoXH$PJhE0}9fzo{bf8h(93Q=O`ho zc0siOE}b@~6ONJb0y{MpJY?@?ZtaPAh(Gp+`=7t=5IqlyZJ>L86aj?g-C-?>c#Gsn ztdMPpgSePbu;W&^hb(^zKWoaUfFd;y4BgEFjY_9&cDR)YDnHtF6veLWeH8hq7MUFQ z#)f(JisRc%D=$AkM?vv(sin2AOuC+)IQjD>BXtqzw9`a^pXrZB}gC}t_aeIWs$QV4a)df^jBa9jz zfyuqV6n2qMFurrg3PvY?cPS_}Ilq2pwNU-wwxtkG3J9qDkMybt|9c#~I~l_BU|zjg!9W<$b&8`-P-_ z56?N)2EO)R1eAF6CK6U@OD+aYnRUF+0~%+EABYuRwCG3E7&Q7{B+_y9pn`e z$jYJXHHtmIQmRED_5J?M~Ok>1?GapG400Abg8aJW}Iw1^BI%w@GCdXp#5V23% zc*u-Lb55{O8vm^7Q%N2xJ&{ z*qD_5mjJ8BGkPtJ4P%v_POnwZHI6E+OJ|s5oRWDd>+azmpm!*N7#@r>1dqi4>d<27 zVhq5qnKa>4WK^5>uuqHH9JgZE%h%huyzgFP^%=HyZ@*vs;gDh%2#p@OBH-`@5g5M{D`~-#oh&^WoPq$>VKJx5zLYIS+30kVZ8rWYHxyZ! zQam+7&j$lIl~OGa#cw-UpSyAAS6&%|+$*}6o<@awu`2yv; zGPP8QPI8Q@9*{X1t;}rxi^THt)_R>hb55T;a|syJsT~If29xCQQ<04L&lO-Dxn%&r z7C5~p8PMMdKi*x{(+vzZO#pcUe|-M`{g3}Y{;lGUH-ozXe9b4gnZZ#+M)x?%cg2Uu zQ@~sm2Vdi{0uTRrqteN%u`KG;T~#%4^?YNj|Cfh zcYmnEzjzS=32&xRh-wgtK(n5~>c2W{X}S6~*S`VY`xBI>WabEbj&7 zeOo3kjjM}2ug&Na9yWPfE7a7^=3cJrRzc(3$B;;r&=5`4%(mi z{6vBf-5gA_W&)A*A1q9!KZjO|;-HLVG+y`zzRrIYA`fnpU~|<+VTv%yf>IP!&=dzd z%s`?>`@7DqzEp7YNL4`Ga2~2oO+77OGhxO`&O(pI(8Lk}%r={%N`Lm+3P`GHPVRCt zj^xv0{zfcDE}>x5UOKGOhKD7|y6;gncJD1&+rya24D2k#24DZ`vpJyNr`_p$P&T{q zM&R|&>s`&w8X6jEvdm?mhR6&@{GD0S$0Q9gQXpalsAYUiq#1+*T?PD&!+kvHd)2XMjaGBd;;<3CxaJ+Z)u-kEda5Dt^zx8 zRt}RN-#J6Q9@c?kL-}v)_iB@ZJBZGTJ(DiU_x*jWf0*aI{A}H9->a#TTI`w#3M0N( zBkS(`#?p+;sl0%t!8#0)p}mye84>0I0;l7r!IioC)4vad@s01~5ccVb5cWf}SE2W+ zxu94HWX$L>^hev83}}$bo4F90)z?E%_59c%sW^{>qtmvLK(~}PA3V^%5g_QFZ$i+3 zwMELH<|$zs0M(cUGM9p(@y&D>BdSzb2bvAKrpmIcFA1Z-MR%OGtqR(nny*jw`qBJM z<*T}-MoYPn`|*s)<}8HYsXwS~5IscZ;@$lxdgmY&JTcU%JNxjMmKFE+EZyb`UQ_QC zk!MV;ZZGT=*REelpS$}?;8YIpmu1CS-ft@m9uz=-`femUT6KcwT;>S+6F5IhEv;$g zT3EI4;FSNNgf20x{H;SGGR#T_sGxsKsEvgs7l5XHeuJn$@BS7x_f+$P$4zqFtK^|@ z#C)28z-4q6B!~jKWZyPsW4a{WI&0vvZpYe5<&lZL2L+4j&JI`A4~fKGYCr9n#}w(3 zOJ#>T+L1L@G$GSDF)}YWjD>DJ4KB0@Y^q$LFM+# zBVJ?Gv3Bm6?8+lOU&^jEHJwgp5q5oQfVzn%{#Vim*#G~qJ}Ag3P~MXvXKr)Sv2WlP z6v1xW`Y_Q}HEyldC(T**>WIGWD@N4STtfQEsm~8RNd!24a9amZ2?x87w4<^BP+ZKs zwJ(l>o3Ptvz*bZb#(*Nlzn!5pvfd)dr9J>X2qb;2h2(tzC7Y=dH&1V^J)s~IaZ#^+8L0T`=_qCFNZ41x-`)Z`U?d=B>anCnoS^I zMFH*zhxc5o6DSXHwOEE>%=4$>YYJCH383H4KuNkuJja8(v*{$r*#Qp*Tua=Y@dTIsO*%x=ug!MPjBrK`CRbXzP zuRInA_<~R;6K%lg0UAw6P~`%04m>N4V89__NuRBbBKaxEG0W+rR@cs;czS}~NK2}a zfzvHN=v$#V8=_f&X>E_)1gSyfhXh*w)a*DtjT1#by<7)xpuab)y#Ri06IKI>7#sO> zd_K*M%#$`3g5sGzGxL00-O~S#n@nXN%RPgo<9XrDX@R|JJr3rH9JC#vZ897NcE_e6 z7@Cay=Mn++5Q)Suoac>#wk6RJ-7J8Eri9zj?Zz|d3D038GEv-HRve3VE=*0sJpkEH z6*n!aNX{|d^=L?kwdutaDXlq|sf0)t^u7ThD zw_^ZB#stKHc{s)0BREBR3jVrk6zXkIXDm&QoZRsc9Z@S$2GzH{i}+gb(&EK)gp0X{ zXX;r5!q2Uii;G9`>oEjwp1|M%KJeLKa2ca}Kmy428bc1kL~K1M`2g>zN)tj}8uHgy z!{`jryh3Nc&Bs_ZWspe2B)0DnLobx)6DQQwg=1~xo*=4Znc*qhfPt)l(}Y3Jp3!%ZkLwu>@yjsz{>xmZ##|=*z|R{# z#}j4^5Xkq+p&-pX0E(~u0c{6D12*WUTgMr3dS2kgB}O-%PY=D<$Zv;u>+{N;KZe}* zt+O=wEU1*=TyMnv-dKG8mGtq_Bm{HZ=HCc_m;}`UDya?gc!!{pA2*LGinVj7dU=#c^St%Q~LdiOhd)f%O60hyO->z?TdH?!_87{#Xbg z_gT{^Tte6WjB*7~^a`B|_S2iesx%A5s2%RFLKQ8Gh5I9rih^h&mFsEULf)%i!|21v;B&K0R2 z&-1^5JPzSx=ZjI|z;0cjDm?u_7^rQ{ssY;-Hb52&*(EUs@ePCT(Ubn$S7PSuqA(r> zR^*MG$5%R&oCs(!|>42-ejOC`el3Lv^QB0yswT^8nE7u+LIeKcK zB~@|1pvr-HF~xVx(JTnwX2m}p;8FcI@&n!?Y#-DyM5u=3*ElKhoecI7Jtq>kgf&I( zOLblIxaN(jx90Ye@k*=RZr|GvEBE}k`c&N%_`CfB@BpA+oQBnAU!%6`hkY=Xh0e=I zLwV1~DHMEVd1UtB4OK7YQwmem-9w6!Injr&F~cihi*W&CP0E%A+aG|0J^gK1*Fem| zVmEk>mIBrqY+LViNf7AYKX6Y-o#hz(m3okN7IvHC z*4Q=~TE(dP`xnc7#ut1y)^n^KsIpZctO=NFw5#{yPo29Cb6m|pP=_-mGWH>%Eq!f_ zAQ`(RNzQX2>m=QmM0L#wQnbpzba;=hMDxDrbHZJzN{HvU!k7Xa)! z0s4I(y>lE|zlDd&06@?k#B3|GnX6jl7ka#8Y{#x=1Wl#kEI>#Ow^6KE9bG&QdiCs4a#hXWH+kf$* z0kZp_F?|2KzH)E_O(>ahc?wXX`v#vtn-h1?ogPpKR+u5~g%Sr^J((Soo7MxlaLxIq zp}K>lr#-)%uAcqUFC26Hye`7@@Naa4z6vnx!?>dqm%?57HU7eh#fl(>w-MC$6M6uL!Vnup*defm8&Z@mB>IwS)sP)-tNZZtD*^6;C!9gtey3fXR7= zje(Twpq?c~mmwmBpPe$#C9kzcXH0suCN%I_`0I@@D@9vhOCJ(Fo)@CVb#{$@^FgZ2 zZ*|-#DO}Q6&Cf0H{fEc92ZVm!5nsP9^zbFZ@i?lk`)|Y%-hjSd) z2om7~^;BT`gtM6IuSUJ1-avQ^K1vSo5| zox;|d(fQF@ozY3RKKRQ?`=v#9!w4|9n@s(FN*l)ZiOmwt1WWlrQYlTP6I~(8f zp$}09nZFUUZ!*IIxH^b9ZS}{}t{q9`0SRWlj@Is96b%N~qcj^JFn1NUnwzdno-K=S zNVBXO)jP2RW!gR!OiwbmU0_MbOYAVYGu--{!45*;1Nq=I1O$B$V;aMaBtjgPoQ)2) zbn)fs=`{6KIqCOWTtn&8h?)CkiMK!t>D4RZCDwU)nj|PbDvZNF|3=*ArJHW{!~5_h zC|D_oCQbk)hz?7K_9N)&!h;N9BbckF<8pH$##3%@FF0M@?Z&7+Um-PDlMvL;un zQ{&ndl$#ee*udM8Pe(PxZT4vqCIxyn&LyoM9X(u4&g;0R_3q9OC2OsZrKiB8*z|B& z0X9GR2)q$zhw#$P8+@`gC=2yQi=!9D~WrUhAbbHbM*b+dj%x%}LRKHdI82d?r?s!RO?UgZf0hqKj%0XTUcRV+ z_l{>g{AZsaQsda$iy7mhw9RD3t`L))(32JE1^W|Q2VRb=<+qiwcZT@9_|Jkgw z)8f5zspaSg>0w1eh~?vI0b-~X*?)5^1GVmC)_Vr~g^`n_zY*F0_?7HR$-DU{wuBJ1 znQuk+iDky*Ge<&D$91soF;ZZGBMQSl>F*~y<6_ro+?5#RE*G({QZga_+|Bv?mfT&s zPd|vCl<8!(wKKSD=z#(UGKG2L32Jd3g8~Q^43_r7Ko~`Q3jHeaTmw(gbsBh7OL5~c z5xZ?usxZ{S$B<(5Bgq!Y{ESZxjy*V%vLcBEvil3h+1XL`8g?2k8FN*)ho^AedV03_ zg`dp+cnvkw^Ou>lgG-hAdAEH4MX2Jau7=D(5q~zbU&lQcxZbbs`_^NI5d3zDbpS>7 zjG+}!8S{xwU@%WT4mdSpIk4zw0Cfb$fE0lgV{y-ZC;OZ5tQ27hq&RVro77`E6*G}cBe{c^x_WTrrX+&O&gd$(Y!J5bmd+rw{{*U`* zqj7r9O`r7*vK|gNBV4DQeED7aLSw4JIZBG)hwk?$41F#li(TG{2L zAFSI8X;}eIwr`XMf?}x38l+zq`yUpqE`iPRmqW13TugKTRuXf!fnxP|v-xZDsAi&S z=SI==Eshr+v7R~t53ap)?o!YVLR@*;xsHNW>}Y&wU^)@qQ!vW5o{~SqLMe(IvpA_b zUjqgxirm%Si6!#H1Y`JSV{RMH0fd)L*ajGldgVSvlpY18{f+ob_hB02@Lf>DJ9Hxk zJW3xt?FG8zr7DeIW43KCr;L6_f9n1<`L}wM!=aDO*qN``?Yqsxq-O`3jDzqI-Zf~J zSNN9W1#HbmSJ{$gbBV7Msnu=cAD&A1SNI2M~~11BxFrt!~XWbWVVDz7}T%vr1pxd-?GypJ0^ zM8W|Av3+v9xYZM_X(O8vJ<$-JwA@Kc`SttF=o%x1B-w+UnSs|IR_UcZN-l$llW16m zcTIGUMxQ*3io^2hA6mik&}7InBbTz*(fm$lo;5wod>JtP?9ybeSW2hW54BkvnNz-> z;=VV~4v7L=z6|t*pa&prISWV4f&VaS;P~L>pqG6A^kJ8!$*5C)|?~WY>Kst}(z=FAo#~tMNnagn* zi5v{)e`uup!^W|KduT%>avvr%7xJDis2UAORJYLOJluQ0p}_UA@u9S}cKLmHPvGMg z2Vpt89@WXI?Taw;10XQ;vj;eo1H`5gtqK{7>7%)KRY)TwhqX78Tvd4!Z*QL=zSoL?e!yF zH|n`JdZ}*Kr3E&8s`I9_Q@q>2jpf~gL81}Cj(?B8k0)jk$fg?!E-c{V?u02|cqzh3 znX{mKixhRbG`2Z?e4Ci(RdxsRzwvZx`f0wgM(DHDmet=vr)+sjnP@#|v}MW6z`c8# z4Wgnv!ylC@ofiU6DRuoamzT}1qcVfra--d(*n%`?zYq`|S*x2qi#o89&DZi3SDUP% zT{GRGRcfq1vg2-(K7N<@)Zi4tjRZ~u_(<`ShRO3|=2#&t*^|#L7VSRxa}gtIQc@IcJ%gYfo8WRNSLH zZ|hT)+-|LZxC+&uaHxSVqlbEFl8wM7s{vfnjb;kFQY2`_|r=_*u zjaC$RW#51Pp7q4j^Xt}uZ#$TeTh00ZkJSetL0(vUZq78WO&{tIXYj%CJrCD&NtVHd z1x2m&&1{dZx0B0uwqF$Sc90_1<;ou)B^zZ^&8aWjjKT903hQk7=U{*enKwXGOC6`@LSBSJ zaA(gX_);r@74hy)XTpqmn|a`=M4|f}PvVM;a{zIz7{P>|h6*4_d>!Cy3vSoKZd5aP z(4-Dln^QgLFtwa!?pqFrSfx@eN2OC*x>6O+CPb90vvCz2Z$EJu@n8cS-*Z;~h?m&4 zA*O?^fXh*6F`8=??R<%y$bY{q9{EiY%@B{mFJi_S?7-$&lidkr;0=C1-ZKFq3~i!Z;KL*Vue*G53&5EOoMXF}e#h4=e_b4m^qDai|V_Q=>YL zBnMz=0y=Z!Yp19_us^8*^q!Uhr$)B~3ig{a%;4rY(9}Azt~{T8_Qn&gC(}11?A<-g zPM-Y0N_-tf+jZX5v8}Q6zvnOgQK;&axpcrOTiZhsY%KgAwCFgfLqcB(PMa&D2@r@u z@fhc_Y3~`?aXn@KGmPB(hA-C*8-p4V9#1Z(jb>joKZ<`A!V4e4uZ0v7jQsF55Wd@N z@`1>SUyFBp#Z&KOTtAT$BpvY$$X)*MCgPPRAs7g=o#;C-_Pubr?=H;gLYeXz-fYkwYqHzVKdljsXwVuFyN@ zHn(+}6>iau!7~P)fEDI0y#ntA#X!POMw0tj~tFdgOq#9f>d7u99#%2(VHoyOJSMOq$tYA)LGuC zM2&2$np6bca7*>^zu_J1W|F4+eV;?@4f#LQxFzZwh?NKDgjB)IyA0z^i4Rzp4tjj-X zVxpA1ZpS5fmZ&=&bg799#jI@({AD>@kcJTv)e*Q_KFFt6NfB^ot_3`~RERxHlXi=O zUK6D_xlnqOwAPHeW|SyeQ!dMn>8n{k8lO>Cf$&Xf+`c!ZrQxLrMk0zh1RVWQNVNy` zE!hgLne&fp6K)}&bJ0%fy)bgW!ctt)ysiItOPby}ObYhoe%Xb%w1jKeUSuMc?2lUu zAiI&;T4w`bg*ht`(>K#~Ee{7{HkbTs2P{u2Y6?V2q&$m^U{;WGuz7^oho|w0?)h0R zU?W%CSh1WdP0?{}$0%nqVB||JBEH*tZ_PDbbRp}E6eTv@)P~}p^BM;)y!%T z8;Tf^k%f_`fKL*FvF-X|WS&M}lCRZY6uDO!+CL<`E`L%JU{CctX=5ze zCuidj(+29ZShDHhB&uN9SdO=jmMI`KoLZ+h*Y#HK|LN!y#$x!n9rUuB1u*G z2f60CcxXx3`{t!e@9@v25VAx*xyM_l~1R`N#hlq{Xd>wLW zgr+>_Yh{CV8L60`5GfzK-o5tt+rqHU3+Ve1-JsR$oYKOKJ1e_!V2SI1!^pyl;Qa{U zSWW<3FX~JS?wOly%L~&wBi=5ib&Av`@+7r`57}xz$@_}=9V?emX(PG=58p+ufjHLW zCk2FoEP_xJ0x}?wJ%9EPL@>Y&M_eR01(yM8^kR$PoHv1c@y0JJwM{ytt;nr^v%)6% zlUw4t+_!_T<7`CP;>vb&(Xa7jAJnQaMUcS`0!WsxX`KLJQv0nIB@jH=t>dC>Dck5> zp)wg_p%_eRJox@%SpVlsf^W@o-_E8$=ZXae|mp!2P6iYOkEm%^Lq=+xAg1V(n}vROmBwQ`B%qmP9j9yqAq=EN3Pz@ zboJdAka_7blTcbUMeQPQFB(2`k8wGdJjdwwEn>W^JI`?Q-pvLNw_Be*IFm%62{fAn z^DG=JafEq$GF~!j=1YR=tEqAeHI}ve&X;GZWtYu_$;x!ZmCgZF%8)EA z7PA{{wEGnLde)hSD4ONb^!%mFR0W%s>@U+TRCRq1bghXFSK?HR4N&$Q`WG;!|2?*; zhcN_+=Fzc`CqXiJ+d|~44B0fMjyrTU11*Q!Ur?L&Jd(Jkys$PMnedaf08(*L0{7Ge z`S^NwJEs50BJmzrYzmkjt#MqUz_%vdt72b0i*hk=ey_kG7VDOfY6)FHk}o9+>~s;{ z>hyB!Sk^R+0$tX2@vXXwidsu0Xk5MR_3aXI>4;CCx)3=6=Af|LvpCdfYh7Lh;v7vT z`hKvtz!trW>h+(pVsP;1V#MK=Ce}Dj0@Be%B@6O!70J?$z2oJ1=b`P| z&C>c+?LJxjEn}`;wo*_izF})MQo7f4j-21*q+;0_1CIHDI;mcPL18j=?^^8br#duC z3TiJm&)sS+*WeVkzWZKY=Ey~azGS9hH&E=U#=o_qi5hjI521)rQhM})7Q0zX?VTPh zhsch_sPbBo_I4BfYpwR@2Hmra`IC*7luDEN&;q05xL+Y?(4slN!L*V|Am8j=Q{_6X zQlU9HR#q|XMJdQn7$;ZMzf6=h;QH_=H{C+TK<(t`D@5k`tPNQA-u7T)>q|3?6A_{Y z=LfpSVw7*bmy?fTDUp{{2Q)eytUbcu*$4N5yi1f`0)X9E(x#(SqsGm9!I2lnl$7R1 zWg(Gxg#^5lSzoeb4r!m7EGEm0cZPX>=WJqOH11$2G8NQMy2-k%39Glau`Ep7jHSbDouv!s!4 zec|WpN&-C{_4$^PJ z*EkTi8K^du~Me zRQ%=K+Ehc)iM;<&H+DXQ^CV5pTiQ`ZO;Nli;*$^`({=M6klN*g@q3k#p#U{~R28&5 zy#nnXW)u`-v?$zl6;MIBvgVF@Th@PdvfV(+o{269bu+OWDf~7l(L@-OU$M*n(`Me@NJ=W`8wwKweDeJOwu(69cqzcYeOCJwO+?SZ?tsCL09Qg%#9 z@Ovsy{2NwM4Low=nXqjC_I_Rgo`C7{l%jY(U^Yjzam6=-`uGR? zN6wdVwnyG{KY@L1_Td`ahd#Rf-n?V3;4l7Zy?`5ew0O)8@-FNvS{y40boD>-!$D;t zk2kU4eRXQ%2QJ<_=r48ujL#MZL#7Vt0B>)wN1DTJ3)fMSAKjng929|T&@r8-eWw?H zA0BqG2#JrpdtFB0Y5i(0JM+DWyG(!WM1=uDVaL)W_jU!iuRgNe*NXJ%f%;5}_-E2D z`&`c?Xjpq>i!FWr+KD*#6h4u5rRCF4fyCLz|6HN}SRfF8c5nh^9Vu;kBc7NI9R&Pn z4Wp4QUSSMx{kzuYrjh=MUedvgv5;D{l7jtJwhJE5_zeGQIQ;XE1#ZS4)@tZ*94i7> z@swW67|37FL{7bQ9)Jq7795*FFX@72e7Bm>BWf?#qITs)XOmMxOuq&}9=VR6;&i5j zm4xK4%y56%Y;JD6Q8SX(nK91kCa9*smTKUVMV6}M2 zJ019VEIAp*o~LY8EZe(Gcu~A+ZZ?Qy*ohWU#7S3D4&Q|O>)GTE@-ZtyD2=n*Bbs+* zOl5s+(xb)9mf5y&U%+Lj8Mci9bsO}7b!DaX%8Jc*0aIlYu2K0GCojfAe5q#-@N=bh zWo0qnf6{8dkBR{Y&bD{@B&%kZiY4(I(F)?K8ER0MBcUEk=LrMn-Fxf7J-rFAKt#Cp!0t}~M_DQ9 z^$0qeaHZJPTdGK(CeYyC_X-b{Ln8NvzQ?3#zuZxg_k_gBiO}+@SM<@<|(;w7a zyU>$H-3+0T4LuKX5Y_0y=(RfR>NeDN?b3&?luhtfh#{aY3aSgZ+ac`=P(Si6f zOYPP&IO6ApM5Xx~qt~4^F2v?PW_lrH&h^3JF7vs$(x<5X#mu`#z^VGZwE*ktD4tY5a3G`Tgg&4wKsaz+NwdI!4j8 zF@WE?LkfBnCfYZPI{;hE&fr64%_`RK7YFFZ|0(zLXvkPwND2E(G53fN;t3z-lh{;0 zpvL+h*rhT!h5`U@;5#&S8j7fsYicxe$gNj5_v6!Lic+$vO4W!=N~groL^wNvx(7ro z!`7zP7q~qr`fg6q=o0l3{-i9}qH#-OfGhk2OMZ%3s^We30N1C6IlLf$yfFytHf$Gv zVdw?h&h?6B3sb;g^jevCo1wfRK>t&{2Bkz;iLWP)-}dl<_!~`#AD=?!!4v+G3@d{f z8oaRFk23N`2gJU1KN00Z2`!)O9&3dBGP5u6n=gdEK76M^srSykv~ZwI3BBNb+{C>T z7p4GvTMnUR!-cT$HKRLUN!@PA>S%J$ zJ~2{pAbavsqet~i+H!0GXHwQ<7oP|>zEhJNtYY6$Y2x%vGCq6@_`1j=7{eQy&i2!N zC%=j?dC5N?izj_h;cmam>)Fqzcpkm}cz;dpjd1;w|HR#1r6{sbF9^Izz_+jT68Q9` zXpRj-Tb6E}4{yN{y;t7?^@8tW*B@25ut-R7&3^4+`}1fzeTvpTz6K}|GC2fu>J5zR zSCRg;+HV+9$@0wWBEMu_Y~*+(v}wn9m>x0jW9zNT;ayTs8)||eULwz4&tOI3UkFv= z)xd)pdLaycCd^1&>UY3cl=~rmp@wzGm%0(gp>8}yp31GC)UG;xE|85g=1Nr{M)0yA zjnF6%HFU^A9fFv)fPFOG!8;EEI-!2Hb#c*#HGs}CK&t4sGfl(bzrkxwy4vS4Z znMI3DCRo1cvHqTE<+Zghw}tr)Y?!}r3A2aJ$M^f&lbGmu<=QC#l-W6I#f z9`6R8DV+XBN4;``h43gXA@voIx>mZGN!4OH0pXY_V45Co+k-L~yQS>yNXP1PE>HGU z1t_m=tq~&c|EW}!Px;uVlr$vbdElM4ml7XVWvIn+9V>XBd7bIA4$^V66i6Jl>4FVC%pm-Ys3EK>Q`IMvCn4U> zK!{HheQ>6t?KQD?xMIsZVEhp&W!_$eEykm5ao#iKnSxrlMK$xbTo^N0=-(*H9QDNC zh;b{SzY)q(;is@Lq)HP+G74X%FZ8`4G1Fn!c=djxLjboAMJ&Nq zo32IS;F=CJ=+brSq9S(bO)6y%a)75tx#pXfpKZSMy}0HEWX|G|esQibNcG_1@n1&h zGGG)p2($t`$iAwfQfSFC5fN4$_8s&zidHwM{EXo?Y1lR34g5EvSG@Oek+-DF@6P(M zWDn1W2_Lm8|NMMXo1(rJnp5mk8KlTG5cU~O4i;VgN}?N*$C6$9hvu@Gs#?f#zbEo0*BlWs2;g?govDw z-}a{X10b)9lLL%2tX3}h^X1R`1$Sw!aMQ>inw2Z!MHEry)-Vxp$AA7I1#Y^H;bDJ0 zEpW9fo4T14gMFgVHrQe;r@(rqGC!ExaA_?EHG%Eo?%d| zdM>cwO~CbiB3$;uh)3yzv5X;bJMya|WWhPlCoK}5;CA#+#&Ew;~RAcpaZzIxXL(F1^i8N{&8qJrN2gbf+bc)jMmHVErQGOT~P0aPbXCIFI()-rA0hy+2- zDp)ZO$CV80UaGLuQnl&n)6>Z@*`VfUeB_Y;rbgt?Lt0Mu$shO-eRkS4@*r@i?;I~3yPcX| zZ%;SQ?^c>Iw~WR2ed^jq?6Vp!w-ck9LVYdpZ$|etIMB`Wb9lS_Cy@-!CGarEzUrTZ z3q$|fzET6H6!#T>+{OyYhi@uIjUPHB@kp+#Oe4C-=yZgs)+QK}1FDshR32Iwy%^h$ zJrKt6jmAlVdSlB!*{0#6YFp*0(2c8;<@J8^Z!cX{B=QxBN#iq`sCEqQQ9v~KiUIwh zx1gx7Yjh7N@0Gxvjk2MPRPhi~FJ^q4Ax^SWL}Ri?~t@Yq|GEC-D%*GuT$bn}u+ zAN06lE4#}<dBOI5xebCn<%D519odh^`>5|N z7-ZHzZP$2=!$6(91NeXiVRUT~J@s#d0JKy|!ao5Wkhc&q25eee|8IotQn;C9#e`vn zFME(t-Az%YF7??Mogy$d!z`D_KF^`1JdsT4h&3ZQ3gK-q|C z1-)w|x0l*UFsyf|#?ZIl#n57>0kV=%4dxSo+iWNb8w%r3J4b)Ms?WbXfRV1FU}jjq ztj>LTJ(~04M}10RVi@FG@b%)Cu2;uo!`4it`{Vk96uBaF;}}fd)XiaNUKn(mucrOs zKveyEjp&~`Q4f-AgM(F)JFu{<^G@`AoRld^l^+np__LZrSv$$bTJs_XO`(vu!QhP%*O9ez$sbXf0;TI`KC63I{Cl@m5lQ0h9AG64Mr^5a!+FNTEylC- z>j&!bAa%deRNE5Ul%I0)_;M(lbUlbPK>6|ux;m3C;opv@9oZx>G+gRt(FcMBx2@q* za!FBnuzx#1{y_!#JBi*!j>bhxlkS{NUithz!u+P6R}HqbzKy3a2!vnPkrJrYv*XYv z;S3_|1Pwq3BdAb^WYQ7c<1$On-?u3nV;eg ztB5Yy)%OT~~09S{AZX=n15ULs>Q zL#|Z49Z^zFof-bj;#?wh=P(D@avTVR-IHaAY>uArcO0!$RYiHCB6RG6tTjh0RU_{I zNVU6ep&`fv09oeCmZ9v%2+*x*_n$DDg;r}^iydhia1aq2vprcGue|Ca2?F&|s5#)E z?bEYy2OqV|pDafw0w05F-j z<@34<9b#OjM(Dkq{Wi#r@{u*FtEK=%^bn9hh*_vOJUP{wj9ETQ1#SH=DFr zZdQD8x;s*!3-O?LI-N!Y@7+NC3P6FdX!Q!HW(;>2NV3vUKujzd82=ftX`_F4qG@6H z>tI0Rv7L4cP^s#rsjh^2Th_xzyx;suBDhUvP@+W$kSp~has@+D?mB#T{@4n-Yv8{i z71RC!Dk4y`G(7nNw5zs)rs-q@n$FzTt_5)Z;P8q_XMDuh<0EKxv-pXh?SVjPR#O5H zH7wq{ovRu4czww9AfEl^fKZ&7{YSs=9L3j|fGb6&Q^DI0oQL{0X?rx^Dt-vy77dWY zt$SmTK}OtR0})&CK%OZzTM-nkqG|sAB!5AmK?lMMKWB-~FGk7(uL~V)Np6}|Z}OVd zi@!U>)gKT_`Q=)BshfqiW_#Vh^c0h=U z3XGT^j{#JQH-cbXh14`e5rVq>2?D6a$1oCj+XghvIPB1n0rDMzBC*qzyn95S6@2j; zh@$yT|Hp^^|M*Pk|M)|NDL-k#WegtqJ(-pLPGrgCh%IFq6EZS9_hzN4rn8FOyVI{f z`F*Ok)CfP1>#IM3;L57L_vp``JrO!P1)Ab)sA;d%v61`EQlIB8f31rzF88T8>^wYe zrT92WQ*z=(ICCTDan)uLM43b&oGLw5pTh-+HdSHYqhlVXkBUylgwrJQgC8!FgcT3{ znshK`Zi=2)xY>vJ4am&~`yW_6!GqewBL**J|8`m##smJVV7bY;D7t1Znu@ zxBL`yiSuHT_QJ>8fz8}_L&sw*OM(@*`Qh>M#hsm(wzzzl)48>(6{>xa`>rPYjDAp( z{j`~TfCtw>K0`#3H+&sTi`xBRUYKfL_#(bfrs_tZJb;+WW`jd<4s(?Lh9N}@@!lfu z34_g<40qkwym|U-eZe+LBXCfoFL0vSZ7vXm4qwJS&>!$L@sc@^Q=`aE8;^IxvPze= zMWm!_$Cc`wyrGboB+2tELpVPHg>*9k=>Zn?D4%eOOu!}6ts#T`m=y2`fLTT8(e|=8 zJnit87StTK(G~J~>#Sl!)Bi);dqy?YMp1(xO+Z0~{MM8;-^{F;Kl6jNAS+zLz2`Y)pMCZj z4IpbWUL5b~llNnbdu1&8_Hp7l?*|12v=7U}y}*V!kf8RQgP#RWp*F+{a2MuCl3hd+ z@E174l$Y>27zyI2zD>^BuJAfU(h<|uPd>+izuw=*4ztn#7f?6B-I5GkDW}7^vxA z7=;;3x4RZ1((Uoni2(_E#q%51Z!WpD1{83)UN9<4UP(}eSjB`#@Ww%#*7 zeeqZC>)!@ZFfC$;B0t#yql{`ZPSwlYyZtjDN!7M4YwWlWv%MpebMu{3>ggwsHo%-T zd5m<6h{V8O*bCY@_C-kIc3=9(zW2yDfvhL;zcjj9ve(3kKe6kPI@dD56fQLaW-qaD zkcW|Z%%YG)YgZ&G1jz`nks~?YjiD1W67MLyL78<+3D?kjypDE(YAaE8t$gk6+xtwH z%bqzi(R)kb2$AqsXkLI|f-~DDpU4q2F)2y;AbAFjga-X!8Hr0`x2kHPE+39+fwb$% z`os6%^3nP@q@wY(tSIPMcIhhcM_kol1F6FksOw7vxIsR^f{cl3U<@WJrZR|?(D263 z-d1x@?3`PZH|pF|v6snW*B|`(98XVms)O{A6?7d_@^IHc@Bq%yaMO&_9^qv(74`@_ z!tK%2Hg(85wO9bvU*q!_6BDOeV&+y8TlQvvpNg1+;hv3b1%;m6LbUF!5b}OTunv{k zC9G?ThL*4AdM1xw_;x?MIN+VA=Td&(wbnfpEd{Xd+)u6(`^|~Z@wx;R_+_9bQbcwyI}Rol+VoGi@pRpZR>TrT1&8dBi$il& zbTAS#mRPrZDi$}mzvxwTW)~W%EE>giC=7Sz5A&`HpQd<&2bkOe?jtKOG+T0xY`xt* z==PXrRfpIg{;;vx1x;3&hQ0z?o^v&K8IvuY(p2xEL~CE`Ug%f=W<40@SnlS4yFbKy zy?IiZ@gQ8Tlfc6rs9>vMV{I(cJDr8^yEW;&H||5lJ<~EwM->_Y3M2YnO5c(7WR3%F zL6{%8FQI=TZQb!otYwh;SCy>s!XiCHt)?Y)l!9UCuBN5*In))WU`$!JI{0urs729% zc$4|BU)BlU=bbZL3$HS9D%RCd`69STB2sT2e!o$h^FqU7x0=KSVS(CKUV$SlY?3EkV*WBtYN+d|Jl^ z0l8IEP${HXgkyF&5%&m@M7QZI+3HWkw(ba#7w&h0i_qe~K5PXep4`v6czNdGpTI`l zyH_tDUxbiBX&$9Pxj=j|Pq;z;qBL_jWUU295%@hzismX8UI?jg*qGio(Y%ynD4Gx& zKKv0tlmi+TRQyO`@I0Uy_{_WsaR_S($5bi7W836b;kYFGO@8CN>jNXR^TQvU{`#eS z&~xOg&WFGndYle43lVI{Ji3T%6CDvAMfeNy5V$TayWix8q^DQ}0I^r|oZE-G8HbK_6WS6r~Ov;96cpi5?-V@f{}^Yy#F&2LF>M9X zrAcKE=q^^nL`n+IJ!9q*q?E{YAzArNU7I`;m#)}rsv$@Q!UdQmDW`z8skgR_>v2X| zDS8GGIIUMZR#_!0Y*ar3GpK1sgRv^h}OF?D3a9|$uBfPI36eP0Q9iJo#!a4xUw9Y5$BFiaR62!s?CMoY^nnw2|p>IWX4o8#lvk*zkbMr?eSfOU(bEZpSM9YSBL2$ z7C|DGgTxtwg`)EquYc=JJ{%l zPVDkItdo}|Y4FhUG>H#-gmdnMa~^29_nh?L^=U{a^(tges{&7pb}VQO{r>90K#-=8#N+tv@4#hywLOEbNpup z9Zu@b&v{>XfB{oe7T~8Fv{>&?D1{rqcwLP$%`I4Izf}jZUu7vg$7l5>agR~xEERMI z0?R&Ah@zRDK%4F{QLi5;yo+0t`0(`4lbBdJYyk&B*aQKMw#VOv^m@her!Rj;i8x)+I{QC}nkl*RlU14yCps<) zZdot0geo&q!&!2tSv;V}&B3y%zbyCbH7;Mspb@(eHEMy8O9Zioq=$c6RPBCToL{+_ zHS{C1={TB&^a`LrWe3xbg`Vl~EJ^;-OcT0(Z2_CK=AZba4I9OtUiH%=? za&9b^+4LG#Qt$G6zVoT?M4WY-6Q2ORis~>+@<&Yu9RE@uQ3_wPgKB&ea{y&%pgonQRK#dWnm4^KVaj2Z1w&X z6?K{s*X9G#sa`3)3$2Pk@8BP4Z^43q2gy+`{%P3?l6X;K>!i{>>n)+C!VSslZ&E(L zIXf4c(30DCF<2}$?kRBiJ!v2h`B*Lo^-<1@QACMymbjfJm#mnCYAe$chq*R|%Axje zKL_|O>r$t_b2=BiY5n%se^iZYuNqiSQ~F>$4IY;XyQY&$NIG}bz@UuUdd0aRp~T@d zgAF~Yaw;&XCc>rzdj<3#kjpc$8MGm>04sIx#L*>!MdkQy4>V)Z+Lt`7t|A6KB3PYC zFW*hs5!w)t`o<9d52pg?9WCtNk#I!Lxou)IBehjr(L$Pbm{65iGZpK%Z)%9%3;ogx zH`m>zeyzq8C58Qvq`!H}+v2Y)5}9|uBWq*B)l*hha6r1YCMY#-;HgBK@oKt5z5~t; z4g3QrM?f@Tn{frU1vxrmnZrE??=>|wBd`+>?YNHv;-4_Ei)8USG0jqA4K_YTvToJ_ zpzGBC_SFoFhc|e2qn|IGPcXt)_vBH50npq<&z0IKlAbH*95X9k{wA8G|2R6e^Wl{n zU$)#ks<93fW>HXs^#j4`e|r3zHj~Oik2j{VkoMOSJ>_lGFmB$)jibx!REF3qY|Pb_ z>w=drlu}WFZ6?AvhV5(|gzZo8(1c?g;?ZZrOsXArPh_);5V_y?cGey%rPNCPa`q^# z=%>QmPw!OT1JJK3h4B&&h*rrho1H;KVUaz|(95jMmuo8K+LxL;nwp~XGo?KQZPhid z_r4Fh^zBjop%6UH>;*Wc9k8%RxW<-`z}F^Jx~f+;Oui*2q2?FnZq$}piGqUSVKD7E z{fP1Zs4}npra6HA%~Vt4ct~6)l(RVE6X3D^@w>5Mk1B{Mwm6`=s$sS@5MZjp)yf~8 zR@d&`e)r1lJF9_bn&BEWx$^kQ5)vOXZdj^x&MVpeYL>XS9etcz+q7ol?uvH9Ni^Wr zL3MdsWDGQWZ`muFB7~6bJwP)f40@Hw>WjZz?L=8^bLv#3fArL7S9nul+iqmm_|f>L zFSpV{tixVJSss!AdKpcW;?873hCMnU<!)vch3JJSC;%#_$w4m61wddk^8To^ZtF&xv>{HvCw z`u@S%yTnIxO0#LF()TSkqyNVwKXngf41L;$5Jk@E5JAu>(~UXw3kyv6(|vUb!#ssN zx1FMfF0}-0@2h$&C?+}uRR5CbMdRH1w|g;Yu8pkllx_{7i0^$nHlg9S278aC1LqQ5 zewiaOLYseEE>L|JPZFdTcza8S()MrAu#Iwd5u|jWP;l9uRirSfVnAQPfwx_wgrqSV zA#82`wWu(wA-OyN%bYMMkS!+*!D~h?cc)5ki?+XHzwS~xzlf#0bopPUOaIB2f-m?x zBvFciBme_EyL}217v84N;S7ndgk*$+dgR#Mgr(~yDjT*{Z@TM20WyQLgJ*yJ@cMmi z*dm19&@PcF!)1IJ>4YXK-y`(kS+`-!e>7c}W#AZewBrO8YMCJE27_t;JIUR=Dero0 z<`#{WlIH_2M^n|w{?;E%+4v2XKYM}ysQ}a=@ecg(Smcn%DZoQMVY#}B>O5P<@19}L{N;LgdRZ7o$xoHBzIPyK{RXHQf; zY?s6JfJlSIUfhW}u!_$PjS`@AcNyWeA52c(;-%7^^S7d1W0mo@Zd{c$tUW0MxGDq(}Z`!2?Uyend>re#W(r61& zhV5CW>+RgJJl$-L<(*Z?Ck=QpLgUe$0awnnE9dGGiuYDCPg^UoJhpI7IW)-JU8muK z$Zx^p^n#W_=>T8OL)NDBz0;P}&c;-8l3e4@YHKkqAsQI{d1V~GMqkWzadL4AHl4qk zEb*{3E^}#H(4i9WzxUu^Ugzu$azi^0p4Dp1UoJ*P_{z@;+0~dqn_9eiiaUMTX}H)R zSMP$KVe`Xne_Fxhki2sEaugcp38`)?pQuA!hMHV>L+o^PJ)pbT^AUr^TGRf{4d;C} ztGJ09V44ScKKoYxJF8sNk91V@5-ESi0B;wp?834VRYM8MZ;zAn&os0=?Ws^@AL&)m;%cMX~=AzZrr6N2$3ZjM>TN*w0jyx88a-0sUrZVwNG~Ai%yEI|? z4A?x95D&;$REWSO(6JNo%dmJP-NBGXUl6~2WG?3fNpmi`F+0eu2Gre#A!&;yf(izG zm;$GTs>GVN>{2ZWdu^E?aWp(oJrNDAqlBk`T8@5$7(}=?)^3U+8)7A1KDM*+ z5t!CQNFSUyu%LqJNQBpO^+RjWpjLaD-U&*cn(JL;ix@|J#I3M@{=p1-hKzOKZrJAO z1{`Kk`PlP~@62f__5|9CjrkKBT&`Y{o874zb8%oelUOJf!!+DxhpbeWr1=Xx%L}mI zpdM;a?S_bwO(+N9u)DaJJRG_k&Z_L9X`;2=*HYq^=EiPwBRK=|?sfK)VtD-@mB9A( z2Ywk3YmiOg(+?nVI&yOyglTfiBVrrQU3Z)mUQA3eZXDk|n4FY8XshJ6vFv&DZgGVD zMjFlQ+ZlIx6`nU_A2kB2((?zSrLlzHlZS&EK&U*h)$Z%WI%j5+z5M7Lv;I@6L6La; zs^rI=^K98v7|3!N2ACk1$$4w!H*hCd4E&5|gJeK?UMz)Yyd5q>y8Z4he!Fsv<3hez z>EriQ>50=NfPvr(tA}+@5r5%0`h^*7evRAo@zYUW$(^zHu2)sr)BFAG`_4Qva3OJp zPsqI*`218a0jTRUvhgTJ6SM|fy}I3D?5+e=-+5@>tEf%$poB#E}6*~nBsI0{r% zxhC10v-c`PIp^ERdA_;$&-3pKDAA?=M;7}pn?(ZQZlaVKNkHht9K;`pi4T9T4>}VD zD}buvY!%1s{Z`#S?=1CIRkXP8d!6PTptt$YYXqGr zLS1E0qiO<0n`mZ+&y4g3A1Z=7{dThHy&j7W(XoKj_1`#-61$hCc8rVsnjT*$rCSVA za&6TogEr{@h(Vck`!ui8pLQM8b^geqPw%v(gKU?Y@XorpwPbKL7JG-&%64~{y9CZN zEi9j1ROmTJA1o^TeqAaGtodZJAvxoGp5K_?1-xP*q`NPf?<&TBxkqJGk))e~8)cXxsMffd-(Jw!6jME2Yp_yqiH`wKUsRoQIP$^-w{~yzK#g)4 z(K_t$pkt_HuJg%dm1Qa}*e`n>LF&I2P;HnbF+2!|Mvc!bxzFXjD$H6QmX*t?Pkr&Z zjaoY5PJLcEKh=h0^(A7_i3*xzfd9}}$&h%M#IN>~(|SH;Z`>_}oH&*+XRmUz6)i_AD48}urHQ(| zX`#q(F!%k)4XQ&o>{f^3{01Z-+s^0isN9pG^5XYzykM}m{F)6@rHUAg7a#N1DFVNNcgG*b33`%YoxlRsCe}@pzXH&L2#w4z>-Ia z99)2SdSVlml(p{c*4Q-~cmT zSJf=d3yNBwt*hUQAz+e_q~OUahc&uBFREQW0F(;m(}#<)OUG+Bzog2nrq14|`KC|H zHaQczuOpC$gF3=EfUJ^pEO6l3aC|B00sf?m!Z^(kk6}p)u5!sZn#;)NnOjXW4~zDD!#l!jicEIbTVWg90&9Kb^Z`N48O5JM zz$12kcZ^R3Pf2!1&-EAcc&};EQTfAglMpNN5A+m-@PA}Pyz+ue>rm6H`UIO3wd;hS zRO!)g)TgdLdL6u$YAou7Bvrez%F7-*S-uzyOd8^Aum8`*G z--O#V<(ETkGyORGY1*tQRyf=74q1}sY5X3XJEGs3HGyObeKxB^g+nCS^ zv@G(R_~2{MvRM?#tb7D2ML=e4?X8f<*w6*`;Si_&Z5=iUb&Uw4#8iTU?Nmkz(ZzuV z&}=93L5g>Gn)&29i2)FZI+)#$=AhVv6O%zpE^LQ#|?ziw6DPA68)#psI zEi-PGbzvTg9##dd0JdZ_<)Z7fJx^@UlGs>)v2265=~Z3!>U8b<49~8$2&Om`QHA^H zBfLmbyUEA{O+?UtRMfzgq<%(k&9CI~Y#`agtUe;jEn@FyU$cB`%`2LGrKi#8a=aNY z@X52AZ*hjbWH~G>uE9gqIG%vj;v<#qHcB?Q!kd%Yqz$!rAi;oD=NxjSwzw&74Z}4u zvzUH&N{pMkH)79|wvE;MV`M116trA|D%EBWyE5)y^uRbjVYnUI{ehA-8^vk8xy@j` zQ&aZz{C!_ztxM-W$a465*2;a(iT|6|9da@M#0`P8m{9h(dyMMlx5WyWNIhkGqGjNF zda;Q%jqz_JFC+xRTPzn%Z#5P7udvc?^1lou_!%@!Stu|h-+#)a7)f)s=hW$uq1U>c z)CXUEk?h(^L@$h3hI{qK*~imeKd5X5WO2=lN1Q7Hu@2S~IEjP{_~|I6ua7z9`^#*^ zJFodVN>g8sh!w?E3z5qxgP;r5h6$SQ7$N9|ie!%TLzzuKVm0dewCkxSAk-6*I*=4U z$RkpNMraR+yk|3Uo0OHh!tW1DnT=;IVHFzqW>| zGW|MOiwzx2wKVPHTCP_aBa1S@#j;n^s=%#BSEp8r=4Sl zo9(g`8J2B@*_e*CTVY1#kX?90ou!A#oe9wsn|^`3iwmkbpOgjV1Qa`IB#?Kh5AB?Y zZT%te$V#x>I6%Nog)eU~X)XNocH(jbSEW7INx-!xZC+f3ucXWB)Ez4`?>o;qA?iOc z+|LdG@8p*l0#Aqq#~ytO9Q$#Q*`o1xAe174nGzkYNcjt2F(fIfMk4=MjKg=+DN5Ak zn7qUY9(-;^qq=!Z33^FA!}3w4-#Il2;mNV-kr0L`f8?o%dGkPZtvBl7c?3n{U)Qzm zdTk^QxV9b6k;GM#HO2j)Vw;W3Iq7M-)9rP>>G};(-nv3QQ6tw}PMHZK^4qzV$$I8C zml{~Nh&Q*r7^xQxWfG6iGJJ%?b^!+ z$78;LzfdO&n7`?GiXebK7{w|smw)(H6e)Wbyo^^7gg z$x!{N#BjoE#-0143?T{cY+3=Z$*wWOiWk|9V2^TiPmyG;fme0psT1dS6s<#t6Qtt z#}<@kxt93MbMk=UC78fOV)(vLX7->4&6g(*>6GfB{-LIMwO%S0fL;F&{5o?5#<$)< zg(J-s*Z50s$$XC!mmI0Dftaq|HWPb3)@%$o8nN>KQT^x~I{4B8*pC5acwLf*Ul}=D z?`N3YXIEDCOv-p_6Q6>yfFUpMU9HWGgqy^4h6R9Oy!nsH|I0__^oTh$tK&&)(L~

psLZKx({vO@40@v%OzJL-I{-duttwNP<~*cMcV{j7S@rRwLk>r z-S4jt-Sc=r2mLkA3y=_AcEEpR+3)O-2TfwG`pL@5U71JiXUqkf+gsAzxim5*gcPUq zxr$`I_f_s2HnVUxD!3;Rp)w+6*Io+ip@?eC@_N&wmAI$bM(jihI##dSTb$UC*haSw zD!IDywt8p6RokK+eoX&B*2A`X0r$BJ(gJ&aB=z&$MO-{=t*T&>sTz%8$ZL zZ*4B8Fe)hL#T=j&v>&YLc`qMFop-c*<5MC38E6?L($N53@md%k-uW7y&;o>MpR9cd z{`4Js)ROe3;M7%d-HrU(mP6s?>?L))Z3I2ev#7#^dBL!F{wvC$Xo?S|dgyY!c*abX(eT1GpxA5AG8bV_`#whxxi$PR;5{AP0>6%sFcf zdKwCkt0RV5d!4a!xibGBmFqXzDPq{~6}KHxx70+(xk&ArDl;t9H^LeQQ-RMz)Awv3 zm*8B3Z^EYYDmy11S9YtdH$E-?k|lX|E~Z$QR1lp_@jRL2{E``Tu-YEw=c&QEtHY`^ ziOc>6A6ewOkPUJ#G=r*8#WS0{Ut2xC^eeZgGTPL}JW9+(scDsumjC*XYTk4FI*$G; zMj%PzX3dw0{4A4_P|LgQ-;((3pn08ef(RCx@rz}r;(kFb(lr4`r5Tj&4U?9nlC>oGV_}--Ucs* z%A5h8!a4DITOI+iGon{%svrdXawZ@U^?p!Y-i<1rRAw+wr<18fJ=!VUf(njvyW8il zXYH}hv-YRC;!r+et+XCqoYl2zFprK=8wRS}!!-cCh0$SO;sdj>DRFZ)o2@c;3nrJ`?47kMlT7B}Idp{uF;g@UnXqZt%AXkB0($2HY%D(8z zUtA{smHhsW^y4uL<`zAc2S%D#{>Nx=@M~JowmUICdK2gt8UuWi0+7I*Hq7cvq$jumnbg*AqoBl5QwsF zs;_mIpS_4SZ%-ap%iWOp`|kOu;f&vW5ml(_^RVC|&@WlFZcOo>R#8;EZ1l4}_%ir+ z7)n1QIW>bMP@pK;u%2wl+DLQ#8QiyUOH{gGS8(}&Pa-*TYC`?N*M=x-)ERT!6;UT6 zHiHnH)bE6p*{JST!|h9u7iRF4oQ?iueN1cpJzqWPKQoWo$^;%Hv;Omh`Iku2S=RI<8n6)YuNHA0Hli%EDb5I@xHM+%J za!~YgIc1F?xynDV3hYaZGN)%82#c-iX0uCI9}dCfuGv z&lCRvgNH4bb;&L?9eCAkhh;eP7xU?l9jVEutS&W!DZ9*HWcQzUrmmac<+F4GqyPpaCej28-z>&;M-A~KEA`=0cOixUK zwH@`HwPvi6>2NnyJpZvl>l&?~hR>0OXkyx8qzKVoR>{UaCBcDxCwtrUT&ziC57)<} zx)-nH_pgXaZr0ba`=#z4t5ni|Eu2GJLv>mWt0qb}zum6=^6jb!Ti1@%+6Yy9UZZMH z;8cKlGAk3p^R|pq$`z&bmQ~Hq8o~-8f*yQ-lfromTsoIosDw(X{#T_6n)k#aT(&;3 ztzh12Gu1@=YYuR(4zl<7zM0#ytMv&SaAET3fUcldA56##An9f?gWcL}&;@2!e)E8i zR@=NQb?Mz}LG~};)<>|!2&M_`^CYc~F{L5n`ttIj^0gN}b>2%Qg3DEyl#sm-8U%Ha z555f3IAebqA5);6?P>D&tZ$>ao0-`5W9W+RM;7>sSC>=@HAxS*l<)jmX-NdH8kcJm zI@=M@L~oR+@n%(J>RIj;D8!K~i0~el*^?xSdC;?T0dJmJ?kc5u%?9yoj7O~LyDnQ4 z$}urz-3FM6DVF%@2op*Sy&Y}6X}-55XTjc=(jjiR!%N?_!y=VOF4knjbLbHt`1oL1 zd960BtrxC~H=2k=$I0xyen5?>=`#fTR_+ijFt9l5!0Muf%_T^B?C`A6;~xa?*-si# zd`vX2CUCvC5W{e$Wu%u^9G!mxoM|>ASb*m^W;m?$`8==J&6bhnj)tW$DV%DLos`ws z&fy53)rPgY{`-o8n(NN=3Y-k5S-=0ghLkGdD^5Gx6_?$cD`FBrnP_jUtxt0}ewyi% zno?yD9w=8t1^FCJL+3j9Qcf+}BL=+c7egpx(_1xjhQi;_u5-7RvJs#laFKDmkDK0> z!g#o$QOCm#k3Ow${gB@NfHIcIX9(X2F<~=K@KxS^qH?zAmYgCMASw9|p}yV0YT^C7nKhrXuL@52Bdso70Rv z#2$YCjPg8n`~@Bz%F~HXIY+YVm{cMj*pP+o&N^B-wu*YoG}FyT>I$%07hqtC_Ucn4 zQ%Bd*F!k24El+Vq<#`8&bC)(|<`m=m$#2t7^so}WXYcLGOf^5!qkPvaj7XNGnJZsI z$N$@l!#qeR8g0Lr*FF{|i}aYFI;r@NO6rO>a&JhhsUg%kx!uNV8P4El-8+4pq`p`z zda}yG8JZrcH59D6bq`|_N7AT34Bq{FD95?j7As1zDi=&$ZyAhWs;AzrEmyk9x&8Pw z!{~GpNY8VLpb@2iS&qb0LbRy6&!1Zcd3{>nw_7qhTUptfnPFvl*`T0I3>Zlk+$E0= zaRLGwlQMIXUw_fn(z`vynE44k`Cv?=}pjwkx}e&O3`mx+r)^sLo41DGJo9PU%a zSro5X4(#MYx}GS4q4(b3{FQ4c)0{(Ip|5RK@@a*)>DNe?cH_|pcRzt{R0tFuG7<-0 zwVQx%&Y6*ECtA%b-K*9|m#h}_0^Ggxu^AFT9AM{kX}j@%R2OytJZL`#{ka&a+4OJ3 zIE~T$KPm?Rk^Onl8Rk9(9j2tf1UxV}+uIw=%e1O54AtV@BR6^tpAP8XW>IWt=~p(x z$VRWE@K8otWfL$~G8G5gw+O$rZX9&XdQ5wh0#oza_Br~g_751kVJ|}*x~drsi`L6q z-tD$D`}k}E@oKiRW?6*<0&HR6T5jTq8eNV?1W(@8{t zG)t&~eJibcWOuOm6twlj%MsuEdza3(0-lqS;)MnFKvV#q2)EC|^9kr2?`HgWBCVLi zoDKG2>f6Fa`n>|T^kHuWZ#ahzC(`{rO99>w6OjeV)?MvV=IW=5O3Wjp&il%{s^?@T z+f$nsQ*Sguqj-C^6Sw&f%oaFU+&3=?{Yjm7&vTv>yoUZ1lQx$wymgMV<5g%rQ|K=; z%J2iCZV&Z3kw&2GaZp0orhrYtb0636Uy6nzeWMGf0&AmJ{Y9dZrvVvoPEQZN`vSsn z;C>8M*t>J|q-oyya-Dket*(gY(}x8#iO#u`JLq$e5{z|Yj&@pDtY+MEOIp-*aR+0? z%|#>rTr`=%4VO0?g&E2N;!&+B$$jbWMSlVnd6nFa^4Y2|{qJ-(?gy>?b|nAQ5&ng% z?lxP7$MRoH5v9*It4qFQ^v40FqSf@YJ%Lb<+pbR3mM1BV?cOP@|Js^y>8E9xzZ{E! zp{OV%Dw=gP_jir(61PEmOfCtN!60^IWit8(cB(vOzOmal`y~nw?OkL&c^s)>Ms$S7 zLI*U!dT77`g4Wo(A`nIRrhWX-0*obCKMxQa?ewWv3(8J zi+zTuq1!rfDD`UlS;gppbjm2}hM5k7Sk7%a*em2Ye{_s6d+U_4aaddElfU<2xy}^M zJBum6z&|Sgg4KXO*C1qb2hWhw;6jQc-H{To{I<^K+lyS3a1kEJ#eQj*f%4rPK&+*% z-jpA5UiJ{a+o}>3sMFvarY=1vEhJpSV(buegxrhZ zh6pt9fz1U;nib@!P4wA%BsZcPE@)}&hnVXbdWKblc!NXdMOQ_ecQmGenraHo2D<*A z*)r6H+E|j%^=ZlZ3k55?X$YG2ZzP#8Zam1PHN~-T0lfmmMK4aniw%kc>bK6H`|#tG zZ&L32Wqe*pn_OZ5a2FF4nsCSNTB5a!`6EFu$;@r0JAW6r3>gyrrk5f$JOn6av4CO& z)wn_jP7fHmNB(M>y@C$cLA(E>zPXaiJVhYX+@rFI!q&jo&$oZ!Gl&ylzp3_lAiYt9 zq&(-M3VncpZDrYj(FMGrsRQvS<*szRy6nXq=mPxmQJI*f`SkcGS}( zy_H=~+1BV@@EFnZKr*_h?25CFvCdn8WBq72hCiwh$ z)Hdk0&F0lrlMZ4jfJYBou_e6@(`q7;apbU5-+z+ZRNJgEEemb%DNq?vc30-U(opzW zQSgJC?I0t4A&Km3MxP^|nBb?k>2TZi=rds7tIBktk?hvp$ts0`L=OpXUab(%qmQh` znr7J3?Uyj~+cywvP=)}nb&!#xf>D=&H`RuFALJ}2)d%BH2GswUjX5~y1}-=&;+@lv z3>ArJjZkjhqq!hRD6|NeZuUp82Y2*Yn4NPFIUPHsg1~;s5Eg3_h?F!8|7^L6Pb|aK z=(XSGg$m7b9ML=ly0g?KRI?K^&W?*w%fjzaHn4{Ybz0=GdACrO|i4PP~g zq2`Ddkg2J%px?=qW=iM#Sd9Ioz0e6WdVP$@a6e%r!;DS&Wax9n>fiae?(|%vPR-mP z?rNX^ox#B3xMIUaRK<2*6*Xw<5<>bO?j3c1W$tb*iL3|N$SKmbDX$!Oc5sQl)ZCBj ziS*rqEL2qZ3I5|u>|$Rw3>E%V$Pv1oa^nlKA?>?KiNH&3w|8Cr5jI5#!t|x4q*uI%&Qg z*?)j-h>&y<<{l4YEHy1AvSKf$hQ)V0_+7wR=0R_#D{GQE`ugf>>g|pu$P-=p+U#+*?WxnltfidP zJ1`pFrmhsm^Z|AKD&yUgqG~yiNur}wDdH0X7I@!zw2AM&t<7L13iei%-DqZ^N&%rI zH5tJ=QiMc#2%`-r3ZTgf{}ZQLlR~>s_t;-b2(|iGYvV^^HKjgNMyvBN%gZ%3{*GU_gZIN@b74Oz+z1O3+*IYo{1Ueg@w>9UbBi8(C^kjmMF$yRn0iGP&K^#5H0%1B)J_y0|j0O=ggw$(kU4 zX%DaKzKu{d8JhZveJ>hwRrFVy%!+ilDqFpCn=D3zU$VD?Md_K}Th8t4|e&eYl2I z+_UT@xmo&OJ?a0lvMWrbHV8375Xzi*)Neh-Z7SQ>TkpX&|JF73hZrnn=-`b-e#8Lm z92}!SH1TlswhAy~$=lj2&5qH@F#u!j{?v!&RVZM9gYf7GhMJ$!5iZ*tYp)ZP_JplQ zCMbF|4Z*Cvi*mQgIoi_TtcDlN5BOook6v0Y7DAmp|4~&O@M*q)pahW|Wvb9aUnC!e z-x}~c!-EU#ZNrn*_{42K_6bNLb`6ya_&E2s(LS~%ggE=ksy z@Ui;2VEXp86$g|9PnbRaLq~kp3s-x_J#L$UhgWld?diTcKA?q858pC66?ql?B4jbQ z{KS-4Yby63`FLp>1>xuVO9x)m)|+_7=#8?q7>jv(`}Ep-F(S&b=jOqd9&3XSPZt_; z6EuBRLxS@zk&c*RF>utPL^rwtoiufxtKDWLdm&+#EJMUg8cr8ESs2SA>F_>Xb;6N@pkMU5;qASg1 zAMrx&(%s*|V;fkzu}{b_^4KjNnndhZ8_Zg=ho%E2>9W^0vp~X0xn$MFWWKlW5WHx;2S*4I>gt7fj%Hl0tRB?@6uXk9|I zzZS0#d6h%<>pqMns_Y%)KZ$55`8yeW%V%ygTsjFxiT6ipY6&U?-!rc0D?N$dauj^0 zCYUIE7)s%q?qNNa%0LpOqe$+pKpke;s4w4Iw5C~ma6$ke49u%44T zA~-#=#5&r+Hc|jq!xUW0L)nY`TzT61rKRJ8<``4}ZxyFJBiOg&Q-`Q@Pp%Rzdij6k zN>B#y+a0m18{)Ezjmc0WeCXt7Gmf>%p}U&ofoqDZmvQf^bVniY})UAuRZ<4<^GOF z3Jr}Wp*lO_0*Rk!QH+o?M~;ymH>rR9P?$C+GA;6UHTC}6+aH})@y|yz{Ty!FTN{>I zn|LWmzaEZJJ9qFzls4i*hu0*iGvM4rcP!-Ry2Li-@vL2n_KTGpo9Djx2hqnfLfYoG ztASEFyZ~#XoC}kiYNso9d#4%R29cG?#wy#Lu5xQXO zv~^cKtmpZ3akx|+?y8+1VH7XexPf#W^>^9QWr*zmE)oF)li;A?$L5Z+KK7csxYjka zV#Z01kx_D}=-zJMA{Dr<2^p*Pe>|!#B#C1HYz+*~v6%cBd$e`C!ai)f)yju-Gc?`) zo+Bs^27`gN)*paPr+v+hO!r0DLHl4J2e#bX-v@T3yB^zGXA;6xjs|?LUTPcCVakE( zY1lP5rfYlbJeVEn9${$A`> zi39GZ*}5X4KEFN5PgaM$ZCl>AtB#2&!kaB0rDLV|lr$FI`bn~`F+xlB>h@|&Qjn<+li`;$FPa@AeNXT9`enhl9%1gc@mMboH(+hrUpOJ#pTWAk zeYT9N?SM$2At99rRJarnr~129%$$YOOWD(PS^gaSY}Kk|`%}ohR(1A1J`>uS#WQ1B zV&*(+@TZmteTEh;S)G#N!Gcw+T)RrQ8B#$EQ|IjGYpfBHE!K6-2Nh+LeXL5d)+{rs!+R?X?cacF#CUbJB@<-{O%Fe3BGm%1PZDh9H{y z;S0LlKcOlh$L{@#I^=BFe^hD~Vi#4b=(j*&>Th8Nq>^liCMYOF-7agpVSL#H(SLH8 zYHeXf8R%HY|1Gw;l+@aHM9te%6)hK4)eH-hM#!j;&m1}r@1&hzmOjokrikvQ8?-;1_aiU7 z7T92FqS>`*w$~2`!Wb81U8T&+|DT)*giX`T#|mxw+ra);Ov_qc40UHoD2fYD*rn?|*kQ1>Nm zyym!3$WOkw8RYfO2jkFqL&COx2k^~~lc&e!S=dA(jeC$flj=#yC%GoHsGthk#oX2! zP1M%bt$A?_Oz~T*1v4-i3JKt3@snOj`!M`B))QRUynk=vt9Q58wdq}E@ zI*gT22lr5V*RoMS1$eKBymRrAc<-!j{1c?NnZ6jV*6jk7IRV0%c;lQCUe!eaa38Zr z!u#KZe+sR1y-A`s&oq}4EN!`CEJhbS@Bq(j=dW6W_J#Nlb@AU%)2!3P5@`08Zt_bA zwe!~Hj(^ak{OgUlY<2s`rFu0K0pPlC?1AE~8p11EKZ@QUwY7nEM+`Ad9C5FpQE`Ah zg>_#Fz65xWz#cA-QPjR5>s}(YSe*@>$d}7vM-B5#<0I)_p+x`{=l(AE8UFwYSH}$t r#(&zMXa2>v5PgsA+q~=szh|6x-1Fg%d(VgS<#`raW5C#JXRo=}TyxIfZ_dTppR)zPWnE1j zO#l@Y72p}=2XHnG&;XpHqWb&!pN;xF_1|`ZhKBk)?FCxeznzZ$B0U`)105|b10w^& z#Y>b%OV7m2c!}xn?|;AK@9lr@qI@pV(bD}L@&B)#{Q$7gT|m%OQ&aH+&aqHYvrwIN z0e}Di)p^Qn|4#USHmY;fH0Li+ro?cOGNAS{W%|_A=O~jsPeVf)9YVPdpkXg+0JG(j1iQ6?j1(YaE+aVlS}BP@U7b-GO}{=3W`b^4ET{;{+D3?E!Py_5;YZN@u*n<>Hre)OY|+k zzy05L@UIU1s{{Y)z`r{1uMYgH1OMv4zdG=*4*aVF|LVa1T^*<}M8d8(U=C`_D<&oG z>iq00aD&uak~@Bji_^}Vb_8b}s7JoooKuLMiD)#}3pL&{Wt&Nt7B<4Z%qvsd z*v%rm4r0jN9g(xF0?#6&$^d781a3pZK4MXeD1ZYVEIFM4Ld}V{R+)_3*Wr*ez}wJH zGW69MV6X2C5COy_{DBR{<&aK5nEX$QhP3&K8z<&1unu+mGk}^L?05y%k$~6QR$uEG zA`{{Grvp4iVd2f6#b=fi??G=4HwOTCP;mU7Ms+R5?=W#{T0k@X{PU+X!RGxrwY)nM znV85;J>~TBn}mlVRbAa2jz3gI@||*TWkC?{%d%@4!9BQx; z+jqZlrRH||0iB_12VeHOgFaTf+3G!c^g}Qv>L15t9eVF4 zYC1JlSkM?x#V$dXjVgnL+a^MXcCw7vc)k3D2ULTGw?}`g=4fo2fMndJ^nA}-_-LAx z)P-Vx@b-uy7%Pg@YmxEr-R5v}??R1d3$~I6N9AD*P2W_JYPJ3G zrAMyrCF?Qyk!v4stU6OQ{|;rP&R!;Sl!mg7eMvF*?wMiw(j0QTT(HM&5eE9e!)$e; zfQ?}1vlxX1ZlI-tK8~*l40&5cy1nD}E+LKhXM_7-?m)a~R6F8N_>+Xk6X~^aUZ3di z)bvAEgQ2nEH!b~fN1LLVq(sC*J&XncM!fqD(a@_Ms{ygtX{`z>#eIF+^=F&uU}(`U z)~7L8-DTU_t|*H~86lFmch95abpVrBY?69`z8g%oZSAg)^=Y;amDT*nTVR z4myaVw@q6HmA>_SU;Y6Lbb1r^bdDWg8Cakj6Chx#kyR#nATiRap*@~Au`ST>re5p% zP{S$eK;jHQ11EhU!q9uu{fmWZrI7yajMSXqVkvzOpR$>L*4+CyIz<&}b5)|r@6?zf z@p}3m+wPN*OA0iH9^}>#UXi{pjmzT&+p`HqU!p}duKz=A9CZ?YP%@??5mxVA0^3Wg zP+nv?14K{FlN%A_FmM@u_6)#Z3nzt2kY6A5LIuwNi4*^I%{8^_NEPPy1VuD1z#H|7 z`|0a;>enZkyb)?@(0=;WpW)rb%NM0D12%O4CI{y`Z!ZN<(^9nI3~>G;J_b$5nQ}to ziqVWgr|JMaxQl!(e(H9Cy@wl4-_|MZMc{O^J*R!|lInRKtsVY>SSe2JqYH2NbI_n) zZP&;BS{-m-&j4YN94@HVgvuk#xNnxhmPK9b@C#>QFvn`DvX-M+LQwC;9)iU}5@VSOFz zQ_efN!5RMGRm2}BoEDaSd;a`#8jyqd2rF=JJfEFKqtr6r?PaIv*X~)|LQYzd+W6rM zZ^lVg9p!tfTIB#oq6Xz<0gcW86|gb$-ZMb0ZwD|u?NKApCHxGa!FCkTWq5MM4|uvl zdIYs5gc9|Bldp#U;!3SV^Xj6sM;%&f<0>k+T?J}h-y#~a6{vPdjbX+3L38JjD;`!G zZhrjp{wj`dh>TFe0%BTY5&kaG4#?%~bJwle4D>PMD1nw}_8V&cD%om_3;F@!C9;6h zlGg|>;}c*8z;Q4pMv z#zU4J*le2Wc1t1j_t08-W)*nGn2$~2xmZHwAwI5g{eJO=1Tl~y)NMRIVIeXHF=Y=L z`~goQ(`;dT3HKmo>XJS}d)JRoxMNYH^R=w=j(B5H`%K5##=Z*JN-lBh*I){!R$!KE zp@2L}ADGq##FrO|CibX;$F8>~n~W6J*C*y?I8T})F`qqQ3`7EcjiPopFlpPhn2e<^ zJ;AAE+_7F!K@0u<$%@|2htF=n+siLo)E6a)JQx8itG}8M-luS*$t8IJ+^nyuvnyac zt}FHC2j|utLbAhAp`;S`F$ie<$-()yx1RBeI{uaf@#-3SEeht0VAc7i+cZ{X*V+{T z_RQ6e54bgDpvm_7H2P#nizkVc+gcV)U;Hs`?XO9vxC1Ig0M7ufXMiX~ z9W2&E#dMJoW)`M566aMkOT{|OcCXRYT)+cx#OoT>S!#F*(%GXVo>>Am!y)yvUgj7v zt~B+$<>VIN0Z31@{#k9vkZBE1(32eP#T*5?J}}#~BNeHb2Iu0f^|MdWdhZVj+_DWH z6EcIQ`Q>0ds1-zy^|w7W8P_D_^(1L!x9e+-@L}MT13ErpDTfwcB6iiRej=~tQkX@K zKc)=e(a&}g_iZ!B6+5`}jP0KEks|zRjlhfObu_YdbcioF@T3>niE69z;Hj zmbXo0 z!`K%6Usm*PFCHfxmmWF=sF1Tt4-j2Runul4xN9n9uyv~7xQE_tyf3N1X0jx4TW8D2 zvXNjD+M~`(^ueT|*~5|Dh)A3|6VV#S_|cC>fJ5Uw(b;~W4fdwQM_NFKNxI)SJbey1Yf zRqARbIPw5>?Kt($!$j@Zi9hTc_&y1wxWQEIw zoWFd%5tz&UPsNO&6X`}^uRxJq?WyB`ZqWMA4YZW4VI;F53z`%#id2=Y4P?r5s+%>> zIVwG1F0yhZT{>3Jo)yLgO~4CUA_<1WdDbQp{P&Q~fH%j-iFS>cEu1Jm01gq>AM=x7 zpLPi_kLKMu1Jr1)59O}ahUhH9X~8j9y4gFWHU!-oS{%TOP4Yb9sEiF6b?QwcOUlp) z$vkm+H;KXrgM5qeQaJ^x2H9E;t1uRFXU-*$YxBRbg$2S%xfQdfYW6!@Cc7{P?gZw6 zYKS^IZX)wix@76hNyNQV^RXbUa1}j9;)hO5r(W{*3;Ih_G7N0ii8VWFxn{dX3&3lx za_>G6l}8Rw1Tsq|wGM6}pVr4781vdj{l+hnxuBjSF^|C5zLOhv>@LW~!DrwJAfL>Z z?I5<^$nxU?`bCebAsa3$r-PZ@#cHMJ8~J>t|EIXp8+|V>y|_7J&bx$olP!+T8td8X zRN|4dkL-?=jhngsi=%{zcFXGERx`-RR#Elo;u!#=mk@(stM$O&LUj%Tt?=(d(E4(` zUeB8!@_B1Vaxx2PDlGtYNt^h_Vx4@*mDAms?W9@iT;9v^_t)kS{OZqOlimjlYdeGm z+viBF?;|_Aa7*#Uu#IJ!s%j6;r{T!^TcD~Sc;^v5uSn+$sSs?j`@w!hQb1yiYpv%| z-^Znic8!_R3xu4{h}z>gqAWoR$}v$R@eY;YJi6fJXYbkWkTGG(tp`8O6p8|1dfIp} z5)p`(plj`XnQ)hYCz}kLPB}VN848O=o`VOX&z^Z3J~m&6?Nk!jHbK?r%d?IKvVmB; zOs}B+OF?fhuGV1&<1#It`Thb^0*Gnk6s%q1Evi1X4LB+jV%Rki*~7wEbsh4D<@!&+ zvgjqg7F5@>m3Kv;Nxg?5H~D3X_*}JmLRIj|1w=)BW62Cu+V(;J{pzf|dcBjkNi&*E zYR~iTpO3sp{U&;ad^7Yn`7ZvAf9SC11lW_!T_sXOp^RB7er7B#l&DJ(rzec2(R?Yz6>&(9@7e_sc#wsM(LM=ai;2vq&5&CFH`~CzT_Z8&q(accE`uMJjD&`@G4k+-W_$TrZ}jDw?Qwa z;6zq|&K)m~X$NZEx6J4s^+fnOto;nIFXI9_I$|eBielx-bSIcN=y`(8hI|~imI~^^ z%(=U|Cyt1M&`HR{weAcjQs|>IKwflg8q1E5(C=xvf?hi5S&ARSF$?|^z1d%x|96iD zqIfjF|MqC13Oo>ld2$;Hc-)BsVm{{J^~-UjdGW_*fccxnuQXmKQ1?%`)b(?+gURkKw&P;4Th4tcE z`|Q<4*RA&j_kJVMj0XBvE8cP^Bj?Q+Bs&&q*-)>)#EJ(!VZSnaPPh+X_nIl}tLp6% zlP4C>U*9G3!+s`zI|HC3Hbv1O{H?k+6=}lVO>;*?WdnC zykdEvZMEaVJDiSrmPnc3n-64dQj>0)Ch&YtlC)zkfQAf`DG(kFVatgQ7E`xMd7I$7SUS1N6YuvN5q!`EdWctbx@twtnO3Q4qM75rq zD$kra*4ACe>$^dh*2ydhdI{#WgoiJ+f*Za_a_Uf1_dl%-@;@pzib?r>TU6R_l3XK8kox&Yn~ti&9rvzl~;;!2R+L*|$Ro_gvbP zNpKZfCyxsakq1WS&Ll% z?N3aK6@JKt$=ypA^xXn!*u_24)s--*SDBw}>ht6MkW-87QD-EEU_XACn$^Ln(^~D0 z)08+^>sUUF3G1Z26&xaZvQAN?4@9>KH5I6Pj>-1~8#Vf{RqaEw#_A;QfLK4jJNc^u zoTm%~<#r&(JcpThZ$d@OsoW)14VZ@Hif9}!$X@KC_ux^ppwl!UzFh)1sG!OrOAE&x z5>x7(@?3lcOU{eK+pF~rwKd;I7c5?xawG%O_Z#Ofj07#1r==};F&J{nl{IHIWE=_^ z$SlFqHbVtm}F* zPRTV?DqquPhdeQPcE$5+l5_^%N}1a!!XeH=-k-KSsHU()BkM@P%fwXPebp@aTS%$J z2WhvCX9-Z$CWy$rhM5{d`@*i4i|5DImdUwu#AjxJFz^A(7Ik-SzqUnwaf^5&YAb=r zUA~az)1h_&{~lf1e=SOH;(TpNc{}jV$b%PBREKkjtU$}22Ay_+{5vZaA45Vg$~ZSH zZKuRkeP_UgSH~1^nxk)XvjHSV!Dhv4BZzxY?nHE)uxjY|HP~Jyh?gjXPy9v+$7>{_ z+$^N5Tc>Y>_UyTdqT3+a`gC-wSaW8>`_nzPC`QDrYniVT3YS>jxLfn|Q7|%|_30zP zKe2lucoZ-xjfc%fu)Qd_^!G7wnf}_;x%-#;iYS)Kyb$WS)<_vNp8|25Fe{EMva#>JQ>=Yh+W$_g@9RQqoyOW)FC6kg*Hy zZmLsVtp%zv)(vT#04s;q6P{Dz>BmrJf|L!+aq`a;KHTcSOxlC>*_+f(CsF|q*|heK z%9deLbh9QT_VF!o*GykkXZ(T1C0bkxN_}Sp2;8E&NO>l)VV#I-4^rRRS%OcABIPrH zr{c-^11f*Ao>^4b+o^#*dI#dCQWP3tfm`dJm3D(L{f6C@K6#CwJmYIEF%pv6l79+wv({}bz0!%MytJ^GC&ks zZBD%!DA-x$&VNZOf?N-TyJF)azF`Kp-=A-)D!gwI1s_2uz!qBO#3x#qDOi6Udm%4( z(<(c;SgT~CFHO#z*-fOl;mIH0>r1?kJjY6VX^lj1s$9N8BGH-tjfaB3D)UZdg%0FG z`NiiiU`PaFqDZ|8-udxPQTr(nc$c5)h&bv3tex+2xo@cwvJ}$a!_@3Yq0rhke-#jP zh}td9H88xzoKihO`Px*|_|VX2&rq1_zF8n2SfP44^}$$RzWNVxiU8@88sBu@kzbyz$VhFkwPN@zbKAbHNOY_F^;K9QvYg6iNB*pbQ)WPzf|_<-1U_y z!o9@M8J_`IQ^RJSy_{D#-U1CiI^L1TUgLGx0bwdzJhzrV{5}J?*(@4p5&Sf_aEby}tX)8%JMW;y1xdloP!-9}HQfGi}C-g4jkiFrMl8$#g6lgq!Me zV;^!1z`U@951?wU3&?iRxoia$=LxM)|Puyk=du!l6YyzamRnz_pIM@Z*Y)A!=H170g`u8GvrKd-aRJfm*wx z%A%HWm7E8P?P=U@?22P>w*Fs6nN(wsS6-ZiN_@3tMeV(=A9z;TAJ^_5@0EaFr$Yu2 zmtxxGd(r(dgx(4XD`HyWa_mVSEcVjTQUlv4R0q2ZtdZ}@aXTDYn*Z{e`>FlJ6 zwCEf5$n+Z{5`gj9a%JBhKwIhq_O-r4VCZX!3-ecPbqczLR@UD@XZLBMIZW>Uh45xe6j^IbNI{Z%lAKYgCsB&O@VxO6(L#!8QtRAJrMOK#CXqC9q3AVoB zNc3>+6#9HD6wQ-;Pqr1E@Y-KuLAYtWmUIPH>Eff~cp)WFT4w)?^lfib-Gzl%C$9yp z4P<+&1o3^^^fyqC$Nh9)N zNS`Fh_;LC6oY*;j(g&y^Q|q+?X%{#7bQ!?Hy%#@RG<@x2J)ZW9yq2738A2a-^tR9z zLCg-c#w4yGZED`EIIlT5uiYxB>RKjFw;BH+iWXo})4j0CkH^cQj|jEHt4X&XL#A)v zqq_m2Rtsx=w|bEyoBQhdUlUu#1O1Y$AI+p`<$r!fj}8}wej~-<-LTo6OKCqo9rbg{ zbH_XiynKmGm?vd3ujzGQI~xW;WOrxFB-5kNG|vjSO-HL#VI$XSDYQW^ zAYSj>1hKf`@=HH1OqvSiVcI7gn{9?0pKQrr}q32KCxj5-8^Xqti;%yRAs8oZ^r(jMihnbQ5 z*bFuI@>e#WU7w6guPA5Tq+jsiapB)t@-RNCka${UishTlz57TdXc<+CFY27)OGb-5 z80~EOG+xten46uF%sB&aE@9enq*-(I6DD0!)L+~^9fO#>YC*|os++Zh2y(!G8sY`xMslXB7fW`{ z1Wb3Pj}1YUK+Mx=P;F>mMpMwNHEg}zE{d^q#9yUC_^JDETjt*PDk{*~f!4rp)(?dK zh+szx0_wShKg*RIe%tX?7|1l3IffUj@yl4Pdm7`-f5r^_X)LS9ka5qoWH+pOqWLvF zP8>OAkl)$f6q2I4%=T?82GXvrugdwVoKqs3xGEZlJniaTO!)}>XljN@+U{vT;PHju zcVm38JN<^ZNq<;i(Y9kK$rFuJX2#z{MYSYM2mdf)seO)k1)EP7#cj1C)ozWns*PV$ zr(lyf7{;*lK5!>_!68;;UQ+#dQ9b)|u_NvFcsyeC(K;MSkD|o#mQ`!FR8dBU}5I0x5}cqxrcV zEwaMVN|iGh%`0qDOHiNs>ZY{N4Lts8@d>{Ar`%;X$056QsynP%v|^aD*ld$^Mw0v2}rQl)cQD6gvY{&v8128TV;cxw|VTQK4I_wMT08x%#`T zZkD?b&#!aUq+NMnFk7iRbQ5wolP`-t190r_-oTH#1PUtOsTGsy#3hA3N}*1oCu4q+ zL=T}uECv;(^rHkh#9XMJ7Oh0Zc%HR^bKktz_Hak#W~j)-%{EZYd{;?gFXv@a`j2OE zIZaHg`@yDrEYMfsKiWj_!G-fKF#DuUy0+mGD@)bx*0s-bpj4s9Uv@Ti6mB$|ss9L0 ziLpZ6QF%OZJRA}|Yc%a2AkY&FZ!TWq;Y)|>;xi~&cGyB{!jq=@g!`jQ^v~TdjNJ?R z-ecIO$cIL#|B{*-_sX8s?VUnSdCKv#dMiQ9(qxRu>1bG;OVna_{`9(+(xi<|O=YeT z=7B}hFUe8CWK^rSdNq+9#YfW;7piG_FR}kcACsCXB}5Avhfhf?Rpg)(CN~FP>9>xd z%*)xQSd;kJ#B;Y$q$mkl-~}>j!j3yVLpk_c#6edDhx-o(1LYl8N>PhA zkM2Cia3{Essn|X?fOI;+OuXvTse!KVX1)UXvOjAEzGj^X{&>Gm`vYCyg{`sh1Xl{X z9ky&ER!_V-YM0pT8WgC89Ad+MHhmZ8Vi4ljczR4XnxZh8Jhe_V!Rx=MZyG0Cs{8*~ zQVlXM-8?l3P~vv;O8^H(H7+IEoyU}R7XmrwMM_cxgjJQIPx}`lj^DmI>~}M?NI8nr z7>ZY@$@FTBkp^6R5xf&6 zq}S9I*VmxqSa^CyApE8aI{GdZJw5bu$X+HOncq!tI5*q40@H8(d8dGId$v!p}pW`&-@j4_Ann1HPHei4#zB^ z%bG%Ki)LJoAJ8Pprs!6SUZYk|2_?D$nZ%$hL72_UQ+K>7Mg?^*3@}msSZ$*e6@`6) zULR?bCfr#tfA&+01qvwfN)XXT1 zF{o44QOyCY@;Y>t|AYBv6p_6jxlg=HnS?t&+&o_Anzo;%+b;dmnpJmRTrv%j?)p8SllnBBzI_l%1>8v&8tmwmwzY3iZw6uE26$eA zC*HC}%Y(0C>aABPfTnsQf^gwOO##zJyvS$gWj-k{Y!^*;+%OFamB6^Vqd-vekpcdPKOr$fsgi zNU`wCpi>U`YBQM*OE8{+Y>^opFx7hC;CD;S<9up{v?M;`tsT%v!dp$sa6d}vK|>RT zbPy>JP4Cmo+x__T(er;IZu%=T1v_;ub@=}`E!UsJe!p2^ili3VZ213lpf;dRaTQBu zLIN?)84lK}SoKDQM)jy;p28+`p5XPc*9>_fp)-R;EmEq3UqugWO250v!$S4veVHDt zg-^6S#!5uIQo9XNLi&_7g`i0lY&vF^ucjM4TjF1PU+@CewJBgUb0{2b8rREWcX2Uy zlhShyFOH4JL-TWfAym{|5O|^a7&R_P385>8$=yC<{z=MIFK&r;AFCOv^sZTb^)ksy$D~$KwBFXq%j;TB(WbvWeZ~gsph#EC9WYnc zkHqlg4VxWA4@m=E_uj)nFeo)I0Ln03KJ5|B<45I5CtSGibz=VeIGYD`r%V=9gJ-v> zZJHS_=q@=)tyfTZ*?;+-!V+~}g>k1E)hJXJ4KCfG?P$9(Ce-juoIhJA;pfpG>Z}Id z$n(DA9g8!0N^EX?rLY04C*Fl1w>jE!Dl4dQ3ps8-UbMr=|UpZIUm4bd%vp>J@ zC=NPOnjt^ZNd<6+!Nb5F? zfw_zjg%%%SX8U_r-CFx?4X=6af0nliN5!X{QWIEU?ea44$zF4-35Ly#!&rk8WO^j> zv@-`Sj1Ll5*s?}Ti*N*qr(Uht+5J_Y-T0941kcoRl`h0_weFk`0TSMTH zGoG3L8b>E^d5z}#p+r^}-~F3y6pkc3Le?jbLdi-)b>vph03X^i(l>SF8P12MLx!an z2m}0QfD#d@5F{vwh@gy)dnQd8ETN)PBL}Oi`9i5@hNldzfLLOS2Mr8D41eF13FNU+ zdKG$+co*xik$od^*-{Yvt*@laG2UBrD`#^544{X%&(}lmxGZPv@K1vodR`+*ruH2g9=(FR%*R|0^~T3jxMxC539d0OrgKrVw#;VTyv&!N zaa&x?;)Gh5R+V7h2m0=`BjAStqucUwU4uev;p=TO2|t@E-=z!dgnEcr9JADh%junC zrXGpqj=x=|ud1krG>cPEkr_X!Z6s`91G~Z`N zh0ZX;9>=$_5WTNhZlWgIfLOLLL4JNSP8;?+lB8`Z^7EForM>r|%0$|RvFkaeSk1IW zRo8_~Lp7Y`*Ej>+LYz=ab9#!cu@bCSkVPLnA#sO&Wguv?2^Abtfa znUhn*{8uDx;^h?&N=-6cat6yX64cvK@vdE5`<+VlDfCZw2EA6zg_Gv9)aF+5L*+wd zw|i`wq`Psd2@k{3m;8~d)phoViex`?Q+&+Y?MV-J#lA46t7@N7Y(fXgPcsWs4InZ( z{t;Yat(LzH(bp)HVyfi9%(!{m^M+}TQ-;B$#k!+Y6E*NN;-E1n=s79c;8^V^#l>Hn zBjuABsxh_}6?;3y>W3MwXD!s^|h&;9W53-Y8I=6H$F>U3gBUN=h;#-AECpQn@wL+<<-Z)p>uG`zr zCj7f$kjoS)KEsgPb-F&J=~edI`iu9lktNn1(-kUJcTK+nmqTUucZ-%7UjnbrZYZ^q zx_c7KBuCf2*SY09*EH(C`}z4jgWo~eX6k6X1n%3oM;ns`1ZFs*@^p#LPxKpO_jc7* zni2iJHpc~x_^k8L9_H8^RtXcVJl5v5dqu^kHw_mn1bAak^ULxlcAP^SAk7^a9tqKP zwN4)2wiWn1XCya}R-GPi_%1W2V91}!=-eH4b_PQgi7b=^_9ILJdM^s9Q4XHFIey2t zIZv=o0(K|Hzd^vgC!`1t)5U>eff!TJyKGSx^+m?|71Qf73TGUeiW@HtiFaE3{;#|{L<;~Kg7 za$`|plwG8vMh~8+gR+53AVqc)ey}T$6f>8+mX=gBx3ag?x02=819FlIG~UwDr=6WQ z-2G(F2Z$zol81HH23wkWXmJ}D6|tTK?5hj>sxnV|Zp&eq4wEu_ygx)m?#*FrV~2Rp z>riq7;WA>v`Zag#>%z{JGNCP&vX_(b69gy+rdopdy5)d(yvuBAFCoVys^# z67Jp6g54h*>sS;b^criacv*aGWF5NEh^JrJT5-(j40_SGRa5#Xt8>Ff7EZj{zmNit zM$|jt!HcF7ie~^_0_PR+poP`87}Qa#L`Rb#>NKT0F7EU7vX^QcjWhvp5tBd$mHx?k)imJ=ZcgvpQx|Z^!K0fa)&36WUW36IP&(AmQz+lPjaN8s+6e(L9gPC zm-OYk?06|c>9bpnI3ASkiS6L@-0=e5pD~_>Yqmp3;A=HT($^Mh`L`{Bi?iz8e7r5G z%I|AP_n|)>A(`G+E-2C_+QG}LH8UP>S3W4O4>K|UnS(JMdS^_$LyA2rcjew{rgr&e z*^oN*GTJ6Riu+@D`0rSBJ#M(*QJ>_Wj+RWTsvM%C2hNnHo_GsEb^aHmN+nd8;_aiv z$$jNN@sxVeOG`1*pJDT+`2OYw{%IQ(FXgnv&(E->u_7`Zf*=B?Je_K>5&7pm zaUTcw= z8K3+xbEs;pT&%AD)ru~$NCjG^wFJ_JJaOora=u~Co_LA+GZc+rfGuRg>dpZ1ZMVsD zKt^9%J+pD4QCTsarHtFMx0f}~accKfwy!IpcI)3VLCp(csF}Y2(3tt1AAev&tWuHa zU-xNx<)2uF@mFR_DF!{|$ONZYFB0;nYOh}O{2LDVVtqT z$L}0JT!%he77w_5L7RgIHF5^vretbOD1~}(N~+^GU)xp6R<5rERx-5&lj@KZ!*r#= z$5k@>LYXU6OqJOu5;Jc#=!y4OO#%h>;d77~%<+3>2OVjHf@;#0(qf0MW~x}Wdlj&V z%AAiMAZjTq**EdbxkMP2aV_#l&&<}!?HqOfrVh0Xi%ouCkE9p5JYdlzEQsqsA`(_} zl(pA!F+@P8(qycx@1iIyRb9lslMf+x21sulvcl~}rOi9x7YkZfEB#MZvK&%nBDtNA zeW3<-k$~X%@e@#W1dqBHUj{T`{f>d(4_3Zj{tg3U|C1;m_q}$*?CWNCx9pdYCRTOlp(*+H z`Kq;!5RD!_>2bHiR$=`I{>=$WSb`Qq$Og|D!0)~S%~Cih8CEMyyq17QAoZVXVsr;F z(?9Bi#a=DNv%@ai4Gfq*xGP}qZ9+HTSiO}&sqYP=RI`n#y8x&AZ878u_Zyg@PUA~T zs^N}i`?+x^@zs^M zMr*L}0NqY^NQ;0$3Y$nKkUw`*M^o4_1Pew2YgEgRBgul0W^~~w|Lbn2qM@SfVaag;NQ;RB6^HJypVgM;UlplHr z)3?Q?9_J^)vN2JaqwJpZiiuV^)0q&?@pDbk%Xvsd+M#>FO|g-uA@k$jo7A%9C)U7{}7$5IPT17u!qy`sf9$tzi`Y zA(lTqJyh&^b*B49nwR94QYTlr0QpdhAoKTOU-F);6tSx&_Cm}%pWwOCEAg#@!*w}< z3V}rUe4uihTcPPmKmhV7$jn(yXq@Orp#mf~ymg9^9&Ri(CLcv++$2(Hup2|Ikc3VG zFR-)HM*H{BMeW9T#otmkZIZ8BUTalQ{U{;RV%2dM3S?Kho9BfFf}!u+S8Ay%c%;}~ zU0SW`vNwIDn{3S&&%|c0=A3=JpgzUF*F7hZ&=He6w{aRV5O+}zc2%~C<#Wh%Z@rIaF{4upgRCxR3WD*h0il=^iFc%~wW0$n=EPlOqu}Q2u7< zf6>J@L6j&%M%`YIB+*5w{i2l8OqU(SiSPA~$!sNYOd~%25%Jb}aQFUJnDk(18eV0s z+saZ{aJKvB-a+(XRiEzN;cPoc)4rhmzd6=jg$L2Ma?kw(JCpoU6d!9HuEi$S?}V?- zm{dzXr%x&51B^jMuT^^FQS_;09-KNeBNFIMJM~oebKPCVP4FfnUPJ9_JrYz~p-!n3 z<;2=}+d2=F9@H(qrZvvkhlG7Cc1$TXHcv9>_Zj#k7!go+c9f~S5fF&?YK z*uB^`D%Xs)xU!}3BfVB@!_k!lqhSfwE4v@ysi_f-Cm-(u{Pu^TLGjTc%^pP}n9pyJ*7SdDN0g3tVM+A>8qX54vve(5{cYkJSJw{iX7 zKl3p$dEZ~lSSHhue>dfHU9`;LwpM+h?I1NM)Lr79!1>dO+Up0eq3o#{EemLpFkCsy z;aGX><(u4E9>Rq%6qrYrfE{)JS6GC{4LWNX6m|N^Y{3K zMXfIQ)PoMx4=mmeeb8;?7S|H=tF~Bgd2&|e@c@m_T?PxDwN*tL--_@q^%~8BFU0$p z@t6&z$8x8!(-XnY8ZFUOm)L{C6qi)Yzv&RXw?NdtvSrTz08-SCfV<`0Sx0YoynGJ4 z>h9m59$0WnQrxnc?0#ij*!tEO{0o`$uGEk@Hm1LLve(aC>o(ypZzk8`g(+G>MzrXl z-ydd)68H!^K&MFGmID(3acY5QPKrZ|ph7hk4}7M?oW(fg#L&?Gk9iLzA%9F{jal65 znU9?^j5y2$Ua_2WHF_1MhktI8{@viWY3~uSbUWb^a&5D3u-Sd$In}aPcaRFV%ERwI zkZ<3bH=Ke`b5ikfvB>9=-;3iBQfB66qr*8oP!Cd;uLS;9J&{!t-Er_}r0g)~b>Z!! z;8l@~_*%gUwDXRD!?9c6&x=>eLPcQ z^X%RBKa~@^07`jcVRzLWK=!b9T|d1uz@G}-5H1PxS=~9Ov#tl(&1Lqxa*NC7`;1{( z-rHxA?F@9k6yRajxc=WjwER1jM@;j0^R_!<)5hOhwI9g6ztcrsWb-iJ6M|{UefBsf zsP)yNK_TY@O6|Yl9Z^LSrM>rR*D5ivs9JEjYUrRJTw85a;g$01x0J+3IJx`u4m%kA z&cjd4$*8hz%>58upZ4Gpi`x^xg3ke8FzPp-%$P4CI-5DJh7_!8@{v96>!+KL>pXm# z(xYrEplbCW6u|pd2|B!GcYFBgHusYdoz;URsrEW-?LA=UDXQ<=FBwP9xy@PTG8`o? z$^-MwKn1`Hdl4ul&14p9NZHDx*ovJ=5Vr3GM?H~SHZO#n0dSu&Lt*fxQ8#x)9q|#Q-K}O?vQ4i}AcC>W z0c-!GVifm$Ld5}_v)GMlS3kG;{wsdk3d@NVPNTbYjHEyg>vkhc~U>JGa8BidN zi@SbD)U%j0&VJrox|g#AfdS3(@fwrqi_zA>ku&AR8(UXz%~lzzJX;!7fCu0~1Qtjp z<#0|}+S{E!m@N5U=ui|zie|{>MbPVjYf3i$ zO$9M_$lSOg$@e#2T!H3F+PKqtKv8da9>g?hRle<;H+OiCi3%q(;QQ?+;I)Q~_;TO$ zT^3&pW6PL<+IQC(S3}oQ#W}*?v~SF{$m2fX8H+X$mya$EN*K1b{XU{!AopCg6*YA3 zTSl8r|6K9*Z2k`3Hl)XM`UQ8_R*t$Yf4U?h2oA@O@u}jkQf#hg->HU%8#sWcYCsEc z;g=+t&G&#dgk|##;Gpaxu$*v?NISZBy`^cqo79ti`NVUh#7uVS^Ch`Z*hDK5caIby zGZe$eT{~6*)h4rJ*-{$E%2G-GW$?>$&F@8Pl4$z^*-$f3p^VK{7nv+wx0_~<`>xmP zuvWktG-(wS&%yD=KkBaq9c8UnmfN~wEM|rJre9fHPI>l6=&rUx5A%f*i$8k4J}MbC zAQ^*X9R2uB?9k3eAwHSUX5xI!?HAhbtbemM%LTDCDxtQV!r!ThjLKmU;23H`J_Ow+ zTZy@rxCoreT%2gkwU{gnNBdr#zGynhV?>w!!ss4Q4zt?Uj|qE)HA^kpyVy{jV2f^z zZMv zQ43>ot)CJv@CbIHF`lZ7Thh!j{RrIsF37SA-i5ta(g5>aZFtC?tv+{)_a_a^3hzN? zFvqOf;7bX&!CZXdg)^>J1%o^9v>mSVT#V>lj?V%o+>L?gMoTt3A49BsUDh~UKP9}; zwcx(gc<<2RTQ#c*FziuyvfnkeZ@Dzb~clgu*D(0dcv z+Whte#n7M-FQsw!)q=uLnzhg7qd8Uc#H-w$K)QK-6wyx>@eZo*4m})wo3rk3NEB-F zc%~roqe5rvHF%vlE?dFy{QqF@&Euhd-@f4yQc1FB9Yu)Zn?2i5*^(q7#8k3}BwL0N zA?rjTlBuj&n(UHHc9DHwXNF4l8O!)EOZVCDzV7R}p5N>F(OPIbH(GCf5D0IFXz7KNoe!E$%ognP1W_P{AOQdsY*T0JEJ4>HW z_E;ShWJb1VQi|*@m|Dt#4{om??bAWJ4}OrUbNnpSZF=|o#nNNXA17DuE`L}`DqXf6 zyWGRf)c*Eq^HajvI~Jq0I!KD3D&#^$9-6udvOr$;$fe1`36hb2Kn?kei^z& z`8nLh2CH-pd>>81Yh5{Fa>Q{5;uIXA^lhya``p8?_bM z#|`;9RoOJu)tYNLTJ*R&1X*mmd5OQNR2Poug)PTtQ>%?-Mf^q)ar1X~9Z~9bhkMn; z(nps{cm4_I2=8)|V`KxY%zDh~dygISw_xv+_mpm58+fPc(=C?B!YOH$RLBw47G42) zaep|@svX4}W~9)gqg%b5pgU{+byI~W;0T}8@gd!V@YZupAc`~AQ>M708 zo+*kQz5hr#YN$H=EVCEE@(|7MJ~mlZ`>Mg>ooB7Opx)`h*aYcj$5v5AuPn@!TIF)+C4>O|{{ud2^P~EO&q%;TtKim9#n~(6+`x(M=hy7Ni=g9kO z)lv%;#rLzPigt9?I|yWRGf>pN;hAO5jDGJpb2in3Ze6Hes9Ji()YTRvomKb;c4Uom za^->HSx8Cx@IC>FqnQtMF_Xc#5Bv3}FujrJKte#U=3Haoxtdcl*dL+YiY?ttF)|_s zU(SwcL34PD1oT#i`4-F?HG|;i0Cwe({JPPtt9pBTfoeC7`-vIFl(S#(;@I7a@i1vOn47)-jz4ib{wY4T?M3a{`f#EmbK zJ4&JFj`cZvzhSNqr+6-&4w26*XVnzet6p`{jIX|#f6g&b2Q~Ow{=Vgtd+VR51r4;S z(4^8Sy9XB(=Q{-E5mlHr(?%}3^oUyb4G}ZkncInRu!nA4G4QnoO`Pjk6SKYk_5{n$ zDpHp@C2?Yqu0TV0X4rptnKe?}&F^tAE!S~n+;}<<{(?B4AEslwCQZxXLhMBQ|Ik&K z7mB8zoIBQF;6`zY?a{lL8PyQlsShs5?H+WbOf0CzN3Zri&KE@pn< z-Omm$>Eqtzkrk(4`!6s2;E)mzDSia+fJdfGID9B0H+mFUv;2s`=MV5=b;OaZlf*X~y%SHr+|<@?1wWJY22$n>>NnHLwt zYM$mkIy}FwIW2^{QvF*(3d;4GSK{B*dgx~#jIj8uAs4{G^Ns2vS#dEbD%~>(9xi87;qg6?Y=lD~k zZXUm7{BDgcYugZVpy@)<7tqa36oKTmNHdcO%%L=S&GDeQnGaRODIPSqvZQ zUZVJ{riocy>RQMTZ!>j%zgP*uTeVVWJR4aWn8Snir#LcB-^({y&vB7;Kevv&2Fa`Y z3N8}q-)k*X{SxVOji(tYOu=AQV@hvEnFpRvVgZ>Ia>rZ#;_e4)GdAH4ik$5-@&R3D zFbh;n{qqfz%JPeEaT+bZC$>LD&-ZL)PAqP~P+<6e^nX*-0gGsiu2)z1X_Q!km0g}- z%IuAA-TKR)u2VIsIGZ<#@tH-*huO&UEb zBee0QTuf`r%mS*gGYI~%uZh8e z9n-&Lz=gT0kBAO)1onKPK70c>3@MQ>tX6gX}z9mP_b|0L zIW_#AI~uF5^cBwjxG$zE6}|lQ4IM^yLi}(T4E_O^%?~^IFB<=Wftq`a>>ro_@W>C9 z82mp!+cliwbVv~m^XC~f?y+L_P+y%)x|Sb#uP(rxWepUT930FXLdM5F0w+pjd0E&WpPT;qijNI)YqW1>gM-MqNiY z7x!EIFKSQC{P-cGoSJcV<0a#n-m&yP`f&+3-0ZdmqQbc4pKv7WZ{eLECznhJ%0M8bSXs>o*>pzWN7NfoXBub!f_xAnXwp zK>5ypVXXyZt;N9rVfuev1BA1wH*Z+A#V@z+Gc;Mcx)huoA`=y#BuAzdUEo-QvwVjjB17x1K*RZuH%Cjzy+SWkWw*DsK{L zcQ#n@M3mS04#{UT$p>+HbRp8Q2HAoTug%N!PG&*XO4V|hmhQ(3VrD}>1=Q?K6;8Ow z$_9c=D%*coQTX2xboPJu{NEF7_`jX+0=q*XBiX=60DdZB5<3TvSVy?QpH7ZF-g+%{ zr9Rd^zPmjrYu`-o#@nQLef?+Ahm~hk9*^Qhx5bfE;t;~OjqysF@itr@-Q48`ZS|K> zl&sYyyr#}{GIm}Kb>Dtoi{ocuSXn!I*;kThQ~ZlOg#v2?1^raVgfx^(uQrJ!R(>vL zxqie{b{{?y4#~lr`=-{J2Xl$$9IZ$`ggYY6qX0@~R81jP@Ns)((>b2)ceP_*)`=P{ zjN%{Y3D8-=`iXC`WW)8?@PnX#T9Fr=T;3??ZeFyMVI}%%K2+wlw=zLkdfMq>9jE%j zH@BNL4<4{=aeQ6JFb^{F>HBl!$Q*H7XrbUhv&_=5t3!%x)Ya39*9;G9l46~n9Bq+| z3`R}@*esmvOO>IZpFFgPchVd}iHFQb-B(R-NWQ?iD}j7D*csn_h)K~8x%B>=a3`)4 z38|yo&;kf&_|)cH-^%=)?-8%DoMSPenAOk{uXEDnWafL%1IbIKfvd@)MON;vQXzsr z{IV$U}jcU|KB5UIbmJr5#(Hv7$?J_p<+v6>U5j915^u zyw%=D05+)=7&qTzBv*Et6#V!C+XMOY9N8}o=W&;>a%?>|M4u(pBARfJ9LHbOpi=^c zZ%%fVgK`q3jaz!=PAxQt-WJbTEcE#5Hv3qdz2E5Aka>Vk(Z|Q0G+{dP|c*DK)0@RoBoZ+7|@+q}9*r!#~;gR=bgU zZg^YaHRt8} zWMb~(p!8K~a?j86sW+!TI<8D))~diB>mve-aK8xD#{lz@jc&S%aPf_PWRb)?9{k*% zjE`39etB<2O#y!py5Hun`th5N`)6;{Ze#vLWvx@VStUr;q78Et6~u@F&8COiK_y#>{mile&ik&bWb5M+Em|76T$FBG-8fCu1sVj+UOO_(UW z3B`dm8q+0_W=m6H+<&&4py(w69JM|q6>6` z;0oEsB5Ly2r^4S}L;^qLL|dt^jfGo@JV_tWl;Mod=t@96VE20+#U+ZiFp_!jIqi%ml80{S#i$Hf=(zpvcU%C~BcP2mGkaNVTH-F!wd5+OZS<=;RqjgD=2-=Wf5h0F$F}iQ zUq*orC+#L3#CQV8>#a6=vgGy;uikvFkyiePHtDem8n2!si)}DBt1F8k!gC1dMhcH$UUXyZ`4o!ReiV zdPu(G+HDzlnDb~plrP>BF*s&gk-YKyZNpQGnsWmE80l}>@4tQ%8=E&DTjc<}eo$VW z>QTy6G(IfvIMZ;vfOP$EDth2Ekc^IoXhn480$qT{JS0G~g?g7;@LY8CJY2*pJ6Dqtn{6a1AaWlmEFdB2UtcJp0DwZa0O@c7 zDj0JEGY7EblZdH=V!IBvz_afi&on~0)}SWD=&17{KLveG=6kq)pf`1n7t zqXKD)4JxFBsaI!us?;XpQfK_99`=YZRaus^d%u$ZDx`2z4mo(U{Z5!*oV@>X^{n#8qEy_Ajqo+KcHlXP@$yH*r!a<>tPpK=6;`9 zF-W<)YhC{9wxb!_X{-BOQEkBs^rJ1L%V-V+F%`RSt16*OgDGYn0G;zMGCfJ1PqjEf z3_sSsA4UfWzsMqVQ-zb3Z#-z&sx&d*Hrf6Kw^CcT7G4tbV%aO zZ0Y=)=sshes~l@;0=pm-uzwYh<~-UHG08cHOVs0`yOXznl@!o5%B_lm6C5MWxcp5& z-C zY)Yr}$2HGX3OBWT0sJ0QN}h~h99h$=Tt$R+GFp5=smT#1| z+V0&bi}a0ic0Nx3lrHI1`RKdo%$Eh?4>Cd1;CXCIm}nc5Z$xsmL1T{7rC|t1vbnML z)+hOGgr*+Rb`!ud*gpe$%(x@SHb@qT!+3zsS{kfwv7azAYj*Niu6X|^K4TGf7R$hS z*N{H@Jq@ZL5K@IGDcW&nC!AYP*ol6eB2gVcd@Ulb=XI^h(Y#ukZr&KDNNf4~RMVW3(fmiv zx*{9k%!6rQaViW@9#)7Y+pb!)IUKr5$y$LsNWmww)mCQ4^o4rH$dOTcIRCjiY3fbS)JXIZp-B&KsmrYgJHY@l-zo8;yqTcklKEQ43J7P@x|u6IPHUMB=C(4<)DEWsd8C z1s7qFr~TQ@-#e74gxCwxFO%<%Uyn_CsvUE_=I2%JfwF%&MM*(LOG5wH-MDdnM?-^8 z(}WnuU-|hgQNBp}dHy znS;(tkIOF4YWVkw37lr-A?IPq!5z(~EYli!P#0ZpM{gNE%g>puaxJx3vHZGRj;8M$pty)qSPqunDxXmkmKwC2%N)d3DSiT1+Xfw%uE#N=gm-t}&cwZHxewi0GRtp>A4$Ki z^#G%N>X3GYv;5I=q6lLYFslz7{~$FuB)#5_JOJ7Bk@7Q=r6@-6la)7nKB(t-$M;h1 z9g=(=G}HS)YYW!`YSckdW(b*sc_?8{X)^RG*&H8Vo#byd8epOtvTb$fhG~(ot!M9@ zz$tz!9k1wXwyclSB*pP`0W5Wj@c|5Ru!wFH7y@YYnEJ8DbVPwCE+B8(36j|{aCnu= zsUVqppDp&btos6(Y=Pg0xQ&=}W^%>?tRgTU&?E~%#aHAp14u&x>qVUE32mvZM?IoG z+q>s?K~?$Qi68tiGRnBF3;Xo8Q3ODVP_Vz1B{30`p+(oG)}TXk}pz&khx~KH`qU=@`rzkIvN%rdlV<^^`SeA2WHg_FbCN zegM~}1nG=1jjf7klWnISPY4rJu6xFC1Nqf=~Y|Q z%FeZ^s9!J6Pc?iYJapDv5XN=E|LrCEaBXlPZzG5eOmV=C?C_2pKInEHU3;`aHh_fZ zv0L^<2fc6e7h7gen<-B`@$Q-FC$4NVfAaQr23;e7FaUR_zq;~+|5@+6` z7Za}G@MvNRb_2BHqA^>=fzM|ug^R!we_QOHUuc5UM5~k-XFaWV*6F+|Ov@k|0zDoD z4#uMNLzXSU!eR&aEbq3TKU)O@1Iw!_EmR{lNru9%t~zq8sTTV_>hYsI$-W&}PP9+3 zXj`K&r6AtL-RSLG&u5Ksf{1G{ z@bwwR==2=c2si(3Rb7?bP~27ia>D@r&$RllMl)cI_@ar=2|{>cwji8$+c&=A#_x^b z?iHI+d#-{v8E^G^UC;O{7+t=N$r3o9dS5h9lb>ZT?mxg$`z_*8=r@pPCv$SN>NR}T zp#fX#WbC`(-M;#{vsXjWZ%Asi%w6yJ5SXXb5tdJ|Te7jbYwYKx+`~tFcE8I+Q!@FB zA?hq3sR*;Rg9_}ANUz|mb@QrFtN`k)^W2P_k&1Nr*F;4UYl%;d$1%B|IdWj6nFU>h zo#v(iEMSO6yiyj@xX|((-E8^BP>N2to0+aq*N+a6*}yJ89A=?e-DFl%oSPd(M@6p; zqs&PR#MSd}%X|Wk-N_7s*~)$AmW0vFgL~fG`N;pubDQulRuqvn|p1pJ-wqVb>@l8L?6TMBZYF>4y^L5nz|R zxK_(Z)lNm1VE;Z_|=rv2<||&XLRarHMVnGxpDY zLP9D#Ow0Uv4DA%~MfW@hy@~-_KO%nemjOIqat4^9GHt5>td4MSo>d$C{=i(+|A1>^ zZf_0?FC8%CXW4xVSa1v7b;^O+Ar9OfgF0MuACM}bXax?*xUD`WS;PraJb8m%q1T93 zLAlSic7rp`-`{8N51zRA)ZkdE02}cmq8|Nk8-=8(CV*6qV9bC8lj}#uHS$cvtGU<~ ztdN%$F7~^L&gvB<#p!8}-FAztVqp&Wm!$pruALlP;^dvV{RtH0*DORFUHd&4_Qd`4 z^gB?Vu71{&bKG*KpFB!ig~D3S|e!8{t5=QG3_JS<4x?@$GTBmaR2VPq~MNBovGvzY1|eI?G&J+DBcxF zMerA@y>ARt0J6P4e8$;2V+o9g6{pSE?=W0xGt$8*wXhWc=^@M$>tJa%G?Izk3hV`L+A zw28g~rhJU3)^L%bcS9e_3a6fgTss=~AeZ%&oX+8Wg>2!H%w|GWY%7SJ5}1e5ii0^i ztftY;Cjz7g=ni+2=en3(=eZh7_SLZ5H8Qn)I(_w&3%ZuVzYO{iIJ#fl=o4)Z(_ne@b?2*vCye>{cBan>J1Ljcb1)Z|bBDRlb2%>} zc^kBGUKJjJTFiuy_sutS^MnRWwFETpy(_;Bn$akfrU@Cl$|(~*PPN} zj|cYcPAjHiMLn>iha+T8M&^pG-Z}7N9}Mw+Jq}sD(2S0fP^j)l-KJb15uR*!k8=7J zdg)13mVeHE{e~^xPGJ0XcR%{*w!o~1NL%4k)z6{JoF$>lqMf5R6wmzNHB|nm)5S$v zIi*vQIJUV1ZWkXVstte0r@{OK%7NTPH^N7+{~U0w%vw72AzSEdPf^98U?=Q$H=bY_E6uL(>w@GnSl>Z;V@WxVtySp+UKJJj;>6S|Fo6^j z=IN{irD~(?-SVV5PWq6<#fmSC zila${ubFC`Xi>U*ZBz@V9`<=WL*2X3tH7IdC%mzI3LJIuP*1{Kmg=wgsC zu+?3lsky!u1c#>^;&;|U{OO0g5Dc5@VO@9l3;yY&u`p5}pGV`_0q^ZIQdi4IRG0H- zxJI&m15TRkiO_il@WftRoP;EqXGELUsQdV#HrViB{2j*D&M44Nyk!=1j8R*+D2imzCG&M`3i<6NUV#ey#G zNK^2V5FRc&j($2DFjD%lV4b^2H1mKkW9JJj9H`tTEPf-EAcme$#{Yp0eRz%>V4#7p zB>Rqm={uwyI-eW0i=vU-C@t>&P}iX;!=4{L5qM|Gx9VHiZLObXib5_8*$W@`pTmtS zLBLw2USs}&wPNV9K(Ax2aexGK}APC&DOaBq4tE>+u6g{5lW4$gf79 zM>hehkZeVWL6EJG@o=sHbd2}h{wnf^j1adlM#-%+qNs}YOoCRt+=HIO!4$RQRFW8A`l3rlPSSi+!I zrO6@Y!Lm6t3thQ?4OA(4Bd-r3Xti1}m;Px$2I{VK=W{gx8nwSE4;uI(-+>v}Ej!I+*1SHw7IOcg*A z4wk32;5q3^^W6J?*QafI+S=6E*Tma+#Ux27^KPE)?d*5fJhu3)uNz)l`}7aY1Dk}b zmU&Jdns3kj5S%a<6Mb65Fy?uPzlZ8*h12de!$VaZT{E(o2Y*&>D*DAU>rnyojvA!e zNHq+Pm#yXa?>+X!y)?%6SWaKIvsJ0*sMgoehnN?uK*;O`2!yhm`c5KSbA*w{ZW_-m zqO6a^`G-@@Y z_O3}?+w&8?hfR1@Z+^>`FJfJ^1>Vm=@KGQO`+T3GA+>%y^fLfY5w7_J;^H)5T#|0f z{(&7I0MV0IUf3<)dJCo?FM2~9=2wo=1`py5{7-)aX{Zfzy#eWj_rH(lO%2I@|JHs* zWe{(LCqA7-uAa;b?+Ta1P~UlbG)ihDC^}eZy>83fM7xG1cN#|BYg6ByTZjyD(=B&$ z0C(aRcSAUZkAIW%lhKmmbMfmrb^)dZ?P6DE-(uF-b8gi#aAgUyc(2 z4e1NP7lG^3KS4VU<|5f6AX{ws44Uj+*laPiPA#U`w&D07!kjwhLvZB^)}w0hSA~N~ zc%K{lyL-G2pEsVJI_JE}mw1>30JZVNVfaKWaV$y+WjME<1m{7iyjd)|TR*I3Zednk z^H{&{`}#fphl@GxtP5!9T=yJdWr3nVvTY7#?q7xY7`I=Bj3-v*g9 z*U37`;!e#s+$R~p2k_Jcw-e}wnm|mU`<^d2OA!QaQd(J6(TB-CBQVshPI-oKQ@#+Z z&CYRW;n+cu=bx}U+RTF;lmsepy(2xCy@7fxWVepKCqXLOoL{``u&p&_0cEoSECtWG zH;USa6M1>`#a@4(1#<0Z@ZT8syEIA=2;!-(GYxY=n^lMth+&1fbtan!fu5ibFG8?4 z>O!*AWU?Dw+KGi?6E%N^W0l*B@iuq1BgFex0jJECw2x?vQE(@YZy2l|gs>~xE8pLA zazq`1Lj3{)6wFm|h5H6S^V-GrNvE1Wgp;ozA=&k4@YyVMWi4F+d0<(D`Y8|XdDHXV z!VUMVh3_vVUA~(nRNXta>kk0*GyJP&s;;F+%j-Oo-z?L>x2p8$@Z^p|3wvNI--tDh-I?T6+rixB-df}j9hVT zl*oC_@?N}-5$}$g((ctv`fRxe9HZ7%AZt?JMc3$s032s&{Z;bAdVa$^i%Nw?L>>WFRNQe4LJ zqb}Cj*oLfGcK$Q@b4bSag3Q{bQ%IG)YmC>-<6s74$5-3|P{6t}iJCnyvQ-kYb;0>) z{H*r12xF0U7|#mp?aTY3FFiYOA_8#=P5gqaR06}#Wicz=`Cb5xiN8zOx2%zs{4Svg zO7&*(RmC51XzyzCk+Pi?L2rAm2*8v!iy)?B4*}NZ>FpcU(+Y3aivIB^=dpB<(+S91 zmD?pGuXiHFJb|wo`SVV>RE2n9*M0rOi11zyEZGP>fixN}TR#)dj?%5IuTFD~xjqzG z<&!6H%YXPNS0Eg;9h%B*0_yP!VIsi^%A7x?GZm(GQ6%t{*keVtk}-~N)d`Zsa<=rQ z1MzDHhrBjwnSW=YaIgk1w;=?#;1h)Tn%8sOJWxD!lO}3NIW+N2ri|_PrB`8>PW#?^ z52un7 z0s%Ua+^Q^_wJRjro$y{H*->dC%BEjXT%$EmgZdU_KxBCH1n3^`n|W67K)v?tbhF7} zmct(oh_@#JhLk|;w`pn+f^=r5C5!_}Mm+8|%b)Np56dfSS?Ku`IR9`^@J{cebj`0{ z;qgeaB{y_jqa6dCv}pH&C&DYQl+S$$ZKekOrmA=tzH%ZDCzEVW4eVF+9)FhKoF;#% zD+1L=RiX=8Q9m)N5N;rtMVw8_UEI5wqYkCz-dtpAej(>1`TM>1u5!))lj4eF4mcAT(W2;0r zI^ptN$xGMWwC@PZ=o+*QwnLhoe#Q z8QOlIN|x}DS|xQF`nX4kN4Vm}gDu*Gz&NzdpP{-=p>A+;TxY64nDQz%+Qz$`pVoEF z5cWpm2FEJb|A8t9I1KQ%6A*rY;PC)og%h8 z@m1*xLCITBF#Th{K|l)5JYYfLr2aiq-yw7xQ5`fDQ zVx8Cl8p~TezilD5LBp-!rp3XJ4bGwsJSECL*4}yZ^4wPyGRd`c(UJFnNY~JJ@e_nO zLu{KFoW0ErOt0-{HmOJKeE%ivb9Ymx$q$ay77e7-;&fmSvE%^Tgm4}x3e4O^3$@Hx zklLt4IqFFt_*%vUuiZS2xYfW`7m+?uhjUU%hHzyFVoHu4#8QZEQfWM7FlX*XTy6X! z6uoM*r+aDB+@UIizwoXm$NRFf z_L_8$%*tB(KWCG-cP)TQP(?QxH*nG(f_8BNe=D-{QvGGb}C-awTSq^?2r==a%HxDinHsHH{e_Fd%8nAwq> zRfXc@z&8r%I{Ug*(yxZg0eZhhH~BGv>(OL^c6-h7but6^h1`(ESEkVDYJl8YnaHoD zMVwUir%1V3QhY0`v1-ppqhl}Fj`BxQK}S0vx2H*b`j=XrNyp~n4!srZ5U9zvA%JEC{?yeIrL%t| zkJKc95sGDvSdSu47LmDm>tWt=4evJ64(@vR3820MyQl zp6^d5$JELUetsb4aaa_->UM5Ddj|(Rlc+QtQw$|C6wb};1kKBo@)ptKwoy*QRWtHO z2g52zdVMbZnC@?WC%&naKg&2_a^(ui3f(k_-)~EOzyQwl9dzun49T+i2pQktPj*x+RqH$Ss$}Ze!)3T$yB@0j?z7W+r~&o|ff4tQ z7$835E4aT>8JTz&r5OB?A9e9-pgPIw%Dg%-&@}} zk&;p4Qq|Y}7_+#C2d&D8YklLHs8AXr7%^Y)+!H_P#PUf*u>Ln*H7j5^sxwe?QQ+GF zjxy34?l<7WKYgb}5EJ0U1qlwF!xeC0%>zqR|;dD{r-^>wtJq;g1}@DBv60z zIEBD;Jw@lN*D%$H+NwjrsqdH%Q(YhU1njerE{pzBerw`nS@w@L+)@dN5b+Qjtq1F; zE4{J1Q1IJH>|@mBp|TT>tVA|3pmx&i{m0Lbtf{%EsK_j-Hca zneiaYPd^196F;5EW2?sGzV=6}A2M?3hGfOE_p|s6Uk^9jWS`FJnCFR3I}AsJ94iqA zX|pmdFA=Dki|9SBwg-5!+d4hhtOHK90&m?K$r}U4&%Y&wfC!1Y(ufG=BO&R0CGA=0 zcpWbioNI{>W!zYoJav8a%*wCe9%~`Drow?bnmrBM`TXaKJ{2`37A8UZ5=r z&*UFP(jVZ0y{Ng6+q`bznArgH1SDiLypHUyt*RVbb;}7{uI9OWx+Cf`FENlh(N{e8p?wJ*nBbixB3CCtpENdFeCq8u=Ny z7y5M^%VdgVE3bmSI7i!2zhFmFrtf*Y5Q}UMAf1l_O1?=3?}`vcHwB@IMF1U#Nj#L5 z@uWc|^ho*KNgl27nXdh(t6mCZ)-nz)(xI|)01hLGam(v=I$50xdki<;9xiPJS|5zi8M;f7Z;aN<6*2Y?p1A1^7_AL zhraYNj{qdX3k{ShNQpB<|G=)KkC3orHw^uyTPuzR&j;bifHW9TQ)CD3C-_=?G6)sS z9NDBg)>BmB+_~hkDvj>KV2s(&JFS=XpOv4B4^uiFzIpk+iKA|6JUf)3L3&`ow#ppj_~XB{Qx?R*R<{@YH>8h_rN#sBo4VXN!> zgZzO&2^)GeJ4XK?t`1H7g^3N)LjZBFkJ8T?>^1w=pdD7=`QyZX_gu5i+*FIJg6D3g zp5x4!z%?#sLIdPi#?Ur7woT#yMF2lpr6lBGtEEd1n-rdNCva54TBF7zfUNvsOBcpd zkQR6bmU%W6?b%Ux!}Ys3kL|A9&_X4V?|AJpO4wI0r1P1Q)m;*kTsD>xubPC;;*E2KVl zgUsFP%jIeG&?w_dU4xn9rJk3p6%HN;Qf{~jH(fJmlFEYZXfe@f;s!iHU1WWgo9x!~ z&$O@__u=0{dtZk2Q$OZ!w_0{gv0%n;rex21NbWnxaAnjXh&aSP^eudB?W2x(x?k>6 zR&#rMM@#*gY1{e_xstyVcx%_jbnm_>J`&Q)LT3lTTP}#kJo4Aqs`Si}rl@KNhKn_P zF?VhqG8#1Xe)F61kI6$QZ=P36E=Di<)%OgK$t@NJ3%W@>T-*@_?oQ-Q)S&^Kw`)a% z#Q{4fbGG7cX^Z<<_v29y$}&T&`hNuCn&FUa<1611bW>l6!1uF#B6Lg2q;V`gUrHFFIo%ZrY|j(lZFwB1wzhR*^3XuL(N6+pFr zVEV(qL23Z^i<^ogH)7F10A09{-{W>TEBq~`uw~KT$M(ld^N>3(Z0IV-xIeJqv8`R? zckC+Ayb9MmuvB~c**~!SM`19uc^Z8T@UmTtU(1*-Jeu-ov-zj=W{RVrzCq`$@^+SE?E|y;!b8lv zKp`&UcQ0*$@W0$5y7v*tT{MLLft^sV2EZAh|9;ot7Fd~FAinV8ofncSxrErWbw&<= z48ecRE)mg9xuEX!g}E2NsQ{W?T5)IHvj!y_tP9AgAm%|!vN5n-=EOO!L*=B#T#LE@@J}SrfiwYC2gLlD*%OBOt z4T`33?Ai(`_r#pxe2X~wIgu~9P+)Ha1R*<^?Ldi{SNqcCt1XxaI|TXO*oyZ1@JqoonO; zziWpppp*uge9)ui7OIoXps!Cg+bgS2jVh*M1I0&OgA^h<9!Gd>EmQ9SEqNb(*Y_#D z5*L9oC_!nDD5^OIwPXe6hR8$MYl)_+qXXZh*zLaFtmKP;2~cY35>|g;CkaFhBO*YF zAr}K?-#J_Cs&YposIngLJwqB2e65w>QvDMz%!bm9auP7rD1kELbu%7boZ{55YZ6@mR(M~j5 zywmiqG~Emk;2dDIK2w6OLcWQuENl_yYhGxd98N(b;HntPEHW66mi*LuMVmS2@|QLbIq1|2z1(DXRH4FFmyFVs54g`7I zj2~k~8ljsSkB0Zo$}lR|bo^)+A*MD`croZ(5#%$}it@-RqkVF1`1B z(&pNoe_#kl1aSrl$*^JnmQ0K@P#5OPH0?ocOqSGe^}`f&Xue&9d0AX%r;ClUiMHt} zs}zwe#-XB8+}q*O`MTOHg@lk7IuF{LQVn$?4-li4++IreEZzox2gI4-F{yjXBbnvb=4%3MNw&;T<6xGauk$*u=XyQoc|CtU_jBFXec$IFiQzbo z@9|we>-+QmNE!?myV7J`OM>%3f!O*kzWHaS<}AOR@gS;vM;Q{Tw8ep+fYC>MNcnd` zCF^sz%HBz|fm|Z@nEsE1r|v?Ll_x%ToVxFHh6A>($l?Wyz-c@T3itlNDorElS`6N( zwGLT{!GNbTYKzH4mM|0b*F{M#v~ycs*gajRNQ)OS$*6J!#exJWzDomBIK)g~=u);v zKL8TYHu|OZ@dS6)8L7XD2;N-mCp&j_$QtS%v<X%m6x!O2}esdZ zyA;!*=7w6Tu5TB%R**j^ zt$mu)*KO0P`q^t&9jB7>1RMVdz}ZZX94!qM0;0m`ILKUn+R=>OIDyq~NgHjMS`D5v zjrcXLuNrG=bGEY4!o2IccQ-ZZPr%Y>F<<_|k0C&dL@KJ&VSQq}D<8V?bJRL9#w2?^ zC_(&vc;x1rRLvRbT^CLqVF$?-J__|`-5g398!Yv?s@SJ>y6fyo-if;7dL`WuNOiuHx|0fpxAiZv zy5OR~UQ4ri$xnh#!N~%|*)>R2qNP;VdQ>xB#(><`W{N}xcOQ&$9k1%C{$|Ve0VbVf zh@XBT(K?Qb`l-oNxvJIL>g)V%8l_)Ve|9wO$z`<;Kg7{2&8cRSO2&Y#Ju+<_wgV3`w3vDt;vx=ccWN=~r>YKcRL*D{#BlFAp zK1I2fo`5UyZM~p2b1iy^MN@j+8mSI#7{R5?hiEZc{+kI`mshza7P5yt*{tRHthw@U z{ns~1NA-dZqSJ_p3>FXQg(~br;sZ%AVqENfam@}2wm zr3Y~^mir$V28e1B$j(7Wt=stFtq5VFCoxv_wbz4J%i?X85UOarcOm9=%q@&cd4tET zlcQqvZg868D5h1EQ5H`Y#oCA}n_Luxa~Z&@pfSM(KK7~3Q)4}BJr<&2IwyYa=wWbh zQsI=KK;RI!q3G~eI}Y(`NDjCSh<*HJ!{7vAJ=4K0scqeM7GFQZ*IR~A^?I?ssZsU2 zLSS{_?U=9fw*|PlI1=c?rHba2Uo^hnT=Df?DOnx8Zez4B=~PC)ogplE53CHO`vbtq za>iXsGt!y-|pCfv!arZr4X*Lye2zv>GZST0! zSwPFwRr135a4o6~PWLxs&fc`!Sj+d0tE;(=*Aj3ZK9ig_Fou#bxnOmIs?Fm00|Nmi z8e|>9u+j6m&p<5ghTrq+e$NGa!cO;vA=c06iGZ@={cxsYQvK3Rud;(hMMZC$VOQ8L zyc7hh5s}Gi2PdCXfElOfP06DBHJZ*gM96U!aX_&6`zxr^@7(#&JW4Q>|qDH(9 zhwLA8#cBZFD%ZqYYjzSN_xdoBP7kjf9#QyqQ*WUu#M%V5(>DgY2v)$68B+yFLNZp& z$jJ}4-I_dRvq*bENAt`_Qw`Rl2d94?{cz;9NUf;@Uh*9{5_w>5n??{bKxI^n#_5Ta z231&5L%L%0&W?|}Iyvi?pLsLRjV*uIp7zw@GcJzM^B?WV_2g-rV2N)3E!2{NRZVU~ zh-hflqs7&{Gn>*%q_Q8cCU)|~RkL5YWy+Tr{zPFH;s8Mt=-({XiOP`x9s_B{HOh2u zL4|l(4fXH~&F4KQZl_wpo>*|P-|194J+9-h_DGZsx(LcmTAoXgC@IK3>ZLic0@^vg zK95-|eM0}qxkt5i`!5j>i*GxuSg$10Ww5%;Xaj+4%P@9=Lj|^hqZ@|9_o#dyxSL3FBq+S zx7Sr`axugCX2quiRVmM3EH(3wHP6K{5wF6Ig3@YVOxxy&Pw##=WHR6a;%hKyb4B7H{cbolD-QIjnlXeoXcYt~xddze1?z;qKRZ#n*U4!y zb;7(B5!dwJR^n} zM0=Edta(4-&OB&NIq_^O5Thh#8@;}t#SK=7!-}f7I!iSr(c&PeK0SBlnd*dHS=ETl z_Xg^x@9QM5r*)GfuP#{J+-moo*A`(4UcQ(9c#)QZoam4~aVFx_T;iG6KD&2?oq`FR zIKuX&wVJBzSJWka1wfuT*!2pr{D$!xSuXt!>ddTyA% zF}O8CeZTeSP5NsMv-wN!Uvu9Pcrg0wk$;}k^paLhn)3vRMM@(-QvZ@8nd8oRlX zTpgF#EN~LnAkBcNIO^tJ)t5^DAURJi+!}kw~ZIYH<0i0Jb8q&GLp zqb^I`dghzw*oF4QaF$le3w7{aJM|@C)@)apBK|w5U;2M~>8OFyi$hq|H)sjjfbTEd z1CDF0WlkyFd~kgyHvn=$K+8u|flT{_Q6i%_3k6`rrVKOL-te)Oa>$ z#ea_9f36^@D#UZvHTf&C8UQt~OZ0ytY)`3d?0H^&+DFIyKMN?&-a)7a?I4#MWW5W{`81ClpwNX zD9ZMPdV$`qMLoOt{1~YESjYs_K2Bg8 zn*DtxE;L%-pCz(?g?i0u)Ggct3m`$GtaCMGPwvPZZn`#TXf}1w@)3+Ww82yzx_8fJ zk#gUj(#3Hlokxn(CDr1T?CaDJ3k`6D=wPZ^@)5!WfcK-MhJKkwB)2XGAE~V$yc--| zjCMFbNzA&t6`|cnRn5u9!NGO{`o zoTQSH34*Nl@#vT;m1EYDSeq}{gQFQz_a0sAV-}YjdN;@4xW}ows;j4nSEis4@AxVV z!H%05SsN4S7{{$9QX|QmAO!fyhxIFl+W=L^Jbt?Z^yX~>M$0qFfo6G_k zXq_~XMJs^TfTZGQO-kWftl1!I6eTCTQN2irnfMmjsAztq^JsuRDtW42>6B1TT|=+c z`e_yBL3?;KODZU^HB@O-b4t~y|JN$Iy-zf=iCOj9&t=J1_`aXyHyf!lvFq>1w?V<= z0`Yyu3T8Iv>wcOeC8u?54{N)(t@>+ylUMtlrHosF)&X*oTdyCdCA=uEGHQ)ugfO20 zuvW~t8AAqL7&N}tEUQ&TaSy6!uXsn)ioTnPG&`Ieno?VLbl%s~!%p|TCG7M;V$hM7 zf-P`J$C4HRy_g|Rtj%_HCG$c^+fl4${*MP(r4Wq!q2Z>eT#pdHQ%T>-V@;hKf?fwObI&P)zY; zMr8nrb>^jD-CKMVj;uiXif6@}v7B|l9y{RBPb!yKB~sdGzo4scgL-FGW04)xGJ~2& zHG<=F3tSyV87v9T5jWJ3MAkZOVZdb zgeVj$bC}+heQ@~707Kv7ZRfzmX|7MT>6gvSyaq%)wjB4acl2qn(d@)FFbrWF?F$u% zT`rmmn2Zip3Fvrl^QN0z_N?-%FUF0myBw6$L$?u#>sR_v{^$&9-o&^4MAyLO|6o)? z<=`le){rVH3WWEb8mNQYMBG*7GX{wq+0yFon`Tz#WAD;@nzyXvOKe4yOA}Lkxo$Z= z-gApCPrF7XB@h%DT$H|;M%=6=iXFluE0DcYyLi7hv3f`w3aUrXP93prXtoinIikFq z%r^ah6FH%tA*xgv(h{YiEg*o9g^T&WWaRQW}VscSEiXw(FR@%Pp8fPJbnx$q*miYpxI%G-a$*EFkG*4%}f4wj2Kk zw!7N(pS-i61hWAIa}bYkRmA2x0-WL26fxR5UNxj`sx=&yrLb<;GI}#Bd}ijce(}~f zpgMqqdleoIe_xQq+P8?~aJD!wx}y5CDq3fO_A1@tfTN*qQY%}Ae8RT}4yzdqaR(L; z2xbUkC}!cmFpwjPf(PX;yoCCu4dQ!T6=7__Ov-nH7Ng}m_I$n~JA#-4%yMgna15^H zB+(nk#W*q87y_Wy^3K#7$pS^C%vfeJ4WqM!_Hu7?IbT|r8yfgdH|&Ht*N^1YL}0JE zFdv5s<61QIT&cWm+Wg>3b_>d#rRo^0nrju+`p<`~Wx+s37ZCZZ0%(UB z4Orn3FM}@4EwnBg9AjwOW%gYf{i#4NIv>z|EN1iZa%tmf zk3(g2NpPtmxKuT?mk8*ma)^=)r`eoH!u|n0lUWj%&x5F}5WdmFmX4A2miNu+&7)VjIMRnp5?G}_)0*W+)^Ppbaeukykizgl5(v*;`z88B0RfgRC6f884`vYq{?{VUG z`}5y)TB8Kb8;tM@mYtczI?zk*i^B2WwdRK6&_^{kn<~@qrIb#kWb5fAoN(6AS6`3Y zEQ0M2{3(6J-6V}?88Tx*l_g1vVN2EeK7fT6XpTKHcW8VTlf)aK7&h>c+Q~B`AxUo1 zAsz%mZ;dp53Hj6M*iEFRgp`^t8AP!9S@E@lu10!3-Tj zy!T{??QbVlcC}kswq!1J#(+l!<=EDJNgbTyY|5{9)E_FnCI4{g+%+Si^9St4PnYGl;>lz8+F223FsTK-X3Z=qR+p}Tp&og1H z3x8lY5plR0>lNO)7uszIal%Q)P0G8b;;$f8=RBs`VNu27{PjbBtfTXEZ<%FYIB|cR zmU<+rUz;K1)7v!yI7n{bh;(%^PCjp53|s^Cw?&V)=FE9l+Nh59Pt-jm*v>Em7^ ze4#eEL9jgx$>CcVP}BmaxaFI(`u%4eXPc$qr-t3OgAGgOpeTJ3yERZS>o$=K9XE*WS0Mu!fLZYEb0a zXYQZcM}E{bR@$CEXr2+|7HGk7rV&6ez^p_(B6ecN)?6S#vX#+0(7=93#>`ONqTcd_ zloQOW)54wRJ*j%7&~H!pBmC(7Tx3k%*#DtjirD_)y}5c36i3yuc=F~dhJFCo zsz|Z!$e~09P0()*{PICeIOtqZ$#dd}T|gFFY}`U+>^rrPkJHjZmDh9$jsw8%C3s89 zTqkw3;|2Uuj%j?mwt-wDp}Hcq-BRO`s=_6;E?%?a`;N0m*1T2R-f>-j;#w}P;AS#O z<*wsiH7PY6sxl_Wvu1-m0ydOmT`b?8I=v;wx96M8KqY$6549jd15MSs*0@=tI#$Uo z=uiF@P}7_je}0KM#q?Tg@4OV*&#auO!#C2RV zan~q0FcI?`@C4Rui{o{Sy$^L}c>1#MkCeuU#%nf6ztuGz5D0?vU|k1sGw%?)3?zag zJF}HW^^l<1a%;7SPiFoCLm^Z;XQjITzG)s-WzX?3aP!S2`q+&d>?JDT26Uy*d!`*G{S&9Im9?zcfN*qt}qp&K2?(+H3| zIRr_=QKvQ^{*qC=Y-wf7t%nMD#q;z0x!I4mfr)v5-EW~PO=}puxtYd?;JcbTa7{} zqtG3x3DZYLE?rIv)<36fS7R4`?W?i8@SVC3F9khuetU7-h|Cx}&^Ib@#% z5Yi`LZ7(*L+O_-Q5km)1L-KX;;IW)?PlT#{s zLS7Ep>IAKmq)D{RCz(Q_hMz~uA|?;);|RbHE&ZGNp?IF#rriZHjN}dK)hc}wPpz_Lze;y8OI=rxtJ7_r1)Z=da=QktxU^W--U;s z`j7r%xYsRCvi}NZ@vs%Ny}`Ap;5?xj*SFIg1Nm~op`O)dw+JO&6N7QvIs|R7STL+} zIRspPjs!?E;Ks`J{=0P=``>!5087 z7^dpK%L*DH9D|;cD6xyCYG+ek|LFW_W2RL7)jT))z_(-kI4`~szJY|7!a{(;+a1p8AW{;sC#h>) ziOz$*3$X z_1e7$Zg>5h?{gc&FAfl&HXU%Ffzub6yZXJ~F#QwZcqutG``V=QkB$4~)#{sW`?*f{ zu~j4iBI&{W|5aV$|H0#i>*HFCfP%w;81YQLfiUHB236!jrUT^-es`hAeIJ%xS9$mbgLED{lPyts)~NrY zDD35JSTp{I8u2J`<~2CSHv<`ZsFaFZ^Wu2bExZ9w+XMlEis?_-Fl&Q0a##ADS10{mvI+)m|2UyiUiu0W<;M5K z3d(Awwzz+OWkRayQdqGP%t$&yi41rI(6nQaG2JHgV~$9bDv*v?k-0=;E2^ITNQTsO z3ja!a|NgVrVg`;LYluB1^C`3`l!W8_jUQL1DM0ucIDcs|nst)$E;nzH{2g^nM3189 z-9)+f?z`lvQ)25mjQs#0!Z8Xx0blze{<$>)?D^y~R%Y$a4Cf~~G}q1XxsxjOoq;mN zQIiD;cV7y=AxMK3(;rv)#vWzTBr^|Mbf$+ryctJ{a3YYl`I1(bi zWYU)Y-gKJa3TB~))iH{y0cKMTnJet>(-CVC-n%0jV#jz=xBJ`s+tt7H6j2sxUSw{^ zf#|a;*g(cOU&p(;04oNs_}rxKJneE}`WS`bFU~h`OhJ1D(%*N?NY0RreCL6~&qtxa zjBjzo%_N{CaGu181)v=&b5_!Fr+jqYlu?VJWJV})Wdl{}}<7?LsDuH46H~NIIDOQE@<67{p+joFPcR2nn); za4~L@gL=qWQf2i*e+h+l9{ z?*}hNCvL6&YPh<`#j5O5@PijcT(ba$CIWa zFxJNiS-QmfIp#N1RJXBTdL^;Xhmi~x;VjNw=SRFk&dWHjU3*w>GRj4&kgDwcx~Z%t$^j!Ich(6;qrrSrJ7MDfLo8sv9<0z!k&*)-r`Tc6a|DPfqj!g{No4Et~eMlhFJ#_nGJ{^b-4Lwv*0JAynp-v9|n?~|NU?H z#R|a^NN3h7ugzv!C*Gdw5{+BsAn<=V>kqqQ7w4pO<(ICbv@EyY=n8H7`_iK4vfwr0FL#HWY`sY1DcKz)x!fn^DBenRryISG;Q8Eozt5(po-+Tt$6=Go(nMYHtLB{~G-R z`wbr|-e$bq)#ynBcZvvO906bb8S=@UB=N?JTv5M*vtE~1CSF6w_ijB+PUV*0ZL#N_ zQ+0%xk)qhjFeC^L!N)(auH+Rk`L24G7?lrTbZgeZZe|?fM<~G53O|Sjy#nxcIFGgQ z2ls|8f`L_ewwaUgHTXBfFY@aKG#-&C1VQ*(-= zqC{vnPUSWm)|VsvdFB@8pU3nG2kCdQtkq{mnX&$0V((T#Ifn45$k4AuJ_p)W(B}a7 zVmxmi0am9jx+>kBB`jEpA2axr@NrC=-);tGJgZMI;bqk)e!Ly)tA@t53Och4Y6<1R=G&tTuj zdm4<60Gz)9S0@jL=fOtRo!(QCz!AI-AH&bdHqD@fu%?Z*qk7tYG|RbsuMFqJqVq82 zcn$Lp9Twvg7HvubuZYKRGavf4^iT=Fvpzhr^HKrO1QpFXhmR$$r~r|bmmP+9Z#^c$ z5)frR!TFMsbTDoBT@1rooEZAZf_3$Zk- zESk?6I>k1NF_avB<}k!8S$3k#zi=zM0HD4oP-+N&nL`z~Ub*-OX1p0oq@)cYpy3c4 zT>u-!%WQ$GTfjm`%YWgRdcs@3l}SOA5Wc z8!V0_5XmC+3%F4XA5O(B3=#4cWV^O>X&{ymF~RzbpjeS&>M;*Pk3zG*)`XYr-~93( zzn5c>?}rvJw108c{v_txTdepV^O+ls{;dc~fa(fz4lcHsc@`n`JNR^!Wo@Z&@ESh} z>D)YU%5c4>|54J(=M|uJngT_Mf@BFKftg%pfUR50#md|Ez=A4wQgzPTiKT6<#J#k( z0Ca#U9LUgC|6UpT8~DJX6IUGw zY7A*^%m{*pHFy!SM6zdOhy87Z&3l<{Py7l);_vxMm+mQl{S>yMc$<$$dUWgw!yeqz zk2tT2s{|CwI?=3zmx7CzW=0v(^K@y3KU9zlAX+QP9QjlMBZ4|1BBN<^a)br50LLsWr9 zJu@;7pB7Kik5T8l62fx_Pbz;XY^pz&mKr3+co{kkU&o>(i7O9AKmjb_K`|Z3>h3u# zIN)WyW`LpOhZj5j{lA0T|NVcwCbI!E-E0^d?{n4+gu!9CR6@?WzS; z8JfRgq4CftOCX0C-^f45f61g4N-lPD@z%B-x6n_HR=rhv!|Q9ISH5}AhldGKwnqC7 zHU-m!zy_WU{R~>Nt;IU91R>;HtfrudqcgLHu6uE?V9>KZMp!1-=?Qk~DQCp>w{qD^ zPz|0)oEb%JU8rg9d)M>__Lh}z2p|3Z_W-D@1N}zmST>9q0ah3#Y@mR)`N4ws&&^!d z?D*9#lTJ}|I^^Gc@%JehKm>xMCmzzeC(Y*{<+}3 zB#g~p<=Y+OA@N=N+?8r0KEpJZNR(jJRYkl{PL*P7&cg!<2K)N;P681ge9)Y+${P^u~Om&AJMux~-q37+-#i*?<)6KrT|D1e#%CJ12NDSfN#lr|G;PGmqa zd@tnLZeMu6U`hNA&w<*+Xbz#tyK?K?ax-y@#Z)+pFKfpTjzLricETUkr-kI1K~?u44#ciUix)sDYa8thaX{8#-w zER!%}7_n33S@FaO1jPjC?T*o3Xp{eje~FR-BX%avOy$JqP|hLPzCwex!@n@HL3?h* zO>a04kc0i-xO`YCmVRHSpoGo3&Jod9!!C&!{e;Z}_C#MINJd2%hW77&g2Y5%K79Dh zeSsf_6W0J!kHf$z7Kkkq%nF!~>oskNZJ%mq1K73c-oS{5p{9o77Va~}R( z>hlx1OkeIDJ6zjT`=qy9IjOy(p+WqHLaN2#FvJ^c2t>*mQgez7kxR7J3aDca6DuC7 zHllaZThC3VVnF$IMTF!#-|yAg$jAfNW{lK}j?H-VimZp>P)b`nX&oETVMsbiki{!G zW4(LP?eMTKRtZ3iFPBqg(>*p+0vuQG7Ekov`OzG%aH%z+{J4yINNe0Qp;YHE+|()Jxl zLZlFZK`jjm!jNRycyiw!D9!OpquK0o?ONqB9hvS_dc*k1760UG0!ADW+zJkB+6-AP zx(&l{0O5`sOK#ILI673K0gtWD7_W)*RoOG~rpC{EcLrcku-WBqar{kM4~HqL_fGz+ zgqK`dxc0q(k>*FkP9X({1B&6g(`~**4ma~`9H8BzX(tCmB^I~m`(iaE*YZbxe>!Kf zK36|sWvlS8IJqN$S*(k&{SAjT2q+N3R{EpMRy+k*WFht(|+~W%f<`I2I?S z@oKS8V1>)%Qs$G6j<(A3UX|X4o|QXJv7i};g+BT>oyFJ3eA$Qw2ZWOs)TRzo55Jw> zGm4cDEOGm>Zt)#yc^{)$ZgPr0HA&dP)+VtzY-~?0jFoi+Po7}$deTjT=H>uQ2{Sj< z6^Ff8j1~^UbV>C^RpzJrmMxC^CNgNW>i&bX7lk!sV9F~$J6=p>btBv{5bZTe<;w4ZcL&zM8v+zm4`PFJTV;T%Bt=MrZCx4a=xAnyYS z{M5s-b6yOIIcx6LZLf6N*t?QjC9aug^7lqLaYP>oZ{dhICAQTMXhsYfz&(1kHj~T{ za>Y>tm{R?Qa3LZ?h`qTKtQ7`cw#iV?SzyR7yLZ++pwJ z=Js;KF}pB%Xao3e@iw5H=W-vNL7Fp!j}$QW2C}#%L0#g{kOZ=7?K^RqLzasq7^Z5v%WYDJ@1`Rqw4b;flsfC!|N_`yTUXRfu+vnHs zKSrW7Th9!3)rz0#pLIOSR$VCf#1@eY+NpqZA8`k4!<}Sf;VeacG&SqQ+5!{P$%=4gH;5l6oQsW{EPUO-Zx-h$$ z86Q!)Jzw5oKRw&3o~J(dl$_JtPTRqc@H6()!2$tLoi<%{5x}M^k*H|O_+${uF-Jrr zaHQW#8gnO%pL#0Wrl(f!l-#^&=)%tZ@&BYahd{qE1VhNVi^V-l@$abCc2mOag9f`s z(>InKzO?$TxW|q5Ast`FmbdZhC9HHQ_a(sc5ceF^srYv7Y#~USJcYE*TG@`ucs1=h zVOfUiNk1ji|Gqgg3N{lV_yD%}7C%!%4(tR_XRIgms06C$a8nr3Bnd$gKP^{HSw9!E z(S^_FpC#esPJkLd*e=2}SmhfyFVJ5?11v%dR=HmTA>vc5gGf4FFrc}7;5a?Rb_i*8 z>~+kTq?5e8JxuNKZPN>e2v_1v1<0(UALdYO@e=}@bMS~zg)vCmk!q6P=ytGTc;$6t zgm-D_Y;^?EU(zPVw@J6eNI@r4QiuY~O%W56z<+~&F zuBQi&)6ZsU`%ZQGWh>1SA)O6c0<|j{F;h0JPq|l2_O)MbB?TX|tGd;~lDv`DL64jD5g(z|l zSy52Btp0J&Qcuagy+$x_`qx5*uqOvtJd7iNwo}9Cwg!rOQLA1hEZbV`BF1-*u`I25 zazg0NZ8h{Oyc#o(Z~(m%0dof~wi15_C-lu>(#SC=b6wv=f14+@Tgcs@gpOlHCY4bY^GTgEECu&p^yojrTob z!ztDEg>6T=e0fspoMi$DwF{& ztii@ycL1njF#si1Xk=CqZ<55Zmnjk*R@$xA)2C{6-V_#~b+9$QVjHO*UfEixDQylV zz6kT~efUV+%ojjyA<{VwUM!--G3antd;M!S5o&epTdkUW4T47%9Ykj|8xu-jl$&mi zFGsPx@UMCb-vu^NB+)Bb8tVqGtHXLlwr492S6Sk?p@_M*ZDm2N4|jXHrDI;kO6)08 z`)Ce2o836{G?i@|PC>Kyl12cZ&jXpRWCM*K4+F@*fc_wPEr-tZl@jQ|vU)k%xm zyWIWx+hZp7#KMw;Fvs-6WGCo;xR$I38C}HjoK81w2uBUJpln7nTnnvAWOSVdJ?iVK zYF~^y{bjNL$n_1$xNq0llZFLXoWNo7_5sGZ_Zk<}OXlieDxxTBj{^2p-5C7ks)j@s z7ryjUQ2K8EQB1q&=bQrHp1(0BzX|ORJq=NJG%`di4Ql9+BGeT+L->mWtMi$@#Tc}2 zfbt0=spDco@+NGT*@o6F4g>ZVU2qmhn-LSEvj$B7^e5`GEnCD=9Q?I)IKsCTjgF6* z&6F_hX~gU=e{#j_^b#8kJ~1!G079l^8D6QH_=@?SCBaaSZ?DfMhOIF63&swbL|Gpj z`+XFTS2bO6U;~o{a zdKBoz2i%O~)z8;4@9<=zjr*=yIz2v># zt9JfXhA{Y06*dFR*fTH!CvP0Z9`M(55@cM_8<4x>fTsEHRbE@3??#*}ewN%Dc}r9c z{|~QsaVrQ7iZO2d2sIeRo-B1*a=ukn?7f|oqvJ9@g>>x}9T~{mZ6P!i1!QX!$6mvO z#2O6fw$_2#i}eSOxyKNo=EVJ$)Jzr$DsCP(C@v!DdGB{~`Q#~@`oV73Tjh@zMA`bz z!5BIq>>?Mcfa?zg{B7LtK;Kj*VrM^tspr=N-*HSb;~GOKL>H>&XnYkSUZ{^N4zqfB zy!mR-p=!s~zvk0~&W+trNY_8X8M5jYNJj;&*?}14VXO^RD>8s4swdd|PV!tB#bAl^DYa6poRTav0 zr7P5hE%Yt${eqCDJS~b^6Hc5dM+q|Ifk|OkMd@o-va;`LcfOkOv3mGm^#C5DvYB_C zD=Fw=x=GvHAAD*6WcEdmY7MsGca7sp7a-ysXWHdJg?~vU`l`YGKq7568=+?x_`k?QX znOc1B3iOTO#$UdH)J&<#9Yn=uoN6P6BGIBVQivF@QgJtY-_($GtsqMyo!U@6%w`J0 zz%V`vXTT5}g8&v3OxgMqRd%rZR0fH^mWxNfsf0G+{B8fy=$VaMYQuvOEMp-qndb~1#2cKAn^ z|NqTS{_pmeqerPk7GE{s6^n|o zgP(+5_mtnh0*t+IxG;2w-ZCu@X!~tP&@$GRb*Z0!j$f>{yq|Y1DM454iqWQ+DFIk! zkrDEaz<5Gb3KtdO8467>0P#x>LeFKwm@M$yeScuze-M9O{sWsF4h0Co89~Oru0Jqs zlRvOK1+4GIo7;aEVgJD*@oysS|G%FRVLCu)B1$lj2AXBlJ)1+949H;~!4LW{fDamJ zN}T7*mIRUI44{ez;+R@5x$*eiLFM_{x=f|uAkoIu4=uKu(y10RY)j|`5I};I0#_>X zQJ>I&2b>{TQ&*kxD*dRQZ-G+h((N9jED+n58*K;1Wl%xqn#hi^7~$ZpgNpur^w$JR zHfUhba0HuvWK#=uy|KiCJayyNq75$LU}F9lb-Sh95bQCxpFmKT&4&jBqm zHw4;|;5<~fPpgLGAQx{U;{h>d=?FXB?++}R!0JK$x`P5X45fAu2mOXqHc4ZatR%}B z6iY%11aM{XWF250PU2?Qh~*hf>5OI~7(T6~cqk2NFX8BuT`G6)0xv4)fAt(KzAZOF?~RDM4fjAxreI zjwNvRe=n~dZuFi7&8NyUsbebS?4{3+^4^kmPvty0PeDI zt z#-`^t*IRtLN(sXOm~I)>DW{SwUsu(n-b-}+cJ!(wESpV}(E);5f`ZIaO_Z3AgY!AY z>4CcqS#=&2WxgDfn$MNLy*4$LJ8oz0+P=##`xo0ZuB9ABF;5&K0Y_T^*CGq* zkOV>L-=)xlL7)>K55tqn-b&p2?%?pRPsXj8CZF-rN95wfqh^dHPkbzu)OEZPu_t{2 zj>NTqy7PHb0;vYEy7h`73DtqJ3HwUwR6dDzAvbRf=h^*w(%IR??Q}}uiS*O*;c3!h zz7Mb866?dY_#$Q&){+qAJ4JfrwMTHCUwwPL>b1k%^bOYM?BBD#WL>KDN_9WR=#e)p z=j834N%dl1FPOyvzyPfhDx?aUQs#*(Oe)eW%ef(`A5%r-oA#~`$exQ--W2W>J~DJY zUGe3NNEi+}3|?%bET)U`{G-AoGFGsH)h zvKC?$Dm+4njTe(CmkyuG{ZOM5d zY_s4l^jumN{o8>hu;bm3g8*l1F5n83A2N{op@R-;m^)|{>>W~sA6KbFqpw$_+nq?b zUyc}`GrS{qrmjWTi+u~YAJImk!^B^{L>UAndWZ;p3I(x;7EQ!VDqKqM0I+q(uy&9t z-mBU`21=veuaXHg?+;K);}|(Pi5z>q*9}`6G0W<4?7ZW?aq;@e+7x*PNFB#CBW_x6 zu^?X%8TkO@K^0wy&Vy?S(9B6jP>lPeLkHcCQXZ9LAA{^OGH(^TdupelOI;_7Z;^vq zINDP~-l%nd9(c#rd4zq72o7R|Zs>?qz^{?@&`+fOCvw~zR4d#1+h032UV+hvtoU}p$_uc zsCUNmOfKxzx(Y8llw}XF*3nP_w7r& zo>I$K@-Xp3xBOQ0N+gg1I?Q5%O6XuH&zIz8QVdHZYS})*;&-wL6uMzCQ28V8szPd? zpR$piolp$Be_iXV3bjI3kJrmxag3)xYy5?nsRu}T(0iSwNZ|@+i7)CLgyK8wZP9_w zgi@~Y*|;%n;ry6JwEO%bcJrt-4r*UE4A6=isE{w3VDW%JJ)4R=lr z*K#*tXOI1Xsji0UNdV!`v)?%+o-VP9_>pJB8f5iu1iBkozo1VMgFSnnZ=e!!)O}p5 zv43C)kOQSWgEQ`h8GHk8cLq!{klj*54|q(DC_$cafnu-|L^B)Jc#B08d|jrMfAz`m zSLoL5sUDE?+w+z^*xg@CRBJ?uo{LIDQG9U>2?sLZNj7Uo?bQ>ZT3YlY+ucMelAWuD ze?2wm<{Rld*m@=C4oR(SQd-tXE{QWVFmxz~5~R3#80Tdz-UXgzO+iFqWzcWC(fNzx zQE4>#+iMA1_fvhB>%JDchrv|B0j)k$_aX|IV8RP5fhOil1C7vW93O+7YHUTp#po$f zMNs9G;x`#@Kyks8V_&epSflna`ct*epS%a)YUwRpV!!YVX>TAtd9lD3ia@glr)8l_ zl&o`<-}#$>{E1dr>_a2^A|Jk(PdeouuT+8g>!V}p1w&A{K#PcDUA~E!pY)7z4Y zwGhUS?~am%1VcVcZ}bNa=glUXFfN0V6vA}z_$2(b(cUqO~gskqF=5DG6etQAXssiMhv5NcyER8pG0T4Fx1iXfgw9!iC1md&<6DOl zB;VY4EK_ts%;H+3=k{Y@MLoknr|~QvLHb#&F_^MKsDT_E0eFoo$cmAyL(bl#Iw7Nv zG?8C9#7;Uc>u5g=XbD^pfsovj2i2klPUdN z8YuSs5(=_i?;KD}Y2BTEKhH5YK_orZ^zy~jc9>OYE3O5Ln3-QoAieKfWtIXxWd`(= zMI}Wh|K)HM|FvHF87pV<&(G6TK6D)*XOaH#^jihy|Dx_qz@hB_zTuI5-?POiq_UMg zWRz@4wAf=RJ0T)OP7%r)g(69a5|eCMGIk-6H6v#1S%w)j&X}3&`T5`XbKlQ>U)Osa z?{ggQ`yB6k9mkdHn#|03{?6a>`F_8j@8=_*{JKetluDI4xa;8MtT3cp?FQd2^|BcY zFb}&?&pVFW5WeMnQ2Oyx#aHD7cCmhM1pPV66(6gU+MGJ>GM#t-ci8)K4SyE+viShU zt@2KIm`8P00=u}To6Qk>_Dem5Nv@0?cbOwyz<7FtMN{W2M~j09FCcdPBM~~J$!s`h z_Fh@hR3`4N2;Z-ZyUZzs;!}IFgns<@1&AtIpJ*UtR_vOCCP)#wWN4(L zsK&ZfsG2Q_o-4PSELjyIVsoiQEF;4{gYi>0M>`kkt}Pjl0aqZ7@VeA z&6D-Rl#0U9Y*UlXUz4spTS(AMIV*eX*Vp&QuRK1DSy};-UEmQG!ypF2?h4^S$=Tsl zO4dw2`$}MF_?7Igx$5CUyMwtNDf;Kb?4JsE&tK5i*mgreS|INQiZ-_?_xH>&m}P3= z{Mr6EK&=Ux21>8nAFi6u~FI@>fza zv*;Z0WXBFENf3LdO2Ix|xUW+(gVmNCo?{LjZ=*wGjT8*UHj2-1RQWLwRM(7c{Fg_u zFCkcwl`5%Ao~8KGuzKO7v1#X=X-U?$JYLM)K#{NQ1^U3JVZZA)WIs*Toly{sHQ?D7 zeE|T{uXvETA{p}HhGsoFoLE~k0hps74)wFRLsYs@I)mbIUb<`ZqBt{qVa6`HOZVa?i-lIX(dXm2Rt66E@3)=&UXOhS&F(n)cBHCT z3Aw7Jb!Rc&ZeeQT0OMm*8EM8A4VsZdKc&=jt`mR4&}B$Qok1?54UAhhEhEBl*B1yq zL@Gy|AjxWK0RpRL7B9u^^Zb{Zgk?KrLfuA0t7hq!4^J=6Uk+DZFzV$kE^XcGH^bMM zLoF!Ob^nKWQJ3$w?Y2@JLxbV>PQKZ=wBY5DFeOclKEwLlzM|V&`ZfUn2K`h4R!Z6B zRM;J#k_Q_uqp9v#Cr}2rgypbZD&u(+3>;7PnX!!u#>YoW0GM&kpAF{#r+o9(y*ymL z*jUk^YUq22^|*=TTW$qk=6NH_tzFZhno=>7%(8AabhC_DY*whEUcvC` z9bd_-$`c5l!h2r35Li7EP?gDj)J@bwDRJ2M4Xgq(Mx^q}+xfT&ZwE9bcDT$(p>bdCt5fpf{L9~=rmg^AS(el-1ss)Bs5t`u?7L%HDX(pZ4g zloPvLO)}>n0h55(46P%f&(Xa~BG3OY1MTl*Ca}e+QUV z3^Z1#MF{J_=juLxYTy$SnfFNcxDpuq z{QU5&6e95?;9i#Pranw zuWkrOos`c0mplbF%Q6{YfOX=D8Fa$V@9q0RzA;~ zssAav!T%-e#M#)_laIgrIMPg-$av}xwa8!z%ji^e6$pSxyVG5Jr3RhtxX6RvR&W?k zMls~^b7@;k;8;hX#BEkMU7o%2k^6 z$zwx1w7y>z63H|j)gF`2x?_DY-4xKUS|Ol8vhO!$Jf0#3pC&D~2io<{@a}_A<^$Ct zA2W(-d=kU$W}fuATPrp^C1i+6#zLBRkyI&>T_?L0o7;G;w(wnT50jTaLc@D@&tAXp zIwjk)CMFRt1N^vo;bXQ`5s9lbFnWX(EEm>Og$cG9kFX%WmQ+@fnz;DnTg;hsN~|Z1MHmfUcK`U%a3##)9O2(%zunMW^Ag>ngXJ&k7Nv$A5Ezu$tP)8>|X(r zVJ%?+g_FG=iaY+X{wk}BfK3|-QdZm#~lhQ;<7C^<8LA!rd7CeeSe=fv?o{y$OkQ* z5XqDmr`rM4Wq6Sh2Bg$h;M;f_INur93~x=Wb9ia#e4KblyjC0P$8_QeSN>~xuH!vD z);SBDK+`xCK4eMsBCB=PqvDBJ4{V~t)t!!?cOHek552K@;X}1@bJ&xg=D?K8XXN=q z1aB+0+WQFObNDEs0TTycT%*kN#FEnd=Yb%{Q##Xx8{hRFH>XvcX%;Aa+ai77j-1_$jZiktTXnXUXDUt!4mp;|Cn#ny2Weh@xJI0Yv>p- z)QYDT?b7FxLE^-)2DAJwiem#x#@42`AjgZ_Km5Y->nl?=NXi&J=b$?d+m95tcH=MP zM4KO~kb*I$WG8AiI`+3h$%XMUqFV;Lq?lFOP6C}LQk9A(^Mbr!0{tMlHc>+%*|9zH z2sZ2Du<}A+pW??8E=lQUQiN?^mmqF2BOp+#`0q)rg1=|U00j(aVNaw|Pi<2!0eW>h z43=kXXAYQRCI65pgZPK7SofIry!U8NvA^b;z_vOa0QO|^@|OR7Wc7Ay9xmD6aT?Q( z*#}d-3YG9w8QgqL5#>gh1U1XIZW}XwsOc4$wn|h#vP!c*d@Qy3-sRSvr#rr*m{}(U zgU>(OZ2JIOb$vPi0VlmF+xB-PxnYwVi4(;2lb6V1pt!@P?zzY#w4$- zro!54^vhbir=cIcK;x1$*15&7O9RjR#9xTM@6gY;K$!m>jgDzX!Y5Mceu?yXOIn}p z`tSc^Y`p);`ujIl9%P3nd5%%@pR}PUsw=3Si|98af52A4k_`2T>8XLRN4XN%5ee`3 z8X!}!q`QmM0>kMp+S3V3nud1@+(ZB!cWdr0uQA9BS5#UIBYC!*vrpTg!uaNL>w|{sQEMwd6 zK;)Iy2>nxj5%Uu4hw0Ep0)8Ar@Ai*_wnlu>XRKHO)4S!YdT*}(o<=!&@Y}bk@ZR43 zcpUZkX)brO!)a9$Q6=wAy30l(=ANM2)Q%b-YmQ@Ok{Y4W#`^RDQ7Kx3-t$tMY+O|4 zjRzys+Ca0ltIwoR7MxfF<=O{NhDoOARGyIkIl#IJa^@@*ULtq1s3KZ~;ZaJh2;UOP zXH4XhiI`hW(ytJ5d?1VP`E=vw+7h`YmzV>Qq<}XW!A0N#y(0BAMyfmNCw=%V88TI0 z2DyLx$$FJGFUz0S#e#jGI$lfRhEMEDF5fL$Z^g&ybCHA>>V$5g`(H1m;krY_1-_*8 zs9xDGFR2FGE<)7(K6H-!T?AE__VyqN>hgaO7~u`ypWYQt*zi?oyBk@86R4P~sWiWR zn?+VS{zn5PP7G%X#`n|scwfmB5GLZv5)t8%NjNfG`K6t3?0z?Bz*#{z02JMHAL3o@ z-4(m~IeBK{$CAV6vafIUv-_*zJm+O@CnK$MzOb8J&63_oaVChn2OZ+K`;MbVRh>~xAG zLoJFCN8dR?Ew8ava~HkpAEbLbbXK3rR%*zsPxh_qP!a9=mS6B#K$XE*G6g)~wgoKV z37DjN?I^8+Ad@nm5jhOCXacZj308Uy@2c&z3V#G|yk#wXv&PUY#?V~7+Ml<|^R&Wk8Rq-RWmOif~3*mT$jVz_G>pO zEPhcnqs^!8ya>|RUAYHF$cmO8tUF%r2aK&_1^RChWm9lWC-1nVTS|6gtZ6w|7nKKK z$D!oDb%;=c!vS@}aYaCJto!v68Ap3NeBGRQfDqJ&xa$T8swsQu|=N&B+`FO%$QVVgq>ndOP*vu*DI+}D9KLK*awDd@cP)8%$8$Beloie(v zKVBGmY0R`{L}mEU$LjU00L6=40#7Of9`9F*<*u%wyZ`5yP5tMXA-ub%EiTigqf+p@ zT2{LzIQr!Lz1(eJfG04km7CMin=>)gz|<{L%ZDK1E$nd$G@lTss??5u6|4LH5o4UEPs8UtSq< zV8Bu=#q{cVLg~>Krn}Y}RJB~|P%n_{r& zeW>!pFJUvl=-=!5IXqZi2HbEQEao!D_;F65ycVYNo?lYj%@aRQI~EB4tYdsLy@LL3 z(yRg>PF)3FLW(dg!kV-Qszp(tZemm$ZarL8c1R)o=_B%I#`}*YJiqCT^&o-4JA9;n z)n!-NCz6+xN0cW~Bb>=XJv(uE`F5Un?p80PGP*?$$>ZE(51DafSRJSomQ@yd%s@@E zrDT9vdt|V@3I)sZuI~DT`v!Ofi#+hHZPRl zs0{6Kgi-OAEdGy^W`t9GWffc}4o;M)UB63xXL(>>CFVi>w;bo!ll=FZ{g0byZ!nl*O<)Q9>M z9;66a**||oKrqO@jIv_dvYp!oDO^}Dk#Z5fP3E;9@C{N)a#pMzRGzE0a54H`wLMfV z883B{TVv~zx<{qDPujQRr!iN-g@-^by~xTus!-PK9Ev+ag*R0XIa3OYJ1Hag|8Q_H zd9M@gASyMre9Mi4^X9!et2=8g2u6}2ohJv-r}40INM~5NlAqG_@WBB;CD(M-(%Yf~ zRd8><4A*PZx${ZMU)o9zVfFbPj5$D9+l&atEQ1{r><7&HG_c-1@Tohz(I`*LIM!L) z)F7u*Md3~=Z4{jc(@V!Vmzk74IXrGgd*t09_vMd7;D@$rbNo0gE0D_~Z?~m_zz9s_7PmpF zQ6cOPbWj?teVCUGZyyN3MHVqKXNuXrtAA{r5pQ*xHJi)}2O?L&JB$}e0OB$60j!6J z0JSid%U-rNDAI-0M>8ZHhV0KLw#3P@d2q>^<%n@uuKVNoMSOkC_cFsoVq@Vb#WPPM^i+Q)Z+nLMW)L_(PPDqoXGeJS?9SeVD@xwY21Q!}z%DH3GJTAtP${lt?lLl3i-d~S_u z;<;3O7wYV~^l z$Eof{VSb4bz1-Jo3Hh6Fw^0E;zMTG;e+h~oLi1zZ_M&- z+#WujFSwDP!Bmljpo5-@S9h&PAMJ)vJ!=|kK2?4~AMjE9ZXuT3?6dUTjP>P}=u@YG zdvlwE^=R#`7geVEW6h_=x+k3uW_GWVx@B%`9M-;hmR?BxDN|A1O%I zNT_%8zU!APq;;G{)VG|)*yMl}E5U;*g)N*k1k3#(>=8k!<#Y)kjpEqKxm=AD)0VTJ z=qLm!zBcVf%&`mQHLAgQ^ca>Zf~Iru?@H6VrzA-paVW1m<}Nb&(ulI-l|elPj{=9~ zG0F+%^uxnqesT%g*A>O3&)+Zfoe*xn1Xv#0sOe2KFd*Ae^B07s+4G@|o+y+hih#R! zY{eK1wltMBZUZ+n>i%K4_)4~#gy`JJdSP<29?{j!?G(oc8R>-G*PrsnyJkH=bn1z} zA1OkvZF6$87}MjN;bh7)Zv0?)fE&@n#Q(TV;^1Wo-vhO>Q5h=K4!f7RQpoxzz5U9Qkl|grAL|gj8BVV)(4yJ?v)prvT4uhCnF3;^z3JO z?_;?mjAonyFY4#Nb^i6ZJ0R@vK*MB3=FKB>8vzE$Zcv)$JXa6wP|C5D@9$&+ux&K^ zSvADYMO7$S-Zn&<2o2bTfMSB99c@4QnwSmVsQq-j?oxLnSJb)k zfHC2UU9-K(&xX3GRs=K|ZWSS!`5VbFAJu z|7IobX9sQ6SXAczRtExNs-{UmEswLUM`}UQVQ|CT=OIzK(dbQ~pt`Hc33IR7SidCO zz(q;bQ$AZU^aJs`U`zR4M+HEyg$Xt7{D)*q@Q=ZIAmeTH`ANvQAT%57VWI3>)u%k- zTb8ff@$7*Wi@AW=E1q}%i+D5d{5GA}5(If~Ly`1uzg*DK?Bt9+=<|E!pI!#->gK?b z4g5M5`k4murguhY*0vj^;H(IRdeGQ4Q&6$j^`^foVU+L6+mKm7|DHwB6TRsyKD!I# zMe4&H;Difc2XMl@pZzZ#!+%NXPJ%=$b|W3o*zoio|2WX8y5j3B;=BZg_%ryfEc!RF z&M!xR-KW^m8$qAhd(&Qel0fMHa~@A*2C=aUeWDF5F*c6;dnVu)<2L9!arq08^BP3_ zmx1||^)3w^g(A5isgokKp1j>`4F2C#756^)|J1*9#IL>EsK64wWBwS*$3d%ZqtMoC ztVQXjJe?bC@=sMmJ?NcI9>6L^VGGpjB_YD^{y)a?ZK=CQYw~uyM-=`fW0rd9yZ{nC z^^DL4tz$V}z!0+R680JHKW;-|yMWH+{}&={mI4MCgXD#N^_{~qul&=?QK<53a0>vVbWViCJ*a+)ste^yU+?CL6~}kF8xTk@5N>jr__T31(r-47A~gt}vlRRbyy!11s-fB$)BfBEM;&MIv;v z{9wq?Fj6@ zYe)FZslX2wDFBxi5J<&zh9xZZ17;i#%lT6UH>W}KouDH{3#`?){`k8u=m!Va9C!V8 zoihVesIppF{GC$dGnFO>{Ly{x`w0^h zvRwAZ6zH*9s)!|>t&NmQ=eV68YS})e11IFMQH`c;b$l@o*JK|N0Vr2FSQYA? zd8F6Bx3rW+^p+u)FZCy=u3UgT@H_gv91uYYjAkzfqE?<-vV^3?Vz+N=z=mkt6<{)I z855F6GCu-zxkX=(%)#C}o{yAAY1x0ule;FE;U~K1<9_|wJu&YwKVCVH0l_SMzqbv& z6_CA1=dXY^+z_Eysu&>6$tRNY=PQjdekF-_gt3$-a4>1NU2sVU5Ih@&^OtM# z&yPC9)%p4c?4RdqC{rkxBGp+u38 z9o%kJpdVqPRsYbi%)@V^_3jn^A*!8h%ss|)MCPX+qTN!}e>TBP+p7TxMgVLPolG?X zG-peZ2-|o8S8Y-5Z$yoE8~YZSxyuwvn^Y80?`3>Z zn=-9?8|e4KEO3Fsi}JMZHCqkBdNs-PyGBQVR{# z-5z@FlJVf1pH+A5s~fK#vjm3e&w@u<51_eCB)E>2lUVP2YBQZhWQnVAT3M){mM;g+ z>A07b*S7NeWZl(l!=wC_M=r-oJUaNPnRTX?Dhb|nbPQhn3vr8=X0p`-@58=Lg_Uf9 z;w4FivPdf?z!jJX&uaMKOurybQM?%&s;opkL-}sVq}G`w!XIXNKtbq&jGU;ViFXs_ z$-HGelH}#Fx>7jz8wU;ILhn8O zK_=7d#@X)j9PbkAexwJao0UU!3FI^{P1q6~zwo7Th}?tNN;sAmPcSCy`u zvb1PQGWQShQte}Il}bt}EM;8gQAivHQW2=wyN8PcJi^Fo)3n_A7iC-|ix~y&xi}+P zRn_OJ@7bhdzJQADzJ$d{Bo$4{>Ug%6)7o@|wEifwy`)aZrihzRtCZV7t?{kH{rIL9 zv-Gb}5tLk7$bjLFqJKkDc*(5vq@mKC{g?!B@PEuBo$cQT{{9A)c{J~ zJ?|^yR2fh75fT*IJuBc?`C12?CC2@@=gaY1_I8Z$7-mvs98JGc3xgGMeH;J{a_6^xL`i#Zx(8-f> zJ|ox8wa<9Y1{rX<^jfh!$rXis6=W+piw31Pu)X^RRAKo;S6Xl^g(V*)EX*Bhs$VGF8y%O4 zYgrBGlaB^L0VW=J*$od;b!RawE4UXA2bJFBmN`0I=qXSPXzG=2tGISP`j zhw(hv)o1pA9Xw`@>AFh?P5za6B>(;>&<*?YFi@JpG}^z9K#|Adl%S6(7%v}MJ~ut5 z$=(KM9+D<(e>h{(&6rid^!)+RuMp4Q#`Xx~x(cjmOScLj$B;0@1M%0H2C-jFdP&0Oa7@kZz?16`1C8Dsf0NyD|w|E zE39FEQ%_~qk{wRpVM}h$+;PVEhbaa`3)R_8KD7E;AnT&YdbgU)IJ5qMF4*QHpU`-5 z8qd`g0XUb|!}*nt(Vu?ued%ow6YDykfcRB{IW+`STw--P^ZmU*`Er^#GBTXVb-Zmx zK`Y(ECD!`lVRHckM4(yu1e(koAFJ+IW`3oqawup;^y`?G?PbN|6R+kf8~EN8)XDFQ zQoiU6c53caQanAOxDBEny@`eAL!|bT-LqzM_u*~6Ybc@Y2G2#{YW{hBkdhs1op72AF5V)B*ifFn^VkzGsiKH&%sEeG2zyiUxWL=8jfaQ zvgrZ=Rk9Q>Tq!pgYTNNrjI3rCtEEuneal(m);JE`>pqbSK$o`u)4e31(4PFx|8}QJ zBycHuiQLFq-%Z!>r}5yh@7~d^f$8~wt*9bg;Z6G4dp`6=0Fb}?qMQ4mh$(Ppr6C;v z%mv5ymZL36&EP*gpu&r(0^9UU=v{!s#w|#AdaiG2>&ibx9tOytMtceSPBirZ_?)9T zv{%?wx4#e*pf}_0t_+=r3*hZO`T&{&LJ(r32Kh4r^k+aVNzi5xdixl(=8c6)NvwYh zIsWrn{#`%jKd$ZnrT<4ZOYYsrs1{XR1GEpglKN>$_5JjxwoBEKu3+yhKySoZf}8o# z`yMQIqsSa9)>#Zyst)|_02~0|lO{-5S~wIzC#r$oRLfly;h9b<=)R}$)sc|#)e25p zol?k=d@?eA3bYBxHzb_AalGe>fbh&qYwnopgc_a`=x^YQsS;Yd{`P1K4l4JsvW=J+ z(|4a8%d8`;5?R3m+NX>Sv`;Z3A`Tke8v(8FgONXxTc>~fQ;|zZYL*Gr&}{b>^qqAC zOM(ia4JVM6{J;BqXyEtZ0#Q?uRavx9I2>d*y{$>S|NR3aiBL-wmSpTN0<|dUbHcrj za9V&XXgbeOdRsqqvs@1jDmA};Fz#@`Mpeg~rxoa@u80e&2B-@P^C(!^gZi440*<2! zfa(827+^~O>D5=fGfg1i<|WS3+2z5|meI2JNCYsLwiY%cTXGS5DUBy!O4-o+Ac}c} z8aTj9y>|qn^ztab$$rHI8%u;;epXCey`2OzycoYj1fEj zf&1vT{l#y8uH4P&n4fPSv&{;u8tPZeO82;7u6AM$v;$kpX4BVEUh5AoM-HGWZQD@{ z6H{Gpeuhr8>15s)o=bZ!Fj8LT8fAKg&+h0G*?np$)zlo&e!-nUt7+7jDKo!N1FHp; zIWBEqnlLJ{_u26n$D|!7$yQOFM)uy)Qr@y@5wi+Ju$~5}>Z{MZ6~e ztz#&3F~)M}(5Ku&bIt%S#)q$&rZ3;FXkP1V0zpko5JZSVR)ws#fbWqf;XBq7odx5k z0q0HbUI>LC z0C`TrwggbmlCbSZ3q}OYm5c|!StuLDX*;mnEO>UK8h(nYviSi4#Tk68p8dN}!S+${ z(ZmvGF9(4Ao$RL^dTwOnuAHxLPnA+2yR`QKVXtm+6opbsdot4?X?}Us zm${cM*t|rk3Uk;iOFQieQ7uF1g_zwo=0PMt>n#{mO?*n^LC+oiGSW*+_NDA(-09%= z=#{%cHMpB+sY5!3w>FMneDdIurR81PeJ~1{C_*xRKgu7V8UjQ;7OL22J);tH;;Pq} zgZzukEb&IIfFJE1(PmCAEQc$(F8*+ny&-W#xEYm|K#xKZGSE^~-#6fUi}I=tzxAIf zyDMqlQj7ccJl{!a_EA+Ly3%WXu6g{cr+C;y>t-nzm$Jwp;9aVoSIbJ{yXOefb~`4Z zwlPM`qYC!Z*>b@#rbU%lC(BU6$nHf2o_YNGO0)*M4ff4f18#NEBQn+5XJvIUspParjP(NCzT|GBw+9#VX@bzn;-OIvk8u*SFHGOHHY!9l?Vtb@S?MO9S1gbRj z(9Xu@1+yQHSNpj;UL}~ECwyc-z?c-Z&=dm5dNCm6_g&E)&qv>{k8N#Vjkc?%$-BNf z;O8`EK=>_&nEikmQh~vCL}v>juTvyxm_{?Q%PedBL6TB3Hhz&rOUXJ#{5(JPdRySM z%z5s{ltV@ zq(p7fSC?;Et5yvaFSu-{KSSntQZVRoEuFQc3Mq7sv^r8?JRxgJ8+Ejo-A9|tucz2|CYU#XNsdV#KyuWVKxzpj8@`&>w zt~*MV9Bwo$^%ft^7ah52v#@a|`yzAG+Yl!HwHp_P6V#0)GYp96aN8abxf!8<5?x$#32Svf{{a`g46{c%u&ncH?3cG+&k>bkEn3dtgP6qbR%6|75p10F)2 zFv=4x4g>ALn*4H~>^IfH8k4_QiR#U({m^ml>Wr*78^nh;KHeONkc{LD?gz1t zHGX__#Q>1drDm<*#8*pNmmDs9*XSr3EnpJi)939wkY@U;C-mS$Of8DA0ftC?;VoGH zrv`w0WU04cC}c{qPM*(vu~|`5o$a;u;Nu|=h8MC~(!2GQ zvZX(H>Y`tgNo+CU5%2T0p$T@GXI{h(~#>~yE5Po)G61Oq1_@D45(=^XmmqG2#VJ9QGaP=mLjI z$6GpxEk5M6;$yQ%^F7{qPq7VL8}e5!yYgPOr%tIdMECR!hEoAAc%Cv4_lQ$_1SCiX zA_d2{TZ9Apoug&KUH!N3y~X1)rp=x9_ZnBeyb~JqMeh51Ki3ptdA2gW!7m2O-^?)I zXbRT#1HW~{Clnr)sVL6d6JClOl~p~;Y;PAR%Eu=2*eLQB)Y5?^fjK~w+_i*S&H)-^ z__WG{8NklPKg`Tj?d2|_mF0HEpURCbdt!8?F7x=zX+)Y6`|j65EL`sYnYr=;tHG@aNkgu!L)g+DdL%@^;bCde7{7Zc691j*0M|`>Mk0ST2@aMMmCBSc zWfVKTtNh$16b+blsW$+M9>Cber|8K5k!G0lY}dN80>54 zFpQd38!0pDm?Nr3AHm3_$TFu+p0&2pK_K2?NtTumbr#!@6M558Wz&<6R^NlQsZMVU zYHe;;sXm;mPdO0H)!sd;|*NX5jr4Yj7ecgI;Sm;rL~@9RM1lXYR6D(O#Vyb6bbsdbC|0LN2S z;<@SVF#5dM$;?al%-H_HzDEad6w8Tn@ed@>%@lS4{EOa-Y-v&gKa;^FYf|k=i_K&E zPr$`QoBECiBEBDWw`*?8kGE|ZN8Zng0;}0%)Aw)C*#MajNI0zR+=098;FMR>EGq#% z&HDXj!%9Jmx!Jg>8GGe)vx{o6?sfwLIz(aEjLbT-zmDwEC0_fe(E=PS+YdE*i$4!* zNG##PbGl!wc}f0A|7y1Ev*QU4TN{9z9_NXo2rnEgJ9tGRoCk>!N^jw%n_DZVn0|Hy z2Q%HlsEN&`i-b1c))Y1rjP35cg_4cbgj&kbH70GSgdIxE^m-gb%KdfMcf#oU47Zl; z5Rq7I7kczq*Qr)>G3lr;y&E5{aUw*c!sMxZi8DNPq-{C}k;G!bYcs5#HPYyEy~sJo z9;Z@`vsCR;=Z)DkUN$`zdZ;M_TwW8HN-w3s2-bN!a8t@koIIP5q=o0S92dpaQHuUe zhWU*lZ6{*4%v+2B_$@W4r4D$+wgH~~b{q0QgemD5RhN|4M!bPuDcIh}Ku6+h51|{YB5A`Xbvgu;^0kVO{9S)kT($OcGq=1hy z3%Do}P&Ro9jyIMs^q#vJT?rc9)x1bvZlf9V%mzr!(PUuFPts4~kCI`7+d}KuD z>L95S^1FUM;v=~)uuSXf@LI8{^YS8kkeYMtD?=LPYNhYhIiig=E8?XBrIOCE0_LPM zoV9xf+yPG{Zv-z%av^I*MG`RTOkPaPIm$oG?zPD7G$nY6D*82^c5zpA&@<^ zNE&#OduE=%E=&gHGgnLAUm*Y^E$W*{J7q&I#<^9cb>|=7-mOtJPR{5$#vK{*kGN5V z$|p!?>n591z%@ksQg4y`+bJJFimBJ#Uy-Xl~v}_R>2erJ*1IU6M+@Noy zu$PW|S&5xji-enkSmaXNDXEQJ85jQvWoc26?9mCkgFk$br;T9sj@X&g+rV6YydP*w4Jis8&kfCk^6 zD8-)&W;k8S0sl44X?Z^gM~27li0p+UQ}+9qHoqpVnM_(E*5nLOukZd#iUt%+vhXk< z`dlg^K)Z7FA@tnfyHsEbiTV4BuXsmPV6`^-B^pqJFfIdG`U%_L*61D#RSfJ-M?pTK zP6n*e0b=RSyH{L!QqY&wGH5L6FT{>gDJ+IrAF>`Q7};vbF$)O`xu~v^^6A=XdNB13 z9fu})+Y${);j;^M#Z(p26-c0wQ)9T&6PJ;R@$y%+>3Ec)@SrYJQ{dfWcBh^@i!#~> zWLA~!B$~`EPi*-MaTW=_wuPPsQk+51-y($p*4F=>->qH)bsYR3GiY@kWFxfNpq=0v z%$>iTx|nh>`iUsgP1`OSBZWXhwW}mYN+eF>A4E#5XiiBxhE?l^ml7S8$47l!A3Am) zdVj<1%<3u+U&-YLX~rGfJ<$ZX7d*un)Vvy0B|=NW?)q~__@iB+mX^GY=U5yky%SBV zjcd|+#fxA6uWa#11i0V-HT6@1_@6S$|A&4u!i;-&u+s_d;BOKfm@dG1#ZtOR$VBK` zSQB?wsT?Mida%qH-(P&S-gvuaOl3iKbT+_CILzh?6JMwd6TN-lZ;nRjqu?fQqvrWg zKnwek2$)99_w3Oi=hdB8V-ocj!W{$46}PW*!&x>)Wcl{o{rogyR>y1ZmDXp$fdlkPZl>semPhu~Ht^%c2^)Vp#vrg;|yQkhwZt}U3F+M8VsPy_z zz`onnArAYSf2o?yru%F)(7EHm=PgA<0GA4%Jb? ?^|T^NOU`)4@09W3-mTURdl zgh?mobNUJebW8kc`Pf@`?zc1n(p5Dc_|gOZxz%ILrXPpBiY0u-?^|&Ilt!>mIv=bZ ztY)uLQfS^D&B~58==1eaAKM>yWa_-HO~8TJ(e zVl+}sB4qW_#O)VcW3&E5IdEUoc|^up3)XfAoBaG02lwg_?#bjWZsv~4M>+FB*Vs@F z2%0Gv#|waBnpf?g44gRd?2SmWMAv>ie}oK~wH*gNSPz=$ld$}tiT7_v826jW9G%h< zI6KI~(UCS5_?VpsJMOOJISm3jcc3}*9x+2?iBo=|19TI=2icWyFd%FEKD!OX>CnIq7p;6G~esz+imalyTb>ykM5 zdBn+m_NK05oWJ-dipR6WvVwZPC)f2JmGt>Tud7M#Suk}^F?Ekk(9y|oy>{y~>*Lcz z10)uRClpgf`Y8`-6_EckIzROqSqan5jqGr+9+@m!#reL;3obDC;q2SEmnk}W_VK>= zk*yn9+z*YM&Ey`S3JAw}GO! zj1YInTJm6>0?Jo7glu_cw9$loj91cACm7dT;wiGj^4!b+UEB8=HOC3rigfm4%XkH5 z+2hQXEAdoGB)F+K7!>)KCP}ufc|26Q=7P-Vf;oqdgw;2 z0}V%|4+IAm~=E< z`sn(TWzM%O?!WFVi1ieVh82u3lGNxX)HPi|7`{3j%)7xHHFkdD+dz{b1`_iHCE z`QB@<(W8nXYO?w_`h5KlBrIVAH?E{bvuThpRH0OGnQfSf-C4~X40znhB~YtZ>D7DL z8!U0?V~*Xf$&?SQJ_QPl192=sJVrl)vn;}xRu{>?X)Py>#|X@?VH%6bGH#XMZH8uX)?j%x81w{hHg`SvSmc z8m?5?t;wVxnMGGtCAFAWsv@jp*ipv-@mSt+B8Q$>ji(FXz>GO)JAx*Ki`>^JMNe!_ zLyzVLgPBOg?KU{5;okXXyj%8~q%qz8n*E8YFKPXp!r#vH6z2=WZJ>EDgerid^o|0; zFL;Wnu|WO3(X41`M;|-0!u~Da_nIB<5piv?DkMqM-o`lU3t#3@tlvWC>{f7uIaok* zu!zH635v-U7Ay*lF&kEQG7Fm6t@H_JJbgUj-ldJutEGrL-CPx4itFpsS}-*w2+Ul} ztOr%ncGoS!1h6iAUW&oV{sbiNvX;f+-02l=*7J=q=VGGL*vpk(+Djnz-|>H*$X8Mb z4o%13R4#bEIeV-XFv^x+N|FyhDe-im#0Se|<`?4)2Cj}@NV1+y6J@cv#3r49paH-W ztpTtPvRbQy7d|r9rm~NEo?KSj?$1U&GeRIu_9ZW*UN5dryxLQe6=3GWl;yk2sMD?6MstH6DM=qv4&T>V~9@yUhV(-Y#- zieEz!x6!zN_6UGfKqDYTEi!Kg+z<1hb(+rd??0Sq4Ei5lH9)>Y|E&iV`xt^XFH`DC zwyR*Yne9kZb+ePFtC9uVvn>|mW``xN)qec+>4fOZJMI!Lhffv7nJPEB$U@F{v<8+l zt$t2e2zPvLRZKINU@_wZMjP7FF^`Q9=6^F~6MizADhcuiTvVA7pDF3#)gWJybFt+r zDt-(6hW_Fk_Rl}RXAS=AhOBE@mvpdc3!iOOs2ldc5NbEgbIJE0> z-er>8)``4|qN?Nt3BK_g82?Y5jFMj;Vt&p}FC&IY{>jytczwC*(Z*Ay%h@~?H|!?x zUSpO%xBcwTX?F9)T`g68A$3|NzkHW2ku~c<+mB2@!`67vTZ!Ai9bp{CPS=41%=has zIND;J_9~W=1^++!&2p6)l@7E}1bTFKnCdtKlGWl*k8^si{l0g`acTCU+VLd4_g{t7 zDmZvl+qX4h_7;mX*#4O2!Agw`ouzYF>{@xEcp`+0ozB;W@OQsoHk0}lWG+cCzkN_v zrKb*87)^|o3`NkJ|0yUzk&36^q>|9UW|oSap3tr$GPc$`)ug_`8Rceiw~`7IGrOFcrlk%@mGEPV$I-=OX%WY z*Cq&dL>_^G%v~59@&=eU)O@Ei!woirzH9RdJrsYw(>`)@9CIcLMswKAudhy6^i{p^ z)9#A|5pN5Csq>OpBx>~q4B`-rEMcl9NIQa@AyW>(F1eB7P>U7h1kR>hr@-pGpPcht!w5^j0P3{2iG%sk4&E5{P@x6fVuj;)9zk;%$*cV%nAuf=Y0;S zkI3}5kb7l?hIT83uPdoC3Jg(wvF zMMS1LMYdyE{L5Xl{U1PF?%D1tvZ>QP`#o+1njKb5w>+saLe+*(O(SEn?{ySEh)|gB z=fF*w%Dt3CTEumpe8I)y_UMY5OVY~#CrKfOKO(!~V8kcUoPcVM!eKh{IJYe0nhwFS z^3ndTgBngYk{0SeE`$13%EFX-`e6@#ou{g|#4>84^PTF@+EkR%)4Kj%D%Aid~j+kxB( zsmI8`x}e&6yA?|jsREprZUf|0A0T*0S?mTzjx}-89zbrXA`zG1=rJTIG1Prn;@P!Y zwcaq<@(*o(r9(jXCufkq~~$H9jr!u>EAZ3DSr+yup!j@5uAklB$%-c`uesy$MS{|e?LWMG zry2@0Sg?6WSXGT~DKd+qJ*olu1t&?!>izGQ^V81f1!FDgQKI{BQbMT66l zj%4*yBER%g6>$Y$xhrWe1PTWCx`dOwI+j~2Hl;;_lzR@_nH%s5o)`aBX|IPz>;OktRxiO zDIRe*Ggt3%m8#wa?Vm}i?|Dv3y&|ObK`bM*e5i#PHO~W8_R|tH!L14ou#Q<^9kzz&hx8Oj zmt^k_C)StEuBJ))Hzeh)zqe?#P`JSPlrUW4VVlE5UcHLmM~q_JLxi zCl-jvp2aNgk3v<024he`2MQ8Yp;F4(59Ii^Q{0Vm#lNyxjfVBKeMjX6lwFLV%6ucc zNxeG$i~_{<&NDAcekVIpPQa%?_onR_o}Df}g?H{5`gzrII`OF7OJlj`li7#tpQW40 z%f1+!Po;`!0dmgf-r?vbfTn5w7+dW?1OifBBzvd|;eU|#-ce0{-MVNH5b1&xr3D2= zx*`Y&5>RO(Qlz)2fRs>FL_i^-NN)lP0t$*WX%QiG2^|3u>C%%>l%A--t0}*;_`bc* zz2oe2{@wSEasFU11|#H6)|zXsS)TbkpKap_V3G(zc6gNvpo)fTOBM9{u5Botu%4?A zaMkQtTt}wk$uy@db4)an2U}9s zkNISZS~Z^s)WP(0ED6k)0>>@;c%R`lWDM~e(AQ2VPn9!x%IVKHTeX5cHi5c(D!FyU zq@8PCesr_;&-bt6V{cwcpMNKRdM@VY%)^dvT<9Cr%0(m$3L^E!h?-DN#lD3y3>#$S z?bjY!%*}k#o2_t7yPD1`3F2kKLp2GnmK6a5+dnQ)qS z1$q=wi;Ms`S;SOo7lu#UY4B~ZRGY*7LG6b^xzif$MsFgIW-u#kih1f-v&LL4m40Ns z4_;!E#q=L$T*Q3Sx(+BPDxbqbB`^fYpEas?;}0nb^KF#vADGnS43cax1Zju@s#YPk z?mSs?UKTOgwGFL79xU1Yuqs@CU`P2kfpTTE=FBl>;D+;ss{AuVh0mJWIk$A*qEr+{ux$^{lT*WBq zn|RD!S#37z88A9#PZAiXoh5I#&5Mn{HAq%|@bTcp7ov5{u`6jkM-OO7X-@{|!_Wc1 zCLbaQI0OsuD!Ew_hso+)$}vxAZWQ%60aSJH!EXT)z;G^p&gKFDWW!h(`$)WmH)x1fk#H@QJ1A`v^#Y2^j)=Z;xsgC9VwPuhSr zAE$>Aby?Bu=utd55V=u#kMgSQNZ(CElkg z=(6BT{GIybacTFXx9u;N`N2Zhsb`Vh0P|wTuw0;OVy1LQFu#lBsm`wtwB7{{x(L&L zjm8IJh+=w{pU%nPkxBD$^zfq`o8-bN?i%mCFR3RkNlR9ke2^7?)xZ(Ul*@rm#rA?Z zlE(~-O7aY&-yO@y*|`Awg@jVz?fJj`(WgP_pqQ6oh#di50;pZo>=#my;W@EM+A_s! zF5=@Yh+qY^Ii5c)Q>s&<0at8pzBy^kcmEJ9^623+oEX6!I>7#=F8tKh(ZPIc2giR^ z74XEo+V9>j_tFFb-t!sc1wQaBz0|)|zyV@ILh&T-0|FN|>L%BV+l$RkD8$25uB_Pn zoU{8tG$5#4zehyyU(OFjKe4_;{aCFXAgLyHOco2fs>`;BsO6)%Z7Q z%x`anOUdcA!+zP8N6?>vvdJV`6!WdJ7_%yEyi+-tN|Zc^C~=^* z9iJ!CFdS-3k?L)g9NtRE-Uw7a>^wJf%{r>;$I{dCThGpXB7{FmAd~_+TRC5?>kSiX zoR3XRLaM5&3hi#i2q;T`>E$SUo8I*lCI%>*_ZI&PEEoxxHhcgHg253t+;jp>yy_6e!>(hqk|;T=m~^Y#uwk@>qw0PK&ZYy&$M z0il6>DTo$856T^cOQ(?%AXvLs2nIu%@Hc<_U-}Q9M1f`%e%G<4HO5OjkTSid@yhv4 zmimG zz}4(6={uR0>W4Y2+X#2%r%nS?KK&XbXF|cYrj&(foFcD7HEvur?9POvC{@0+>#1|Q zoItwFjZP_fK}oXS3A~K9J*e5;ME5hKC^ki<4#D+{cVbbPAjPf=zA3m8TvcAma^|6= z36-m^x|iOi8-|{q{e0hS{gnU{89_T@$l&gy=VJ(27%?RssbC$X2%u}>^qhF)3s}vz zef@GT*$bxM!PzMit<&vhd%m>4IJIo~W{G#NTfm%Y-x%7x5Vyem`|b$+h2qdA9P1b? z+{rl7YW#cQXtE`@()OiHk7Nmrl2;OnHJew0=foWkwXr3OqdTzeAi%%j&=S=Nuett8 zd~U>R&(=k{F=yQU7S~8mjzh21sk@&O72en-MMy;?_ezSz(llil>>%^gbSCzQNm7bFCt?@y-rn9412kW18*J+C6i3AcK9l7 zS+n%qN#WFZHSUPooEJ<8Z-4e2@eXjjae`cGHxCM=$^~CtaDkpL`f~*f>o+wDJ)K|A zNwP5?ILrFPd*m?ZE#tP|SgojS$Ww{DoK?ns1J&l1IwF0T_50J7zcSA1%!ZWoZnIv2) z6v?55-E=nDM0hj4{6dc~Jf`qO#WNwIHTlYY^!cIYn&lo1Y};QLqZ2`hMUc%fQ?|QW zPyi`8VAq%O3?Xp+W6}y`I9@>BB;bnTpZ9^MbSmV2^2! z21$(Z1(64o!27CufUAZ{&h#`R+zD-bozDUs`LbyTx)reBdW=HFnNg&z}v5UHRCF8 zJe3?_W5WjYEBUr<02sP29}HXG!EcmPI7tVqtL&k~MYd>P93V|jH5%RtyXd1z?$xn6 zfRdHwOKE=cJNV@$Wl(gRw~^d5s$FR zVX3cs!hSWqIsYM(M!WRn4{HIYWlXKvW;Z!#*I-Pw*(?Emq=lCtTIsEM>V2&9II%i? z^Kj`y{j`XioTsIO)H*?3gX0%f5&GUO;+OcA@fC4sN72!)DWSiAWT9&Ut2XxYjM1V# zZ~Eshm{tbzb1z^*uAE=P1Lm!yx;@)0+cRG8f1IDVHz0pu`pN?>_G|Yd?uyRhQ>jI@h(|ly+eu`@Qa|NG_rc9y!brVr3kLZBB95W=4o_Ga5?-6mqVb`2 zPok1Q$RfQ%-PXVag>rT0&V_!W%1?{Hn;l09cfqTJE4B zgME-9GEG`rc^rlyaZwz)r7QY;l1rwvLFlI8L+JvPa36m?_dTV1xF znRHiS#zKUykbEM&0DTBLxBC=K@A@;-E|2&G;#q@Jx;Jr;wR+EXk;yZOkz78}m&@OA zB_*^)GKm+3f>e}E0=RRauZdV3?@k^ceBM4%K~!m%H(1eBlLzEd@7PdMD(`<-$7vtY zg=YBY1bQqg4kQ zBG0rlq;bh7bFYo#sE-{>pZ8eR`m`mb7@h%#6B87v~cf>b; z56UC!o0_uC-4l8$jN6t-_+(ISI$sllWz54xl(1D*{(kwrBK@_a{Ub*0{>9d;FEYYk zGyQ;xm*sCL0$3krcN{Vdx)m7U8VFCSeVXCeKl~|t{z?f~A>UIYo!4o)@NpTMFtFJ^ zFpf7l4F{vYRAYz)f6|nap*d2E6?9zB%+a6wt?%}1bz{$BKx{Zh`7|r%+kk1hrvV7@q;bGMIEXixn|Vi6vmTzZJEndEHqt$~=;_(30-vamT~!Xy(I* zw~cR``6Ib7?Kp?QmHl4)C5Zdt&iLR`U&)%yy7bT5r0VW+mtK(r##)xDS1g9Db&tY0 zsCnTWG+8Ap*-G?g?D^5Z}7+qnI~=0hm(c=RI|Gb^h8Kr>h(GOF{b+VgN5G4bqdp z`zg9HrX8pW=fgQc6-&K4pAD^s>DU4!m5=S!#LpRzY1PyLW7n z>}+5Zq6j8PLXgc7kjxN3-^yY(*2Bo7q=)%|fp2H8)n^9PSyrh3y*Pc)s;MsR`Exhx zAY`_^z4PwY`?=d z6G7+*_*55g+sUF`7C+D@A=e7$>!ST@nGFnCn+8MYPxJAQ*P00w5wTS2TuSOoFYav-{W`2XTp+m3)S8i1zG z-S3gKU(`1BAer$%NJ~mqQth9li9AV4;E1*kTpj(Nw6Z~ti`d@yBs8-`@7bQG!hVdD*1xRMH-`?h{MeGdV$O|BBi_#46 z_2Vr!{2zxD43$c|IAu|$`<8nZW}l9#EC#@sB>;^CHS#}oHQVwbtyl&}31xKGiC&>2 zL6d?Qn! z!bZzm7^=D4Ug@>&%PC>al@~OnS0=4 z%x1H-F2W1EDS$hIJ1JuqP)kruoG91u;y^vj7uQ_ZC}JCn8+fH|GWv}Cex@z`mXtM& zE4>#;-zxO3`a{P(XhNyGFUA13AYd_|$&&mNUCvOpV?x6!E)(Z&T%%ks(Rt^0{8fy5 zoY9{IH%Y-jsfOdvK0M`JD@$PbBM3jRkbM3oJK2@NVH`wmZbR_57?OW=pgBi;itAH5 z)knQ**7)aUdbQJ?O+>ezMP7PMw6Kky{WADx7~r`e#njDv#L7p#mQlNj&_~tju&$u9 zy+I_$-13lPN6+x^!F3@t zFQvZsu{hUyPG*L{)4}4z!Q7b;ch>!RWA50vDxb|;USf?vNIO&G zEep|ApW`(sPIAB)W>5V2oam2qe1dZE z*wX0y(DmkzF^`RB6LNG!Me}a||t^sQs zjh2|zFI=sY5FYawFa@+Ri>-2|_y~m^%8(;G-KV}?1O#guY!AgMBOge_72bgU#j-Od6k#vA7-P zMHk@X*S8@>;Z0Pj%OHqLYJ+<8L*@$I6T5Bi6{H75Ti-3a2!k=TVaI?OtAg2}M#}nb zGQD2gl#C@Ab}0stgku7gxt-Q210?yzK+a8_m3hbDF}-uu?#UCbjXZjCg}z#eFw7+@ z+nJErQGdN|mo-;`(D zU~G%^<$cGN)H4mnDF%dIuS*vW>uE3ktp6_?^%oDuNX~-+NMt3@g&W4W@j}-&K&|Su zqgg%1C+r3N=^A_*Go5$+Tyx(qd?0??_fOsT1lk=?U%i6x!Tx0mW&>_04iDwf_Yp#o zztEHOs`O}}&IaZuTctxUHr2eRJbaOP{=9n7;gF6`bMx(NI;!B|&W}Oq|1|R+&`Bx+?|F0|;uq-vWsulXbmZ19DB|O^ZQUi} zz216d9r&qHIimk!R=blq;LP-4p78WXw(IWbW5C@P;K?`e$rxO37cv3KgFYQRV@27` zc;jUe9A$=|lu7?(h09(sk~n^M!>+?B=3Cq5t>^NY+|+DPB)9!F4l*-PQ*O6yy@>vj zUkBO{3_$<@iN=OrB!9^!miq=SkvC_X!gk*T+>@rf&vUyuD540Pee^zYRV>>*0(7%c$Vt#Uo7sIz5#9vfn&lUy zPf_k(xCyZ7xjho%k^jIf&zrxv8Ioe1eM81?vr$s)-%MDo|2Yb|KcdbGZloyh*7N9@ z1+w|YfJOBV`33IvI-QHS8XzuOhE&Q|O1-YbsDP-cHy{5YW1Zr%md41TVV+%zAo3#$d4ijE6R`pVp9;v zI&!=Lml&8bhH^>BSTrw$@+xt(bA((NzAC1A+p6m(t0m3p>T{|Q-@m7BGC zNVx++cEnFFw&%%@swHBkrAXbgf30V2)7mj4`P>2H9CaQcyOVcP?D~$n1Nm@5>T#<`iKA(ry? zqCAEAU(Fpu&uEqRhKqwDSZYk3#yGJCXqUPZO_iNL#ta_Y*Df8AU9#&ucSiTirIM8f=jx#*=qip25MixE+t+*b0>=nQ_xDjY&`DQw=Sh zFDzbf$c|5+E%ykvR1uT7%oVODNaq7V_y7W+Y0mZTzNMITU*t8T6y|Y~u}`-A#si&g zX6#K?D=Pms8*+rZ-9ISc#+8OP$NIH+mGrXKy^%iUxw(=MTqQOC`n#*PT!rCryQl0h_(Lp6atR+6a2|<> z+kg$gup6Hs!J1(f+ueF}^f=>l&hE_S;ThwLvqJO!nMO8l_Ss|SZ#=464(xx*`#lg( z_zFmHPQ(8&K`c}dH<=0Vi&X8W%{U}kH#Gr6zUH68PzPe9DD$ryfys_~868b`;)`Yz z$J5#U6#WXD(r<+qs+FC1&GNEQ`c?pR4G>Ggak=n6n!L$y$kqiQ75nMH(?=QN31#`> z2yeu6GfB9^KfYCKLg#pJV5kfNq*OG$0L9&s66f|1SFh!6aM<;vgRa#29A%3?cHS@R zGmL_TPJ>wP7S$+tbhjIdAR%IYL5AHD{1h8pV3mXOg7>?USp&bKkUZh|h~;DbTWl<_ zurl6nml@nL;EqzJItBAOLCJ_G+VVs=_qxz@_vU71MNV{QSdMk}ZKG#gS#F%^y08r7 zs2~tK-|tX?3T!(wLl)C+rlSf8lQ135hL>9=RDJohmEU>h6n*+&9U;@$n6jpsa#?(m zC+W!nolES0s{f>c#48Rxh5XG-=4A--VJK>Jf4pFKwQd)^A_-9uT;Zy&Dr7ZRSE8?ijS_%q=&3~1Q{1oW{5s(G+2sr6$C zB!OgwYC6tO`ZQgy46HaA-M!h^ZjgSErH3G1p<-Y~NQ**_;V?%Lq6orA1Pz4eK(Wd{x~zHCuC-cNhsto6$rDsCuC= zI-VgxGl1R&5RqIElDJjE8>ctABR%G)CYl%`NUA&Z+J`5^Lx#zUz^{w_QH+0<51oM1 zgi3aEa~T{VKquWr3Tj;Bl?e9l&DpGv33Qh4zEITj&f9mZ=!(~Ed-&Nev0Ul0EPPRo z$A&*r89cz~BI^N<4G1g<1sHNUV3jkX2zL%~KxXcv3lxd@c$Ex2?l(kVw*^RtD`)qmb7{|K7fQOk8@6of>De%(q*Nd z-L;OB*Pm5D=5!lr%YQqJ^*eVQ9be>bNj z%Rx(nyLv0(S~>Di1qIBlFz+@{h7xy7y%%?5pi0;Cl|hQLnbSdWem)fhkK5hHBhTBB zk~}!#4wQme659+$@!^mHd{dBuU~s%6EW}7?K=~w0oC`0MiYlE!&IIQ%e2q68MnN2# z0JJZDape*Yed}*BYJc`V`6=_?I|cX`ypRT#!3hTk2hQNgUfhmt<%Zsou(6CI(8Gyo zzgO}1}%dkQx~jD@H9q7)?gubvjoeJ2TnA?}v`AF|9HO>nx|>F;U2)IcDI(ke!~ z48;baN}&(r`PKTY_m`J)hrNOnKbW3aN>!Us^N^2LFy?qo3fJieNIO^oJ|%y~kil6= zwWVF7muUx;=}~HW^Z0O#Zw+;I!!)&u+0`89CubTHR1*X;m^q>px@yE~&PwNOfpY@z zX^z=DD&T$0(?Fa(o$z|R0F|C6uqVe4p7lxoGJheIag{c`;#<_K>$@J{5N{AYtQ_7 z&V!qS{0f}7gTq<3XX2dk&vx_4AW@qh`Ld>&wVML<4<3m2iwwoK+cxh!oZIhr?kg#~aq%s+JEn%4 zD|RvscTU+2R|@nt_#@TEs`W&Ase=FQ7z=nfNI@RQwk-$;M==EQyEw&}&0EO+v6VaG zGB?I=nUd|CW{v&6N;hbllGLZWSfyPxCE?%q2K=|`myrwO0IduTz`>SK9!gLf0+Isz ztvYnfQzSM9G{7wI6P1EAhahA?(-w|EO8v0Y4eQTM^8X z6!E^;7(lGt`8R?r2oj;qma{^2aS~(Rcadwh%HB)-t{Hs|$zvK4H~D?$mW}KGnqxn2 zLrOz5F@y)O^GGO-2xsy{vj9@d`ncNe$*s2^TTt?%iWisj4t+|q&R!0wTo%>sDC0N# z46>FC_vpFEvUBIeU~CSIc)X9#8b^fEZ(3GnEA zveBT-?uY8HfgdklpHmd`4JLY{@(l3%Rvvz-7FDawlYriJGi zJtWy3-sXSQQ-T@VTP)w^T@ym{Qyw_n$tqY2c%34CoisL|8=WXE#P^31N${P$`u+ZU zF=d(|`yzHf}8f2kR@ z+VbC@+;3i%eySGY!# z*snGVmrmc8i4v)0um>NBAOZq zLG8ufHhJ7sW0QMW&5P2abJ{Oxt}UN&99a_2&;IHcBluc3yE@Rs69Dm@Nyq z+??5IF}m$L%~>uQ_p(<(rKe5BD5avJf(a-f=wrgbqkg&qQ+!D~NC9A2NFgXf`&L2l zV~a6uMe@q1RSV@*X`IlcamxDD*4jFLt;5=V9?z|%oSIqFBA+D~xF8rj9{Y`d^SL9byR4WGON| z+i^|GPeP>%)+8hW@U&{A*L1Yvjfz>&xtDUM8E#~l^iw7^`V8&ZS3jCO8PoB|nTGIa z9ese(a(q#L@aK{dpZ$GyHm*aTq-{<1hZx!efc@bdkTgC#$c3J4FOn--RVZ1hRo|no zj26C}v5%_^jS5uw1?zrT*`;S~mFck|m<$UWM(!Jw8=5ybcM(AM_}3RPQ7xix=YpfU z{_e%d0J^@)x<8ZY2e8#ll0wI?3%}FmN@qo%2Q%$=^Dcm^8bi)!2-L%AxAB%3SKujR zK+(ku1DzC(LwJiD>P@I#M6}&?YA1zB9F|GvT=yh}NY$gYKg|rqy?uI!HwO0ZHSb~- zIQmT^*cDt;ectriv_XQWds1QAg_kNH^7L2^>n-si-($#k5Wsz+Kxb)4F!WMAnej#@ z@&WwNyyxC@pK0Cuc_&kpv<)#{v{3d)wRN94lg~k?zwUc9INPG$WA!OC405jAtAFNB zktX(gEYGrNd^!8_UJf;pamaVS!Rv@=cR;HJM-hADTHz%TEsiCt*oxJKn;%i@S_5{| zs#%XeluIX!6dyb!x${Yti3Wi+^FtpDZYF?XwxGp1Y$Diw&l|xRy&^;L53sc|eWk1Q z)6reK^%ffaE&ifBRAfW))ruuFv8MHZhiFCuR)>Gps~X!qCF z0)TKZ^>9VqP8!Zqk*hhz6~%OE5e6zP-~IpXvm_;sfCMiS>~WGg+nNkbES9X*y-(Lb-G!yj_831)9H)Fxh(e>S<&0exqbrtx`jFs62SYy7^<;Zc9IzvQBQ9%B3PP|A^BXC9B8o917(ND!BZ5W}W}k04jI6!8q?wpXNh5@?h=Z=4bqz zrx9Ow_L&OX*@2F)4T5ZkN<9{gu_Bw%BpIsYB;(SHxXMu%yE%I;g%~R|4SjFYWkedqZGbbua zXBF&g4Tv7!HBe`^yDt+y|Fxi z5F~#fKw+{FDZxX@ObrCp;v+wFgzM|a9<0A=4Sbn@&gX7nc-xszU-SrsNcmNVKcJNY z0$kR!!wAANM0HgW<+9DNUyEUw+Qy?(vBj6>9bDGV87ibHddjbpV0L_WiKHR&gm4a+=b`l7$9BL zVK<$_P8B)EE3VW(i}sk!9EmY7x6WvAS7qA!C`f~7*T5}1ADGA+C(fYfF;t&`*OjmA8CfzFt3v)CIDTc*6QE@cCKk)ejI)(5jHLFA!iAmTxaW zdOgZmkeWN88l(AqnB(0QpEE;JY+!&BiE00q&(v~@yvX1@zPNA{g@%uTS|TUSpmqcs zs$O>CAa49EZp6oD>xy4tGR5qOe9yN)Zc)K0X`h9c^E_mmYRZxBX=Lngab@!f+~EB! zd2TJU5A|MK{#7UMrd<{Q*HQC_va(2HNe;$Kz+3^QVnRVNig)_&FT!>dL!@&t8P$3I zsR>E`qnY{RLyEP`!1@x~}O=e>#r}yVzujZ3mKUoeC)-qA`S5i1**a8A@If;vXr|lii-a zRoXHoUrb-#FjVE6-EIgyn|k;Y&sTPZd*5n-0V1TiNy@$Ofs9TvaOAjr;Tr zM-ppF#64n|hrw9B{^Wz+j4v_Mo31Iy*}ibMIf*~*++yQ5~{$SA7VjO3zV6t%hgjGDh`e#%dmyvYjJ>Whz)dXpD+A&}+ zv)Dik(0RwE>aedqQjXg4n^3g*2WC96+}u+0gzTOo6MrT8I=ge*Qy$H3bQCbq{Y$Fx zkw1;^dz%F3c{CqJ1nTS^F+G*-da)f>lQO-~|7$7JnyWqb)lqAX`#=-=E87C@Cuqav z6JirFp4*VLN6!!&B&vs}AQA#iPt7|jDs?l@xr7K$O!6xjzij?vp=ssR{8aGrSq_Xe zf}~+4jv0pi&pXN!PsqfR%|G%oIKrr+&|^|hY>|s6W}4sFwRCllE$G|!t+s5?_DRux zYW-e?S5oXN8%5-6(@VjdhAey|cyBzRUpbs{rjcF;gv(5kylb2UQC=1&=(e-(*OTwY zR`8>KmgDS_eUdTrhBw3i-0a2s=K|E^>ZSh>p}`IUj^Rm~*jabrGr}p~qiDQuAYsN4 ztSFS;hN|atamGDQ&V1>X!{UF|I(sGL&AnPKxi`|w&urTe6fv0;{EzTpvwdpxRK{6d z_^`Dg1>TK>l;E8JtAc4iuDUQy8qk|m)<%#x^-a&|Kfe+rWA3$tqpDW0TkAI@=&Nu= zmN)R|Vh7;=)hDz4wnnyNdX*z|&d_tvN+#k7wH*p?hN5fp$mZG=?U0t|L#SzJ>faSiZ-FE!e&FL*&w&y{WI}J+8F= z{>tXF4es(iOiLnk7`XGb7>Ah1=eqxafl5`?{v-}j5ZvjNB~NH39;AzROdinZvr8PW zC#!w_aQTOJoKN|fGh!mV*rR}7K0%a;;VyFb!5HgMO(-Z-;j$n zzf)}TnMFUj#vjily48ouW}N9-kcGuKJ&^p$_VdVMD$_cWi{?=t9Mt7-{Y?RCT5z^b zBqbjZGX%KbJ0g%S%NOn*2*i8b_>iD9=grnOBpj*_LV$iC8E_|&NWW}O9_pNNw%O}? zlyge4_QlqN+;tPo*)IqG$UGM`cmKwA^L)<+-WeKY@S-UK63XLJpa<9H%d@Cw(r9}< z6U34#y@84iuFJJ$X@a)$k(bh$zOZ9U0TD8i{NNdeY&kT;(M#A)(m6H*N!d};mYg zqw|fLcdWSt2d%rtGgD3;x9PEdPm)2<}#O5?H(S0l@0-qHjA%OPX${}rJg}if;}u6hFgX0O^cwa ztyzWPA2()ayxzBe7mzvhOHcZ?-iy5#WADABxM_d>_|hk6<4LBRDQcM zm-hhCKGc1YS8RR&Rp2$GSywaVzl8f-^O;$2Iz!Uh`eb&mE<1Jr^S?R1ptIx&Z%aKI zy5F`ob9M@D?))1YyGWT#R@OFqqqufxX-u`xM-{r&Wc6h^#jyU_o15yd?S)VO82p=! zR00SwfLCyU(*(12HEN*z!XNc$(J_S`_6H!J%5M3zfjfFT#$M={#;@bD+Se;|H9`+_ zzC z!ntwb(vfVxWs$EyS;RRNQZ`Mc@sX4d{%G=#97?8nEefRXiaO>N*L=ZmF4YBCUW=@`HJhyh?_4`&m@eap1 zmqBf=)t7HP3k{~f-D+lG?ndAK&%~vG>3;@JvzT@bgkLfOiE;>tf zU5KkQ%j`6(X^niGn{BtB0hDA)DOewL9yi-<1cB%|IQfw>IC)-}`tOmgCrn>^a}vTl z+9z@Zi$5F*tjPN6G;gBVc=LK6OMpS!d3EmYARiyJhW+ZENc61DF4*)kcvFPHjySsPF(CwRokyl_&*5^2zoOW zjJeUEC+5(SDRV;@Z6p&Q?Hj;<4E7EA=%Aul1m;0$!G8dZ{u_VGfBFaJgPG|h!f8kL zh`9R1)-%B~*U2u?j~~{Z|1iWgdCvKB+p9Fa1IDT^VCmL}%vJ7cVqaiLr{Uj_jKo3& zWVF6Oi`IpQSvYNXAQREYAPRX@pOP9CSc&Lz5vGWPAh_xIvphH9yA7XiK7Mrgu2tkA z7CxYXBVd7)WQ^d+F4!rV-OEsKc`!4Ht>@x7L(0P^)@q!bzHrO2)}D^A-$;*K(;%(jtDbq#eL1zi&@sJj^D zqBL=t9%a@rqNb=iucAas>kIT3@IA$JoY~Juj5*2$1q_e~`BQnIjE;GhF}#nWWoPHF zuricCkIQs!dxcY}>FAY$eodLj?g~lWSp$Z_KdlSR+u)xrqE{Icn4d-4?PgS-nSDVS z%MyKZ7UNLjN|nxu;PO&xAL%6;@gtBLi$(Wg4kA4Bzm1F%L(!t-Ivm;GZunM|$%j$x z`q9B$NzYTbbQQz7<0Dqm=|OJarI{fQ?0hP87SAwDoi<^#?9X|;Rg zQDPZ_wiudqSRj=ZAVL|3fz7I+MUlZsqH>Jwl8m}&_#1_%qpy$O9!iplm?`&sBkNIWMlrC_y*d}4( zn2W{$f7wq~TREcf>2H*L9`Ll!`R}(AwDnwzXF`KVT^q*yq;Tvs(li8w0ze;Fuw&zCcgtff_XbMLWL=)&a>;&&aXyXrpyE?{U4WHJ2aqo z;nw2(TNLbW|6RSdHa7ACvI#>rN&RD`wOa-NN?Ty-3Uf@uz~&xKXLRF#%RnIRX^q|+~5H84AA3(GUp(}kN2 zBV`!iX=JBVFn)9#>J0+%56QoK4$!@x86_l;ewhM3tUyhL35!M{o7(#4QctVk0}pv| z!*!d1GE%#tAkzxRyYWtEEDFUE;4C!HksW102w;;}GlLwIN{Wk?W{mAJQRSWmzeGhX zOky6&m4SKz6SM6?o$U1d@XTw3SYG1Fp-L&QG+tqwwZMnc0 zGHoSH=kLh!#&%@{i|TDp^1P!P|C2>Y3Ea&;m$t&K2Pb#~!E2rBXtbGWhEcYG}`O^lriXAIzv;mQU~r~V<-jD~(M{k`&qB?`!BIUT zPxxo2y)srWq-38Ybu|~-6eOKW0V{KSavfBCGl-KL zB}Du;k^%kMo#8-i?4#nM0Ev>&)(2U8K?!#$NY3Yu(SgZXZUJ(P)kSw!X`7q8iw9eQ zSw;mSmem0D>xL$R@pN`_wRcmK3^c9#N*G`BRAReb&NqzEOoS>gtYto8pIjRCtqOvl zJ&tnpsGk;<+Ph{YmNNbuJl~LLNIPOMwEL#@BR#7{k1ZZ8W4D?reQzBxeS2+s;V!?i zXVa6jk+b1Fil56u+EkgKL13MRnEpV9=0C7m;5IDiDy=Z=Cg(1|N=1q>z97k#{_zf& zB(w~4p7g!zZ&=B`e1eYHj9t_^a<^`7spN9N^hS_Vp`k>YercjT;XgG_S*Pc zvT_`@aAb?r?1GBa_u zt1}AMMl!f^!9^WHvjn3I@yd)`AJp#cmge1y^kf}zbUR2ae=o%1xk4}c6BoF>iVFG4 zkDY3I-uip=>Vj&z`msCG1g5VXJ0d`_qY^B|BBmbzX*{u1381HV6h*bfCIDd+1f;&W zi==;KoOcSnMA?h<#}nmRBaVM;8c{Q@9dDI?X(D$;H{^`2QpjBXJ9K-s{+5WE4ZBS?9%q|I|}4s~x!p+}qkzP>xY_ep>ImR+p+xSkY- zi=6?T2Eke2FZ{q|_K6`+&;H$vohI#F8*Y(uFEf-#!{kPfy=}Vn`61rju)(*#u&E#o zXL9}k5X=w>q&s2SU*mtqDk?L2HUVbuQbH?mNYSRC2m;(`K_UrH-6Q2QE@NjFzD*+{ zzzvPKjo@5^vj_5WtRh%`e>_6H-fI)Ndg8sK+NGa!X6Mpy;Wl0_mdEUo=#fGguTMei zK?FGfP|PTmz{wv${?}m$mTXDxugvaAu!3p-%88PKgPg{7xPk$*76pfosdB}QOQlS# z^sjXhk8yn`mU>DE@KukaDD!yi z5D263Vxy|Ypk-<-VnK{yqr zz|O_-OO=eTz^5J>M*I$3##~XOn9=yyh+w^9iHNCDyJmWtf9)KqQ~1%_-U@G)*}Bq6 zIVQeY`a6E2m2AI!sbz`$(|{ZQ8ZC^*yXH>~04i7l5Xwc3Y=oK-AJ@GfSM*giy~Bv&W*R@Ex9gc`;qEZ>$fxxY-hjRU0U-+B4W=-v0L_%s>yDAiXlA7rq{{v8D zkFUd)V}0Fwis1Oi3a$;iLW}ZShdr1KVK;5xquY_>%ZRC7275C_huVYT3Kb3XCnXLb zt53Zx^X6YWXMZK2Mb}-XUy0{v6({WAy{Ec5EdZs$OuoV3O$S1J>boFuz%KA%eu%hZ zE{@X+Fzue9>|i9fbN6C8X1ag<>PRcxZtdZM>zKP&mU;chASV8yjkQTdsj|A?(RME~ z4Q}!s%lfIs5LBn}2{RIufg{^qNWlEz3D*HpLzDt$S+G?D_{^eGaQT-TVJFMTa7e}= zmds1^*pMYqQ^VEKVzGp#)yh(@zeR=LVu%Mrsfq90I_a-fG;+Um#Zf+JSVr4+>>wb` z&RvlaHTgvQRxSJ5x(~J2zd}gIJQb^_S3hOCiw6^*fhfypC0yrcK?~=JPpZ;ZUw&mU<+D|xRQdcD5Kk*DNP8>I_tB~<_)9Q04D^*=!l8@b8P`h) z7I!f5^0f zD$5ai4&%rVG#gnb-j%q(2|b+hRZ#GCoe?wA%D0h`u3iE%Ghf;4|HauE8brNgkLmvO5Xp}Rt{^6cuSQ7 z-X8c=cjX!6qIebvwPWzYwS$j)M4nlix)%=>u3d1QKQtLoVHV`iZ!GClpIPBlR>OMS zE}XYwUk82Pd;Wj%!*~chHiY@3fUH@dM*w+mxEJ!t4&4XS-T^0Hs&r*=vJw{#tbqZQ zm|SA2XfW5TLpx6mo+`IHLQE@~9X9w`ST=mzlKFoS_vZ0XzU|-m2$6lwZVD~3m9=b> zgd}OR$0*7clS-6nE)fzEC7)DAAxWC-S(>p^D%qN>Gh<0*o3V^5Gjn~O)BW7{^E}`C z{=T07e$Vd@@sgQq&htFZV?W++ZIzq+Tay4_kPU_d3Z6YjU|>pa))fJV>zL3l5?9$2 znmCdRsvu?Ao&eYE4$qX71D*_;vRo{%@V3=iD+YAy?aUIKv*diOO5EvKbZdmuwToi` z&GIPu1n-2=hg$-2FIqnUa&)VE*Z?^S4!nkU*PMbp{EQua3Rz(QB!p5K%NEcWIZwFA z;SJ%ZbOZPx-lh0+$*|q{wQxp%X~EM^9=aRC{%DnLi@>M~3CG|R~%_Y7%=XADo=_EVzx4VK=vj;AEBGluX zgJ`DwQ+_GFb3$AAzf$Ow9mu|yJfII|TAFEVsgBrNCOkaA^rX-e*qG{RuFL{sVHO0Q zXhol&*OQv+eM}8RLd?m7Tgjn1nrBn+XVdpadnQLe2*d9_u(9RO9^iX|4FlVE(s~m_ z4=)C|7^nDtf~|GnKGR{(ftYzS5Z4zfuYgN2Q;%%g1RSm%H9*VV0Jev)4S?SlrGlz| z6)R_Kp4z(VwN-lP7t1+xTb}p!hk9+dPHd*_k~*D!Uv!D!8C`>jJ%U%l#!C-diH_Ga z>wK25sC9RFN4bBzx-NHQ-(s#u=C^})4?K3#=wvuai7M!GHG-C3@ev$_cnv19(adSzh=XLP~y3U$F7ew zELtz7q?}6j9CURm+seOrX!61qbWh~Jm|!6RXftK<75o@!refFj&2UCY)3)qf@Tv5` zbXn5G>YXZsM2Z|){R2JC=~(8Hllu0Di)wB=CoMi*PKj|(ex&rs_8DJu;!r~MPNAu6 zu1pki;xPwgX$M`I6y)h#CP112;N!G$cHCNp(4IO??*oHN$l zT<+r>1KkL;ygYp1$M}%(U49MijKt|(g>u@EA%^CK(pj`cBthcs!7&(%eDnoj%nRO8 ziucSU5!SfQ#X068%n4uTG%|w)0A^ogVs2 z>hZ>VTaz0xDLK48uu3avau5@i&fyD2uQ`@d!NOa?y{?3Ka@}!QSmhN+bP<5nZxY4O zZ-KvW(#4!lPtHXH2|dFyOSRrDP)(Ic$?Cf@d`RYp&Xww5e^2!jH76}TXe4xL=tif0 z*eSGVfenS4>XHyMSHRWVav!)5;-7DnvCP&bcR;9B?F?Rp<<7W4O|Tmw4xX4fLG%7C zjw9ahK2++vXN#&@s*R?mQsbHtB0LHXyC`~zy_xsU(u@@~p493h{M!u?TGC&>foblJ z&>8b-z$vWRp7w6vyv55mI|=aiK8TNYPWZQKoKBnX%ht~OtVw#ArLPV zmIl{;%U47dYLbtn45Hmw?Ip&!`qh~;t}OAwhI3(&J92dHeLdJ1UhSBddf}4#LCK|d z4>`cdmi1C>JX%w^&J-BP! zz12W9h^)za_`bj_@yd*(q>}fM>bIGpF^|qan!dBo_v*u;ZLY^&-hShFuTx+@y5t|! zg?k|8bsD#T%AoS@i(&xyc@?~X9dO{PYJwwKgh9{DD9fS=k|v|0X?9FnY6GS)T>|?( zI_a~I;qRt%x#RnuX9^c&?eGlD@bvW5ayTJ&USk{g=RbF5;s4wi)UrW(2eLO0b_aT5 z;BQoa(lJu=M}$Jw2ht~Oj0XO>Y$b6}$cgo(REhcd)yS_g&AOU2LSJ#|r<#Nx2{ywQ zmOgJ;97riRXa7gZ2{W1f4^IL5uZ}_?)O44&2-5aJIF3LNam{>@vPoxc2<%o2kbs1o z4!{-PVr>q{&nCppS`ky$g!k}i+k71IvAnX{(AZnqqt`tv$6fc3D}NB^FXhtK=V>=O z<2{qczcg5@0?@vw2)bMS`5y1q`@Xg%+iIIm)k=9a*dP3G;^+Pw^=yawv$2FrU_~I* zoYF^EWb+<~10gwlg8Tq{hIf+O@R&9%TR(`lc35W~2XDi2m2FwFY-ip(yfknoWcpA! zh(=pn!0?Cdz>1!^n#Jk9h9i9r7l?t!?tQqxW_HpHC(ZZL-`GkGpW^?MLYgZn?vui8QGkdaz_Cedc|pX{Fs&(M-9z zxc!3J`_Mg@*%qdx6E6$+LIAtpxCdvjBCgqJE}nZc@bdLzBO`WUA_vm9+PRJ$*l_iw zlGRTBx_`9auead;Ujm#PVB;EzOuJHBL{yq-(ZnVzvd$3`8X}qOwU{)mkehe!L*mqX z*?YeDJ3KTK6~L+FoyiIau=dNue89qS5IY^DHy$yaK*lH11v!u>3DWy#IsgL8>9>e zb(8OlIfe;uJYAWpg|}DpPhW&ifV13z2?|F0fVQd?LR(X-f?Qqy$bx5zfZvB(x`Ll7c==V7J46}FNCq#a z6+L(Y^KZ0|ipW>=qOT>0xX*AFcsn!S z15_a(>Ds&3Q}eSgU>tyUtsYowTc6f%`U@j(>Cdr`uJ?q4S+)!gY~a~{j?zbM?Hz4- zwSSwnzr)U&s_4J+jdn(AUPR3uxO%q)AjAOqC*TjSmvF))lYk3E-ibfLv>a1rN|c3- z?**)d+xn?g>cjh5Y?f?9n6lRPgC}Oq_PFlK+dLMFru>T~bzR;3|Mex$1!gZah=hl-#7=qzNxZMDTzrQ3wpXe(ez|iM8-Xv&Mn6T zF(4%6%4trxg96M>g=VME2IKCj+mSxj(Wm)i-n8HMeTRL#Mt$tWdKyktXn)+-)NX^i zYX=1maKoV{P|2Y`#8&QTg@zXxmV3AwEOv5OK{V0j8;-Q{S!Cu&c~x~)?cHmhFLF#y z&mGth=OlwtCl0FpV}IqbuJ1yf^?#gVOe8;bc_%P%C;n7T=zGI)IZkhDXYC50-#}E! zhEUT%WG_A%E_@T+(u$R9h}d2=u2S7lXIj_R_FAOld52uMjPbDz3CB<<_kJRCKW=z> zLSa1&DWM51|AlBolaBS`L*D=t`?Fmq`wPi6^_5Q(SG09N-tvT^UmV{?@--MpjsTID zPZq+Tpjv8{;1%GEn3q9jVZkO}vKYgk542CHX(dMI2V|OpuH&EAP?M5HfVqaqJOZVA z0o+&@hWh&$CaHZ)qN^0Ve z5rs)b%Ts<9;eobjzPq%JWw8t#X`x`G10#(IrNBGVlOTpqboeTmc@boty3Q^X|r$WT2a#ti{lN<(r?8`w#j4tY^$M?`pJ)R0RlRiyAS;gWrns9Yb4uU zQ~|HSwdS~aohie^f>)-#J7wCaGl1yS1XzMsxzaQ2 zbpo+FmIjv@&bW#5PhWb%AswFLftT8rKg`KbrzOItK2E44Ufk5a2UH0+uTuvEr3=6% zCwfT8My?0v>isCatd0zZ_wen8&XN2J5BCmgosLhm-_>&O5NKl>1yt1rb2L^YXeZ0v z#1x-SZ*VM)b{aQ%Q50PIYp8TNyXTYMmHOdMHOp`Qs>!-KOWQZOMeY+n4JP&DIoI%h zfY+v!DBG|H-D)5gOh7Z$`n~Ma&J_NZ`Q8%7S98bhV{2D!N@r_DK(3Vz2<52h@b-=l zanMNEO{d%j_PhnVIohEaBcW%aQ|iTOH9x)P;UD$F*V9H$)KELeR;%$|Uc$k%hyK`J zA~GY=0lwr!8L$N;t`m!p#iX?R?qt%^-#j0vOtanJt8a41cPma#w_NdfLoI*%pj%Uu zK)^AfW{@9(a0_WfVe73y$j+~bcJ!i}_knt+F89aYweM2RRnvOvvm&iT-foG<)>4>P zyXV?i4h-DvO^)eLf&weNkWp`N=V479DH-c@Z>Ss6D=-hZSCsSDl~l(_O}P{NQjfm$ zVDF&Cp!Mq>*hakVg-N_ZiGyTISHv&P>21kA#vGW+dVQorr7*{1K+~qbLJB3%fq-b9 z0i;IOdiV%l1}4zTzsXv>c{qpw%tM*yUOb|7gpLHo1((*oLxXJuSBYfj7%q zVY;XInF@&V1ju{HE!@V!$HySnMHDpNQb(I9_d z2{zlS-tEfNM6S-upz@j3t7|@8UXGDhYh(_*UbMH(I;E33+$C&^nOh66}Dbm**FAd-r_I7AO01D7z1#?5ZxrAg6UEp={b21yI{$${H= zQ>9%lS|5Fdvxr>*J0$o}aT0sG$04o-bz+q#wUCO2?SEsq=VLMW0I-~l^?z^MxlcH7 zw<|!*iX2e5BK-+vlQa<-{uj8>^cC9qiBFtia(DxE2_*Z)BUs+Zov2m=?#=}aJO$Y< zLQMwNMj#N|HpA3Jq#?L2MUwsND}H5(OPQN^bxrBlG>@3}=x>zg;`Dn6Lt`#r(`n1L z9274PXgDK)GLcN&OxOkJqNB!1ihvw~S^YTQmgL#=>U7mebxqjCdlg#c zG7Qf>SziqK+CvvL+4bP_C*#B6Db5j?)~37)a#+exaEMm}r?<;J_sfu${hsFzau;6S zet9Qv=T9B>4n*w7Kemn^CbF>O26r34)|s``Ri{hb45%A@@=zh$UB~>$lwX!Bx9T*0 z66J^rd_r(UsA%S4%#abCGG{P>6aS6gh-idFZC0+FH82j_uy6vf_1YUcXq{|P(Jav{ zrcVY6Y!F7NO>?DV5D`HZk+F^zP7g4JWGsEI1&((s2dgU#yk1c|8afj?EII@0&1(it zoj?5mU&5Dl&P_;zQuYeL0GpHs1_3BT{OwrgNw7|ZUi=OgHEEe56y4>>aw~l~GqmkX zc3|kaGdBcVJ3i%@h;KbLbi4Nac~!nl%v%wZ&XNNn(ZDOVjN<4*O*+amjn-Qln4f{# zZg3A|q}eZinyodJ_44h~jrg#dukUB4TVUGTq8UNz{mlizK|DB`0L*Q4H9)Nx1$9Pq z@s8xGbfM_yeQgRkk$JQzs|TBFG7kHtUe--vC*br5=Uv0{o@k>GFNpERS_kSm3WE$) zf*NqhRY^ZB5U~$D5?K?BW(!&yN1L^e?HaH={+x>;_v#yK79R=Mzd)w3+3#HTC0|Xg zs@v_Rr5)IArp%T6@5~N$0gMm zbI!d;E(&(<_2>`=vv8yxQHzg`WLox_efqQ2RnAY35Cc1~$KftIR{mGcST^hAxIDQq zf3z(oSerUuAnl}*WIy!f6yt%al+C=(GK5HgO^CaIkN5`8v((xYPRZk%uy832nk;py zb6u}NgH>VoD?!foM}g-oPrlySt|V|zh99M*c4)^B(LWpp8s3hZlb(RU!}O9G7gEsS z1M}nMH5`WoT-dWEt&1sSPt8+lm+U>i@4D-V0?3i?Xv|Ayb1N+GW*Q)993;B-Sh`=! z3&AQa&shUI?}c5{qNZNyim5(zrR88^psuU1BI=4f404(Oey{ff15?D6O&VWlb%C|e zv!lUr_|Wp6fPs1> z;3hB|6At42c7PxV9ENy1%(@rT2fP>;Y1@fzX>4?k3x9OAkEmlU88ZE0fhqSsTGo*C z<9_HTSl5xCfA0p0Yv|%IBCf}~OV%KCf){p6LCx*Vsf7d_{_=QjB+tTyY%4$RrQ#GM z{~|dy=uwWsR_`A}Hd_Ma*Y5w{Cv+ik8t59&$Es3-&;Z{c&C-G8)ZQEFt5F|XPb&Mp z1iG*tAx;vZK}wcDq6NKRqoO93LR z*NOfY*da5nY|MnjdyWkVUN8g^c&1Ik@y57|V)PEjZzfK+ zQtc~7x`X91_qsAJeL5<%fs%GZ64_MCIRK}4RrJjfHgZpgmi+1+liqQQ#5geDdz$xz;+Gz8&Bro!W5Ov-xuU(5K1cU%Tk^mAL-f$ZJsn6& ztRV0`D#iqu)VV;G&`H>dqk})fa!b*=IAg8Ed{JED;iz@i?(VqgzH@1JautR4UyL^S zDc^_;5GAS%QqkoLt;9`1^IY zxz6S%`!K*WKg6&H@b2+o>bnk%%9 zbY#qmy2aIG&Qnj;50JPvg)W7~9Jfc;0>v&SXB;V43OsROI;2F#0TmkoK*?xx3|0)t z6vVy5A7Z9rq%4mYz?T~hy?};F%zo+QFJ5`Vb9iK*#&)0i-b29A1x+aKEkKAM9F&F? z!CxzgzHbK?O?;0GlbE;Ya>zbin`4q8ta-^I>)gb>#yFh?2;QfIY${>(P9q{6EQR(u zObc2mx}*gyr9OPJ*MN&D7zxQRju^5D|7DPwHr1)V{g1bjfOKvI%dCY6Z|j~@0=VTl z4oO#Naag&nly7e)D&m|L#Ks1P?Y}f`%nSw!bK@n?+o)qg>{6ymTQ=DWR-wO-IuUvQ z__s6V2!Yw@pejrsU?j|J&7tMJ$V6)nHKD;Sc-;G`IrT@~ zyULolscnn?s$Cj?J`OAV@GFfy%7^%Z^4u5&P5GIG_XE+n8K^i-81r#@z&4*jfKwuk zV7S|12A!e|mz-kRCo_=+yLxbAD_@u?YzDS~b8w{S$z;8~*Eds<)`V@H?GxDQEVu@| z97kv&%0cJ4Zvj`B>(7#GHm`}V1P*@ob*neYp(Xp-*&q`T2olLD$?Btx(G#PQ%Ef-7 z@|QNEb^z*Jk$=`Y|6hH_7(?`NNm!qlOrhU17m8iR@xczZfJNFeI$VFej{42hxL(Ek zYFdD}s$rYX$IrVGT3*E}K4BIuQV>m797%rT^B>@OdB}S=a+Zehx{q{1?sbx?1#-|ZyE4+t42)rF}Hcm0` zGB>QNb$9o3Gg362Lz)tRxf z{}%elO|j7JT|^dqchpwj5k6a!tEB9|tupn?7kT+(f|n*cxC}yLMkS>N+jP$LJPtm( zY!qtPw=+=rTevvz}bE#P8ql$Md@fBjZbl`*;DugUD>wxKk@Q-4#ccAo(Nr>Miz z_l%uS1L5A~$^)vf2}J7`LecO}$yqEDTri2J&qv{_v<~GMAJfN|4L`9@bQs?|e3Q!8 z?0$t6{D999}grTo2+@(TtRUfR;fKAM)2-caqU>t^%wwgLRI@(PFe(_j0s7XvS7^0Iv zs~dUX{P0Am%tfPE^!@$lYBZQ9hWF)}KpOvm8DwQRi^y8XRjgjB3TcBN29)=J}Z{7>Nsu?8{Y!YU}g7Ld~&CsqGXCv3D zcezq6*Bx}Qon^JBV}F>59NOj~hE@ZjU*Qij3vmPRuScG)%g%DenhQ|g!M&u<*HH)n zBrb-SK9cZRU|ygWC>;@FLTqvwD~IZQzh$#BP=xkUPu_g^SL5+x7rrFCs zNs|!QzJvn;;;->cesDzBDSxAMXTAexhH`ar2Uq1JM?5jDzwHdek4^%g{LXos?jlQf z`U(mIK_PYe977{KaH zzYRVW0MSc#f(=;yjkzLkv)iv`J2QAEo84ATO9CT8L#;-NBh@~C8lK2hJgw7yD*sj7 z8!bI`z$K4GB;LT!IFdkgcES#TWyX3#Z5dPhZkH>Cb*^-Dq=t%9K72$pEK>OHHuvL8 z^XKIJLYy3C<|a z(z?N5s~HCn92Nl|;}H8Dd_)BK(quBk&)(a&aeRVEt|n zKq1hY6anNS)7?QO_i2pk!(OOWqKEX5L#bSy2>wEYk0f{ZsZRG&gkJKmRfAg`1)wqr z!2OA<3spB~rnkBX!7(LSeC zGLmzF!y7BO6qmF1&7WXYU(8)c(vuyHSX76h9 zZZ!80QdK?edp_1YR1&uGXXo8r+gu07FgJ&~W*q5hD+uZ7{6+jrN6juS(~xghqxwaF zAL~BX2IO+AW}VKEo{C8FUN)G}ACanX$vs0JaFl)jkBu}Ja$Vvz{?W<{O7z>_9}l;? z7BuUh9rM9tDD%F7U3ul*Nnr-n%)SuA{k8I0{_zwif7a{mhvnZGbJQXPr@6C%)IAfT z3n9WQ5aLOliQ=lPUxrdDh7Q)Sk$Z8>keQpsAgSTa8miJyX4@?(IQHi?aBRIT_jGAx zwrQjU%Nz{6c(|kl{{yspWSApy7<&8d??IHxzPYecOd{&S#X~;cnk(q1D8d?O8R3J0 zXb64*NB5kd31F3>a36O%r5V}~y&Z889P+BB6vNeHKYro#^-KNgllAO}(0%p};{euT3BJ4^O;zA-)@?@Oske$<;B^+$5mWj+-?1%@T{!jmF0`%oi}w3~@s za7_Op!q+>vsx3GTsA)6i%#EJ9KiActkg_;Q9Fe~GU)Tg=J|6HOgdutfmBw^;`93z? zp*#(6>5RrHVl$F?q4&OTM!IX?d7S;r`GVS*U0VJlYlri@ni7AY_Heu?$`Jj-&p`}5 z7bB;X6<5>hcsSD>B)Y3xBe3QFrDFZRqCc?z|7_WDWdU|)UwA4CYx9)tlRtWNsu4|&hB37*yyF}mmjPztUAR8zon7No9A@)*>lNjcle zz406lr?w*6tj!e)Nmq6ncqS5Yky=JB)2DBp%XscFH~B}??^B!mr42rWucQta*ecu< zg0~m7B)XnI1!sqN++1~F`xfj2-Wgu$zt_l#U^Ng#FS}0R%5&H_*7k~VlIXAXY{t#1 zoy$CrUlvp+vaHSrv3r1jL^s-XSLgglYU0*v5^ELPlf+d7cn=1%e^|fnymezqJml*T z+T8ae7IXO@)lw6%ls+JY{g5o>40<@f6)S|u!IF*HSI2}r!dB93+a6M*F8Pi+iCs>w zD6Q((*!ysw+__aF)97Ox7akE*!AoZ-m3l?##>tjG%whTfP+*(qjb3^GQYI7zLwaO; zmOwepN!hVxuYro#iu|GrS2aV$0c>Iwyi)5K74lHGR?UHZu1JHd?yBE(N~V?0Hx`Q_ zW0D9aHSm3)=MJ#spPaM}CjDZsJ&^t+^hnuRx%1U`t#H2By;_0qFSrYGb``q}LX5h% zySf-*yQN~F=XO4GsB6TzN1P)bgmy|fDV@A4Bq!`gnLpiX^Nq9KDsHYhKkaLUUyq1A z5zCFwz3#qLR17uEU@DjE(po&ZRU`DQL*s3qHxRZL*f>jU>$$ zN0w@yU{;!KW!WK#M+Ko>^7L{SnS{`jjEV& za58}J21U0OfK?P(eaKZLoxFHf=$4W<($ zK}qNatFM_&lOd>Zb;wMU(KYWGX*i@A>2~wLYLmyDQU=eCrGDh7?PeSfPmlVG(rX9t zLy{B04KdGml^B_DQ7A*lrOg!b;s5AB{{wWGH(21F$H*oJj_7ZSEQqlsI>vEQJ|R*+ z&D9E9vMTPbMNdWls2%$FLr&VlaI4tlu0-c{J1#JY|DnSO=}^;TPbHD*sHO{=-{JVe z3cnjYY;VT>KX(w1pZse-z(z(x^y-{d?q*GfS^E9U7fFIC8P()moktPcRHG@_(4VZZ&ZlF(6AzW+}1zNf&In`YW1f>E`5p@w-i;) z=zMwe?Yx@3qV^X6HX{jNSb#MRcv8g85WVIY=t}&ZjjXrsWOi_@fJdm}1={P`S+L7~ z^hnL@wQ)*3rZQ^w4C6lY6}82yrfC#wk9;ccrB~UPc^^pJr~BjPO(j&hqe`5kUj4og zOd^~l%iaOfXhHz7RtXeHGkl$ZeNKZuDp?x+jq=y8_#0IyzZSM%-G;M;&`1%2{BUA3 z%9y6V0UHllsEj4pI~IcVU((89A9POs;iO}k`qHRMH0;Fue!=fV0AHrNMycj#I z*~a90z*omC;p9-m;hCKFn`bU>d%Y_B%+u>o_*utoN5Za(IG*W}(P4l2&!+zU=&vqZ zSq1hdpyxbXIOn6?hC%MIrJbqJnlc^ZBk4wc9j;O1I1(mZaR2njWTWX2wY*9C?j1tE z*C+hvG&!m1R$_*_Ti66sa?ZH-9KS@Uyd1|qvJ`|~Y$sv69ro^zw>mmnej|xGzu}UC@9KHmM0GIzyAdRr)wI<|IBA{dqE2q{N@S|C#@xjQ!HD1P7<+ za=T8&p1R3Zq_Tky-#7$)ltwmya|I~{(CwA|fNxbEbU~vvxqe{fQiRUGQR?I!p}vMJ z_a;fP5MM2pe4p=c_uz=APAYos|ne>oJ+>a-6dZARGOSI#2taGVUGS z#lST~lF=@0;vsCOmu-4&pL!0$V^ZCkx4Zb_x)OIC`m#yrm&i{3X6OPa!C@Pr;ThK+ z9Abbs`N|E27;r(&+q&9odAY(*5`(y%K1%m5c7l?n?ta0?)=nUFht1;I22ApM|6)%r zKa(_6$+-VX@{a_syYm{Cv|qM)>noK21iv*dy3tI6GuNDD?V0Ovr{0w9fB{={_RSvVGzl4px5>W zUI%pkCZUnXxX?dO0sgTCvdM~Li+mzm@sj3NaD)9nedS?qM)kXxnN52F^>5A8Ir=Ui zn(b_}U(&j3v)}rHzR=W;HsUa#jG$d(xwKYb@k#?GF@urbMeoSHl_}p}8w41Z1&8(M zPq-$0@s#)9FBdECzz;+`J_Iwh30s5X#a3Z>yGOx5EDCZ1u>LMXBxKXv1Tmek9SGp; zd&W*M7kyyBxN7O_KNF`dXo`)T<9#L8soHy7XFO+inUSDAy_?s6{qI zReEny01mab^%nf$-pa{4<37q?f~WhQ@3?aD^alYJ9=3=h|8Na&9H%G)mmWCie~#8P zg7!#|C3uMxo=)egf=jJk1h31~nV-<|ouOFFQObGTFc?(^eaVEc){?{);-&>l15(GU z14CC!OMhhty7av8z_l@Db925jz@aEIvApqcDqE&M`;*aY44W z@{&O&^Co;5$o=*7zfoH5--v=_(dd)xNZ9pRG`ew3g2*6lg3k|Js2nwHx4JjfZWulP z>64RDneHt3wr9sb;ZLmjdF2@T6mAd)00KC#@unOw5sm{I9lUQ8G)fOq+c64TpT^r zQnbeWg;!1GWK5?L4=HK>c>=ECTj=9SOfZhKi@Xz1iN7DXhHj$$(aDtuU8(Pjkh_Ob zNF7A4!&YKuFEj5m+;1&YhPn<`uoB+S=TXhoQ*hy{$^{psUS|tyWbV^UdR{gr|2U{c z{tvV#X$~}c3Y_ca-XnrH`=`2}#?V*Me=NJ>vFvT&t`zoC$8VsZ@JvDvrGm({h(cuZ zi(vyqsRNNF>hMlL=b*p~i-eBJ36%;CQB!4=``2h2!2oIrgqoC~pSNHVu;6vTfb4Ls z#Y{OI0x}6A`1(hrX5jd=fN2_9rP*v03&+6RsiX{A94u$vD7~AfANhH6OjON)+MR^E zjMj2Tt$Qo379Wd8b{CRy^(B9*fc$G>CRqevGrm8@flt{(UC!Ds7)_eH30CwA-od%y z9n^!z!corw3$`@-YF}Y}@R=NHG~Yl#LJ#}NhEtifs0mGzIRGu&W{PHpr_Z{-z}jAx zY{ZEDsb17gCjX5(mJ!9Dkrf)J81Xeu^7=|TytBI%x`{<$&X5#7v&M&Qu(vdiDxdq2 z9ufXJ=eFjm>D9OTqDK$&OQna}`%5SZoTJi8)+aQq|7}7e)*56z6F-B41xmMP6TtNo zjg&-=lD=VqmYLhc(czxprH6a7eXDvcb?6w$5Cxf#HKX=0gUXWZZ8@QB8AG-|G&VZ= zd{>+hQ53$jgO$FHy%9*_KHf1`24<%8KpMu;qz=WI7!TyCJ>#IcMsWP==o?iOmC*Ej zAFF)qFLSAPWpx324a5@d#%|qJL8-sj>4*F%@ZOW_Q=@TrxWS;SG`0 zMaLO(QHIzL$GrBc#0cHf_YM?@J@iE!QjKJjxGJDRp#2Kw2bcW=NDC!`DB&s)r&??k zm_ng8(m;Y1sOzEl=H1-=a8lu~uqlUvA*YBN`Cf;*v?6>gWe=X(cRq9PNEzRw#^XB; z7SYWv>k5%^4qzJ@dEk`?jAMRQdArw82Ck$GUeGO&DfSGPo}Y9}kWgx|r@!1%A|crX z(8r>Lt@brwqMCxq9t38w_MLSghW>Lvb2m^ zbT0^ZZh4*YcAE#C{)7v}U22@r3f8cnyY5w6fq)}ZXB13!LyPH{SkRhwbbXfW#qf1@ z2x;%%N6B;TKsj|isGkFe^c(4h042|DWzsD~WoE64l3ZCIakcJzeSogM!T5B{ zaN4ZaV5|rKJw;WiP3W`ki=II5fuzO9kOoO0C05{|4+W5Bj{eVb)U>q$-j~_kh)MJU zlN#qre<3PNN|K(}>WO|*z#^k6`FiH{%%iAdqDE4AGW)I6@ZjbNX+s~S5G&}r#ZjP_ zO0!VW;P~U4G5%j6Djxi*IS69?UIWfMunU6e>+LUsF6vw28NL+!P)Fj`4>9?_+0VO< zb1cF(Olgb^Ge(+3i!HqQ;^v)cZO#OM?3aaEd_Y4(gS+~n<%N8}zfGER2ah#JK0@UN z9afha^Y6zcWkYkOlg@FhxS-&Zu<}GUog(@8$Id|87QTBD_Ln4-H~=rLn^K-W=g;b} zg`wv#Kug|SF|JMi>~Bpe6jjmQj=7|1@K=*(_s3e4%>MNck=x=!1O7)E=uZ^1s#;Nd3oAH;FM@tQe+yE1zCqcd~#;AqFi#<J9ep8_O`}7n}-Ac0LZqoFf+` zf33O{g=G8j>6TfV9q-ugq1LN$X2j@R4GuS&gP=e}1gvPBH31}bj(on%?&@2n#99lW zCCNKOzR3;+yRXeXVDd^LeAwT5T_Se9Uew<&BLCDj&%fTE(Cp3vcSqnPP<`(RU$E?B z#zda>VT8I9^~QFhqY<+a>7$95r1OIZwAA75H0&T9+)D1l-K#wLa_h~`tyx*I=gw_? z5NjiFzHV(N|2XD9!Un7pUVxWtt%S5d{M6<#{;$Ru6s%z5xsaXWatDkrLDOe zojO|+ju~!`dnH6y$#>_sNh6TwPC?aE$K!QIQ}>=1bPDs#K79Pwl2~JycKojF=n%E2u`nZ z?s{YJVo1iz?JgDgS*@V{@N;WsLIe;+U9>4KxygF~ z0zfK31^Qn40{W%pN;@WC4)Rnbzm|r4n#2c6`@`hEe!Mx6rFbUm0rOOxjZmC89t*5N zQysC)m-KSP$C)YFq8{0wrkm@-Iu|@1#-%54+d<4D>7#^8t*gFmhpI%B`!ajzW?&x7zWK+@$!cfkSr>x<| zP(XnG@P4S^-J0Sc)03axI;u~3hbH^IOZFY#CQ^FuQ>@mxo9Sl>o$l*>8-nM7Ko%kE z^QV?Bz!MF9h2p*D-Gai1uvTf(+%@_>Yc$vRej{Rv^ol;4A#drM_wQRCE>G2ny7s2< zuIb2rN2T4FhNhD`@1dgvKp#Bp|4VPGHS_T&4UCS; zr(0VDZ_9HH-RD;mp(zI!Y3RErX2qx`PqGpk-%b3!ip2bH#SE||z~&2HU7wW~bnc0q zNl)-vnH)+=oN}e~MtYwKeAy<;S53F8tvTS4HMXfmO!eGlt_yX}fUBep+_^s1Wuz5& zGJOeEC_)}^YMEDxX0tiQ_+V;-XfYh$X1%4q!R%s~{(`Z%7S=;4R;71eLz`q-dBM%p z;wH@;DlZ*42~!J$f$V;GH6Op1Nxg09_#94YJ#wZ`rB&JEZB+v??ct2UUoD^U47rI2 z6`PU$m=8&FU<Fnewy?c@Y7;m#e+Gbx)Z1}+#(EK#abQVT2hOd5?C?c6iFHS59NfbTvKNCh)0 zAbKig`2Zdm&A3rpa6~1u5Nw%`9Nz|##DfQV!DAjGl!sm@vxqoh*Gvp!R1 zi;IHCRGEadS5Ca*yxYpO<8SL9WNxc-h{{NP6(5i~tE4F`y!dg)_pEh~Y2%*1QJ;ID z5pZd0U-6Lt$Go+GaRQ98)exDsk`mz71_076B`sFMbS&4OUc4>+=zv{Dzh3Up*~FlL znkCH<$25_57f+Ec9wL3g!Z%5$0OsGrG9Z$hKm>RO`}Le0A01oF2V$MULCol4!kAR)V+ zMCb!tH|~;Z(GDF>EE$(XUsL8>?AP;@cV$RVdHS7?OowBD~2AMD7pnEgzEortEY+-`0R3-<6SNz{q zoPtTyOZwZp8p^vg{lvR9>+hkC*8|QebaPIby;wn_(3ZOcJTc@uwraI?&IJ({ghMAp z(^Wzjqe*RGr+7}QZ72)4ql)?p1XXko>`Q%G5zJQ z>jCAsc$eHM(!m8OOVXps_Cu;Nwe5~t{Hx8N6-;C;u`IYCf#W*xzUI+W@`amb{%;IJ z8hmwor?ZZm=vj?+y1KqQ;E@|EMqLEosHj_b-4;ILlh+FH4`| z@~XJs(x3T<6&cAssUll1)k=KMJUSms2!y*3v8TO{+t!MEz|EM$IQKeFwtocbr#Unt zGIZk-V1Q1!q1=Wh)zEdZwtqtgI+~ve$=*nSM@l>!?>s>IPT^<&5=(G>-isU?Z-u8 z2}?jM{0pG3C!GOiH*npQza|E}O9)7NcxaL@>%o+V`YmOIFLVtg`r;3ecfpf&9au=tyq}<|ir}yKb(u{O#FXP@Yfv2wyq=MFb$^JW_1p}i1CSkG6 zz;8($EJDpShM&=CsrREdQaf@|hkg_mV56gdBfj|&I$f8wPTqJWQ}#hKTRO|aX}5Sy z2x%+F2w0gcWrJo~y1F83$6dnSt_ieQP}GN*sOP^B#jssrwq1WRF_pYGCzZ2B>o>i7 z=_it@NcaV^w?qu{v?RPGFo{Ah%Y+?Hr?*qFlH5>+;#jlZ!S8B^zx`+%EC}+YeNK_k zvdC(=E4o1?jY7n*_J zlo0*UF)6Bvq{Pg@;GAf3`Tee^TWb4DOQ&54C1h>2OYgRr+6Zl=++9N<4G=vM%d`O( z2BgaKS%mkJ=>c%YgrtJTcjs5OfoxHWP_1(rXATLJbAPxzdivx^7|LKDXKkr9=PL2y z*AeWX&8`V;t>3SWx(|f@S+7`WfpLK66BF2W1HKP$LS&+?X*Tg*Nsiadylt0Gg}obibwDxE0ppT}(Y^(DK8%p(&~ByI$>j z^^(u^6tFVKgjT)%N#B3pzjl>coh`%n?)6EXYO9Ol)_ zc-yiogFVev4`ju+DR*<_pcaBq9o1wLUh^4j-F338_sCyWHFde6Sr7f<<7Uh9MlS7F zR1}!xU%P-+2-a_MKh-z*7oG3Y6QpX1#+8pjHWe|5@Gjcy&0?-DYkSEJ#-E4tR{jSE zO5H!1`{v7CDmc;Lxgm(%rP1&|$a~MQrq;Dx6r?vP(pwY|q$@>GkcdbV5k%=FBGRQw zCngBeJ4`@9L7LJ-M`@uWA|gdfr~y=@NP-ZeDeJt<^c7LpWptme~4bpnUOKZ z`}A_(4^lsTXHVgqB^ESx6_|F$tu1>h{#0Bv=51)OPJQy;RWd_Yh&~9l8%bpz1C@Cg zSYyeB(P>+*joIxv;rUyz_Ti`&w%1GQQ>mj};~{-1g>NcB-<1_!6-vZo*my$6AmD~i;{|P zt{U-s-ICjP!rTd2O}M@IZylKj6y@8vp04)#r!Q(OZYsO{VolR{XyLD-Iu^M{<^l17 z3l6ozKyvbg^0%y{nBX<;;IgzaXXA>>kgockftwT0UrT&*InBItf?O&iqY0qRSGaW~ zyo>nB!PQ{}1DF=VyETrGAA7-Uq>-CVOCP^0>~CtK-a=RBQ7K?p&K<%q50?>1=g||hkC!&I^m?|WrRXpejU=|K!X^d@P5|TBPWa9w8 zhua(2PLHEX;2Fx!ho!N-6a2q`2)GT)+Eg}rT<>Q0*mKejhhA_w)rW+~6#`hkC0v}HpK>9;=3J+cxppo}a^ zbj7YPZ=N6VDJEK*v`!psIvM1Yw6-NWCz@O|;Yui#+E|xv0}X&n$W#!j9GGDQR)>?d z8P+`9-s1LBgCe`w?k36F$x5}ni^b+x&gPsIi;~QoK{Iq3g9kGG0pj9J80d^qr|nz| ziwMtkM*YIEeEDn4ZDNU%wfpj#4vhDbSmTDNh{< zz}r9xy*s&AOf%yC=wy+0=8cXgLe`d&r&E%s;=x_MVDENpxxCVG$xAq{Saos9(c#17 z*dzG{_spe4H@Ex7_l>2pJm~%u0vVJ?_O>{IcqsZFb>axkrw#$a)2wakdY>SyGVfV5 zAJTsBrI1WGm)V7!`IBEpEN{QXrm<4XVcV^j?^cU_`C8jizsPTE$s+z%8l8h8swo3Q zsC@>Uv->1j=Jj@dY*8HCxp>Rac8bUBz~|}$W<-H~Gi*sDrJ9{PSE!#0PDee-O6*+& zJGwaQ^wg0x(?{krlDOf7pFvlowGuU?U%Jp`@^UNd>yK#CWqrJ5EER=(z~3!BN$WN+!}&9oo4T92 zLJYSo{{lsg7|NDsV>9M)N5kTnaMt5n=T3P2f-GndZsE|o<0NIkqY6zaFi-C*K<{hpVo3YsP8&U!k(kX8`9oNSh$X*3Ke+-jbSb{x+BKJ62YY4Ndm$Cx0DOD1; z%iQjj8#}B^xo`nKBb=LCGD(LzxeT}sQoGg1)_nx}+VER3hV{4wPZh_-YkfwcW`s&3 zU5Uzjy!>=Re5A*K5(8M4UVr53a{6)Z=S+6TLWJyy`@gU&X`QOBlw`$UKZfr%p{Ml5 zRIblZAL^eqejQ@V_uE#?)_y}4=SRLeTZNn$wE@X8Jy@%$M{TMk(d@6x6%q7tZ95_) zg)j{)BznO-ufLG3$=lu&_3c-kujAM7*ZBrC9FwLBY1g8vJ6KOdF;xO|4?dy@eu2tf zhbwHsBL@Tu*(k-eog!HKj@AOx&xdB<20a};1I{yzc&E2sf52+Y`teEY?xJnu@X-f! zbb1O@u=6(QD%pHi+jAYr5?&^$(>Q3hJ-0kk=DGTG?P6gR&55cC)~c#1;M6!bzf)O~ zfmY^&7*w6opsk2E89<9j68~CedJf2sji!QNJGE{2_GfPSW-T|J)Gq1Ru)S+_ZH=xa zU4acj$y~?suET}wb)Yf=y5uoEPz8CKJ$TXXWL4{=kLk6(wIJi#uTE(x=lu=YtNPhr zrSGgovAPZpgw!`@wHt~8gtU;VbPLbJ-^#$0WROAD(+6-U0ha|BZ`sa|OZ8q?7Pa@t zpR^nK3U)~5MMQpJmfn9K^*obMPw&eX13(lYRV%n>KOhIZQ%Kxd=ujn*)N2PFLdHdW zk2};Q>@p1s`wZno)PJ=8)K+*g$Mh*rc5q$!s`i$AB ztyF=2j>}u5zDoH-YrUN2(T-?m_LID0!rfCjOWEnqHs~aW0m2V8T{$(_4J3`8qZ~2@ z-Xo@|IUrZH+S;1r`2Ha9Lt8_lm^DrPb^7R#w44K=107V3SR9Ic>Jg!(m58n;J3x4~ zb+H|`#h(>-f>1jhk1yIMoNdeZjC%X*<4u+J;}l^&;XW^FpR%?m;8m!Dzg1r6oge$H zwq{#E#R=xL^l{3;aa@u=Tq;DWNtMS~JInjPgvw!XY)NR5z3x;IqMsp%IH3L1%Pn7O zS|~aoD>$@yHK|-z%p-HnC7(CguM>TqiM=NZ0b<}@P}^DkBF02zy+m^GLJMfC;;q)P z%hhWMr9Kg#w)HL#XoOxE%9NUbG+(>_(W7eg*NMB_DKE$*+;=K_Wb-80vU3b8B|P=< z?r%N1JyJ1eL%ce-=ya7?S@_m$7V~_(Nb5z_Ni=-(6?+l;rm%az8U2rZ(_;H8GI~vM zLkpksD0gp{3`I4GhtWKEk!^we-$Qq@35e|$bwHHf2LnWuVOl1CKbiQ2=|k_-N4KI) zGq!3yS`~f1yX4tm-0zj`F0~{mBpLn!TL>%Ds^uT$r5-jKo@TFkBiD0q`ilTQBwHlM z=Z=}xooJ2^4LLkUmu>(Xl#@vWHiDOne*wN25x z!Q(dbZH*~PPgbud_~g2OZ*T8=)Ru4*lm+PyW~2S|iIMoWe!K5!-Iqq!!0Mi-Ky_$^ zM>{fh!Q}09j>;(4JCZ}UT^`%prwVucEm8omH4JADLfO;9@nCqewUGM@Q>qe~cHNAr z8e{JTbH1EYLkPTh6uV2F>-HhK`md(xvNI2FGJu8m-(YN&BjOE3v?ZO&Wd_FLStmtDa?{a z`J?>-Fw0~WN+o)ARP30|a+h{3oi_}!{L0VDj2L!B6>is7newqrp^?y6LY{WK~vcNaL=EV%S&mQF4(CuA-~-XW1m9f?k}TPVJC*$ZQZ^81B! zh4p1IzYZHYVbhb@KVjRf8`@%lbCCeccgIL+@`? z&XVzgCFIkIRrVL@@6CwGOau%V%}YIg<5JBlRMo4HnXkvAx7=(1VgsN#0F$&@sB?-~ zy!KLh*cdN5D{%Q%nZn{TrEjfB#X3HJZM5p%!wp#(CPBJwFEGxD&I8*v{nV51D3#4< z%-%4@te;}v-_3-}i$5%QDba6oS@CtbL~U~|OXl}{iH|?c8|YY7qclNw7lrCY6SX9% zEPVJWLNBGV8vJARaiP6@NmR1~D3RP43wu5+8vg1JETPhAUCZM-hrs&Gf^3g zqnw#IgeHkI*aUaf>R>h+9UC1Tmadd3r)c;3o;aEG@u|FMB{bU`JcTo2;YL4INhZFa zFuu>u5Bi%C{hE__c5gTJF*s%#$`h(E{ND6|J&e=KNo~q0a+J+B-{VZa?ANx8Asb)j zF_vJf5%XS?3nB?j`ydrNDM_gZQ&3A(Y2<3t@x8UdaGus=9@09~&KWE;*-43r!oKSi zrt0`y9yXoW>vq)VAFE3qlL)&0$p6H~#*PfYvE2#QM+Ot*gjdn3swlh~6NfA+XK zXY^#gsHFSV-w*WpDth76z~#fVLNs^_18i2^Z_^gdLfy`^J?&svl57Gv>}S(yPI!~9 zltgKx!I+n1B?8QXHl+8#;x7-5z5Iq+g)sgARG@F06@0!IPJR#x7gy=I*f#;zUV24Y z4_S#!)RX7I6P1QdQCVOlZqIIMVx-HP(X~Ax?-A^_*tNs!4ITr%YP$g&jX623eRt)Z zb^VR*T`cN@{F5F|AMM)Xa*x)(@4Umy_q6+r)fSeno3U=?bJPxg+LLDStv&lhj1aw` zXT;JYL$fMO)rA8eMt=fCw`bRCO~4Q-LVH@|W=HPlB^wM|vONd}eo18B62CZP8Bfm@ zIlBi~LD;Bl$aT~L7-4*l&t7vM7~?dX(YD0NO%nFE)&^$=;l`p?7Rx1#3a+Q&1}+L! z))@WL0Vg@fCK0|U0cl8$aQBlALai5hssW6q(=W&>;{gZ15QQ>&9d z)#Q`4nJL`QbT5)Zwb zTOXTPwUDl)CN=ctda8t!q-6PK>*l0OiTN->b(6mrY_wjn9YUhLVww6(698fzsAPfLU*q;E_YoUKMkcmY^@S?Y8 z4-&yGOi+7~T+r1b2+mX<$xb$Dh&RM*Bz5Wquf3k^|M_;Z_-$KDlBSYo-&k_s$5$^q z`mzMiLtg#$zW_2Lt^#NJ35j!nB*SnwQ37Ob9B1Fz$sk9UL4F?ZHq7DIn$S;YjGg#e z(pb!g-raM2vGL-Zh`$NK9H@&8v~`xk>j%nN*t4(y2A85HLcs)L13DS)2X;;O2FBx@ z=u)k_{V0k~Px_KpRU#iO+$qj8nUi z*8-8fkF;$7I&C?n4qTrTv^}tz-ex;{J1foTIq)_KD1tSnp)EA;8v@WJeG1t3qh3D7 zXiTU_`k?-tTLvA(J`+#?2dn;n^ufEzL-xud)DeR^P~uDcskOGJKRY=**Em!Z>C7e% zFMQonZNHhVwhYs{Ps2c;VGNXKeSe$)8X7SXNN;33q&YsN`K8wCQIfXnoBk?O#k!VP z3k5Hs=1B&>c#KwJj%-|w7;Lz%o}S&YCW{CGliWjW0Y5L|tq06#$qtS(wlmuDF5>T8 zQMQ=j`@3u(>I*ziC(^sKhsvjj+i)%NQQ)sQzXDIk0R)AQQVGf^(OGaYJZDljA2FuK zNf-x7IvD5oia;ftzBseS#`fOVj#u*Rw`{io^UJNqtTl!+Xk4%%fTmra-MvHd3VQ?) zzrm%vAY5u_rwPk7gs@E!Q^~p+o8t}aT zfT1Bf%3g-A3zSiWk*A=m%X;&7v1z#WXd7&prEbQ)Wo=jq@TK9*^?`z`b+Yp zZjbh_vByA%8A(8JIgy@bacJC~m3R;)c81~>?1&&Bi*GWiT6XlFcwDqw@G6%tjz-1l z+9cf^O<9&E>PO%I%W$y2vXH)tC?}ax2meLG-!1Afdod&eo!%){9TxE+vaxSj=;zMY z=!+Q@GWXo+S5KwjynplJnLyf47{C#_9=~bZ0WyE$M~$ykn<~bZrS{vo!BB&(xggE= z&b{=))Nv1h&jCuS;NGHDsGp0#*e|aICje>fq^p8oIqqo|jJgSLAoPdcs4sG|o^59b z@Q)vdX=oS&uj+t`00TI_XkO}6p)~mWIoSdqAA=CwUg6SKUZ*(uaeMwrd%}sipLA)l zNZMhw?7c37(ZT5qUDTEZI86-%2eP*V-EuTIo)iIz`w1v?n)*Wi3}P0>YP@q-+@=nu zwVS7)QgRhCH2g9&Cye+g>9$vWYOtyAJU+`xb~DP|_#0*S7-{$4f6zk6)$ag20yN*0 zk6{#*X~~L2opxya-k}eRwq@n^QEsLe&p2{=qp@O_3WEd&n$klJAaV@sk(ne zmFT6kqBl#!kYF=O5yV3P(em(}KEjy{!kg$VVn=@w&d(3LEN>3^-aqPjY1;OzKLX}N z;J6rRVaRkQJ24Ymr!*rEjh`+@=l43j5mfREv_rjW!j8l-&=*^m4X z=xqqndnIOt?{JTOfeb*zlWUtJTcbI=ueEAi4tWD%kd7suwKdYbCv@v>lrA4B28a|8 zgI)mAi0r-HY?i1aKx{afQyCW1b8V8!5%FPJOEsKf>PCr>;=1{h((_qj0U15qnQy5m zz{d+dNDuP7Xdc8ApmOIO#+oH-JLAw!IHovly4r=Xg4@%bx}5XHVYTIk`w~MYAHTC! zF;z;*WkD-3=Kz_6CUSLCnJBY2;72@XNInGuE~=ns!*;g>&MfIPUiF9ET>JE@iG6iw z+G3EWAx1A&%9pw4!75~8>m1!N5J@~Ybq;v|Le^7IoGlnCLM`kOHIf0RF+i8_0r&L) z{unntV5+jxTPQ)qx<7I$_sowu|8B}~uFi+wzT7;7F6pel*ap=fwK$z|2A z2PGUfUvukqH|J-+mL@mpebc$twk$t+#LPm6+tR`<)feJ_(UKzoMKA{i#1}$qlsjOC zAeLjj_6|wur#-CIlMuZ&x$VODOCm#5+b5;^gP^r+&rU6z2{MbrsZAjuYW-p0REC#1IZQ46#XeD zeU{mN3%&Xir6~P9sL30 z3!EcRqLwn1r5Nu*_(Nm?$4b2O3Yv`S7TMBRH=j^ExohxEey>qt^h8jG!DslYSg*r* zdUOiOE1@YKeeAdmQUi_uPAJJkk0npspj5$B{#VO}R-sS*)Yd zSlq8!?kD`IC0y3c%4`t*PZ41a1ZRRGN}>i=V$nE1fx+=;aoM&>@)i84b?yAuP0p{E zKI%>~e93v)dEPOD+1PBJO{H!(+LB$5kq-C$4v-C2>Ri5CdBDS!rtbr``O5|gBOLM1 zMGK6DteDu}2+;@gQGm|(b`|07CLW)7=mv(ilhi zl+4%NSZMmKRDJY%Pru-&+w|^_co=PHv-7FZ|FvkrdJa{9Tm@l3E1?ox2Xs55gUL6D zgC!m_5rh`3iMP#;!q2Z477CI(<`hb^OiL9qTL5-zuYYW)(rypk4>|eYlK&7GsQ;-2 zN_7JE*zazIRptEKm}568nfK=0qeTT#P+I?z0wk2z;Ez&5U?zY&4~ z@Ky@`Rp*)9!4vcX^8^AWPXM6Wme5a!z2(~NaahXsaB^sD(`uV^Yn ziQmoS(!8PxfO&uJkwptg%z!|El?B}G_%vJ|hxx=|KR^|o^?~2achQ+|F{;JBHHbjEm-1HmzSREdvVjR!%T2b2jK z4}ZK03?y;AfB!~y_jA+_b1y3&3DMGLZdTVmgfj{8Z2*_~1590lktzak$^GGF$o(rc z@M@qE&xxVt$IYv3Ag369ja;bjY_NBSU-Q+8$0Y)9f1fq&@ipk}XOBtaK;&4zBu5f| ztU~782@hvF{nkdSm08G8+{qW|7us4E_+GGDopX*89JX>xQ%f~e_)vmAhhDo_6oo!8 z?pC7Og@W}fbEaJmBpbBNIzE__g^EG*A5biYewlUy{?5uQ%ocx~$C2ZQ(VTnCA2JLK zPshM!;sGWkKsCU{Q(5;a!U_y=jztOD%4_3nL&`FrTi&nR7bHvE)|owZ{b#CR=DG)0U=@y~XkS$AyqlAZ_sTyI~y3aiGsB+o8`AWrvMCaea-0I(? z1GpJf5KbpGUvc(=0WULfb9~elH2qFchsf9sf|=;%q+(YMhgwZ#bFF^Zp-Roc1vs}Ejy|vjtu*gZL zWYEkh1HXZ~b8xOAleCS`Oa|?p8eUxzFk4U?kgyl3`6ZWV z6HEP}$%uC9zKHmQ1c9*MNz5hy0O({|T}co&O?}_l+K87<^r>1{_{5iSBlA;Q70EHr zU#Oq^#i^^1-df~?6`XG#cq%Y|!bhc@-^$g<_nDZetE1O2S5pxzpH?bQB-Ar`4{dHm zsI&t}dXx=}c$R6WF+~XJhKiE=5Xel7EDHdOF^78!WDf}$Il77mppWd-A7i>yz^H{X)u8HN0XK zNo=a(yeEuLph5EIg(7OYP+725h|g4oY%5C60VKISYempk`yQ3?GYF8{p7Y)L#ItnY zQ{XG6hGMBJdcP}%$|3?pqlc2vtxlcg*I^<$;y4Y79P3blg^MqptNO85WK5{Dft!5; zV=j54IqZD6*o>2JUvbW{%Krk5Ue*Foi4cSzmz1sGu2WP9G2Iw%okF%(f7Lau5X6lA zlr9qM0sr=VBGjm?-cl--=7q1&?`q%?aS^jSq^j+ii)}?5Itrd4p#t24S2i+j-sfgE z(M&d$eVsdZHo@F@h(*x$(qd3^fqfs7e5QWQrz?W@*v-CS!jDmlS_c`sE|jQ zMV7(lhq-&Y^Z`y_GU4P)N-lB}g9L7htdVc{FB-|uOzVDR5k(vc2JqHF)uHu8NQ8Jp zs7#_-*R63Qn}-`2TiIZ0-NEqdRT|3no~vVJ_PeMdBUEpwLE=1MbOFGUH<}nH-90cDC z$5B)p8DwnHM({HeWK%pF>CWA0ZIh6heCe2S0~E4K@Pp8jMq|k^&-;{*9~m8aAJp^G z9-ieRb^(QLCNqFyyD2_&e+vCu0&R+1ErPYpMgc%rNRv`Jz=B`;GLV1k)vjl_dr+a{ za%MzQtzy}hiuSDD7BxRZc`c>+X6!0UZ9>o5n|J)2sK3IYW38$9~NP^ zzD+3KR3wBx_Y=dMk-nJ6`eE=oqaty%j*FO_cqhCKX*Y@v=5WVKfskbPFZh%arv@S-Lv=_9EdHvY16cGVOa3~Pk# z^5KpM7)qYD>N<+X#Irh#ZO|^<-Y9=g%2PNu&fYyCMK(2O{lO6JwaeBZWIS2;pUMC* z-GunCOTe3W!BXuS=a$9fpsW~-`sS1FPfhBW_JIdOOJ6OYDFt53PGO~{duI_Ir4kKHO6FlSdVo(}d>myl3i#-fQj%;;%v+Zo}UD?dE{v z*QU;_v=mlX>b)jkeM__kT(2-mBb*n9j(4!tS+j*6l5-$4yVYHTVdm1L-?LQlLDC&C zzk&D+A^HQ*BqOlBzMn)r4a#jh0L`Bw+IGBX5zlE2Hh_|#|L{8+>dHCP8s~ppr}RI5 z^MCV|YtivAoSiae0UMg4a5$zxNF$B}N8Vp3O}EBdE;qOhzZv%l+3!6~2y(s(S>WI! z0;50=1_oog4;7Jv!tPIyUh%a(zsPYSk!2JVSwJ$NdS#+1uwhlU@R z`i>9eU#Tm8v)CW*@|b7bPg_a+QmJ!}t7J6u>2Fbr0Aqxf3@k+Nj+5xLB676^&8{A< zpe#A**Db2&?i2C(eB0YXF*Eb<65|Ukh3bWju1tT&CrU`+ZPF-lmZU|7Qp(|~ekEiB zfLYmzEVW%?TZrW3@o2@=)(FrF(T|F5ZZ$H5n)f_la(;2j$@!mEXdSuB1P1!Rszbv(;1bMMr|OC4jw<5mNsRI|iVVqr9Nwn5 zW)lt9WYW}@gm-J5XM>*Y@()`qL`~{_`0eN}mP zn{?Boh^>~aweEMq+eun_#i)I@wDn@!?1PH5X*HF1wh$T5d%T%#mrg^b{$0eBA4U3J z;rTUyZZI3-&@gRkpH}D4>lOaeQd`?>sxYh4_4Gq#$vK6OAFr~~N@Xxa9QP7VjML{f zby}1fN0WU+UR4V|5B6RWendtc)iGZl(2pFNPU zHKLaga!Wct_N4o+!`CS2v`jL3^?6YfGM#*QnqI9)XnL7Zr1<;33@@uQD!FI&& zO(_?9qrxNA9EXivivUa%*zDp=0x}Yc>mY=DDQN%^M#C5N1^a@5C;7|Ep7$F(N<)Xf zgxa{tl|}c-NNU}J779~80&qDb&Z203k8qxn(cERfuJhy#+zf{vFm`r?^!b%9Y$dK) z8pI*h^a8bGM9-Au1wEtf?1T?`@1v*^OyG(9002Nj7)XXVR$@2lwxL8-f&PfJcwhpG z@akU(MUn14HDu%IJ+=-UUo;`%8)HP^h;+PB_)$bi~=e_tU=3+$HT ziU_+DFuJ#{9CNW11C4p;;>hK3!Pf8r{DOLCz;`nbNpS~+g?8k1NMw&i*maO3Em|psTn5&SUD~NeyZnF~2pPDmLFPS9+m`es*_wcL9hbxW`^-4WL>OupyQb+&*`! zZ(|(a*LUIJg{%|G=;pFmeieT2PE(2Dwo!(+eH>^dc~JR@2IM=TNM}d_x>X>KECZ4~ zLdXPyf!HxPOcJ<>=XK(oXkKYHQ!N>O;MVxD`n%2-8gguO%m^Uhuux#Z}VeQ5Kz&Ngu7l3yx4R)SEYx?`s2-Y|M3k|WY~%b+N~%a+B}r% z{u(fhf*3PF81eGLyXF?WQoMJyp0}*U>+8l7U+g_4{b^P+ogR4t*{h`FJya?!{q3+fsGK6s zLlnWvf;$%pEA0nW`I5{tk4@~Si>G$E7qNFoRpl66dW1$$z@@|&P{+^*{zV{c#&t}B z;R*^3{5KF*O#u7AfzU$EcBg}_=Pn&|x)*g9>d863-Ggj%-x))Zp?ElVSaJ{I%2d{+ z)bbg6JWsZswo37k<;!6y%;N<>`qPTU0-z{l?_bE{?b`$( z$~Qr&_F%hw_UqAZ9@r2@0uG~LF_2rDxS_T+s4T!c(DM_4uc42b(Y}oebGMm=JFi1{ zuPI2)wA{-y_H=H)#8DG4{FJsl)$!Ta#Ji#l-@ijKU@n+SlqB6md@hv6f6G}%3lUF3 zc^zAe14@73bRYR}3_NYFOI31U3bgl>I&te9-S1D7Xz0ct1RHYhs2K3A_OVhoxA(0H z;4Q@H>99~gpqu^T5NEucF5&aip;+DyQ-g8i>3h+b`MT13pFiyK3-0&}D)}1;p=PT7 z2Bnj$0}0z{t?@sr0RDlc`tE$OeS)2aL+iaYPCQdYE+-L!!Hmdq|G#MT z0GbOXoyt*5iEBn-mSHXj&oX+N9x$smxc{&vS# z_U#uhPSE}aX*-iEIA!nu4{dfk8nEu+e2)5HM7g42Hb?^IO`^8(DvZCjWZo;(?qUTIYS02nj^o{!4u#fpn7YyoEqMuiGW4Lq9JYf8 zHfE`=PWwnji9XEPr1QUS^vFrbrG8p^r~Vy$-7n;P8#V3JDv~M?0YnA>6!@1kkQwky zy2NTM8-yzi<}g7vezoQCplQ0f8`~|slWHJoB=Kj4qjRslE>hB^V zp7IR2{U!;Sn+`+s?IX#2Z>TJExD_hvFI@a}u(9#8#gxq!SG%y}VMoo^bNJEinc`sB zKs8LWOr&L3VfW}sxfg}x52hF)d!McBC+_ipV= zP}62a84dEO?_k06X3_j3BclDP;GWdf&R6eUL-UTuvzG@Cl9JV@i~4H%)M)QASnAVh zj?f(09k0|SC==zG$mObpTreS8y_QbVTC&b019i)KhlP>YFv!gf)F!xdpu3DD#1KFV zM?XyG37Z;%G+U+Oln0hITu@#J@y^+UQNQNJaj8`GeZx?R_tk1o=GHfsA%YtxH3&d$ zWBk$gwUC(PK8*<2u5K~hZCY)i3CZ+SQQNTW2kZFR&O5@~_-v%X z+5GwS><#(OpB&SohfL}5&^6|wr|1KF9SLwYt>ff#bX#bmx;NE40Mv53TUIDnVAh+# z;h{UJB>6B1ftm7*tbv>Az4-kGEFsTY@jhfKf`fR34Vl$uTiatNPU#zCzO-<#-Ho~3V*8h@ z$H)36C7h74!v>>xNxK&fs>&<(kM#3^@NhM&ncM)iDN}g{qfA{lH<>7qi$&Z=C!mS4 zE4z4rIxQ=+`+S`*AmOEf9WQ)rSQsQ;9@0g=L6W$O+b8kuFiVxXc$J181EK|vBe*ybjq)$if?^=Slm5( z=>7O2#&+Iu=4_U`r3zo_+__6<1lMV1Ux2p3*fGrCA5F6-z{@!PK#~zrSNf7rHQxLP z@kyl?j%!=XLw!R6Z{NLp`nGU7g1WX3^5g2)*L$S0+jTQ;jAB0%dzr2Bi9i zHK7BRi^Q2BiRv<|<2w+N`6iV2H9jvHaoP*Phius-*u6NvU_ZZHc}HwyzuvL5SFUV# zQTNgi=jQN`8O8*=va{{utT6xq5R5(bTQxsiWfFord=&kJPULYvU((BAf5!$&zo2K| zT+-=&9T#JJAcnTH2U8tcQoFZIl)rA@r;?fpbN8^Oj_{oPbO@*9+_4DX|aH2lMVQ&e=m?Cm5cCYw{yTR+{&A4bVI~Qt+V3Zu+o}!Q5ihtInAN zlKw2_bOT$C$q4V)+qYTF@?t-6!d@VNi1fA;)qiTi7xZ6j1?Yc-7(A>zB4hnpa0`eb4W)Uo@N=lmg@?RtKbO z*wAt?`16NQ%|=wW7xb-li~Bb7O+8nr#2^15MWv=j^K>+`WHbkjsXj00ces$QB1gjT zZIt*u^g+pu13=ilM=nDXK(^OdfJIX-FIex+?v4Wq-arhr=k1!UE$P`%-ZVpFgmY_y z)Y;ORmq0@y_F(w=^`}qU|GHW6wR4dDvSzO&W4j+wTA3`eLJ|gZ{SBa`$;pD4XjdLq z<_Wtn^zKbq$oPP;V{hCV3wKFp*shePkaM;RJGmpt{c{vz`nVSKTMI-Be<43eH5ZZ2 z*}YN(1D|6DK z6dmYKqh|H<`>ngOdK@%cvanLE`p<@6zSg%6W;eL%Ce*7=REAB*lf_PfXLS4b{NDss zIIL-vd^t2heUCP|>+*Jgn}GS;SSZ63FaRD&YvDW!L-%P)`J^uDO8G$3R`6`f$HH*k zX$AN)MGz`ypQ8gTq#T#Y^e`O8`{jl7Gu=4NGB z&ohSc{V{``0UFH#y$1oINv}F`^#zri=(^5fVm7UPTWs(Xs!}oIBfqL6^qEi2K73sH zXC4jX=~WkcO8EcojVx!%Jes{*tz2DW-=AKMY5r0bN1femRn_b8dnqEl<1QrligOkA zH)w?S+7$pA5zMx0Z~dk2L1%Y&*ULI3Kbg$r5rwP^3O@mM<^`}5bO}`r$e9slesGXf)lz6{ig?+C-e(n9dttw*tf?dp&u%fIcfgyGHSo~ecw*A{jn2$_jt`lopVZETx*9!L zxDf3A+t{@1l~(*}Jo(xv3B*_v0&Tz$j)HIA_&`I^2bS^iivX?8-ey3gnx0Su=7Z){2QQ*|A586# z?eTOrrWQUtlMmXdgK2^vjj|jJc15!^N4bUMUR5j(WT4CbBfX_1gBnb)v$kbb7>M=S zNJfrce{!gm!JGqVE}LQXul?gvA%s6EI9+Bpsf4vS&2@cI346{8nA+UsJwd&&O|(CHs~ir zR>jrj;RSk+XFCm%Yok{T#Qe-AH=4XjmA(&}az(12B2@%3>JApuIv{c%C*HiIjuhd% z%yu8Iabom(>j7e3KCFM7Ww3Ja5g1R@IX$I*wmwe!mtbvI?dYgNYTNyDv?eM)gB;${ zP-{GnxxXpkH5Ze?W9m_v)A7;k#?z>Sj^wz>hvm7G*wyfRMR%K%JZviIhoyWsx58ZA zESRON_t{{W?__>^g77Thw*;J-(1appvTC)B2lWh|-m1+mXtrX^Pl1qUzEq~dcS%39 znEN>}qengueV}93|8&{_@{v6WZ&gv$d^A+c5iE}{8No){f3FwdBv`$*Z#lL&-S-S! zIZ~nq-05czf?;U@Qw~k$)>?;txw}upU0SB#%2i<{LyC<3{q@Vb6%OIXavx_c z%={z=s*IQ`*e>+*&&P6*KpPJV&ejI~l^@pJM#)0)B7c-@p%X4sSD>pku>E9Zj6HS4 zIth9(5eysnH#YMcm>S#6Ul*v(b~A>TD6?~bWtkp5lz zT;M$cJGB=a*rS*bFu)GLOmbBgJ)yL4+c#x0?Cs>4cJ}zqw#L3OBAsI<_@8XizW&gS zWDw0`U5b(Wt~3bv1k8XO3AE`tbqdyWHn%){-7gx#t(K~6Jd8P1pOk8!S{vyBNT&&E zhEh-Z#an*0eEka(Je}d7Cw1z2NN1p|mENZAhB;L}4_Xiv|<6V#|-+?~{ zF?IEAfMM314$Dyg=l^gOGIe;o$@X!UEIz!}WdVRW0W_vwp2f7(L=XsM%~$Mx?D_9{ z`vmB`L1f>|boUyOu?0aQ+>8`V1cc(8a4(wLka4j6V;=9khE7s=v%I1EOdJ6f8 zq(0acLNLxA`rtdL%&_%ZB$ZU%{_^FUhg##>4e_| zOe>vQe!(0EGyZnXwcker=+pThsB&?$adGkJY8h@@((sW>z{i)JGYz(N7_cqi!MX5=uokm+ z)Y~O#ik~T3yCS_{Hm%lQY3CVgr`Fn#@Fm&NRM=EBDUem{gWic(K5jJBh$ompc@mh< zi8dv^xr;}WPcaeTGwQVlB%3dXX7TN3$G#r!_AfdteIT4~GL}j+(Yhd&d*4_Z{roW- zmBW(qw&k(!r(R8*WcpJ`Nk+TR zBOmO|3t;p5g$|7lN5H70iy4lb|I}ZQ|FE`p^3=JMCO-B%nhF7h&9p!M`Y(w_iot{;cED?cNOvg*}DwoQs31h2vgD? z{#5UBiZqSr+q0Nn&Du9ezulhJZ@T|e+9E)znvh87aE*T{&AdEM=0k0HQ%{})ha^kB zMd<>#Nzmv&us-ivjRAxlQRbBAcY}9f2S<2@P8`O*^YfTO^trMGO!c#kdqUzjjk~zd z9uJNu*dl$AetuBp2v~b%;QF2wUUspRx)?mX*(v0A!Bxk(PmHQb)D(F<=z5+^0B=E#d=LE0*@keC zgK;4Xrg$Y3R>Zj3e&ytNIAf94phBTXds<{7k2I^njDg%@v_>Ee(=R62Y&=GZw5r@3g%1A8rOu zWCxxIwbJ4q{Vy6nRBX2bF$e3YjGc`kGp&h0c%8iNc;_Q3ju?l&X+3R860wheBvtby z^y0L>!a!D48ZG7W#Bl;<4-&9=U^sCrWF84kwOwEq&EjqFLwPP0hCcpy#Y1# z3;N98nb%47l0hB;IrK=m7D@)Pw-vN2xa*@(Sh9fIs+g20RD-xK_bM=Oue3mY<3IW;8 ztYk-E7-Eq=oPRVSLjR&!>3~tyVdY3%_2;FB0Wa3NKchZ@dekFqXhJv0DtpDW1w7XrMa)0{Y`J1@@kP``@ZV*BUtIbdsEg}7+DOo@fb3RKS(^(sbyNC87^;-Cyc-EdmoP23ZT2{ zQ-%Nhv8~(lv08l2mvdjgWhF{wOh$%P{>TEI4t@=-fALH)O9q;#RyRO%Ho!5ENpM%9Q$aDd$ zVtEv4W75Vp)@s;X1HHU=(7!%j^}D;>AEP7qNI~CBHI+H^xg1w4iJ$EmJ$>z{=5zd# zCB21}hgAuh0~lzVh@c~B28-B(_P-kv=`t)g+B|yMH*6zd>wCt=fyCEI{nqD=lrNDu zH#gWSlB{aP5(xuF0lSWaMzJGc1m7a4uJBZN-9U7+E*eGduWdfq1l_DPiHolZmn&Xd zz7<+$d9Xz?N9f}L3m)O){oJ_2E({ikV5gb;Xc`6c_(}PWNV$VAc;cDI!(J`dhZd<> zg=|zP8+8U!eEuZ!74;i)x2!&aVQ({iLAvp9r*AVEDN!>Z8bpvbEaX|*GpnR)P z@x4GZ?x@)VWT6tk&0>?m0I)c{E1MCDjm{YD3e_i5Uc?afrp_^;4_142ef~z2mF3#h zngTuJggDUg1loe8E`lb8qF^a3oaUPF`bbu(?H;M^gsdQibG&me_L0hji_+TzwnXfP ze%8#KneFVHByVnRy2Q-R^7ds9sEWQ#E>eG#4dcvqfZ=X}LY!Nd{KjM${V3#2Q$1Hz zJN#QWQAuP7V-tb8Y8yZ4**~k`!u=Zx_bGms&}C_d)+%ZRG-MPt#|r)d1SE@v-@eT#cxUgUKbQpI%9vc*7cC5 z4JuHv_qOx?)SR^Nn&WEGJQ{>|gsCdx9ejJnw^)j1!c9}znqy(!)FX8zLG){>-I3P~ z@6$TufLnEcPto)^wF@8UHY+u*Yw*RcA84UgFVBvxj4h~2j13h@0_iPo4|h?ZXSc@! zUWcHDYyd-2dMXRFVZUychU{skUDC)sW`}8VLu@IN>g!3rZeLwHh(8-BXNV+xmznTX zPT+}oNhnHY+agFK$T)K~S{G9)Z@wR+AH!#R~b#0brH?`taf%ZAT|yWoCas}8_{rW<^oX88w`RaOHD%dp3NB1R*}Hz+a3<2}!%=arh#9BR?}7Vw$ToqM7Nc)DvKU`XPK2mgWX(9p46qo!uq3rGC*+<(o47dRp6B&qWHsFZu6C0WPeN z>Z<$KDsNq9=Wg8r-wv$DlZF3{->#SoV#H}->*6#mdursp#$#{MB^`1j%`UkGB^LhA zZgo0xzKrJ6z6MzRF#k79Jtzo|@r7cG&~B~UUJEs35mB!$pOSBroxZsX0md#cY{_%j zgBJYiFq;A8xO!MFtg!bJO7<$!1f1U(9mUeAiw$;c0I`gyHV z?GjuV%+o;}X~0nc(zhSRcrff&9LWyX8yg-{b91iw46pMZ{%D_+@a6FQV9udeM~R^4 z!mpW&{RJ*U@Q3quya27dVF_l6#T{~1-Nfv9?fh>{6;rl@ z>Jg=YW^qIH$HK{3R30#qEaJ2E7oOs9 zMpCqyscyl+SCyxV6%`}302wQCH{lTiI$Fi%weevHqQ}|NMyL~COfOp@eFK1mbiks}aNAa$UFK2+z{1WW8>o zcluJfSBjpNHfEa$qy8#B7g=s&3mpf$9nZWBzD5TXN}4IsP~M<+(M)Ms_Iu-way$GM zd8*D)uUhXahsXK09AlmZ)n{6CJwo9eq$dsXt?H(8l)BV|H@!~GwQe_6>-#Fel+)yCp!lRH-Y zt4_|>xu*Y(&d8M$enFD)BQ2o+_eOX=`+nI@?JJgmm4Z@$xRtB@oJ6qPkr$}>IdSw| z2QMw>Y!#i2eBP7c2M{MdPwyG=y`>UuCIZ0a*kr_2miUFLMud{tIkfhoZCX>G6kd?w zu9)8J)br$0YATSrn8HKazyemET+%{-_8Jj0(v*1FU#R2h&Bm5Ro-^JS@s4YV+MoyC zJ#^xxQLJ#sL1E4ZF&6kA;|{~2hEJ_$`}=26WesnnWR&9svrn&CxVR`(iCMsI*y5K= z*#bm{5qir3Jc~chZ$>GMAEoscW8f(4NG@>Z2)b!Z{`xcI)$Nx4QQP7)nP1$WT7ck3 zFYdn|Z;SwtTtQ%STx!*;2=0)}Cxk;{)% zB~LA1(Z8fpeEWtNI?Rb0mKaw3^g`#V!>B>`?AXtt{~L`xbNntH=gG9fCg6x~3PDar z34?Up+7$2F4%P^;q=~hY9Y5=$1g)MA%@^1-*Cy>m)1aMb8oHZ9nR5yFG#Rl4Iv92s zfCa|nLzMZ#whXv8IFJZ!ZDA!u;Rg~{zom00<>T}Hvt)An8Z2aBx^E0=e*i1?C=b1x z3J0uM|5XZbEwc=H?bUVt;Nv&PL+N=-@A#6PD&r^HVq-)MD?zSeg3+N#MTmH6=p^#BqeN_2~m`J|t*-68j=JWU9aktKMkPE5K zbCWi%J9ZM6b(81BtK~gS*HJ_zh5FH}&l|Yf)Jh2lsYjK?pG0=x_Q@486|5M6!57jx zx|q5IDNoi*>jypY@-FGx*|$nJ(iq~6}6)vKREeyU^(Br0^N-PjMrgNbKF01tL|Gy+E;jUyl( zqppp~Rr%YJQCjS9fMd1tdWR@G0v?OgS1c1q^nIKKnN2b^8+ zpmSu$`#35%ktX@a+b0jcc6Dqqfd zRz8!zRoBDqnwTJ{r^jnHU-q{iKMhPJ&VfN9%)e&qqK6sMPDnm{+yiU&=;1xV<7tsx zFJUZ!zkz)g047XLjmDHP{BT47-Pa7HB7$`lyaM5phGCpb?^m`aJp%L1oKoy%UWALq z>8a|AiEFoSwm;JXxPQz!9I*{q;{r$my7RgTOdp#5S=yg?Ob9s&(Yl?U^Ug}-!4IUD zRj9wEoR)Uw#%zhPv2v;$bDLF#l}4wKy4F2u0&UN<1|x(xFTNk&Fk@(aOTE!4=O1Bw z%u;JN@EZbsh+qqI-G)5Lx}Z969vOYhTyC&ND0tO@wvbXb(-pKQbIN0Kj(DtnW5Orh zz+W3pQ9Q-P$q3p34FLjZyhvvNqkC~Yn-Bj6-8SDhR%?EMw!1E5z97o_{OtDGX{&1+ z*To7S_PkWsf5}4Ui5NG|81dgH;!ho0cw;B;NFq>Ab-$w0Q%X@<^M%qu+hcF&r&8m4 zdW;Va*(Y#j$r-%)NEdn#%?Ix>gyEd$r~q(aN}2*$eF@8XqYk_<=E9hC$A!E$)24$c z#an@`NN=mssE^h9AsjL%#|pJz!fM3+uc+Vk1mKu^jB)%PLa??jy?@sDa+*rKI&Wf1jk665 zX2P-TxWhaHAFZu$wg5=>xPL;CVD3w3fZi|g7IPTri8N7=qLXR2f88ys8D6{DWY+lA z=Ek^~XTO|V{(fhlQ{4gzscVG|V(hH(ElzM;mI0GY4ETd@8Nfw{n=37pD&b6}bEIpl zh2QF0SM*hbT zQgFaYm5+2??S@Tp3m!BNv;ZUTGX|5RbW%wgoH*f&vzqvLUAM%TlytJ-USSk5%d zZcwPwm+V&CwEdpHyzjXnUJw$uqgMoZ|Pd{$7;vd)os32~5VV~?U4K8fN z;i|z#7ngTsmAtMQYc!h2^Guy-)*5=YyS{#KqnwhqP?MD75_us zvjj9wPo=(2A%{&oYn#tadFL%4%XBD4Wd+yC**=u8HSf>I2-EWX+$J!>Az80#s$%<}Z3p?fU#dI)RaQMqJY2M_K^r6i^8}{#xt$%k^Yi{|M}jKw z8!!NhOi;hUe-QxnmLo3ABp^ivCWT8Ij41_!0AJrPd+SiyNz~2h_IIRU<`0x&Y&D7#jbJt7{s-2@zs^Dyn21;dkRufyDHLngg9sE*dboOWteIP& z2p!pZF=AOIrB=+Fun3sqNk#Uy6X)M5tPZsH$v4oOQu9ULPVP;ys5T!VQI$^(g_xAK^o6*!x%sGF;|ag_{4q{;;Y=Y4-2aFU$=y+kM_jJMOM*h;F6N=ZWY??eyR@P34<|v!& zCw5>%Q3z`Ikv9x`boy2sV3i`*=4h6bsF52{aS?ylo^@GOHC~zG`pMx=_T#9t*Vg38 zJkSo8#gX2@;in4hAy29isi*a0pkdQ_q`~5dOg;tv=()&(DLo64eTVJCRm0Ojul_wB zFH}qY`a;~DE9ynJLtdjFf6Ww74Vf;#m!nL{CzsLn0e}mJb&)-btO@8U6k}eY`R7o( zI=v6kJZg>Tvr=PHBZp(g=`PWrYv!2U zaMmiG53D$F=lY>BCrns-7kQbT}-o&Nbl(u@6Ju-@^{U$%VjW48`ob;ebp_iDw3UvE+u!q$ zQHbHa7$b7sYj^SMsnsLm7iK-%B2d?g)#C%(vp$!U`k9GF34u_5gFt=B1aEgDG7gv( zgAlvNYj?)W3Fxjv@_=J?O=VT|9tA0ywsR+6cWWuYmndW9&K3%@FF<8huIAUco3|u& zW^}aev3kCjo^-aWruyUL@7cH2MmL1eB9>Io)05^Yuap@tH{Z9Dqz{&?Q zu7BE*+q)kl2O?lNe-Jn7m{N7~0qY^)#26eS1m)I@Vf5$cxQI8N=?)d&tEjkXKNWZ- z_bz(I#i}PKUnAXxU(wH@S3DsNvA>1@nc-W`qKra97;(Yb9uea5w z?QgA=QRBHPW=$E{$soY{^b{A?@~{#>4$LwLl7u`g;G=(-qChN990K3tHirPt9#6;d z<&id0*4F^}Bft-npYseKqBNQ_8Mt1~{3KeBd3Rs0*3wrJKjgRvbA-B)s++s&`qjYe z><>}vDz9%HX35dn68ig%yET;-2)Tezqu<6p3nJ$YYk*6y;XTlqcmN>4%a@)D zg+ZxB;ELjhl>{u?9uf9=x1Z8$;06hSqNpN64RR#c#prD4U=U<8Xs@sMB$adhO{2Tb<UCz5o>0NKsbrqX1u1fW^hIeN_y9g$KsY#adYIurx7B2`t;Bqfrt zozRX~K)I^AHB`#s>BlMxuB@wFKNTv}$o9=pJIf(g!XZyX7d%g@Bpn>+}Bt1Gr;hTwRtst*b>MX+|n5;=uW-_ z&TXv)IB@r5J9RV6!4%a14JdK92?>sKmS43ohhN+Hy`H!LDyiRJNqv2HSi@G6|2apt z?RI|M|8|QnpaU%;s|VJ@sN-})DiHLSe*T8!$3$vwh}T?15YCtAzm}njV=QAJ?=ssyQl$jnx8X z3A{_3jPDSR#(5GdBR$n^e@bz&LvHz1K;~ft9uxUS@qnc{1}0IIB)= zXf%SIr-ydy&$Y$i_n`K?rt4H%b@A(e6|rOK{ku5SYzzLEd>l~%%MD_a+fOr79L>~z zOuNx)J`Q!);jYD}^xm`CBe`K8>{YvzvpG~{B|U6y0`|G}Wg^Ca|MPc%xdFEJhYRl! zQ%KnYo)21%7+K$xep=`Qg#ww1-#YsbMA>l0E?d9_ zac$KE5&)Zih0kh%gCJX%?KZwr$lL>`pJC zNtR-m_dlGnpYL-jTSHQ93jb9@!x7=N+$D+TT+Gu5Nw#{RtKMu-uG}nMacTI}0lA;N z&0lzfG8K0BN(j2c&g=aHJM<^=#@vjp*`}?*@X>5rwJTp##drK!9FH3Q4r7YGrK5Bq zZJssi*IqRp%{1+zE2(ZbsxH0hx?-%_1jAj#)dVmlY)Bbv$8bfHObhwwmeR#}~?B zI9%ZB0WGcFuCiPTG5b-AAdslBzE5T_#VGW!j~{bwe`-RP=Rc`sTWfC~9$HUQduFTo zB!7F?zpKqu+$2y6jHL3=b!a(-3E2IZ^HhAy@7~q0T1%fO)7$I9*A(7rNpi)eed{;I zqwy3#A=PQj-(U+k(AFHZv;Ue*@9Anvx$5i}dLrX^H|fq-o?m~d6t=)!h6d=Tn0+u` zC=49f4+4{-y!`<_!F6KRPz2c_g3oEgona`!(Oz`Wv2CUv)rmzSq(s}MON|5*x938 zO6_ic9y82)hMC;dliB*F&G#<#_1$nUKyLz1);=5*Ye}~UM!t)kVa$CX75yvRyaIv9 z#I*oA=e4lyRT$wRjA{|Sb-3gUP>gtcKaYXBBU0@(XJ@gcaqW3^LZC&WsET!R^Is%* zSmNyse~_+_(gC|fVQjNouy_6_z459`ycY*iX@ujiK0ri>`UFzuq18x0H?6Kqx;hgO z5xAFT;Gm#(xM%+nLA&(mTr=U4au8O)Lujy)aaR!cOz^3V4vd3PD?I6$ieGxy`LPHg zE8QB!1H0RJRof5Wc0Sc3o9OuDId?jcR<=UE1GZ*O-v@C-kD5SN4=LG3jyDEb@p~_p zf5{!Pu}pn($M-AG8e#`T-%Z5!&H*~bx=;n?^!QaE%8wL<91ACqFmaV+X0DtvApbs~9@L~?uZ5**2G?Cw&OTtmD2I8ge;2tkP z=2145P1A@U9~-hslqtVHTsbz*ZL0(fE+`gQ=wO6sMxBLKe-aJh8gm`ds9Aa2*Eb4L zNw=mgHQQ_+b;ZF^mVj3I3Z&QI3j3}xs-;o0r%7KwrtF4)^kiKmz+ra z`n{|WrPJ@uUF80o{PVlAa%PEkk+zvV-4q=75r&|HPV5CV!?=mA|3eAr(E3O@!dA>* zdMbYM{_g8nd+xc6s%1U06g|Zms%E=j7zozyCpgt{K-3EgIM2eHA3e?1fW~R>Yl|*N zlui{^2Tr0`1KLNpus5PzoKL(vap={X7MAOCK>`z65%t3uhx=x zs~%op@4J^zHpC$P;h2f+;^%}n!+RT&i4A$2(TLdf@Emr&&_ z%l(rrdQTK2FY#v{hsEsUEP6eCTya`NqViqV=VC>LP&b|!-k(A?r7K6atqYIMj0{dq zPIeAEXs6YvSjqaU{etbX#rMu=vW361@*@sftwJts{Ji zTTEK0BkfVuXY~_5qCu0jBv7dD7y538oJ4_d{PB9NNBy@P1u}9^T@))Stq-{@+y~q)3#qRQ91Rz0nBan4^#HsKd}7{ zv_Guua*!yU**=D~82_KX4Z}H(0L&2!HV+W#@p1e%K)**Fh{0Gf&OtB7=lfsy-Q2#F zL^eHkT-(roskSl0*EShxGUL(Mrw89Z^}~1TiS1u-tO$eDB?e%Lg9B(#+2B7#_>1(u z4idDA+0$S>oW;k}*R|8aq}xnA#8kSj=m+ON^5=Z6b|B{Nx>5{50pw{hG!drwD#@)g z?3>*H>8g?lb#0FOr$HjHw-z#tJs!rV58 zF+55u0KZlM!M4QT0Ye(Fq0tN;NQudl9pw9XtoMAXe&rR(h?E~L-D?sso=Y%%!#>W1 z4501L6!RegK6**Zoz4_$mFUpaSZVW%jt5?HNZbx7eGAiW|DNrD zO9Ah5!Epjff&+zAV)H?0p7d0@eHQ-lNPXJ{A189T4gY3x z&U!nP&e~tUuKi&8%z~|7R*G5b8GEtYAa)yYuo?irit2)#n4)(=sk5Mal@{J=eu#Fc zCi_>=Sor&=})b?&dG7eI4SQ|r7-4GLlFLFyq38Z|4FeTkXHJw6scxm^xUHh%U7u)E|)1gk+?nd8r;`vrc z6BStt0GKro3BmXa=3O5D3K6g|soyeJ^#_APP{QfMo$fg&#)f_GeGT9=|0u!1c(Bu& z8iySLK#F5DuF^5(>P>iZp#)Ql)Py=)$M9;(VXiMW9-V63NouG9KE}rVT8aNEK(({7 zoA;yV?5PW}G)hZ-d`)u6p3}|9{n!Icz zp(R5R?r5xGf2lULdETwbIWpe%zK_27F00B$%~$Wk6G_zj!=XF2l9AG^_CAj@$X{ND5T*~x!pIbtXMd@o6?&t}zQz`-pVi0#JxQh>mH zBaXZ9AK2dUWw*5|(`j}ePS`gq-&T#<{0O8ueGacY)V9ZOTPX1{j)Bc-%;sZv!z-`O zaT&mQmj$Su_WLR0I=`)K3%vA{rKR15GR-Eu&WDCaCnuKfQME#L?5MX8^%11ZTWr-4 z)d%#4G{yT_)1S*jc#_uh+1I+a;pDVnm& z`Rq5PFnk`U&4A*wNX#771uwXC3n_84+)l^7?MB{kO_S8Tj-fqzYGdkexqN(Fx=ayp z!?NZ(x*ZQN@y7s9G(h+gwuLbOSZa+K06hHFUwC=|D<31yP{xRWsM3Z{sF%biU^M5; zu7?E#?fU&aFWY)$vq38IRMfLqt3RkbS_f2^U6-w$Ql z-UvH*6W9{x((m}5GurM7kpr-AnX6o=ao-m*Bq1-}deE zQQaHM*?_gZ(vhT=eJ2^(VhjBPYXs;Xn|GXnus8eQpoR9`wRzZI+3}D7pt$X zFLb6as~FhdtB&DQxQ5!y+leXVh^3=caKP~k41kaV(rgknW2aFYB+zH!_sw@hC`w3p z%bl>4dBSLf{S5!6r;uG_Em~|}N8jlWYRIN5J+DHQeY1$nBLFge%^oKFx=v66Oz%+w zkSD`k1ovAcqzUfp)YFJ>My&+?u$m5rM3|&pLd$%PHzYgtg>@soSh_w`nvXXU#mRRr+Cjb z!qqSAkr646dCZ9Z6tJw*AR0Jx*dPz?0ulHAO9Bi*0!q01jXDP=;q*c%i&{n2Y=?_s zEU0~0sRtddAE0b@$T_F?jmbZIM&7yUo1WT>UQ~anRvI|^aBq+i`#O$z68MD{Rxh?Y zPNB$tBRKhHb;x_bktXyrFQ6Dgwdd8FKMe>AF>RB3C2Cjp6!c>E0&#>uPKLruLWVA+ z;mDMdr{`03*Vn19rZHEc2jp_Z1)5HL3a%Jftw+{(7Dk3dge!=}$KUp_+#$Fz&=N8??9rfz?jX<16=v?peULggDjn7 z5OuB-E+$m-j($O_ddwuM24DC?GWVp-ZH4KhVGm)u?mZF16~Li8Ve23s*vQi;AW*VA zft1tuPXbdGKz`g@&?klmE#=5OzN*m(2-U42}g52WRT&RtuGHrYGjUgx*G)>|1{- zb^K`wO7zrcPE(st;ILD{kw!97=RVY}Q%JZZa7$|RAl>!inP;t_Tx;2iP5Il4P7R3} z2c@viuWoiOG%UGCb&{q%Cg|Ng5 zzba3o+9Z`E68hzXJM|)u-FLlv6=QiBs|pTfrfde|oHuIh(2BPPgt+|r%XK-gi*4^_ zrftLaTv8m<4{n%nQ#<%$+FDOPpZ#*_xv(_-G9(|sIK=)^D8iNx`co&Dv!mN4w9Hnh zYxw9?AHYxpOaWo()&o=-6RzM_-*iwrv73z@e_XgkF{0auC2XQ^&{F? zE&A|wmjyQ#teqXs#@J3CS&Lg=B;!+%{AMnsPt|iu29c9#K(bTgLxE$x`;bc3qZgUB zvhJfFZe`jkjBgm?vKzp-L;AWwq{jJ#ZjeuBQ*6>+p zYpdGX@Pk&f$*vZE!s{?jw5#qeNXzCcDcknN;M(QNblb!AR>_7IXNB<2KAsS5yAb?=diYR~G1^kH{e=!?>C7juqoF4X0_rn^tQaOF{^8yhlE;7$Mv zd_>#k!`S}94~R8dF|hBs^@8MKuIuTw@=ST8djG4+k)-OBGr7i>K3a=eire1$(JI=f zYNhkO{v*$N&TYi{AVbF!JB9cl(yn{|&z)D33g#w*KZd?Aos;AS3%HR%^U;2?ik-J5 znL!}#595M>WY=5<973+lF|mZ#1julJ-VY@(u7c;#y*!51Kvkg&y{Q>Vp-i2yx6?oH zVIj)3)wvp7h&*Tpdmv=b;o}Dqa&nB!(tscJXQ#wwBx5glH-A8Mybnpmld+UgqXccFIJ4dCVNKK?!6!kQPrabYqC9sF?n!1F?#Q{3sEx1__!1OI=) zrWx#-7;fw++kgc!$%jY*>|V@?5tY$K&o?$WpSzeW-u$h#ygyP&TKBphb@NwcMY~tR z0yx)#|CKKM1sMiEh`#$iu!=!D2D<1Un1i$n^Y}O&{LTLa2QBS5xpBUPY-yfINS0=L zoR$$+qm`!G?k(4p;xxRWd<8F9`QFbgf=5R61y?K)zq-ig%>qBG49EgYHjskvXYW`@ zNHyvP7g8Uh{l5RDWfq3k4CqZZi1NV=&Jf`n)|)tmoGErcvf3!^ADCZQ>NY%`t%#b; zh-V_74=j9}oiN(qP(|(awk&#c;B!maZe^g8g+o`ni8Bl#`~}+F+7>A!ygflZ=)Ix84J+)XvID6?trpi;!Pcuqu=IFjjv;AiBau*pV@hbYBlcjpDI5F~XkD^WMTFO-V6sc|ppvW7wRq#N<*n%+$I$~8 zj(1p3SZM@*A!)h}@F5VPR<#XZyclH8H$)jGKMQ|+P{!|7=5|=t)Az>m?~NswA3E(0 zNW-WHP{0AZGAaD3VjElIXBanB6yl$JK5&}z+q%yeoVgqVP>w0(M)WUrSy(l6Fpz$p z%=a!QYnN9DL|g+BeWH?^e<W%w5aux_U?j)5BKe(7yiNy&#U zRn+^v_cQ*9NsmK6}SQDCHn1TKhVRK2Q zDHKO`Okm66_&dlAZ!kus4#g|df3jz0C*yJ#q&)=pD}K}x-D`iUkJ!^yMcZO46)?a| z4|XyE4-M3ATtTtz7Y7SK-b6#hm2WuGGx+c}#h5Gv$h+Ni)Utf!m999>q~~xqdvgG4hR7X|PjJT%qG| zE;pKd>Xn~c=7&oED{u8IIAa`_)UV-cjM|@ClE&A%X!4zMc2nl%^GJb}J0026d(ju} z-nI2lHHqQ#`mD#nmAbi)v+*SXOqU#v=#2{wQ-s>ik8L0VajEq23pLYq53EX{E~mqAp_q|t#KX=M8|P@hIVlt^{zk18dV2yB96qhLFO9 zlaTz(V|PQ0?B7#s>mrJ}RJwz>+jOCabW0SIZudr0imcH$EiQ6ho(It z+?l%JE4fZaJ(=y5@mH^BONb{=Rv>|&YbaXl)y#E!fK{Qy!^MtpQNyR> z9MP=#5YvkKF!w!FKxE4_jNF_!czNq(`c222$9hJWo{(r+$gYKA@Grbrr4b?4I}5!;y3#W*LJ!j78KZyd1+K@AOP2iU3GI-(ro2E8HU1P0#P zyt03}Y=TUvP-}=k|D?!Xd&pM8&+c!qq>JNO+jZoJ2<8HqrPlz*L#7vk>?SM})&;{o|j#gZ)(8*V~jtQE1W_Ve3x5cd|oYA-O*#y!D0@ZC}@?1Oc zX3R)4@h$*%-O{l)5kT?tJoQ=DE6B55I^Qs@>1o^=ou+TY^;Ye1YuIh1c3=Ch($ZuO zR`?*YgRll7I`>o78*@SSSMNPELqZoa*<3kF_iR$ZX_|BraCbe(HZu3@I95HZcPcft z?YWzBlVJ?I|4;%&xp3LNN8D%zl z5+t&On)`Xq?{>%tA722Hiq{?5gtS?vd8eMmdRE~D_N&t8B70deI6ryZw;e%&yVTrq z*8=a9mU_FOE~nhR;$F=Qkx@lI)B3+DkbKmhp~nBOgXPQrRq0B^ZUX7o#zZ)or+op^ zr3?r|p049&y7OgGHg=;2*Df3F?je)yBk!9Wyr{+dWOwIuw0<@>mJipOLL;#Fnh$^h zaOeUoJ>)F#3<0kRSXaxmYzA9MmX_E0kttM$Yi*|;4ijB|1`kPziVQSF+0Q3-59E(g z8?Y%hCCv*>V+9#IOef*dm=x3`cnEEdLm8fHhGjd8DxSqIx7t$g40QY6U zdj(I%dvi_%5@{bdlEWoeuX`^SSKdX30F|4U`M14ocfEc4QoQ>+Nidrm7yQue2c9Xt z!4}>-0%}qO$woMJk%-+r%=Dn>+_gpSnY)NOYpe7Wb!qFOXM0xS6Vtw-sh%^}O5YMN zgwGt?ImoAT)SJk)Rf^-J8=Fs)9HP}SDYTem=pf==IUZWB0}JYSN_ga-b?#`?i+cy2 z*K!!Q82-f&sCOef*V`fimJ)Y2yk7$(T$Ism0&S5kP&SW?hf}Sar0F$7Ox1T@$H~2c z>%Y==FPs#TDpLDE{BT?1A6TjF`$WO}2+e3aX9>OczNL-o4S%C_!a{6aOV?fQo!ZR?KahY9kwYP;+* z%>>x0AMlML|lq_yzp%DCjojGt-8|0F565^Hanl&+Adu`5`mq z15dvjSX=99$9~S#tF_@j{& z*Y}5X@oyy-?W3H3%v zHEAlX@-w#>$lUXFGuP$pxo)ir#2csnm>5uxhJ{Se$wN4{$Tf(l)O2MMPb9=w1);9X zHl`1)&`lE(uFb0YXXRY1xaM(+N0#fl1suduA_-X)kd6oht?%&oFpMEHdmbGr(&GrF z)lqP)J>X{||7`Ej1Tvg>Cg!OA8*2j;{5@PHvfbfs2U~um(X<4;+J4A8Fn$}MVZU_w zeur9&{?^-#&)m^Mv8f=*(!p4PQP%>GL;&R-0H_6ynK_iAy)atj0p4J(*Qi@#{l4#W zW=?0>hW2$(&&9ia-$PWU(bs`kUj`V8?zzf8?HfGcPhOw|6~93=SAJkg8=W?3PKrVW z{yBv7OhI4KmOIoUN_&6cT3c$m74D)WR0!5nG??jfj01y*J;_#qY^aO#_=Mj-sII*Z z0Zuf*s{PN_2mRUw1Y`HSHA$^gTa4yWh7XLl!6$gg8J)U`7l34}DG4qWY_iaFE7#Sv z?et@_iHoFNeJ{>Os%65s9wf8AfnU&o@EJ(*HfCku{Bfd7Sn-kH<+}3`#Xob%w+((v z>lQmi3g;l@cw$<|Z%#;_VliA&gOEpP|Oh%-rwi<^B79 zKF|02d_K?dTaM%T{_`9iC^2&{ulstP*L9uOdD1##i46wVr?tc;@Hxk;8&1lu&8@kN zwDT`vCJxVV?~!G9^7ZiDjeS`;L#VPM{3iO3q228(7+7Y%d%qFq|+?xa4o4U;ya&j|8?aa zndhEax4d@7HIFvQMsiO9YXXuCAn4>!h)mJ%V1kh~ck=MxYRYxMCOZ-G{0jvufm9F& z9GlJRsFU{}op+0b&#RHv`d+o;^*xn5}d2ZMUr?`*w^Us26PMj{cu9bx=+( z!#lL`=F{{F*uU}6Fgk`Q*{Jj`%;U@sDWvskxw>&Qw+2c!<@lPR8poyvxGOKWf={C~ z?04Z6UX^~XuY_q`vPZ87?scqEf8p~NYVY{a?M{K~is{E2>MwF}cNy*zgYN=$>D`!4 zqEslp4t^R;yV;pAT~sS?Ry1~H&DBPnV*8=Z)A2xu<(b1`ATdIzw*f3vZz2iasRvg@ zYJqrMz3ZTkacld^N)1OhU(5bKA=O>VD>ZeIRyjnOC**6|>YuQF;KFI(I)R4dlSaM= zl!Dk<>vIu;ic;^9D9M>dcR-7-UV&9qmwEv!`55V9 z07vMlzL|1^(NB@bN=98%w>^Dr&i=i%JmeuxTm0}x*PdId58PR}P7-qm3E{g=!NJPa z9j`V2>}sTOx8OuRMzu*RpWx42`rQ((Yf$c${K~a znq#H*jR^?&(PIDH-yVbY5Rd!@ODPp1M-s|c60x2nSUIQ};(*tWf+s6M_o?A($db;F z6qAqV<|Uj?Do2P^+3oS}*}0K+-;wPu$HM#u5T&jK*N3+sWZI9gWfu1vH732YOcV(= z_+s6K?T-}gf{m`me!BYeqCk&WS5uShk8nAt&2Q@o|CPc7SeG;m`Yv}OurCbN*9wvC zeBj0Uz$-?!;ypSxgYFw+e5bI|u+X!bdg-f4Dmg*E*DjrP=1;uj%yz$Z_lIAmvmC!& z0K6xDfZ)9;OxgAl@K9(0e^4XrJ%kO0^5L9wW&O_w0_(@`Se+U9ymjpwV|6 zxd+U=Ql#ycEEvWxhGHHa5aRnB1fufc;tLH9u5o#_(VoXrcPnT+0V}WE-)v*Uxxi7H zeb8hN(v1;m(8?PtnW&`L@m#5MByd=k-7i%w&s0k--%*t1DRmL`*^2$WASey)F!o_+ zFpq-)Bi3otp$0C>687v@|DR>=s~%Vp!WzDoS(rt4n6YDdSv+8aa>8E1-p?Lpic2(^ z=;F8*h9<<}5Z^cMzM~oAMJ~-3vkV=Uo^a@JO|0mCVBhP?HYRg24k-cm{YN<>=)+0@ zv^!UX_YZA<+qygtjdKD+`%zW*G8Xdn9@fSk&b&=MC?I*h_?>xE(S=xUNd-boqb$;e z6+txMHA`x5ERzGcCb$m!PrrLt<+j&z|M4Gt65mA}JY3iNTCsr)0mSpol^v~| zpXt5npEx?NiwKZ)GA|ib#Re3=;@AKWTNFsiE=|VKvDA}pz7)CHs6U^Kn!$elekLpA zwfmS}h0o1vUmiU_slrz#RJPN!N$S7{VlqUo!j?hpcJ}~XLMxdt=FUZvwTH_)vPto} z2?-%!&jxHduxW?*O#i)UpHKWL=O($$k%OT@Yyn<5PE5H!l-?1DK0upowTprHn0iIt z5IWE`MbI4wCp!<#uh~~-Tl(kfMjkFgrSZPJ@lhgpto~v|xDwJsDzj)$0SFTKVK6aM zI21nLez@|Jty;+iUmgX!hbuwB%DidM70b)Zc;6JHqlhB`@F7*YGaN^+WsFe2;6V>K~@Iu0XKXEK{O%4cU$W!M{mf zcyDiet}gB2sVjR%i&Nxte7r=0{Ro=(lGG1M%BrDqZ$_{&-~4NmJ?zAkff$GIJ#;|g zjXV;%k4j4ZZA27+?-J|>Mq3ux@fYgeXH$gd1y3g4*tzGM=jZL*@yv(dILEfY;#Uj| zas`M8V(Fjo@?Q}3TrclDmU_&ZI3^sheZji$+eq+De3-JGUZqy`-Y1IA`%@Ig4H0Tm zi*Z$pc^})lrxz(E+TN6==pP5|G%Q^dV-g5e@~^Mz(0Xt66gsWa|! zE%K#yHMuk0l>e^8&7vLo2RPZSX6nF%gBMd?U>ZG&-O?v&CEdsHg=)R{C2u@f(r}1t zJ@#^qZ@tmj>vWkQ&+H>55|79SB+tuS`#fIE#_0V|+JgVk82s=3dXae)PI$&tqV6rB zK%&0VP%P8v+35S_{8y8o-?};+fhkO7y#U1G8T#uT z_1Vd91_*B&6-=Qy=lf<6xI_H{D2~+cJ%TU1^ooZJs+tExQ=1J7VpM+curFyO)^?}i zOo8gbM>JC&Y&yGbsAxcLHQFM#wnKJdN>1V zrF-Y0DeK-c2BW4ug)av zo(OhD$?*EAoe@)r@aBMz5XV0;nAs61u;Jn$m*i`UYdo8tQz<3X+3MOZU$B8(14JuXf#yFgVDI;@|)q zm(;(r;hU?UR%OaWF{DE8z%NSiJ#B<-C6|U9tZomF-pe#ZUcYSC?d_?FoR~R~TGpfA zCItjP((NuV%HC)N}N&nZ8HjL(v%1o6EHqJL6&hEZEmh_1rzQqcWVPgowzV2lMB?wQa?q*n9 z(7w(fO48-|@exUFVY6Lh7K*|LEgZpw%nr+E=ubd&ua{YLk^!L1aBMW`k7PP&yYE(a zKaxOPT$t+pmeM9ZS&>1CorcZ$Pnur$yPdw9zhL{R!Sih`=MRtkxX^K+J^TRh`+=GT zpPPAM=aG6ER&Jab^3umXey)MJa#HE zNQ5i7|1;y%xn&ct(Z(+}olo#H7hV0oht?1aK^u442@%`{FlYZ(LFVmm?B+2)+=xRc zkpbw6&7#=LcCea6+ZZGj>Ih%(qF_W#uV{SE_j@zxbJI@v$*YH>pLRSfy~ek!aBaJ+ zJvYq>H{rEN`wOMa(KJcd`UcJE2CizgVqkqDT^&9~1667B1giTa!_~CjdTf6fMbp*3 z+S&f6b@QQ~b_H&3!D1b5vELR;g$#VCD!53uLoq2G1o2P@uPH3L>!w`vt5tUjW%cqL z7*bZZ+Ud}veay(gT9z#vOq6c$A$;Iss0Fa{pvBWUV1nDi9l&)Vm+=4mcoWMJ?$DdC zA0rs~d(JdJ_XumN?mwG^NDa?a1ZnwR0|d`lE$6^Q;+nFU@= zxRI)X4ZHL)Pp&Z~j-59S>aC+ANaP>5Z=IdtQdI*v(~b#wl`Sz0mDfHuxdz-nKi5Yg zddyp(R4}qtvE>#+I-~_bEy%hAYUo%5X67LXrb(^FFvYLdbr1(A;oO zqWJBZZ4D?Hot#^5^lT`)NSo-9@teIKFD~zj3KIJu#td|N6aER%s@B_6M#KfeF|qPS z!%THwN|K?Ux_>s|BSf2X#@VGyYckyAO- zD~w$(Dag~`cbIqo%{Ng9;njc8Sec;E990MX?Xkitc)4bgKlJw&z<_#_6OKX{G*AVUeYOp1 zPM%x^B`0r)!RbxdG+ZS;mxim6?GIidA0!Cnjc-Zs_+4yB(N3p5irPqEa8yE>UxrZi z;FkGhZNt*oI=&d2M4;X%T)a-@pftWGlXirj@zLuQ7}rZ$r~5}kuf~*nl75vp^L6~R zb*L(9Rppi)7Tz!+11*){<_qcb%OXIZM$tggoBHZV& zr+F?lWQk|G#5@Ro8CHJTK-uQHv_|^*%x?u{4i~qHPIo}$+RbXF$k1dbB9KCFtmt^y^!>9>^)l_$s&4Z%o?99F4EW{NzKkp2mLCKI=BnV)XQ+&#kyJBW zD#ASk4j@i_+6}Xq?LZk1tHLPHd$pHb!hCdV#OIpvBpeu@iOjiKpE$hI88y@AZ@l|u z*R#q5?-XCDx4w#`%i4bOxUA}*Bh`|flIdnDbk z`1ld^=NbAQiM*k;Ap%_{ZL(+*$laGtYvZ_4EvC5jSyskry5Snu7qe*H8?!E*RQ|2_ z&@;O@>#W_PinM$QQsX9dd@%%vPkYA=pa4^ivaG&)yN z_|Y|DEV1L}%*yw3naabTT?kA$YlhI=os9M(=9z&4-(YDKe5m>xYncpXPdC30o(IU0 zrm|nSe?>z(NK{8?<2baXy6kzn4cHXlt|Knw)3>@aHB2_}`62K(z-MK^ukTp=l+gJ8W=~FzW<4s%TRUCAPoDvo&u77!VeULK5qi$rC~&CY%jcTy z3K$JHKq}H{^FQLp{qbYH@v~0v!zk?IN$?II)}8~<@f)+b&zGK$po0NgcnxEKSLBNo zB$Ic&_nQ&;)kq@?1`g(G6rZEhNkZ6)PNvd3IJ)8JRpE7y;A5l;+0l(0(xB9~qm9ld zOj>boCab_?NfEi1DeBs67yuO5@QOVh(%~sL>Gk5ij~ACql$*IP(N|>O6q~HZj((Hk z<^V7@EPMs;g`BxXlVC+Z_zQK~=nf0y&cNQ~LdvG%ItF$Iu33+9j!ZMylNvyWSV@0U zc;OriYD#Q44;*wjz2{^ zLBFT`+JOpF!K@bSOyp z3SF-tC6G4`Hprr5LF-<7A>5dj)y#116~)H~RLmRb8!i{VJa%(0N=MSKR4LBLL!+7P z@-8r*aEs)LgB7wC2@Emj<$go~Frn;x3*dSy3u_t1C$E!(CLQN4A;mbdJg z1}I)W*rK4cITFgyEHOEczfe2^u79D@esTQ17*Cu0mP6lw?LcJM0o$o2nFzt}hT9G0 z-j977zbK_ydi5oaoXz8KDKRa^t=nKX=Ox+UvVW?^|956|+|2~)%@pRd;pu8t2xud0 zhV(DMbmH<&I+K-jOAA)hS`3>eL^Uc8=n1G=(JG%@)BBUzf9Oe>om`u6#CL7yyu^$H z-=CE_9mGVDXqxt7rP}pSVXPswo;acy0$P8?>PT8K>RE zu#%RQ6lzuYl>Mu2jr^otdH<-T>WvoD_cK0uws?m#h~jL(+~7aG&9-_C|M&<8^!MgP z(rOAZ8lu}z1A!91C1`&oD;H|^g~$&xaa1R}5fuy@LC9(<3^E-QIroJ)V%&44?(C0* ztG~9*B;{vxcAIiWeb$CB5HN5~zySki07S@DYrw$ymlhDqMr?wnvsg*`*iFkp#$B8! zctb%j6@Ghba%qItt8?b&eyf9k(^>~rQmwy!N!Nuc>ipvDa=z-p^}^*_Ky z69k)hA;!>U%;p@0@H2mSrF~1=m6(;FfySZ=k(fPf%#h(fJHcb>0d=c307UMo3Yt!R0{V@+K@d<3*Mxv; zcMy2tv0f^1q!(5KmR|d9z`yN~_PT%;hTEO#wz?OCzw4IkU%lJ7*89pv_T&cz+aliL zeQ#oRCHHl5RYjm*{{!?OAh#l~V>XO?qj2LP^F-9CP>PwJSEm1{3Nn%?D(L98#&loWg?I`!~R{*l~1Ao$mPvbp@PqHDT&7 z<|QyT9A-Y7o(c)d?(pBU zkG2ZOnhYX12O{5oI05Mh16zniHe#pFxB>3bE2!B-mjK+j7u_fU_v&7j1wJ!*MRT9q zythZv)&20;l#Uk7qCSR$n>z{zcz^%$?15tYKj37*Kj0*f`X;#Ekv!nP_&S+FV=aho z<50r|i9}B309uB#WkZ!vsbZ_J1uyWUW%Q+rueHiBCUn~RgdlDo7--I-3qFM-KI zLhEf4zU-xfKF!J=nfrl!-h$hcy%Mk9K6QK<<&TnOw?tvU0LC%whIQ;BSl6B$U}^R} zB)$OtCmdjCG?(VYO2{@q(&M*C^O2xk(wv_?ych_yi3})NZP<7ht+%csrYhsSeX8-cowc^#V+7Kg6jl$T5Z8r&3pW!CF-$2Tn)ZStG>VgH;H?YFC0x-T zb?ZN&_}ySpAX>BfKo`e>PtajPd2G{rCHi?50ceHIwywI{%u%-iY-`{v^q8@snAKAy;3ltOrqps-^#C)7w z4$a-AA6KZkpZCCH$6>$HYuB!YNQ2?JqLmd`Zx}gULfyYWgyylYokk{fLF%E{xDz8W zPoSoRUZtM>5+D&(?HT&1X1rqhV0##N);YnR6y;i}1*@zqenRa`~mS@xT$%(jkF>rspG;w?&f{LZeM4aM{WznBL7)Stv) z7u7uh!s{`F!5{@1*3e=~zkUPmnkjUZb(69l5x#5AM1s&E)V{a+La9rSVadQr~G-w^0X7^^Do``>@!_641TmpHfCx7*&10dx`zP{jTb z{1sZ*M6)C-S!+&f(+oqX`C+3fwX+k=Kabnds0{<5|I38XESDFhu^z@Q(E~o2LT#%_ zan6yiFA5GSM_8i%LYWSroMva4a^R%YG~h7=i2_)TWk@=_&<>HG6Us~yK)e@aC5GOl z2DC|M&{GQG?jcc$p$I8I)kunZQLr}N6JUyWv|2O%#2+X~;Qz*Uy>5S9>AD5%DRQ6Z% zXUQDD-~EGD*)<6p=OZU~6@U`okv_m8%S1?-E(q|uB?7kwH&C%IV~Agw@_&h0E`E0}Q`7duX1|w5juU$y5!6y7G4EjZ&dSIC260B|lBbzum4`Yz+ zn&4<{fZ5F?mC%to@2-K9#E((r28*-XKNXwHnQYfUzG!+xO9qT5JbpZq0_a&>Jz=!} zYF@pCSpOS7FaQZPhc}wwTj>DgMw&8xXoV4h3oCpu(WZZ~1n6Vyh|hfT4^+2y&Q?*_ zxFLBtRU;wxh}lj@hYQSWfI{*CWGCE}M5aI%aPA~70)n3(Lk-_aCxcgWVV;9MsK{tH zqr*h$ZN-BV6?L_)-Yd$Lq^f%D`!YYA(#~Hx^7C}a-kD=PTa8pXBgSdyYXt%ZSOfH+ zt)9XF&=csPXVI344EIo5*e8Ti;?8KQmPco`bz0Wt^hO>!u0AH?S(_5>qObk&*VuMp zt`xk7YzVXuA{P)Uo!SDDp=z&oo15+KI>VKI$Fahb-z=2ebbSThzi&%36IS#)di1Eh z*kRc_>Artf5P_{N%QN&E3zR{9Qto5K8kIv2(LftVz-H>cv8q=9#h2&zdprC3^!e^z z4UXlCAI@8nTaPz8}VpB6L)M!zAeBEaL^u5vlR9US}uHo#EXnjEhtV= zw?3iUGNotoir}JFCt!x?;7KchSb7(kNa<3|Z*;u7E z+;90I5>+3E>>)j~gLR_lNwm&3tYGLh>fZMhztd-_h4%pSDNs(nW!3KBt-0k<>N&}g zcT>0zi4Vssv^9Ws2ciwq`$8}fv8hiO1>i$q21NT0e;7poHHU+BZPch8{=q~TZk`T^ zm**QY=JuNI=)c8MJu&R(SvFX}ka2r_Ieo~%DPGZEbZ4meXxrZL;wDVRZ;*XJ8W1%I zy<{-W0yY+9ukrE$;y^oJi}P1q@Yh>4UfN?LdV@g2=KxRA=Vf!Zci$Y&Eje(kNwofj zl44$%I&`NKaJ3yMn%GnQ{z`Qfb=3oxlYM>i&vz+mUoBEcnaYB(xfXDQ{RHZ9V#;4A zQHcD)04EB@sYo3dA!DU!KgDhtXQpbwxrO~_Klud6>U0K*3YE85v86q_8W@m~p#BEA zh6hD*Hii?vuub34jkq{Q`>g)-IK}8}?pm{S($hMlw+=csy6U2qJW}cAXHKRA5t9h4 z50*@Vs;b+!5{s4akBD+?hhau;EiK26N~LA|%1tyAdeX&dlC4>h@OQ8fYcibeSTXEdKgy5bI0{hil*d>D)Dn7Tf31;JSBpFwNLsr3VEZM2HB*)P)cBWO@b+N zlHtj;eSmDQW!O)g!`Gi`tb>hSUNOi;r~mP&^AqdHxPFU$BYy0Q?cEluTDA`zl2I_4 z(wiKvOPh?Z)V%3J^8(~#59j*K&CS)>li60Kvl%^I>Sb5U_wT!Cr~Cl-70m2(foHT5 z?gC?xJVFdJylbbtEEtYWn0`rkwiyi4Yt;Z#PP zLMpn5L>?fGrcflnnkK(RIH_(~36(FX#5Px(Q@w32=lVD3o9QldPSu$|GuEmsXQQ1j z*_5)S8C*>ptVP2Iar8VmdU^#`pKZ6{J43r&ZIVO|`EVF;otVO2^cX6jyPj^riZDH?#12T1DMuaZKUp*2&veUo^CsR|&kh`a zsO@55e5Jpu(Qm2ZYLb%Q$KtQbPl!xe0MR&;S@~cpVoz)vP8N@wl1OGkQ@A)Mp(~s2 zc~hFaD=x{(XBDHL6BcRvV(h;@sBHS6@cIJqj=U1Ok5#|~-bNJh6bg|a-`eJajWfpH z&Z?4li?eE+ICIbum5#1ah}7AxkcO+@Dp7YX;b>smr)+q!YwEpLd-lt$=SDiC7yXSt-gT?$KTN*^w-%i9 zb@K7f8lAkSC!l3p@%`GR+!=4j1NkS@ol{IZY}vNPbE4j0$&Vo^tPgHPf*v|(A`9Cx zI?F@ls_Vg&#S3rkp0GCtV~YLHlS6zjnJLpWG7~Z9_dY)wmcyAM$-M^ZasM@hy-9&} zbLdgD!X-9v*`C_LHpa(Z;nVH5!SB{uU5tl|m9zV8V}`md{g@y91Ct6<_pY{|zQQgf z;EsLyM)nE-<|WNhk;040eX0E!R*X6A&t!9m0*f zC|Z)fb+Jq(lLJGcdM_;K4@6zUpL#&o1N{7?RR$d_fl7}?`*?y;Ic66tnRI3HKBgSZ z+8-n1#PMbwxe8wcSz2eS8HWafY*fU)TPW(7{J3@NvbhEDWsC%t#2}!borINAmh`3o z(%oIy@zlu*)}+aFN~4pQZg24AmV->ysj=}ZEg?pXk>!hyncE*(=x~oH+r>DZGXHTY zXm#Zq1KFO>a0tb*3XMgc5t|6{CSu78j*-rlkSHFcljeQ5;O7@K(tB?j?}^J3isRTj z%Nzf(yMTkC`yYw-RxO6T@%2mBg*%?@PrvUwH#%|j%P0O159Fja^e)?*$6 z7sM7MA&?N%462n^dnH4)hW%E;TVJ{xZwlradN6P&ZmW!vcd%#fj$q)O#fB@PfQ(fv zQw#*BpMW{d+e>|;snwpKA4~ZpzZ?6u!Rk`P_qpZP=s;s@o~zhETi+<$}bKSs4R464VP?T*=69TIQ822(odR5f={Iusm1ZvJunTIHitJ0(Hcb`}G{$HrmV|Ho%HBnY`5oehwX2Hx+|jWs za$>u^^0f|N6~GPn%IzqZf7X@85J3oLO=|%^6CR;y8(g!zLL+tJ^_zJ|Y0dVZk9iJO z)IKbr6;1DcAAYLa?0o9>F5-_IMBE~d^y^RwM7~O5IzTua^?J%;;B@a;0#HnuMsH&p z^!1;k9eesGw^}?~c}$P#_)bbWUMSJH{buEX@8)b>AOX~bsB@E}IY?3=9i)`4!2oEj z9zunKN(iU~R|F-8;t4XwTqSVMnZbr#Q#HEdPS*X)Kh`o^z_3>I)BMXludWMwnvSK~ z>R4f8FfyP;bstxOYYz+Ql-H#d#_56jf94V~kZ;6)biZya2k-q#rDOJDC&3?wGh8V2 zFC~N@0U_-c^CBxf59*ubbRz-3OB3J~ymGMu|I^l&H)Qn~vM)79()Q zm!V1}5iKU58vQ5Mm7|%>nEqX^C*AU)V>Vr1v*=AClJ12D|;?Dw3WPX#F5Q+_c z?);*%_`@gJfa}5=<+XcW3wz$*l|0}Sh?M=Uc<^^=c`KZ6fD&TV6VerQoI9wOMm8~R zsWUGqOIQ53o6<+;?g-oq+9ngGKfbs$8+hylyNi|7<1r09XX+bnfHVam?_T*(CE83H z1=zrW?b!VNvvG@%3n5(0NNbPip5vNnyQl$bz#TB)hbCTp>nSsB>vzy z18$gk_@IkP;dLK_#bMgi&&rIv+fhw-3TXRxD`cc8s4aBz7rd5@w9sYSO5*($E?4)JM0wq}teeoo$jQ8)>e_64WmPe*%d6>Rs*W88 zaKdHZs+TkZU>+4`atPD|004*Pa{c%yXaq4LUs&(wsXTcgl`PW4~nQ}`4Z?;})XB{PrF)+xiSPc7P|tMbUNLJf0w zCLFKJ`W(+k3*_8}@9|8_RvB%!6sR;{KWWL#|94P|?u!V3>@d=od4=@|(PrAik7?>n zCOhGfXY<*L@Rg4KwoH+^>=#R&qdRMSgoiFU8%H1G_;_CI@!R^U#wPT{o~?<2Y%nm~ zO3$P%RWGfPNwM1aBBx5K?E~_H+wC!{s+_I}oZZb}(ekQ~5;tY{y_&H7BI9qM56Lf? z?8axe%paVhiP!RfslktH{d6?vR*0&P+)n(S0vkaHPc*H68ept^P%?=p2?DQ@e4rA) zVyqc3)pfdM)3c)_tZ>}*yW5EqbZWlg(K*AVRQ6hP{xcZ@8%ACLYd&PX2`HsRV=og^ z;(@G<4xYmS(5W~yzHptKw?YL0HA=vOlR)&~LOLjVMmhr*#=@k)gp?ild^hU6^{Xoe zY)6Na-&sipjiNLhBs3Gtww5{Q(SRCdq&4y>T$PnVB0s@w50{4fjMTtX$ssPyzyl?! zjHC7+`;4PGA~?h*L^Ha7$O>!Qef?kSuKdl``u{+o-f7K)aT98uAXG?10!3qt93Gl# zj3V2hTTEmIOLY^B!f+=-7P>!HC~zAKGJ-cyhF!#)0Q}8M<4Ak)g_FnQ`+AT zHVh35!L5$tvae1K8>FYD2LIRrNuyXk;F3UKw;RV`5T)>M;V59e^Pan@Zt`BJYw?(O z&_3@?-*4`R*H0Df`J2ngmXC%(czkJ{%~(D}?3W3~FXZf( ze46cxqJ`T_K>>6h|Gg#8tK!9O?OY%Iu=ukqDS{(YX#*hX$WWhQY2wl>)lufGO6}u@ zz~^t1mG`!3GBr9ma#6~gD4Xzfa8wH}er-7F>E76^qyU{!J8E5rd}h2MRX{;lp=hF2JYEgE7=(>t5u9({Ioe;QY=i14S%020GLTZmLtq z2QvK%4}BcHp7U9}{FrX)aWSgF&UBSMmY@6ffA)oM5*MQofmU$ch_E8mWa(crQ{Pp}fh_B+k3T4?@qxi+wT{Y=_ z>_0k1tY?Xj94t>U%qjctK5_aDx*A-fZ)*l2G69nF^Sh8eGe0}EOcK>@i~f-@W>W(! zn9{8jPzRNVsh_(Y-kNVGM*)+F=R{bcnkm~$3!&3k?<=4vtoUyX?}X`ejx}k~&*DhA zuXx|GrsiHP%xA|m=7%5g+fF8moqun7fYnF*p^gZ6_WNs*$hDJPKO-l6^2-z4v#0uW zdf4Uo><=$njJqu-Jm_$~KFNa=M-W%(D;Sn&s?)7|^w#R= z<+&`j)7O9TX?OkpbZ0OWq?}xo{!%?K*)h4q72q47H`tadLb!T;XV6tck%)t2c8C|a zxix|S-a+!cPrMuwf{O<4`bpNSb+tChYRWq7nHNVQ`7+52C|TLT2ZBJ^Mzj}HS0!4A z@uLzkW2bDS-5~41Th^D$lO?C{`RP&@y~5fPyStkb&qo?Mt?k1Vlc<-go?Zt@P{P@| zhjOGjc`!Y5x5T4VF?Q0W8T8NZy1ye0)+E!7%rJ2^GnQ zou+2q2TxsKRTds%jsN+S(U!UPJL1jTc`dM%eZ6xm;!l*UO(0Ie2PFve)?49QbT88_N_U+17yiKKTt)Q4r*~ z=K5h$G}KDl%06E-?9|V%Z=vY4jJvEjOT{yC8 z*>S7MK{h>%3kBIi)Ehu92BHAhubj4AHmY5#0NO04Z9={LflKXR;3J#C3h~!UM|N5| zJk0WvNf2*Bp#*KQd$E*#xY0EvICM8Cj^Dr`mc;YmWWHA{L<4zl{t21j#_kDp%{`m9 z8=mcp_ez&@A{3@ew-ej4QOI{J`2nf=D2aLs`%RJ3{LY&A)%IO$O7WD@ z8CU1y@{O)md^%Ad98qgHbEjDn@J%Z`=R=^HO-9kakA3oAwJAAKv`17UrP%$QHaF~S zsqkANQvJ8$-y~IDu#^4oewcvdGx8(2mq&;Y`4Qo*LA=QxpF7Gf^h#g7#>X#RKQhiI z&rHQzXuKcU!E86@>CiD8CZn0=>I`?NS1TF|+gbxIyx8U(2s-{CjsX`ULOaH*IvY}I zyK}TJ*xR+o^k9$dt&`t#xgg;5!E%F|t0AiEq%4rkXJTo5rf^!`uc2@as>BFU@I}E< zza=fdQWuA#tIm&#^6tM2+F>du=(~mn!b0+{3=3aSn*6gsV)MqJnTN-E?89Q$?8I{L--cnOcE=hgWG z99*MF;?dL8pEe)x*I3bKK_P**0-Gg{4wD?W+lU|y>n%~T=^&Woyr_|oqM_B!3Q^G*HwfHa8EG+=po=|x1 z3baLk8PqJ9slhXYUwm8;me_OOu)M7Em;>j7FDRQZunf6x`3#GbsLpmmcA5=eTLY)w z@61d4^5!E~W8azIS9h3p9#P|X2^a-rKtLbu?r|y4X4t__VIKUy*2TN>tohEFONt*x zf=AC@@^9|A8HFPPQ(NFDZu^>{3{Dx-D2;r{<)c2)0R>cD!%KJMUcOMx33%QU;<0Px z@qttP&ks-nM+^KfHr&ldtbEst5We!M@gx0+Ko~1=I1ziUmiDhl zVN^eNlk^XsH_S$D5D_IY77;s})k0vZ?gCYic-9kev2H_TVxt-Gz;+~(eRln-PE1}ol*%k&;`IQ~|{v`OhT$A$p$*D%5#)nBR!rqQvJPoG?vr=)aXF6}U;W~&A8H;X)gME-?NoGc4k`a?+p_vJSbYTgWp z7?QCuUpx0MWG+zs@%i{LT67wxchR0wL92tt*KvE9fKM6K1qLVulsdDI4d1VTV(NNh z&(|mYk&&^HcC}i5#&7W4dRD(f{mH)5QKOB4;Kv{8E+n8~&8$xI+zD^Tg_C93B8p=6+%jf=_E8Rh+1) z2x#n_@`K~{+?nBHlNBGaWrW2fnh51ZN0v$@CgQhI(z$P)3Q47jzNtAWC!XA#!+pwp z0TxygGBhlm)RtKV@h+7Vq-CqQ9?moq7QX#Dz}Ua3>GL+=&lm>QRgFtv2k6Xj#8P&n@X6*X_ejb@e9Kp+fE-|D`BfQgkZMWB~0s_cS?XCCZ?H2OU zq%R-ZcI|p@jUMY17&6MhQg5g-B~EvesK-H;gTh-(Wya=x$LDv5ug>>h$!=4b85$i` zT}`y+eV4x2X_RoD?aK_wp=gWI6g6>c;V$C@6Wzctp7b=4e7_aILy!T%@o@u>CxPuK ze9FHBJQfestOz(FMk(Nw7Wv3fLUpzx9y&FgNg3Mwe9-tq=FDhNb;k6e+xt75O@FYn z5h?DZQlL0(ZClfiCEv&Ki&G|iZv$ZG{FpP&zD~dJT6gQAGm+a5Z3j7``j(3N8kPb2 z=FL;s3X=?m6~3c&Y9fj($K5>YUZ|P;)ThWNB@8btzAr1?A%o4!Q!}E*?c1zWDGA}( zx_Lg#J8%!Aic_(*fxm;V5@17$a-$O-8kwG(HuOIP9Vk?GDRopLGeR8rMob~ z4Mw{yQ|mcxY{$pn8tp0CcPK)css`)tIldQi`>FR&9 z5VU-MDdMO@I`l0jYflTlYSK7Q@Oj3_&9#GSWK6h6$_ zLj+vEeK|HY3AU`l0O<+Oid;LsR@tB1lPye(CG5xf)Y5o3Qsr(AIIXn*lgjuHO5=aW z*Be;*#zS~49S0wzVB+Ml!;7I>Ai?Av)uX?DJ^tY5SjTptVBu~h7605kzE`R4x%&P+ zgdPY^0y7<3`W@4tt5VU6oE`oG{TmJ`P)_4RctQXot=!^9bNR|ESHQj@L?07nYUv2k z_z=bV?A1R86@GaEik2@jc13**(2TG3p7ahhRW`Q+wL>z<)Ts`jN1$So!nlC1w}n%v zw$qfAq~EofJ!feu_owP*#;lf=WL$O{S6_SVuA{%7OFctj|HVc#aN9uS-*R^o0`$Sh zZd>zx*2JEbF)ajCYAMyMWJ3~?`ROp`ivd#vkUK#Cg7XuNdnskgO?4`z**E2$yz!7* zv2MVQg>cz-WtGl9WE{WSiGA-lxJ&h7`ZMvJLF`SKA@pbfT?cjmp1mP9ec^lXe68We z{l7z^f;lTMGDn#av zEczjHeymqYOz`Qnzfi)x$onwomp+-BdvgN3yaEJXJT36$-(e=1e89t={rRh@Gbj{C z1lN|k^oa7={L7ae9Hf+=pmq;Bf!0T=+$`%F#NzzI{d2Mdf{i)oHn1tp$({QCJ+U1v zZSidUJ%K(Mdn-phN!|X^RPKaxfTE9*Kq}_E(^#$bPDhhy1}6v?pGf6SOi>wnbUz{q z2O%N^ZaBv!iS#ZdL6_lbfpAp>BP zZz~tT8Ua31AUEX#k55!2{O|)ccx^ zcigY0{Cu0&Xnp*fcj6nvt&_knZ!xFDm#!aVA_or_jE#S{;IAIO!5&bWNhu}SKv9|iQud&;}=R88W-gX)CmeL=AH&~ORzfNqj{Zz3n{J%q>&Uin$2KGQK0hZU{;id@R-xh{o-}^ zJsaaqWec+GY35DZ-MhP!Aw=L6V2wK)LmJu^GQ+rRN-O?|i=9SXzzQb5B zMePeQ3th=jKkhHo0CcL<;$Nm84`b?Z;Kj)_tq}u?MVpPp;oAOi4ZNe=#bqP&!9VZC zyAK)JOz^LVv!x1T|02XyvQR@+jT*2BwS>|OOQl*&(rDD;gFkdFJ@1S3qq z6A$FHJ+Yf2;YJWS;RJ{lN#P0M?S!hVRvfIhzSvv9w1jWI?kvSSwFy zo`FUdrV1XMGqO^W(-!#Hsr=`U-U6;1vuH|9p=;362z)K>N-G$Nv*}lB zSvh3OZTiIN&B-k*kRwp@79vdnC3B73y?_7&H6hwG-67g&*Xt|N_>lWbCl(zqBz@zI zushK7i_0q%OzD7P^2f>Y&bAV*4ig?&m5dg|=O4SB*?Ov?|46U*&*Pa#ZpMsu=H?%i z)cfrToEfndg^Nkl9(p@9tn@C;q6v|seNOSXp`}?8=zG%8RQpu@aNQGOK{Pgk2m)EG z5YkEd?f(~wA4F+DV;m-ys69g)z&~lEso_NoL;p&yHq8^s0(Hyl+pEnJJB{!ZEwRzDoaG@%5fP#u_7T z-9L)>wcUE^=SLAgi2R{8e zX?Y^A>y=H7du4IGyhg+QED~M51_TNiwcK_bCgv$cdpj<@tcmtWKNGOe zgX@V!<%q+j`-EoZPJk>k!@C+)X_vEJj(bB=6*6u|meqYK51E<1)6c$R6utlB&iCm! zO*V~*bN=w@%ClrlM*wVSgq$YQdJ!^aGi5G=3JNg(gzBIzBc4enpfY=x!9d(!NOL_l za8i}QF2M7=v7V-t{pQ!ztUCp-f4)uq;O4;-RX<}N4p#5KJo}-K!9jsr*e#B|fMu`& z@%W1sfDw#Beqh1Ots15koU+~Gx(86<`@eX5^LVJ+zioVkC|mXrLllLSJ=>(R<&q_l zbt=0FDa#aQMD{I&qDY9cOxCew?2;|}ZWh`^W?IH4Gv<0u-{<#w{qEGVaR8{>(rm zBk8sv5Sa`}!wiE3SBSs^+zgSS3I5-R;2U5@&PXJX^=Ju6A-c1SYml{dV|(Mom$M^P zbyjx;wx+Jvy|`WwlXXNL_catNN}#$1{k0@+ih^c{(PAL(0jR~#6>${^{x4M^ql;i)7Gt{tSX|hyV*>FTtYd)m2A^bbwH1XB zKnKo9ja=3@S>juJtx~ICb1&snrAB<_Hix6&C6B+t|I_ZK0xkzU@Y324NNYX2$3QUD zT8+{F$VDg9R_VWY3qfnf^PP--oG9SWXB~sYDbgY%^lb5kqq9fzlvIVd*cQmo4yW>p zUCeZHl`~z0`hXuLZA%e~e=`ef;O9yE# zWI|kI`0;q@TME-v56ArBQZaZgF4(FrZWTag%E6Ea?LH%HyzoQp*yZfjeB}Dj3TGbL zrkhe{@IDoDtT6&#z*flLiO(yxWicL^ht-=7$j7jKNnQqH`-(L|*aw<91m^I)`LmmB zsPc6zKh8E};G;r=7>8qdA5GX-`4J!cR(#NPsP^GPV ziUNHE?=D9h-2@(_3EQQB{W{+6_%G}s2JF9~leEyO~=Qq#Eq{lHi+x}mXca#HaxoJZbJWFS&q^Oj)F zmA;4~rtn=B|H<9wz$cUwr2>tl(Wk5rqmlh^Pt=IMHNbW>W}}iHgO`QQ?(-zGxsPY8TWDTY)VJ?bsg`y0vTe=Ok=T3O!~T+jR*XH4wP4ev4|{Of2M1 zJlrAo^*-7qcS2Ly>XOxd9l41Wic`(?Mfnr_@|?+)`@(*>G6AVOFc$8hlT9ytK<&Ut zDQi>@|3J(9k2zm?Vkm-pPPH)Y6s?KkkWP2BzYyc(ag$4zLeSo8<7QGx zG7h1sCyblhGv%2e{1Pn1e8B!b&IBvDOKfk-B>!r#_&S${c?rbw=NlTVD-dO^pB<=MV*_j5d1G76$8)lb6@@|3fxYSL3_ z@6%WOUW04%Av_H|Y@6F#b_-dOmcgO&R?5#+_@}?Cpy$YgsOu5B-#Pavg^t~ z82qi$95?$0_Ooa{g)RvhGsGADA2Clt;biQ-<4y0us+iZlDb}iVx-G+QzqLePWH{qC zwrjiygkZ`HQxuk0s|d0-@S_OuX6~;=JZ$)GIV7#OA*4`$&$~4tgm(RaD7(t7e{OYO z%oP6ZPe=ra0Uc%;f}RWR4tTJ4_1$c^eh2JnZ+9&t>M{n1X-3jimoA!emny}wu{On- zrk&@7DxJU7`4{V(t32qe;ImM8!z$5LqfXWb`4PTgM*jywF|I?Wh&==m17R{?1MqFi zJNkpi%6aPYqYrFsdy$u}!7R0fg!n5Hg523_t1wl?MozK)L^5?V=!U=*ghRkv$AIK9 z0^1KyA@UP$sLkW-X+Ftn4MprgmJ<7_(o#0nEXZqLK)TY+G|=yyY<9$#N5A!C%rM(; zXC@O2Y8m1aXOR?^_!0XX^ljL1#bbN@;B`g22ME%jq3ad)XMaa8xXE>FmoMf3cwOfZ zf*?FTJpXe+4ccvr_*F)(DSp!4C9Kn>)cjM9LUgT?lWWd(_PCX}e_<)kyJg@O#}I2f zm{ZufM8eMPrWXQB#55$;LJbI@r$RZe(LDojPC!&)ErR030lF@~uJTqdo8o7wo+P%y z@^@~#MFdhVvL3flNt1+&gY!pyQV=(kzaoxLv77u&%@dR-={lTUe019Ui-hz@#OD82 zcOzq|cGe4Oq^*rEh9orETO`Ke{=p?7rWr+krQ0L?|B;9 z)Q@W5D2`0Sde#B3Urf^|1g8A&F_x3#++NIQ#(ZgM0VY2VywEq*U91W?k|@A-34_u+ zVm^UZ4eX9a zr6Hpc|7VNH{}Z*e|4J@0LUT-Q!0z+ot3c(~6qH<3=0FvEe~f>4<||JNyW-oCJK(*i z5MVcdeTfzavC<5Id56G@AZ?E&NwSJ4YSAo_*lSaA)Ri)yH|sk0{r_=T_)IyWnswGq zPvc};w6#?~fgM!zs&;D#UKt!LBC>8QxkDU1<+cLp41^uv_~nN${D(%KK%bI*>=izs z5`3vGw!aS+8lqCYhfzJG^$q*`&MZg(j1@r^uDpqWw5%zn1ey+udlf|QhrUb3HwM1~ zMT`>|otWGc7?%+q0_bNKwo%utn`= zwKmX^^I;2N1sx}xSM;yu+KefT%eU8T1RwEFkx@`*-ev6Byb3IPhazi=r3z#m7#4n@ z=ewAuwd!97q3rANQ^)EGZ+hlrIR)Lw9t!vNanEXeHZ+flY$)LHWG+lG>>zm{^FcKTcl8IDn(D;WSVZoa(3;RqO!Tk5Ri6_cQV zz%-S*4`lfy5wE9Ds4!TTHO!HEh<`E@`O$0>%P@oJKmmw%6LO(Qw7cPkaOil-?8T`+ z``ho{jq=punY$=t@qD<@{)`;Q&qJW&au;~r_~Jb~Xd>K6fQ%Xer+CgJI=oL3(W}b7 znQv59j&kAj;&K-I`@esL}L$b#??cMafV9JU<^K;2XAJ zT2lUP>kKMFIf=DYZy{MI`s{*O@a**$JVgx%rlbV*Scf(`%ZiKsYJzp->o0fBnZm=< zDyq024dzSx#XjFtqhy)5GRr10BW)>7Ztp-dfeDBW`kCUpj}f2R6R8E2Cabg z`=)2~G+=*WQ(p21`PL^r;TrEV*T|>exW1x%IH>c>gC;*i1OO(hQrjD|q};Q4xL` z0l5cY4rE+L=*-lWm zn9WiN9C~!5pKQTV>A|&~l?(&)s|UeC;BbR(Fwm^dLP>~Ya^0P893!(hZ$cV5KjWmP z@KJJtxbZf~LILh)E!%Lb&ge6*d1&Nm<~YdPB!3fj&VohDlm*#Jt|)^D7>B{P9K*mX zDMf05&tO)+y9#8K*P54p1Im!@bw4vzvtS9|toR(aIc(c^Bw^#d?44MTtwB|Sei$nh zjR$c#wGZH`8c+PTUQl$M)O4!(l+``D)iu82>g<2sq74oGk;$eF4f0$yw7ZbZHDF{9^OtCG(Hw z<4QLLIkJS(KkKA`q1)5{ZrrOD2EZ3q7?j6E=>QDR_Z77`pBWV`DRdBJ`qUxrzTu;& zW@nX`Z|c&9aySCk;ByB7cL3=VJV_D|og+wf1PPTt&pm9Wstt51y2|}-Fa6_gXf)^C z&9c9QzPhwACp`y9Q-ZpzAx(Kw2%!5y?A9`z{M>aWX6V|jAl}oS7~OJAJ@Q1zi;z>I zni_a*K6>=@ZL_anS%kB==czc>K_o*RedPiTk)i-0?@>F>8R$^FvqzgIUKG5KNBNQS z?JY`jZba|*aS}5xk)wL1BUit7b*N8w$e?kTXG2s?e_70g6g+IJd)M%|PphipPEG1r z4#641Jpze_NOeC6cC5v2wx%&|J*5=gJzzKWXGrULuC&AFmBmY^ht$)hA1#eK8nDY5 z21oOwSaQ4&0(8kYlC7DC0RP-5!5hp&4gguIf!w4^o!}9`MIML_wJZ{xTRJ=QNAN0YL2b9Gl0RD7TM9cj;Cn$P}m1p6_9X)()^1KG$A)juO1Hq*t_ zv9v#Y@JlYQQwi7WEnOt+Rt$FMrawQ9`>OLSxDSS@82|r)X>4l&>3@8&v03tXA-BwlpeB>PJr;azRH=Y#Ow=(R^zHoen-%!q+a{TQE zBn&;l_})d3K%NQ$KT;V{s9lRD4m=vX3n%Whz$dwC+nx?uqHa#}wA*rfH$PO!tZ!LL z`%ueUe$eb&BUTMby_eAY9P0p9em8xz|jjKkpzG6SYLWFMZa~G4W39A#}F=yP3v!mfg7% zZbY{aUWMcbUYt*UUKZgkx&&j^{ddc>=Po&bBZ^T+D}xY)dbij4cr1m_y5IdIt!ypa zU6#hVgz$9fjU5qtd%<5JBwjbmk3U?$ZAe~Y^XF>DM*b&^K1$qO29jGP#%HAp2y ze=+4=iLP&{Ww_V%C}SAlSK-Gq&ZD_{^cwq%UZYpLcB1}NKKatrG80goFD898^O>ko zuR}ttkt*a!(Hi~;Mf9$~XJ0kh{4y)|KfSDZ=`!2(RP#OPTo*1i(xWo+Jk6oF29+Q8 zxw^`KRvZ$9e=W{GIE>3Ef(B;yNwD3**g{i`u6 zH#678Ksl$-`sCFk>RTqaWa0ooViMcsk031(VhQ-MO*8Sn%B?i;LmZ&-5p@wiD{K?5 z{5|N9MSzISgw?$%4T0-s9Q>rzx zj|NnPa^TpMZLZ0^}0bWZt}myaz9#VCi;u{;5sgJmOq! z&q0^O(qs#+1)I*~gov#Vl|iGM#DXFkyL%omU8*2Oekg?f*-h{!+^_~L4u`Ne;j#DW z_^9C5-<`iaXs|lZ`Wd#Yy4>4z$U>neqF!G^cpWwE%l=snQUIFu#4ZBWB4|DsbfaCl zGG$;KUV${KuKDg5-MiEH{iUhrwQ%*$7~NG96Gb*qVn3EK*#O$aF(B;dxb2j9Xm>C69ST52p|UNmZ#A za6dH9s0ga~pRU;V;47!zo+4id@X@icECm$t0G5Ei`NBSVwAN0m9LNu8Kp!s@6)MH4Bq;>)&*r#!$~dL`8v1YUsu?U!!xRcK z@f0?Y7{?OGzBVP*-Q>P6!729AuY!%tm5~Z$&#{5L$zfkol8#K+3-%7II7pk$L(jpk zaREoUK&B|jhoRWnN-rQ8iuHL+1daT%s&HPZ!s&K|kv+qjvn<>VD+PZ^T$Q(ZfK?|@ z_PLG903f;9w*G07t+8HLi6u`{FJRaQX=A+;Xe_;W(n#VfUOYYm!Q-E+y%0B= zAQeqn8dLpkXP(lwrha@kj$JdAH;wTL*$LRSijs_Ep}y5kWS`N=87+ZB=oI(Xh-EX| zjwd6*_cWf>x5-!_UXX&%!6y`%LTxO*SV)zj%9>k|hH?})*?edPw$!t%*@*@#^#wiy z7x(^eYsyA7;(_7<|Y!mWAKfxe`;J{=5^nc`*9mUy#><&Qy-BJvO;K`qgF-I0JczkAbRc zB{Q0a-op|K6Kbm#@=mGxIX)MT)Rma_dU+mmI--o51kfB(0`uG~s343FA%V0sM-lTa zZUPP+BQgaw$e@?ZjDC?8)ZN_3XE|U+llIQ&+o>i#_1LWy$+UpQ>v}xeeJ&;;aVOoi zUPGKwY&43~*llRP`Wh_`M7!#1e7Cc|Ei`42mT&I#|Jc8%T1C4yOnYBmudF1y0~WR+ z1WeKXm`ITZD51uroHZwD7Qf*^cpC8gx#gw0mi9?`{;)1hLQqO%;_0dtXU~{0obew; zreF|@ZxBK-I55+yH2~1iCqD9%9}A9j(QGENw}?cDy-9uJzgQ7W<$i*U4m= z3ynk%GiEz#0pL_aheoXJ+M1`dq|f^_15E9sLDk-LON}-erzC}}dkV8(gP4-VAWx=Z znPNTUz?8D@l5z3bl%W9cA*cHIn{U@|_mvP``?GMS2Q<$#kC_^$|7K%=t;q?ZxLAjF zKJb~rX4XBE1gY*b5&#+<3eGt} zv_YyyQZFJC0AalKJg$qy**c@8w$GoyH}BY;KmAhc=eM__$tS8m;qQsKnkm@ytGjtR zCTk1ehb<6i5&wNtZjYaXgJZTX6#^;05L8)W>{x^M!@G(5SNtWrP%T<%cdxy>S@+%f zaqYknyo>|o(Q&Xjk6so%>NcVf_X%6FcU;4~4}R|6zu1g#Ng9xPnXgpF7At>?E$4Av5%~_G*!m!A!pJCv z#alPSwC_-(?Spa`B}JTHv|}3wZbn-3CEdldpwvZgKLLhR zN{!_-`<2K%1cWnEr~n!l;78dKMLSZ?L2VSAsK=!OB;LNMzOMel5rU`r(L&E(ybs)# zmzp?koaX~a&6$$Bz>>u(0H3!?FNNNb;hSj}Sc2iy()+a&!@Uc;)j_g}mxW(s4~4n) zWbqp+vF`}c5CcdKw7gTc)l9n=*-hT&;fbu7Z3jF)|8%F|<8WSeZh84f{?^<~e`-+1 zpDP|OAbgMB^3Q%RZMtNsxcbAxoJ}Rd4>Cv3yvLD;=)F)jj_I=dnk9#h?cM1w#-SW{ zcE2^hd-r5Sv&}oywRS@U#&r%R`jtpE|4(Mw4cyd8#ZIX$eozuh35{Ca6u#7wJK-Va z>g-DF;_DRvIx(4T$|p#%=irM0fP)4ODh^#&z}W_d6AV}$WXGH!uR3q#|AybH`jBP%<|kMBdC^QFZ(7qu=Jo$67Xp5gv=y|M;Y{?+qdM{ed!Vd$Di5f;Gj0kJoF0?kW&{~2P~*+y=Tyh z{Q)#(7uuh-yWcSB5q zxg8DEEDo*!<1?FJa5j|ueOc{t?#-E<$k_|_;%Y*LhX(2BaF5riGL4MB|0fhY|CQzk zo`Il-CNPB%s{%|Bgf--{D)8{Wrn^c9RnOfIzi{O40O#845NdOwc4n@6#m$}**D=>k zgOizFV4Q;y0U0~87*NsfDg*(=peVx)6!px(x1Jqp;8UZ45RVat^zW5JC}-Nj^16<@ zT<$2(-_0A1=0-qalp{*Lr^{8ZGH0i8nj z@}XFK<1~Ryr#EO+&rU*b(1G3s3%ob__wIQrvg>+Ymg$p|ZRJ7@N*|qmd;(XX9Fl#4QZTGs z*QI5{Qaf#GZXY;Q~kHW>Z5Ek zg^8U=aRW12kVQiG{k+~AUgh(?cmBy0=R{^f)Q z!Bg~Mz!iB}Cs6tP8Pqqw`tvX6g?*d{`Yy%Zx%s7TcJ@x*PU5nul}fT}cTq@s?=Ja24u6*FTN5g^ z5PI~cHrnFU@bJMmk$mJWNH3QDX*`r6HB^-9aR1%AvfjIv9^s@-!GJB>Yr-P&%I2TA zNdd=o%)0cQE_1_J?Y_9PyBQe05rz$G5X)o4YcYGEG5LPh?`aFf+NKhdtl zvj?xqWd~CB2S4M_`8ox&b+z79qunF;5a$OdTp1;A-L5hP$nVUMF2@T7#p5StEP73- zi#qor!UJ?iZ{G*zjshO7fKH)_C;u@#KS24E9_tkquG-h})Qij~W*MNTb7zRd%CIzB zT~^M1P`>4$j>eHJw%sttUrr!O`T-mg7VDRh22RF6)3*iWNAfe0LnA376d4$@ab$4xO9IA3;yp7Nm=%r~jwoQc@@b>C;krow2=ETomcDb!WH1 z(z$Pr&YXSE(?+(QO4P?3dwT4eq>crf?2~!7uvUxR17M~81-ux1xx4iS5kRwGJQV&= zFClJ2OSvIt)p06RSsqnxYCWP7bN1zQtSHv|j!XXYTnX6|`emF44@_m({V5ik+ohcgx0`o%@ z%p&~vIi?vHGDH165c!TW_%awlD(nDWF9ty1x^z_|;WUX`n~_eGfcszibewPTD+EXf zs&2k?adH)^KA95P)NL=mz!R{9Ey6;{I&_d!zWf9V*ggsW!v5|jfN~S52uXESOT+^E zNs0z{CHQWM`#PU$r16ktJ0=u85nhcn5&B2x_V7Cck`H(g3wguoDssMaFOpo3!Z5FG zk?#$G_mfs-gt#NEG_g{0jp>d)*KSe+V>})TKYHUlSit^b$)a;DdjyjMp12kixZDLY zuz13z9|g!BNixc>=WGO) z7L@;j@HYkrt}^m3tT7xsEvXnO%9Nsx_2*-Iiny4Dv{S`oeM9X<RhwC64$^=65;(-o?HFXQtKynDG_oQPA?*> zUg%wj%>M&Cp6A=9slb2yl)IZ=^9|dbe`9`X)V{Vm+J67?-(vvIV1`*P`t3^iIZWM4 zkYxc}xd<#5{BClj75JMQT-OJ{KpjboDNp5(txE>2u^3=_>V!F-}RqdZ<{Q|XZdkdf6-GmrQG5xm&M6Hn;iwG_*>A~ z9(oP!3sXdmY!-_-jvljyGD%<-+!ZF7iuYbEebozmLXT&yD|@ylqx5}=!*S2F2UODA zrs7oU6sHK(yFpCxl+`aAq6`;+-B9a#jBv&e^U8<5g0p$XU82 z)(TBm!_VCK1v`BXp{jx)E#ek$?(7_PAzgkKSZ`k8m1jTM8@#r9_a^BXe>v=#OXDEqGNSw(PYPcUY zCAKZ8G7~1&WF<0WkSK+Hq~LJgn71?9QI)#&5KCIXE?()AQJ&PCXC3sUBb2yO-;3fP*dil|%0P`IC$ z^z_tytMtlH8~`RQxQd35P*sij_j(qWEFc@6DP*I37>2~25g#%<_s9|#HOV9r@@1)QK(_e7b0(hnl-e65p7n{_4a# zK3W%dKS&)dcrg+xB$ECK#ETC~uNp)HF?54yR7;IA{T9?rgahQYH4OS`$vHvnUpX?_ zNUw8hCwm80mpto*^Rxv$#fObQ^I}gD22^Tk zN5X{uxnVie(2x}@c>8dLhs<*45g%PG<{KaY!vA8orSbXwou9;2(ZcMY&|rn}YYS-o zi>QcW-U=oXx#Md}(VP3t!aJp&rdxm*qXKYVLE1pEn|=2s;ub(9j$jTkJ!$(vjoE_c zkQaK-qrxc1t8CdfYoYs>bj$>ImVJ)tj1!ye!;_u+9e=_hRlqdxnqe0y4J=#=(-gvj zIsq^gMF*IMQ=h26wR3l5p1-lE800y$eP`%ikTTaD^u>kbHzmA_&w+;ICFQ>-pfezI zK(gkE0<4q1&SS|&NBvkv)E`A*b^b;&@2E4G9{iPDW+`J$4e7p5Vsk=*a$>S-T1T7ROiz#%R`Low%)rT|}?ZsmYsGMyLssmVBNkfme^D=J4N` zZ{rGNl2%2T1OIoPqQyWmbg^A`jAYUEEWQ<*^`=EN`tG?Y;C*DyhU6vIU+{RdTsHFV zrMrbB%trWjoq_G6!zCvhZp3#AbWiR@(yT?`F9z9Zz*;u{?^J518vJ_|aEf9bXIb4v z4TKDgN~Ak=9WAr6)gz3yv>vYsJPLB|do@)x`Cm@Ay}TGRB|Eq@B1hsJ!gv<}cr$%N zk!6`{NFeh9Zv_7`wqVP;{{nb3Ch&~;N3Up&ZK2o%5(E!+bEGf$xy*%X$-|!A7 zzLDouxdEb8wtWf#^s-FXiV!8zQ}&8PcKRq_Y!LuCv>vEkbkTw7h(5)H}p1DX!~6{8Cj)mq5-QTgrdW8h0Ooj9&}9sw~2T za0t_hrCJf=2nP(VMwGA1$Wx!k4bW4Y6dGGWiPmt@H1(W{{4aW7=W@L+jMDLVPmT32 zHw2Ql{XLjD0#AqX_Z9fKk=Qb{YBup>K}t2kl2D&MX`c|2_0m%PeT$uXRf_|M=c?k) zHf{~K8qSoCV|Yhigz9l9cVQ^DoM}XzDGrxP2lpx z+T_aTix=DDUG#mQwEIBtJL>hR%Az5xJS9MB0RVR38jrA~S4IHP8H2=h(lI9?luVI8 z3iAd=`c~sUV8`_6GPN0bd{a>juW)#dB2Z7W`5XRox zf)vI32~te^#>q%)%GQG>ym-0Y1Nq6t-ry9AALGZ<5BJ)%{ylC;d%=0>dBFn}Ht@dI zkajkGf?AaF$LiD~spEBx%eU=2izXFw>$%5g{sJWJO`>{JKl|;a7o^k&$UtmY3z2G! zT)037gBOzzDrU2N2u^$tk(a68NGJ#_4(IswU1)jE%}K?kO5I5N-e*&vPPTKf-1mQ_ z0BkqZdXqp2I_#gHUO@96=*f@zvsAz6=in|l?hWj!_p$3~UmCK!FLEfUoBudBc=HZw z(;CvrVDZcBmSRyrlv6w&Bfz{tJ5@yE0fFJ6>Y$fRE!pO^6F)*EHdO+HF4h+49u=xe zVN1(87u+J@6jt}1ZI{OuIa%1ovYUP6l5&Q_!cle@57t3(`j}?lN16t+^c$z=AkK-aw zl(;Oe2k+-BzNfk=|CF6KZM^mEIw)AU+B+hs1p|CE9WpXL(tM_Sq}{m8dEI*n@kCbR zeAcby2W?Wd&v;i~)WrOM)C{=KKtFK|SYTbubRjz!-1ta)s!tF6upgkBYKK#|dL3Cp z4U^Rksyg-&GqtU`%Uj3ue=eB)I_{}jtNrqLAKO!B3mBqYa<42DTM3jb;1VE=e}Rz) z^rd4D^^Rap7j!pW`I@(*pfSKn5*qqk-TlNBh+s z&N>4lo>G)X&-FeXlzKVXNYTYJhycq7=TQ@{qLno%aQq1Z&VQ4=M|) ziv39F_)H6x*h$#`hY!)+TAp(2^9-*2!m+fD;@g=fijP%G;y<5|yLvNRBgdP>hR-^O zkib%oSTB0865)v;D)lMH{xrd?LV)L>lBO;IX-C-X=DeC}P{JHvcj$4jaqAe6em^5W zZBk&d9!XFfX7}}I<2cB40XD6K5IDAT3i;O?jRh?(DEev-Xsu9#yMhgZ-v(Jt)rusQ zNfZoOw0#xzxUHs0>)Uk6L>&2?x^PI$U2cQb>$mwyY>>tOlpg*WV)20h3J_!6+I_%M z0{KG*Eu!ZeiLkkREb?qQIQtT`jxai&Os=s=nio6sR+bYl zbwL1>tSte&w4l)uruP(Z0OCY+ge6kYZ$-|MCqQpqcX-^Tv7{$Jr=z2)N^f_P zGDpI|*j^$)I`Z`DaV+Tt@c`3oH6tM$F3_$@$`JNrDNSj)+0kcw=T!nAQ`_ zNbo@y?jOZjN~_lOxhikmQqt{D@cqpLiU z=KtuHMr?KtY4nf4x16^=T~&z7P+-o%tET?n(6|2|wd1iMCY+PB?7rGhfQ3T8-x>eb zcj0sH4=@gxTj;y}(?##%gokgKeb~Xv=5`U1Oi|p18H)$h0XYHW9lYTRWJ_Mwpx_@| zhhE=XRn9*3oNTf2qEKQr8u#?hXR^#O!MtOR6x5m{LO^4Pg@(TruJ_j{Kb6sr>22m*0Lqtx>t+R^J>^YSrTFz+LwE{r-=c z^**qCCKql~p9}zsJ*138z%PTOPCDQnofW&s1inQtFPL4&jK8vyDrD@z@2M(<!VkM#c zS%xD6emN%q(^H<8tI(3=CxRQ(v^)B8tw?onVc!oJVpio5X(s;TEB-o{$0VILn_kSI} zJgw9#cKAKlw1^`azlQ>*erF}>fI1*H;9pq(_uoKJ6ZjL(mrw-wjDeykRu{TkFISVX zf@mZ3>hy0_iZ(IT^I<-+%D&b<@-4Tys-8SDp?UBH$9obKRK!ef5D%n)ra+sF8TP58M}5QF}NcI zs5A9`ul)%7F(Th~;)~Q>7nWVnm%2|>mznZr$yc2Nl^0vr+s%F_KcV>Zm!O2t3 z_@=06I!rO04J3O>L~00PBcdwOg2wrUdFE9D1=r{wu7M3e`Z?<@qom3_=&vs<)+Zes zrSlHYdt_N_ve6WolI5gok6tjvF2p#=^nx;fogHj~s#e=F09P)rC0IT9jam%J^|ra_ zJ88vxqX5=k@$MS?ZiwS2I8#uc#WTpzV$CBPa<;q2*qGvzEKYP=?@W>3WY*B-`gAp+ zf@`1i&D}O{TH!>FI*C~1D1Ro!cShl$6L?UDq?WS|j4-VI5eHYP)=5Tks1w9EU4i01Fl^G7 z7iLBy9!@%QmpESKpr3U1|Xufg)_cC2ZlDLO9ZPik2FrnL#C@ZQ%gN#;p*o&HJ;=BMgSj zGD`downqT+XAD1G+*JQ|Uny%<E1w-7|+n3qt0P-Hf3FEHD(yFh(ny zd<=Pa!FfhY_0Fi2rIb{7b92jG^&7S#VmG5;xF7PAppDnmwpCPq6Hm)42aMkC{H+B=AXkXJV-LY0=_Gl2oasf*5J1;kSmb`X)%aXbiV>P)%p;qA z)oY3UG+%1;z)Z6M6prsy+mQpy6zs#aNdhziiyQzRs2jXLAl?^LCT$irmaR)#tq~Vy zYrF0Rel(|AzgZsPExxwT&iuKwONsJp2@XeMj&yZHn8SDvZ*ULqNg>_w^KX<6!`M1u zVXu{7IMO0LnHq?rZm)hx({ziloYHjju$uWnYOC$QqasrLIL2~9u1j3IFNj>R!U$vL ziPXE{8?iSSR#59`(0qak)OD&LOa<9DyWo(2Y(`S#+K@nlOT(ggVDG6rS4GHXA7BV6 z0x7+aZ<+~CNI6}1p9O@tso6}|RAUsrN~VrGEw*p*H8z!SYzx1&F;7?W$TIWjeyJe@ zJKnJbg0!JXsmgBzs(m$&YxBe(ZcVo?VXICd(?4Uj z$Pf_o0zB3#V%~aH4REq#?E4C@@5rW_xznanUVJN*t)0}Q(NkJJrCqPftiEm`dOZ7P zjD-zZ2WI588wr(Db-EBoelNNpq|k`&l{59Vk4FObx0H5Om<7olK9!<=+Ce@w6~H{G z&WRJcS2>@i!`QP>jJ=Ru7Jb11)CYvDNn+VGe)t`u*FI$B8-~v}_&BIpPVsBk4||y# zfOhF}Crl@DpXT@XfvT(nOa**1K3KtSfV6AquFWO)IPng*>>-o8M5w)1*Am zE7nm>vm8K0h57}l`|Y?Fct35P(K;ZwAmw>0z^YrvJDhy=@Xy;-Y+H~2XTJRof51G{_S!4WH-{v3q*YyW_q*SeO|s(uu=C-ESUhDE{TmW8 zFEwQYcFV351Q4@TAj(0TF-76s_scV|qXEG;d&LYCG5d+nBDkiIFXGTfWhn`NQ!k?Yq9U^>!(>6pcmY73>NRYKmK|Mkzu z-+aziraz>@(6_`a0m>|=vm2a+8sAhOhONy;W4rD%r79r}77rWkuU49x&BUTDEnubC z|J3G!@ZA9Apk&Xb{a=^Oy>4`I$=h`ja>ync-8``8c>rR(;1 z`R#i-6C(buM_6$4>TBE}$m5mGm`5}K=5h?t$f~Ap_rYPBgmia#s_TV+rW#W=wfo|q z>$Pdjo6g#_)fL~$YIYl*8(WcukBC2vl&g$9M{{^y{M9*tRbPz~KQli!=X7pXTqoPS z%T9OI#%cLUw8%aSEe3F+T&-y3Yeg)H0k@1xQ%IA^O|1f{%0KwK0YSlXwZ7zgXBypd zPjK7axp>R$f(xM@&Xne4NY6z^Qri$4)h{EJEz~$vDT}Jg|M)e?uG@?ys?NBTgh}04 zZY+prl>b&_@kD`rATzmI8>gS-Oz zy^7N;`)e3{zy!l!4HEY=-R^;z@r!hvyg?mERG11y4OQR#Fdg#upV=^r*YlUQ&%9b{ zP(GsGa{FLh+UV`_BgFq9H<{f8u+IYe6oh49Fz;DnP!e@!|0Gi@RAA2Z$qZ%x_^arX zhOa*EqZo;Z{Jpe#_$v&N;Jlj!cD7QaXq6iym+JxQtrk` zQZWUgXIt|pFHXu%l5%gpd1Y#T)xO2vwmFr2ZbON^>_42n9Rgs(_hP%EIsj&?OIB-( z2OR{m8~kdJB>E+hVG$CbceRNjR@7?VH=J!=hPmhAY44g{yO-QR1uS_i6#6ghqEtTt zLNtM%y%07=w#mJ%jber^;3u2hecrAN2x)XC*JXrhaYYLoUWM!&gN+MkLg)$;KSC%_wajIFW~IogFKL@tCf(nXQK;w+b4kyV2v3=aefi6t%3q4%*s1& zsgZ>~eBe&wt?Ml)U=K;wd!+nly`Fi6?4~wgbmZOmpLaUeGn#3a#PE9&f@KP-V(vzk z`I2a;d<-`@1g9z@I{=*%2$Y5W0VHs1QMx+Pu-a8EC1kNu`p|%vJfAx`q@a1&0fN7~ z{nNDUisGq~*P;w?zRFLBIKf0BNT6*~ur(a;hz7n~RO^(C5@u-S_N;(pJp2ykaah2OTM6>|n}GZGMbYg455Qk;z9m6$34*_hMD zxQ{qU3$^#^SP~xK*+JFNa!RcVu-{j&tmhqaX~^{ZXp|NeSU21$D5*Jl;jnk7+DSDB zP<6lJuCN==IslFJ>MV5BQX|NL)s_Bt-%Wqo{1T_4*xdGX&?n16@+dIYS;H0zYqQ4G za(W09MLp{B%|L(z6neC+{yDm$9rtq4uYcqCX3T|yr*_*4r0UZKq=8tr0`Qb~;b1Ks z1oP^T(R5wE?_ONK_&q#KXCV@=y(xqQB}bn{v(LTPVBb1)+=enA-)k*+d+SLX1b`LK{NpM8Dlb-LLORrBZ$}gY)lp-(RG_3UD=jGsEd7S_sBu+P+V4X$>l34HW!=oSG ze|Ql}9=bH)ylV0{nsJ`Hdgh;UWM=GMN81`}#Mxvv0X!xGb`9xb?Gxoeey6uWSwRoq z={oA%UzJ*(bKRIj##nUmDx4QeAL=)fn*0-`O%7~Wwp-92a+(qqHaVJL3X3+ikA2UT z{|<<04@TaD$YjDi5H(L~k6PvVJQ;b0EN$dwZPy~GnI>|k=31PMrA#OPXM!ZqCpG*+ ze%?_MERE|VO1UnY$*Ww^#s^5`u>HMO(}S0P=1GZ0+^&Q`P!cl;JOLH(837net znOSQapeypXF^LGI8qDDJ9{yS!Esj2Z_D0rO!2z(*4{%7b^01^wV3vXrJ!-&X^{Jvd zTZ7{MY_UN-R`q&mW$`6?%QfY<9v*O0NIU0nekkXb=V{qOlc?u31lhmfVz>2srNO1hj+wl&$98cZP&iqiyE#p6N@xh$u& zs!C|KuL7 zQKg9N&mKWf>pw3(Tg~5Hql#|?QxnLJDXRaAxGxWfvVZ#?AwpzNmLbYc*(;f(EU8G@ zV=5sesVs%LM998{QlYYDne4kU_N9__?8YEvoso1gOYhhHd*AnX-rw{4`+4u$^^g-Z`yra`iN;cLW%~>M=1=V@9@QSt{9hDm4I_k3z0O){jk;Oyk@ec|;Q6wOj*FffLruOh1Syc5 zK7|3KV7HnCz&&sQ_Nl9qnnvYjj5IOR8*t9593U4C)iy|zP8o=4mD@j7cIi(RsQUVH z;^mYr_hi9D2go0BX;Qs_CpCt(t&_LvU_6O+;P>8dH_3rd0@I&YdzIz%PPZODd%L4w zTIL@9VG?iLvHh40&|K5l#02&8><))My9ZLT9Ue_US2DB(%|bhUG?6&N3H!C2yjWg2 zICiJ*u9nccBCAzGC*Q|3g^!$0GIHW!h)36UvsbEMrJk^b`ULR6@_o;Ier`6N73a4g zv!>+aW^JuhRjV63elpOJA-I3jmjW$dX(BGp&ayH- zwn1tZI%qT2yErki{N&o9XXYj$_5ChZ9HTO~xrm9^9ja;kz=CV0v^_K+o6d@ldyH*ltviN zJKs#~y7A=D&Un;rmGOW47xVT6Kp?Y#c)66(2W{;g!EFZRVE75Uq$I4h;J*4ZzCcy$ zcn2PH`a9Z9H*yc5ZTC2!Zi^3PqvI^{Y5o|B;t=_^kel|w@ z6BJDc16(YSw#p-`Gp?n>@BSigf)aee(M)t$CA};Vv~h)a(GdM$Enk-gvJHU1^g1kE zd9w!}foOb0fsiJKyx zE|dK8Nnhm3ikFnNh(*7Ng$*6XEyy=90gzys2mPF}!T#_7JZAz1 zMA=pTx#xlEL}U!mmjh9p;`8Y{6<>tgM*Ij54s(^We@J{8d@?}?3l3LE5rpT^15b{> z4oJp59Cef1GQ=s?;Zj=W9!(M%2fb93BK?Gtw)>6a`MRZzYaMz%!;+^$m}qW*Olj5i zFG9hQ2Cnj_RUG53kHU^KMC{=iYqa5~sLpvj{5IqE2JA};CNvP3m=J545L`}WX{TyC z_*+h*hOdd$^8?)mIv$T2huXCdAzPtnh+Fq@vE9V$ z?xIILyms!wyk|wOb#?QXvvT#Q$J~xW{Ekxms)xdiLg3+?)pMbiub{UY9nop`ORrh? zeDQ=E&U0gHeS@~>K=mu}esfQvvc(@_aCA~tK@V%`z$R{7l-TsHz}=d69N;0i*DT57 zRrZa7Lvm>#>O-)}B%VR)U$KFBn~wVX0aW$$HxIvmxx;|2_UzPAe(#`yB|kS0H@AaJ z8yDsku(g`AI5La^tblQ*L@$QWI6z?HS=>Oa0~~X?Rd^e;TP?PscJ9_-XUlL6SZ;s7 zQCu01k1VMOy4s3Ck%gmI^2!3DMf0(nQ&!_=`Y%ar9uYq%Y>3#2qLoryX(C=!bvUBk zF7=)BTzrh3_}kx;)d$>0(qls>Mh(|&Z^q_SJm!#;()V+bGPHS<#I(%@?i7G(fC`I^ z4mw75xv=ITA_)*vp*sH}s5}OMFUo3|4&o0rhzRURFpg~67dIGaLT+D=(RH3InO-D>8|-G4*?;I2BU3sb%0yya!hMZXPHY0E2_1zbk2y^{7jW zr3sTfEd(!FHfTQd6gXKCwWnyKh-^x2&Qh_ycSfXJB#UR>1=r9wHpeh#dx~2@e*)B? zca}HSXX9xiZvp24GZ3aMUXoCPv82M&Dyu5a1DdAG$WnKy%Yg+Ah1U|^tqBe#c6g)n zaeQE!R&>h|wi{oJrk$U~F~QpBj6LoCBCMLC%|4Y|Sjl}*Gwz8_vb^S+bVh)QKl_x8BU`sGVE}6M$r>ES@fw);vQ~u2^TBa?F|I*$1<24bb1NYrI}x@e%_h zPGQ-ru2nVm(aLr`kGQa6rEk>A5_*<`0NpnYNmv58W)ZO3@?Jx@mCU8govTjh_u4aL zXYG{0Ktp)gLoDudLJlLn_kgUFzK2>H_nRdD`ShK7#y&Ji19QT%;P;o5j7CyC z%XR+`Z45<%GNYAVN;>uYXiaJ55nJLj0|DcYD<3~!NIhcx1|i41y0klcrWg$`QjKUR z`pX7c?na7e$Fv|>BGUVQRU&$MC;MbjL-H{tX&yy|R4E&x81c{s8V*FRd;sInN66a1 zj(MH}<6Ym_`iHC@Q(RJYRq3_#94C78mRD6!eIL5~v8dnUCKun8iLWbjVuWDe?!i@X zM(&(YeTswa-DYUM!Fq96VmygJOF{EPqjq39jZQ;8NM`8HOa#zZUwrnR&?E*=v^1US zS6(|bU{7bh7GYi;=ZtvS+xGcgu(OV=&S9J=>K@<*MAG=9yvFl7v?UbiA66Lq;iw(5 zBN{$C(^01!V|T9Xeb-riAO0_VP3y!%Wk*sNaehiD(R1M${EWwP}+f@&0 zAQr^4M}U?EY6}5-Qk`I^J(ElCv@Ok>()sM#^2y7gsOZhNSMKJ8eanGEzH0d1u_>@< zdta{T^q>G)_5xu&A3>7O8e!pV7k zKJVQxR%=yTm9ejDM7LS-I!x5EC>1o8)~J=0D=I3gCGgCC#61Fcr4PNWwcA0_Y-hws z^6=U0qv@kSFI@jccnVmTac~Hvtfw-`mr$^bY73Y~t!^$3Xl5{%4QV=49_5-HSlXG7 z^v?H)jug_#m5V$Wr8EEN+Br+nspm>-ZOm{UwsKcEX8V$LGm^>;!rW|aO{yzJox$~m z{wgql!fZ2A@fRUPvmqB}6!`u*?smxGcVyXMsMxmob`|z0epJ1bIfn2F!k@qq$Gn1( z+1mq|ZPi)>?b%(;g_rp7Ea(Qg3k%gZj(6<)g1|&wade|l zx4yzgRQai&si29qN!Fip_IF1`lF3cv^U4WN&2=lacjL#^3-0F(@u|nQ)b#C!Y*vs9 zO5Fd$rG5?cI495wDDn`#GG;So=ik!e5#2q1BeDIEQg(Uyfp=%RWjGLj{yFfwhaSjr z5Nl>O(s1L7jLx^flo=jk3>9v$F*p{$EpnoVVvb`7*3;^6Ev2WA=o;5x>7E|`A*Yq} zMe)L+@>Cz~4$-cu>Jt_@$sFzbI!JB7&@oMFBgl0%0O&kSRsV|^RdfLK)0<11s4*0S zefMq>hJZ{Td*?gGpcW)l>4UTN)RjWTX_$8T5YDSe`Fp-M$wsGnpTAin&!%9E%0-f| z#lzBs7Er%8Yb|CofS;#-P0ROp^i4A+XqbG6F#lvU@x3MY+@s@NyxNLV3j#3*Ktqz? zF#S4&&%$j)D629$R>7e-8`aRqkj$jNLH)=L1l?TCs{kq|6mCvzsnTO<0H&~P(O+wz z!*M(-P+7zA*&pdfsiaGTju#D`>)yiqr4wu85aK|i5^>r+Y8!NQ;BP?Af;y;*wC0ug zowxs(yF9Y!RtD@%wSawp*}xuieeC|!3xjKzP9MKny4`+5QManX?R(^MR?nh5wl@dv zFxTyJ(QZcrtTX}!+m8aqyOo!=yHxL%*v==6MjWvYNxq03Gu!IxZ$OgP$E7<*GTDAB zAM#oW(r$NhDDVG{3rn?y%EwBE-=`QnPRgHXWPgy&c7F!o8Zt>0Pz}>1HX-;v4LGi@ zO{N{N6?;J*sfSyS1tOFb1A@H59coL%A~B-HD*|1ry7J{cY~6DIrz0=yNpw)6zijNX zw?E}PrIg&QUVSQBH(vU<(r1ufNdE=>XX`giVvak5n*2Fs=DE<-#dNv}R)Z>BnFc$0o`VvLBZH9Sl+0WBAT`27F z?YevE_N+;)=_5uqe{MfhIlDGqV~#-3z=;o$2d`ldu&OFm?vNSdMGoXqxm`$Qf$+ts z^1c-ws>z)h0gWm9&I}zwgZ8V6XozphuG=kd>b5}rwUFt<1G(H&^B+^EVtcZKdMSUl zn&7Kc1*6>RwwXu?HW*Wf_r4KmP(1MycUL|Zr019G?=67xYZPMBPI0jjke&8XRREuS z`QRh4Da7Kfclno7N28Drxs{Bsbe7KvcX8EZ{A5RRCV}Ad>fv0_jY>6IA zYnG57syO{{>nW_A?O}KNbE@BncE+(!yG4P~!p_4Q??6p})woMqa5Qiay`t#=W7v8p z1>1^>l3;?L`h4g}nlYGNeKKK@^hGQoRx)ff|5g#x>HBw>x53p3O@tP-W{{{AMNrSl z?L7%&a3y`PJR40EMF#fcDpmgrMZ4`#T5N+b>FrsTXv?%#Yps5x!a78OMz1bk{P8Pi z?>W?(tH0s$?*mj{>e5@zWK_IH>@5)bB83_*x#l`_KGaGsK#xYB*0+4Rk+V^3aVa)N z91zWclkC0AYTsQ#)B1}jwFi}<0pr4^5)=v9w!Vh}PWV1tm9903XFtq|cSJ^Ust0|X zZqNIOEJDOOXd?n_ETgS( zhR$M5aYaufBLh)H)P(F-VBKKJ>+_VpR_ybWCb1dgv-j@$1jrADn;b~!>e1)Q`xc5$ z8PBANT%~GGFKP1|Q26IAy-_#*no61Z`8h#e=9$Ha^K<54o8(TJbOdb%tfMY!L3QSj zjQh-vcqL?MNd^oleGGOcvJqq7skteS_6*KiiglE(ZIOLYq^1||q`t;ZP$}O3=7DxT zxd#HJ3O1|EI;5uO@GKQW+fRSh;F2K@M^i#)u}^)rm6IrsrUKeh(}M=jETr7tSV**Y zlWP{cu?qBbmQ!c|gW%Uu8$Kvz5VBkYc6$Jiu0w(AMa)5$_jRI>ZVb@#(}{e>P^4%T zeN8=Sl}svB?CouNK5#5fZeFSFu(}v?aP}5f4cy2A2rtAqup6p9;^AqfurH6R)XqLW z=(e*KreQwQ&ir6&zE|vdfBDT%<_U)SA}0_GRgWF;1ff@;1@?bf*jHprpq~TZ8&v;; zo&GJmb%Fk(fuG8U=>s^(mHg*|1G1$_B)HXc->r=cE2iBVF^}D39(?`$q@4eZL_AQA zVLI=RZO86n3@VwwGc7}ig6mihbL6F{ye?n0vnnM;E1v50ToU(dmwB3|&>Rz|zqdG0 z5WGof;68eGpu4U%C`Ze-B1Qe5W(6eX{=gh5$JHM{s*Oq$h?y{)O^)G;LTqs;K81f% zuhG=$sTu*|1|-)upfZdq9X_hHR*WSb zDf)*d_AeS=rDOR!(gV~>br3w-NtX~og6D6;xXvf4t?R5y5;6pzAZ;TuBl*ssKjSH@ zJ^F5L_o5sYlB}dG1@(_aHj5EL|1*|+UkinT-e^&(0n9zoF0qGV9ns)ATi=mTiN^{#>;0eV5yvyDV&W=ZKe7j;*KV zZTYu3p!Pf3#bxvm9hyS7fbgGZ0HswV5ayyIXZs?X1gG7!9c9P|>prEOACk2{-w?9! z-sP37Q-#{Wh@LZDPN9jR+rpr~HHk5f+LYZ*tZ@{YF%*LJYGB^E#s^4P?KR1=W%=aC zN&A`Z-1Bnv3A$GC*s-%)OAizI&&xbF56p3R5yPiQcElBBhT_O)C&Q2s9vyc$fYZh>}8vpUl`8P(I3pVJ;(jJS2Pctle_>B0OelV9M8 zf9OGTGb_F~Qm^|1{PLJ5tl{P?^1!n2$;P;Mdf2^hwN#F{rOSpJ_LQlobjgU_mbkh+AZ9xOjs66 zI|Njw{0+7~!vQ4K=igv7p20`Xt@bFHM%youm#!Lnx|k#hgll!XiR@XEVP(l^M*d9y zSKD}DalO&|X|<%HPbXM>w=db{V~d+6oHb5WZFhCQdDPX8fcDr;S3f2>xcsshDp{-C zZ(#mv?OsWVvztq8%4m7lT3f2iEt6X&GD#jwGqB`ZV^9`^zX4!GXK6!SX+Y*dJG;-{ zF7&^;nUIpCrC@mFMZz_*J|3ILPy=8RUFr1<1BSl6LRVyOrV{o3HTNv%2t=~_o2)DJJt28W+ zGqf3cBd7AYjcdMmA&%9gXtKILp8aIXt*agaBLpPUB*OJYAXA$^Td~l|<6XYe;Sz890)C zw^<{{JoQFaYD>uFV@w)MzdLDS;Nm<4RkF~#And0`7du(j)=&h~!sdr_QrC=vF2zSU z-Wp7LD6ba09Y|}|zLHMk!Y%{#6R;L6XDF_#-~zBbg-yw-&3ZDWO%&_qiAzT`2)O5V ztt55lJkh;Yo^5Bo7Tp$6Vq!AEPohBD^gmsoAJsGwsgfBqqZ=tw@k0((D-l&k6beLth2&X|FXJdv3{x4 zTjS(~c9BT_uWWqkJTF%diZIK*!9)W00bB?=1Vs+azy*2DnvLa|#sRT4wT<<*G;5(~ zb#~cwnSl{`5%tQqLT(|YB3Bfj#&)6_hqvRwZ9fj-MQISo<~YEo$cyH%Q@)!gX+7LJ zZQ3z9-&pzmaPc~$tE_DD_=H$uP|86?u9M|7e2Qf-D!Go4iz9vs@}=4UPjtbZ>@e-% zV2~Slc4+nPT)LG{#j|mn3r;co&&6He==&Hv05yID+2ttg zj2SG60{wq&nPgKGEW4T(tn2h~bjeYP;1fwRD>QJ?w*Pr`(!fi2Oz0xx&#j4)MZ})Z z_y<|)h@jGcyyvd=S^#*L+GRo?)duuW9@_ZvvlPBS`svA4Jg-#jzcj$|b!x#<|L*ND z=1R{~>pASwG)S3iIl4FcYb{1F7$qp%V_JB-QRk=n1%4wLxfB)!qzblicaAn%ca; zOa!VD(?{@2`OaAU+~G^3Lg6^zAO6z`>&KCAJJtu{Dy=(;k2yN**D~Xbc1+WQ5_iKxR0N@}Q!HOR5zmOon>_L(jwR7Wo9Co>R;uV(u8e&$ z_jwj&S1*ANMl*&|Q5D+JxLKBJa$Zb>w0D8&!(8`Y9>?4^rp`;&EUF9k%|obc6P!#b=oR&T%=Ob-z6qR@Awm@)FDQjv2R-W;p2k2Y2Gkj&~DO+O+1^fuQq%jg30J>`MUvq)IObgONIW}R!h zlr9{y>7cc?CU@^Q8|6qtuM@ux&+*uOe(c#*!>n<~_TFLM0U3f2jE8lG4CtO@Upg5( zH4OTgh~gaI!GBhuj|qMGX@(KlOm#j}J`mls`i9%^IwYt34!zN@Wq$9`{SB1-;&VN4rmjvImHQ^ zbdpsHw%hsL$YH^3q;qAQ?%KvTm7x2_&+zS9 z9q*3DpxHk7aC~QBvEIf~%Ra4ipx%E78P;lCbPw{qIJi9^nB*lKio8d1^XS6 z!RR2M$mXa~RgxK=s_c&Bp79^N90$~O{4ajPYiRRcHx$?Z+Tm|$6O+Pa%w@>VMy8@m z(PU@jbJbNGF$WZIH*40tqzBgZL>#8B+ZKJ6%(zqOpnNNF@RV)8%S%Px6uz2+To30# z?qe^&g_0mI@{agx9YlI<{soZ6f`i!MXv|`a9|76&7@aF?Z2b0_} zns&$%r{SxZ1D90t@|DUIT;1!7j)mV)Yms3)Wrz~>*Z~eT^d*iM0&Eiv)(_e=Cygx5`ER;xS8H5}c5q;{J`4ak1_H#LYoU`-E zJV>6qwu*+2Xj;|O;-W96U-grE?8)!Y7^ytEc=Sk|&Nr#;v+skCA+Y!-j50lt{l#zO`#VZfo#nE#YPiY04lk^pxL?wvcW04#wE z;QGnr!x2+b8qYzQC^bbN1fgpj zNjs!Yb#^~F^|o;)oF<%SH+dOMsWFaLC3+@&+An2!%ArDB}w4*Bl zo6_r9uVx-*|MZSu?U^jI1G$6FjnONg{G&{j_3V$vzV<#KZ2M~HZL?bN{e3r>Xm1GX z+O$JcbC;aJ*$pwqjAyzGS1Lj)rxm`hU*0Hriz@$q^7Nd@O+0-f_1vj6Hn}`qj%`_n z)De0G?uQu}VhCDg(kwA8yN(n^jF{G>hXI6`k$F&;iK6N-xT0%GX&EPd-I%z{`4|~N z@~~ayZD_{3JUFT*co4xVT6P(MdeNO5@G9LzBs2s%2+i#g2!VmCJ-B`36l6)N- z>bsSeW>p0$DyzjS-WR&k?#J>mu80MTo!oaga;wVlM}J+ubR7i*2~yM>cN~uaE%aD^ z{FCo#ZhlHyPon(erl-8#SxFt?t>u3GV()k}@0&wZX$CtRxXUwSl(_^wF;?xuJ* zpl08yHfvR5^F~LPc-^mhXbbyV3!UE=Lz0P=b>kFEIxEK_tMfz?yo#Gd z7+>q|`lVg@Nl4FV^!9cy`xO58RutaJ_u{IHzS27rkt@=>D6dZ(vdttvSfueTcY;(K z?x`bF4eSc#R7W?7qNlorVj;<@XseJvD?jCJM^|G>8-l?%!!2mx9HNx?e_w6{1`FX| z+R?X40XTj~Ip4UjIyKysLQ%A+EQ@kWKUeRQ_LtL*MabK>>eSUp_pmnz&U)6Kd>1MP zTp1Y@@spk;j-pBve7KcVGlcchK{IQr1n*gPq?_D|IHX7nP+iw1DcLe`5%uV z1nJd;U)})I?gju>e-U@qqfyn1ZR_HJs%S2_wFS+qNu3dF_xAZ{kbB+cQqV)?wyJos z*k6w?&$C$=yffISS}^D7z&-;!nqdri*vn#2Ri$Fp$ZS}-M7gBc`Ho9fqO+B?k%frO z0d`5Ce7ybumx>(&^acm$1(n!%3Bk}#t!aL;*E7R+Vs#~@LX~^pDZkkFD{0pi`Hl-3 zo|$YtVl4w%01}ITk()Eo#2JH6R{unTX9au&?uSwqM&1SakVLm-Yim<-(*f4{rCxd@ zzE<I-KQF_NC{-M+C)$Pi8u#C_MD3aY#tqJF(=i^nG=k$ zO*oHrSR-QV9biaHrM&G7_M|D!A=#G>WfR8l6>635eXB(I+;IuZ5?d>7DyW-qVV;j& za}2^C#}S87<7Ts&s}rCR4&s+4cRtgVC7#Cgw*3*d8m~C{;*W*xbt7-_vh$+1MMUsWvh9L#B`OExQm18!dlzDV<*0#+%T=MqH z^~F87!^=Bhi$Y&O#9-(kikOY8=nTUs1|H5d7`gZE+b#0SjxYMw_I5C z2v;vlA(jkFOTNL}8;9`O$f`$wu4HEjknbL^4|C&*^0YDBH%}CMwn|7;=wX&`BylPV3l(;tU((u&qy_NjF6OZ=FJ-6!Gt1-85kR}YGz+D=)5Mloh zfc5VHgh!Vn!vk&e9Q}J|qwJPawfznp1an>WJ~dOGNn>H?RTUw_uAZaBcadmNf_#>? zH2}0&QFzc+V`&Fbt}$)((GCt`XRhBtSB#_`S)ELGe!~!{tFCZ+iDwxoDJttKc>8hf zQ;SG&FMER6>OuN1>}M3L1^9x;cKL_TlQT(Hd6biX5vFEBftTwChd1r_S$#jPFx|R( zOWPytl*5Ady$ks*3bLPVU#k9G_#YQH#wY_5Jc%Y>&5X$`J&$dHSOOI(&WlfHDcRQJ zxj7$y-uvc$*V9oTPT(anN&S0LGr}fcq(h$uOhLp1hJfDm&J;$G-)HW|fx~_9={{n0 z)bH*lc6h5#e@LYPnKj{VXO;7@?KQEqHpA36{Dkp>_VXE*;OrPz8D<(N*s*PHlfeL; z2qkHeu^(t@u8Q~sd_V_Uk<+!p+&*5n-Huj=-0ajq=-m@}I?$6S>cs?yc!~NazHYN~d{(s$T3`)J z-JZzAiLJYM_}@-+>U;?^T`QMEeVsx;21v4-HcilPE^XC~YC<~?TT+ZV#?#vNb5FEd z_?Prh?NdrXH1K?_Xw^<`QC80&n;PR;qu9Z*It4hZrY|r&am4Pv5R6g{y;M&w(8z)L ziUex~`R=r6kmrYqJL4~^tAsh-V=%OkgASuAzVb0uGis-21}-54VQCQ0dO&!arjwZN znabVfLy*biqhfG=BRY;9z84o3XJobY7-;vFltVSALdr_N|Gk&hkHX zxO6}l3*`4kY^d3hw#9pdP>6)w`b-niu_n8jqBgxKe&43QzP6j(npY!;_Y+gami<%wS{DQpt zO3R?&y1nS&$#W*X?6e>F8(M3(YI{B*kS~!UC?XET^n?y%l-*(^Zq^T8Y~?I=sBCdN z$qkg@U*}29-V71+N)bH6EK}NHQ_g2L0Fez;VP(RqCV>C=%S}SCQG2FfpxN9y7fiK| zZmmmwO)?!*oi)z8N&OM3<@Ac3^*YjAAIM#F!=2=Z?I4xcF8&bjsQK+TC@mGX`guAR zxK)>~q>osx1PDroTUl79q`XLx%mj24)0rXDio3FotbVYncDRPPfo((!MzWev`O0(d z^AqP772V!Q8=^wUHLf@1YWckf*x9y#+xjQ##FohzK@0it1ZS4v+WN%ZthW5@@Y2zq4aDQ6Gt=g&4I_*xtaD%vP1khwVk0bROtrYQ3dUY?r9IE z8XMzEj}#Z}HBzyVOL!J{{EdPfM-~FH2Zx0wBAQR)+Py}xb8KNWl*@Q(Bgu?q`J~nB zJ0{ENc_YV7gtnN#KdyoaTd>*7pxV};MAm^54;X6f40BbEIPgk(B|{%N>dIGnuly4 zd8bT*R9+v+c0>=t73&HT|?3;=vQ5WfX8;AIjpr1Zhg(d0T1L z!wm*mONXq|=IpzCJUP|?hIshVj-GxnCQ+yB?n$-8Hw9u(6BD?FLBEx7>sA8{golhM zduV>}(-O#ODyM%CI7a(k&Tn}~<=;6)#73|L>?V^}0%pCUTtNj6kImIM-qw#-QRl^-6;XVqR5jitVBY+OIWk}YWil4e%_PqP_zU*MCeO9&*MF`mkpMT@s~> zt&)Xr4j9O{%5yPEfIK2=C9Fuog&@h6Gx1eTO5c3;b9fzqm+dGLueM&2>%!j#RLVZpl_I`nGpEQ#YI?d%+Kt+F)|u?lNgQ2d zy25KaHX3GWb?Xd}jsN2Wt`S=oM0V(#;ORhr0!+gMl2;oR@x|4H^hI8&=&s7Wo4-;& zPfg?R3@XesA6EL@U9VVcEfucM06Ev$t(i&!UCgz;k0u9#Ia1-yec-i{nmGNjbAM>O z)uhFUK;f+2%AD`dx~Nf`&hOVvSW7;5K0I~U=RxVY)yfB*_YkssaB+1AMjVO=`0y3X zIB;j_R2(&JR4H~nUUMU< zZ9Y5k&CSH8RxGxszgGXKmX zuKkOcK!Ay^0^#wQWGk{xj}+C_KFV2fH=Nt4{*`t~ZkdfUlP+f-E$~yG zu=_dQkeZs>K3diF{;ZG=LB|sGUooM5fbbidYzvL6lR$+*8-!(?4M7yeHspZ@Y04qO zm8yo&Ha`orNnT^_Kd${NaP19D2Okv0J;+HBfy5fzh+ojNebK+oc$RlOx1vk#q*;tE z2?>9y&=$Jd(Jm8p?gmHQC&>^sc5FY8$>bzrqZm@a?YLmX;k@&}a~QN{AS#Ekfd&Oe z$Q`W|!H(o0H@26H%#V_Lp69E5`SnTkN>4Ckx69`uusiMSMH6VJ0Zg10;(4_O)e zjJ=QLGilfU@v zVgjewef93SVu)`i^|kj0cWoiR)L(XvcE;cOSo&YihYf+&p~qG%2lRIZni^1K_dJ?l z-`wKr5$YokGzS=n1*+kEV4fX|Wk@ zcQt;h#HM8XeiGm}mBEXJ?>)!y&O~O5S#ujWDBd>zOh#-U2tQ%#UQn?};h`aU{vwV< zZvj~KNvje#vQ5C)DS|V6@;+1N0+wdPOM-T^QQ%8w(W-zM$_E(oHmWXi8Z; z%dQM|P_|;@0C`*+Lj88Ta(uDUb5dVggjWk|`ZVpir@ePooO#kKU?tQ^F` z0*m{2$ar}hY*?rq=y5#)-UjrVy#nFlPl74-XHPH(`xgQ8J$KP-@ja}3W2CdjghOp!Bh%m*f49!IRmBtOKw1xr^>&n;#Mfa$?nfY_g}+;wUq9BCYyxKvff z|Gt~WTdiFxC5rF!{XGk3WEPRqU{;0a`A9aQ$l8*g@u=oeRZWbei5saiX;x}d`{i(P zifzk$@aL>Id(M1sIt8Ax3_PV~lP*|T70{TQ&1?Pr3j%tAX4X4?2UnquHv>Jrtx4BS zfZG%OQT1f&<*&9|Da9pOCM2bZlv{DNZ0|^QAGg;T`(#1edoev7ylg7$3L4+7hBuEO z9)j|Y-cdI+Ym<=w>khlkWP^9G*@e0hACbc`V6xr*?BE4UBc-#_2iRmN=ny@S1qvnp z*3q>r$c><${#k9O;LuYSNNAGxzn5Fh9axnjG?sE(zM`z~7@H*D3G6t)^GMo<>bXUF zawc%AKg&qCaG>ar4Rp$IvlFlPU#UD5Fy={2 zCHcE33%E$npIeI!x480xn9{dwYaU>|!3LZ>D5U?_As7Zuj%g3mD8`cfhbGL!;Jo?{ zm;h1m)z%itE!7|urU3^f>N`sS%@r|2a!bxDJ2f8%_HH5v)Z9(nmU|?wi z`)K}*#n?WkU(EFh369ZfN!i%7ke`;AH+qtg3_UFCyUFIah^k zzCFrh#(N7#X*@Hlxa1n-C1$2{47KOZ#MkE^5*m!PS*3y*+8%$*82h}z7+3^m7`-Jv z0&OOnWbGGShxHWOOaqQp(%)`cRdY{YH6{JD>~URCVl6<-i69a)es8VQc$pd8x`QZU z@r+2IK6%{^4xDD9_`c24D$UC&$r_K%`JH_?QspO?_se;!L_Le_1CwGQiraeF;eSva z!(yfY8bb}EHgh|I^_IbXd=A!!s0Ec-uZGQeh3oBuQNpoGfy@PF03es-73-xBEC~%X zXv>S4a=nU_Is2}gH7aTER5oR%T$U!0<3*93BY-e{4j1Cf5GGeu2Buoqm84;t-aT~s z9#|mxG%Nl|gksHJwtL1bL?j@lRAi&&sBfi$u8BfYJpI&^bs|k1aR1lNqN$wYpyM@* zXWeC&S{HOcgoI%xFu;d}_%OQkj>9-o;i**v6LL0hqN_`NZbOEf1^ZzZa{WT0KR;GE5d z@HZDI;h~R-uXNew%-Gw~Hf#<2y^7ZEzD?Ckbdh_elR$ju`0Rk04xI=4oxtFV`HM(y zhx&IYzh#y6`@{yAQig=S z!S=>JR9!+E;rIoL!Kn?Y*mHbof!e%p-$j~3vE%eX@QDJb=RFRLMMXS9)1`*<_bk$k zgh}x`+78SBHSL48`EX_sF#<=n3j&u&itN{J_XZL?dh-`iGJTL?*NDJScm4P_us^Jc;H5uQ z%z#$FRc{|M)Pm!}&ED{s7Of2N;=Mkz6q{>5X6nsieav8=O}1L*;cLEkXWDimZz#tdaAAN*_NLCtq&PrgCe28X;@clCz zZ=pEo58^Lk2)c+327q=7q-td^znK4s zjP}F2d$Jgs2#t>jr(5Koa2bsaKABDf6j(GQZ?O>;t&&F7{rS+%!&|aEX}vvOytTDv zMNMHstuUqI%q^xamN<_spu@EY%=AZk4r(P0P6k%mPmp%rob9&>urqdG@m_1J!wO!j zsy^hxw#@x~!MY}oosMipD5vI*r&Rt36InE|eIL_?M>y+9eQK4QLlp|rQXzZ}mk41mjZL7_sNh@Q!hCf@QM^%YE)M45{Uv8t|6-&;qv| znjSljV}uDJFzwxbb4OS&Y51wn$+c!YlryAejO|s`L z?(_5xd_)x$*Xh=fW8634?D=(;D0!XLP|5I`cJmBH3>}5~G4+UtrUYF7{DqBp__P

dc86txp2#-z{H?PoU%AiBt*!p1!pry9zDMLe=AW(+fgkj1@rnR9KeQmAz?+J zCSdObj_nd1`x|=Z$Oqf$KF;^I(iwDU9MW_dJRI`fc;H&pfl}k!drllk)e*VNMI`j0 zxgp{oYOzsg`gvSb6VsR*Ml5|zV+ae{K=@A&4;Qfn1D#3xEZ-|%2U>NO(QjtR4^(Nn z698a@-NkK=xuSW#v26jj!Wy-Exc$+LeK)Dc{rNjaFd_|M5pOYC7B zC*bD=DWVW*2;a4XrU@p`#Z{6guEW052#Rw_#!l5sWlfPk%4a<}ErxwgyKQ!)U8znX)fvCHx&7JX z@C)uXA%(X6^(HV6W0yMr1rC}_EQEge4pShS;A@6OCe;)gyGB7V-j?JuZyEffa5#FC zLg-lM`q88dOMA^`vM*H^y5%p`p`Rno>jwKb3Kg11e#mevn1V^Sso={@)My6hYw%=0 z22jMfDhAJha-gtwT!Y&jYtja%h}Ztuw~cQ+OXnYloR~%GW+F5K-F_kI*}e zP%s&&j2VLdeOly}VpYZ7r(yHWZJ zWSxG!kD2<~>Zy1U?A&FXUjuMr3Mc)e87W%2y9$(PNddLDsshVb4f;v5(h+$T)&0uJ~d^ES=C0LowM1L$G%124h4kz&Dc`jPHH;1vJaeWs^3h*C^j3IdQ3G|#-FA6&&J zG1n=3E|lr~wz3it6T5xiqrY$tg0%-3=G|01VAVjVSIuudCh%WF_W!BBGyPhp3Gg%a zg^~Rj9Nt86?N$)X1fDorAQAm=>#g3-*M}a^ZL{fXuMCOt9ScA1p9=PvH-n!UY-#>Q zWC_RP$oXl15s}!zaXIGhKw&Be?I;BLjpK%1 zZ@S-3kil>Yk{fC}j4Z{{5e-2gCOC(qUBYAlucQw2`5MT)63!5cyg+DjkdpFe zv_v{{g@BPsHw}o55xpJ#v{v+>w6tV8UG06aJp8b=4hp5AoU(svY1Q3J%bDnOb9M7L z5XEWZWL|LSaIm1|83a=>LhwP9+@7Y=599t3CUt|qIqd_MMSnN~VFZ$tk~VLR!jB)# zpv;H#Lik*OMDTIivaP z;a-x&>m2FpphI*%Jv{C9+LeV$)2>@LiatjwCGMSQ(!BNl4D$C;Tms0_;TTs7$*7-L zXbDwd{d@ZiYhV~y(R}y{Rf%RyPfVTL`qVJ6iDGtpQp9|DL?+3&Qq**ehj*+(;a7Rr z^NG1w`{NyJ%}UOL%i}aQaI?NcbFI6WyCg%NVRstgJwyJABA-X?C}t@BMc~H&?9O?w zU{@S=Q=PcsaG}Tr6jp*+etP1>&vkm;+;|)V<`%!4j4-x zJc`pt5uZUUxSJ@}DR{<;o}@WpCgc%059qS02PV5e+Mb}C`If3Bs$1XNsZLiie}tZ%iRGG_%Rl0YZwvXGuEp{gn8XynJH zOxDfT6E#Viq8Gm;*1}4LmoJ;N(gEcikpO^!=+-^BHw&nakbz*Wdy{CbaH{Xl^rt^t zN6ooA{cSvhiVP5w_soq?D{d-8ApnS+5Cvm8DDZ|1vej!6GDhO$0F6IqS=gngy~HV< zyBg(#SSw~TJO1279WJ?ky zTTCS(Nrj9oBSK^ep@_-88~e^=4@p@P%8Wf*W~6b(%-MbTUiWpq?)&%sJ-c78|AYGh zX_)0a&g1y3@6Y?a1_~pFa{yD=5db=5^XF+eaDTpG=YQ#OV!x)M;8Zho)D{E%bKs+N z63Ei!cg0uhr%$$#@z2`Wp`2O%?+_NjWu(~G?jA`lY%x}RgB3}JZJY2I8ZZh z0O%9DWus&B7lEzafb~njY}U*$YL^}RbJ2qw0YW&CQXlebCw?7&hh%B7DTLoo*2YJb z-qR2S{dCjZO!KQ+ETY;jyXPK!crm9GcHQXod!;LLm3wR<`+v~EZojSid~zv-=1Ivh z7;dqdDS`)zCJge4MD~*v?B|B3+GPU9Pum&r#faD*uwtGA^cf<55wg!IIJj#HOIL{_ zhy^QS`>ZKi^zd@RmN$rbJa@q+Z1*6z3{9bD>G+Uf5t`FyfK+8)^}pxo(^J`6nX-P< zahN6c!|~F%lGLxcZKIERMsMl>=ha9yrt(sNR|@s@nw)a1A!jy9r4bFmGgfr1cJ?oP z{_;g$#*OzCkK>~fO@_X9*cIIFB*NxP;zR-5j|YcpYr_K#q)FzHo2iP(T^=6-U9R;l z_uCreL^I4%zO$a#OL6WxVKVaeN@L?VAM@!5#*}I{fo++P0}qYu>2*-&E>EBsy#(t# z1{s%P9L|oyg2d&D&U}0KnX6)Q$*9jimF6ve-wOdlt-q2+k26Ca_;q%eMINL|-t>Y>v!2Yx*nVI5B=#F4*rIf` zHUub(jF=27>!5Ncmv@57T4Q5lIfH?uyL;bsYz^gOgE}!?!alSI;FlEZkBs zks_ZxiC_BF3p$ssG~PSh604TqXjjI{RQO=u6i$_2_GvxBe419wabt6CPu?3DyC^Vl z82v-O9STaq{R#wki`^rz49JlkyRp1i80a?vRZ`o2Y&%Lv9ss0H)vw`Z6}^DE{l}go zqen*YO`i~nTEvaPFQ&4_qlkJv8PGKMFtP&UX{hL{5s}q<;ZL6={OkD zh;$qOH3F6uTo&kV0$Pr*s4nEriHmPJh9FlYbQnh$90JQuTo2eIU)WZpvXO;5P?4xA z<0l}|ed3W_-?2TCgFDwSg8UFo)fE4*v7L7H?uP2G15C%5u46f@|Dkmt{y%9QLQs}_ zu!g-a8XfF9a-c9hD$tKa+_yrkU97w@W0%vKGSbv<&#`poM~>I)%y?U z56=0;FJ67gO#2FOoDaa0!5;tYKta>tal#+AmJyiS0Yu|zd)@4+8{ixkI9?@4{9xRD z2n~?Us9vcsYr7v17MQ?@4pF_T$^91MFU*is!@Bf^sKZV5uUVO13>-cm{(7Q)|MOo^ zI~a1F1IG6 zU)1w^iYYvpN4uBZYGbW))66fLJw=q-9Plwswsw#!&ALA!^3fyoA5p4m7_7cH3{F%Nh>8 zBc&PxW%jzKHc1Veb8m(`#2rK`xL>UlWNG}V1V`gAIK7^2x*(drIshz2iZ7Xtz}^Re zy)lhHV1}C}U_*)p4n-E@Bw=7Uyc9QgA*HXa&p-E+C-TQ;;fFi#n?;QK2inEukz!!% z?>{2rAaLlw9{vwzgty#CI*xE1573L%vTJ4B?N~O9*7=XHcBO}nVn4;IjcNxiXWJJSbkPq z5(l%tCS$3mp(kW?=c&l`2gSA+FSFXw7*WY0fu!DI^*mF|&D~R`C%TM8-su-VN<1r# z@*e*MBEvljo4!#RHz1g>+6F-3{fIjN{s8QKyTD(>k8%ti{)_M+0&=%H;Q9zI8&9yj zfY*;xcBbGd4K?E!KtexL7pSkStYO#07226(eYw=Raa!DT(G`PDKNR5s`Qle2VNOF7 zB^MikrSgxXwwArmh;PvJ5ABKq_sZT&gszx@8$3P+mKO=h5Zh7Ek)Pq^b~IV?x5=RK zXN!F375nhMR>ErmM8%&EYKN`ZnGT9Y$?7;SE66dT zrr9dmv~+8{qO~${&8Ru!rpGI1oBNHP(Pn~IrEb{9&v{V6=}n}}dF^r&vCHR@0i529 z_{TJIZ37J4YIS0FeD{*>5tzM|Y(W*KRAG=;>E#*&DAtwk{JKy0(!9oa>((fvW$D9T z^imHQQJ7_q$8qv9FjiN_1oV zXp$a)GVB(4o1PSV0^F)vQCjj_?w@Vjv&atI*E#BC3qH-F8SH-AA+pnRiBaLYpC$SB zw)@d?U6=K4)8F$PT};lX?vxF8e!hKK&x(r(n@WP*xE5u zV?}n9D*vqL9xzMG;odq6tHiYb}arPl-xjw2~tr| z36v#Bf^MNjE+~cf5MJE@)j9wF`n&l5@+uyGB7i^n|G^uqaznZXREOra2oa!;JNf0L zgk^rJ(Qv;$Yx`~XLZNpnk3umd{_fz|;2tdbk}5D`z%o=8h zoZp}G=|`>4Tc;d;l*m~4e*OWvQ#YWO#t&!~LG}Y39a=y8mS#kD%JF5Jyk;y}ou8Y2 z~m~u_XLa@unFH$e>9$S9Tyv=HOJy0HcokMej&@+UCm`9t+#()n$)w9&S{=t&V zU~QL~l9dXHOGLutJ#%I<;fy$;2Vm%Tu)X+DSe2|&rm;|-G-uAYt>2n9n0)VW=n8YF zQA|};B-?zE%XJnJrU&T3zli4`CFuK$m{b_)PRq6D?jnXZMMO$`D=KgV!`2xY-v=um z9!_+#>c0|gWS*Dcl3*Rh@Qt%~M~Yt6c65cXAA1WKJ+1aBH$vIF!-;$BcbV>DjH&1~ z2Rn0h!R9`*4AQUe6?PnEKQ znl!UU70)b1@%tZkt&W*Bmk1j^%O%v-(m+JlUIxlf@iF6A^3{#~fx2UQhoGD;K))61 zUNpISD!8_AurMbm=|W3Os7cI+9(f_#Gv6BetBXDwokT&}t1y@+E8d0K0Mk-O9hOhW z;VtC%&c@RHlx{?9(O`iA)6r>p+luhifYgMj+kHPygU(hsIulL4UjQimu7ac5MH3~P zkb%G>2u(4&Gp!4wX;&WiP;yo@w+XHeg-e(m|Lg?A8Vj-PhW;kCV6v(3<{xq|%w|Yg zok0S@qOM6)_r^2a5Q^7n)!##YUS($sKYQl}eYu)?Q zCUT=i)@OHh&T};QKWd&;o&~-;nE-hbW06}nZF1NdYY;5;dTg(Pk|r=#a!_{>c*zRF z!Qi8SGqYh#$B&$CZl+I$s~?936qb?dE8POW+;d4p)iiMZxvOw1QSFSLE&&uma=T|X?~{Jv%}s;eR3K*y!r4V(<^bR;LZ?u?NInTsL8@o@N}CTKtA`44@`ufq`c?`zM7pmNz1Au`wuAHd}s!zVGnV75p~|0FtE zQhj&}kkmsXHsrP>=iEEnqCsUDv9ViJvQ(U~wQ{Z6N5FE{l^No1vj%eC{U@M5z`bXX zQBmsH=aQ!zS}Vp#mIw060Mh)LOq_J$$s5TN5_Td_nEd9L4}_qyRUvI}n1vtA8Y`BU zk-QED>Qdq&XEsDJMiX06r7ilKy?NiiHrKS>>|;{vPg06JDA&L8{StoYPZG$4F2ke9 ztxIiSWk_>iColL0?FTVi5&Jz z^93cN=4Gj%z@-dZ&~}4L7B;uNHXD9v1^)g6QXZbm!T<4*_8KOwfV;Cd1`NzVGT08B zIA-^1kYBe}0suSOY)aXwB#vBatc@Zf{Xg8#UxbhLZJ^yQF)wV>#Ja&JUYuE>`&ckSP0(Sc{~I_)?M_ z>2F)?3Io-kRG^mmZZak7$l}XD3;!UG5@T`P$vB_->XgIpKfISX$jl~o%OzTM%NI04 z5MY{=K>CD_CRE3-x08R(8h3j=BXFm80>O~ntK z2&;z5#*#>2a2$Z{$=2MH{GH%d9E7GiMOB|=4d;BS4btk`9O6uJ=or>H zarQ_HUqg~ld1+jOv^z(`A%bJ!TkgdZ;K=X6#N1!2&N$$~dCIDF7ciHx&n!_&$Pd0* z=<}Q7Odp~u6!U%+{aJR+VOQ0iP-(N4M?!_4GRg}oU51&44-aRzwOHGgQ*YJ86l<2- zCJMMRGRlhR;m~OWvTyu4dX6{&Iz2bXiQIu_eP=Q(dYlw$$gOmG)`U;WiOFd*1w*Mn z*|(1!PQ7-_>K)I;GtA4@f;3=(UA$md0|mQm@AQapfsh)_gSUL%#M9+kf*>O5pwb2-HJ4{0w2%l_n55 zH{K1XYe=g`xn+f_07?uF5?ptaS)B|hsb zx85MM%QFdwQ6w^Iig8wTt$SB2Q13Gs-w87AsDQT7nL{>>BYDRh2RL|2i$pKv>Rczt zx-*ful4%Rrf8Cig$f^2uS+bTny@23~m7r;TxrZc%&6hJ%mLcv^?&8WQgQRN44#<`6Z*p|6vm9P?bP ziuBS>TwaFxx)$;?7&Qu`)$e#8Ue>!7;N_2I)31%VS0H!)JnPB>%=$at1Nsmt@rGRj=V;9 z7r({_Pk~`aQ*SSJxpUg*VlOTRb-z41PEW}u01&*YA5Fkv&M%5UDnSX?_&~4cIK_;C zgx`aib;(ZSG1rl@bx)~>Ja@xKPb8YU-c7W)b{)aJnoZCmkY@1Ep|)mVw8O}#CmP&L zP_rqS7#EB*^XAMIok>tn-*w30ePu%KmZ|K|XPG#tC@=Q%NdqW@oj@M|{1jlCx{WFD z@~4&?S+r}Et+e{cpI?4X>l!#Vwbs6jV>9C@?>l_#$oYzS{!hQq^*AyH7|luwilq89 zm<3X8*Hay%$>(6c=o(-CB6hyP!>pu8+5ByI;QXO6=FIs*DL!FcImQJf6pEvQT1|F1 z87vjEX#HP=V#zi@q0+&UEycsE9=n_vuNgDU6PkIkgSlo=37CLre$?79Ih@NL(jh3N zkoqTGEbbdP@&mA#-<&aBa3dNCS)&2VBORtEqe%H^kw}{8>$#yAK1@>pE=rR9gq!Hq zW*sx-(M#;BW>+35Z!ek@1e#2o1ah)(xhGc?q06Yc?Vbn_e1ro-Fb*4~?->HXK;nvA zpQ5%{ZSH~HPsQEA3sedDT#a7nlsavCMVsj% z1D*FDv=&G#nYpm-EwBPO+_kieWQ#7Z2w0sAKbceIUUu$$9FZ|=lcub;*zVHM5a*Ux zA7N+Zo8m8+kdS!J_fYn6Uaci=+FA46OW4o)5Ifjdv7m67Jhn>wXX!^M&NCM4&gAhK z?9;tZX4AnwBiC|SNmS!o$_YobnSan-Cp%|_wGsg~LDot3#&#KJO8tK>83HbDQw1Rc z^A=VF{KH#=n_xaq`I1|WxjA;PW`pxAqGOkZ zxhR3UGW@aR{wGPd0u0(NDq;SXkdN%I5j`{GZ-B^9+$D zgo8O9px9kZ@V87UsGVHUJIz`d*~==}Q(J(k3Mjc<)f_EHLb?WcxQBIp?#Yb=Xr0aW z2;%EtF+h(9?906tY)TaS?w_P-k+~vgrbJYZR)6w=)K~pw59MKIVx6tJ5bQiT+6dQ*vgidpCQRL4dxw!G zhQFgMwj%l^+y(3NzS%l2<-5#7<{ORi5+x;Ikn7JLYOjnk)P=1EbZzl-TJbH?{*t&^ zFE+>?Y!>rIOjjstvYyQ(M*E{tz&NG42>wnKs|`9o`*UuNxofGRDLFqOGHT<>iQ|3R ze29~}d+BM@9mPN3{)HK$ypYc`fw>9|TSarW7|spxLb#4o+Y$GEdptM4z2iWQUi_ZL z+hU=3!niAjg_$`{8^5d>!bSt}XILbP6xs$lOrb4=CxCbudC(1rZ`6m~OugV?!J}Te z?3ZFics00aAW)UtPr82B#l>qzT6xD89ZDek8?V<>%t-NZZM-yHvR6zkfeqtrKceD0 z@OxM#^W*Hw;>niW@9MYox~+n*&B@BDGyICDowfK+`?b3O&NQAL>Iq~Dkrl*Hn0sn< zT=HmYXXT*o&m0;3BMnUrua|HGXSZ+inY1(VwV&yeJxYkh{VU7cVIkMk1bXLyd2P=< zg$Gp#$n3UP^r=AUvyo`Gm)}Zd(a+6(nf`DU`uO0`Qm|MVB6!#m-GQTNZuG?OC=vxw z^C! z`&~sY-RFA9&;ySn$+uN|RE6MSN!2fU2WjUa_Xng5XUh~YZ4Zhf*_+gH>!;@fd!k)H zaDG&7tzW{{3!ynG>$Zv_rQ^tM%-!)5yVNnDBFQ|9@A3kW;pY&uZ5>6_0b{kYxq7Ls zrRKG4{Zo4Wu-x3!_aV&->30CT(sscDdc78$isR64N;v0lnmrOr5eoe*dT)`nSn_>8npnmhH>5D$i)9Biq^@&b z96y4ORt)oZ?v!WD`pZUinAzs0x!n#m!yI{WUK04S(UN_+YRQNGBJ}Qpw65P`((~u6 zL;N$sKIbQISS59BtP?(a_p$kqSs5#(^>WIv7hhgtUg}*B2YckuNU#q0v?&gX zdd9yzj_hnu#_Q;o-#Jzov1?gxSD`YjaYT7Ri?3*NlXAd_{<9-7*cjU-04M~Q=>n=D z1W0>~@W&B7(hXV-AdNoo3%nlRZ@@i|?@*P{iqWylhthj-eC6w6SK^ISCXF&YU8o6n zUbIy9swZkq=KD*G*({#BLZ3)Pe-;J9nR^48|ASxLBrfo7oIvPJ)*gzTb#N5!oHQ!J zdkA-jx5j1M*|yQzf2H|uEW_=NlH>14`yR7x5N0CUG_jv0p**nrz!zYb!2gJZIX38V zI?OYwUnRX_5|-tC=KwRqs!bjLik4iWZgOq)A9v2iNl}e^c2%Y_nP=W1D}itaI8~jT zf#9ZP>ry}@!wSA&4P)A|@t&a_X`%pu{1W))IcxeD=(Tm=|MT%H*ndA6{T+0C{%2qJ z8oCP|?j#4DE$hw4)OSAK4*xCxR`~ejot?nC6Mek5Eu~DYZHsPQKgi)?vgWZnjSoSQ z=7+u#-(hR;|BRSN>6qo_V3KC)Vzr*BOI&3#Y_3h?PVb-Y@9lfx;;8-JIkbQYftC74 zrppUSHbwn0^jaGSf;RP?sGU1?_$W%L90wh?OSV72O*>9G2p^SOe65;odh5!h$QQY>SA4t-lZcb(J&YLZ0wo(15T{7^XZH1= zLI8Pnptw9WPM6*Oi?}wDKVlg?p6_JWF0Su#St(ZO-TJi|%f6li{-D}92eVF4szzY8 z0m_N)6yVGpXSul`>WVVjnDE)B@%O>oPaV4XnU9r#+_ysEt9JV$0dRo=<3RQ??gIK% zk?XAmfx=KV3B~L_E~!@2h~-+zM)x1$8xj>GTK)xotSP(xdBQmBV4JpA)8ah!(%9=Rr+43{TqU4a0^WbO(c2p*;&(0ay^4?8E1|S$fvHNbv5*gHeGV^nyU{nO4~u z)p%tu)mn_$x@dls_YK#32@Q#+N7AlbF}z|VlQhK$=oM_hTjmqUet6&&2dFbpBzbPe zqLUEmx>u}9$8!w%&JLIjG}A$`%3rGSb?~4c$AgYr`qKI$(S?s(R2O#E1BPE6FtzJ2L_s|bnk{5O5o>Q_AJ?&)$v9yI)DPoxCR^9L%SNu5szpW|yC!`|JJ4AwEbhc! zrIAB}QVl=b?pCE80jU%;hP^-}gVsMyc)SaI7?W0qchlUEKZ(xV1qSGRt5<$GK~YCi z&1x!*-CAm+S+kQD-48r;v=)4P(ktl~5|*~17omRO_S8$-Rq%fi0}2RwDG1yGvA5A8 z!GHq86dA6^Oo0<8kF5X>OmdW5f0dKvwt0AkT7n-`}sYWZ|%YVB|$Iq)>q;6N96Pd8m7U*}e)JmC|mVM$}tTgyJtW%vDZLY6;o}sc*YdTJOxyC%&_um)qOq zYbb{F@V7+3bRkrqaQ;)z$>h;*(s!>_Uh3Ls#Jy}7_#DJXV6TQ|YSH*ysiLELM{Oa` zK+9QkGppyFZOEDos%VX?>E; zgFmU3#nlWAB`=OKppIa>?tqv?md+2f#6Zh`>@VVv3%VRh_5u?S&0yPyb_EJiN-&)r zkG&!SgUHjHlgd}kk&l1uycEW{dDuyr-b!@Kwe&h{68BTBaOE$e<@h-U8eva9J=B91 z1epN|Bl1T)89RG?wB~eub%SQ6WOGx!t3jx9igy30%&eC$&f=@V>H4M$$$Iq+;2}lT zPLvQjP?O>pNjphWd|m5#^Dn}FcxfanLa07Xkma`ZFY%V>edQO`Fx^<2(f>tZG;cwZ zWZadc8PN+-&R&s|y&-|e3!!2ZgU7GwJJd-Zqk#8WnSxO=CRc=CFq77qV!~oMg zzP;JelxAJ<GY`AVQ>3&?rK(DFJb|mU z5PpQ!1j`2l{8>FJR!@i~3k8uxYS0MBy64z{*(rzS<}&>q&BN}8g~sof-hU`)U0*qA zoBX}<4C)}viETfIbwkDx1-4MkKd>ypQijEt&~;Dc`B5Lm5-}6g`Ic%k(OkE;LQ1_M z+?SuA>#=0a(wZs7U~c279_Ti;OY(pcOb&EZ*Z9I0e$)DZ<`O%V*sSoDSM84sfJ!z4 zB6;ZA_($A+6zTV0L@bpcxFrn5l80ujx1Y3qDJe6YX%^cxcUbF<@NmuCv2E#of4%%w zP_y6B_0=6<8kpIz?cTpPAep@Z2E+s;)^W<>8BjD6{Hda}G=9~1xMtVPnfuMHI1%R* zQ`;fYw>H+u6MT=JUguue`1et#@6{1sqk+g4rgnpwchI+a|Z!;Xk-tjxh z|M-<);7;1#5Gt_U_;6&MY?oKGlQLP4FoR<5Bu&M$_3G8S9(lTL^J{u}`+TLLLP?X- z)sJ@XCfTL>Waf6H{%<^PCcu7pQ+c6UiU+j<@}L-iD*ialZbpdw(n%;KNa&;4^15nv zOIxF&((wbnSEGFupvm*v@1>VVF>!$+aZuFo_i4v3_fy#hL8O>J&3RSYgve0&}I z(C>g{q9PC|`qu|6f{$U_F+R8h*mhnY+~lq)`DpG4CUeZwXM@Y}yL4H8{t%v&N@ZuPh zEjcrC_Pq~X)y&a;VqK*?uPz`W`KcOWhB@u+aRmPxAAA=mT>$k+AmJw59Td>swA7)B zG%to8ub0WDntv+Cmo7;&nF>&h`eg3r9&^!U0b_@73aEanLMrjk0H3R*Fy%|q9G>L| zEKcbSNt<4rUDCIc23%(<_J+@Id7FV^ei?tn{t1)Gvmy$p^FYXn2kbb0Y`b=_HiU=S zipbtqz-YSKXCK^tD0-wTtrf1$9Rey66H^!B4-wVmTu-tUHr`g!?lFOzf9!1&d6{jF2H1?_$^CoGm6JzWyAO-L`+Vj><*HoxFo zVW7GQ?jf%Jo4YH;k&RH0N*94_Y4O!y77bYNAbnltF5H0|4UH6;Q*24jMI9->9Hl36 zhpTg5W#2m9cStPc5V=c=`A4@P?E_|RFF^q&RXcEE{AXp)>`w*}+z)BeI_Q!#l>a|T?+$+KS z^^J87r_&wmelD05f@Fax*2x}@>3`{a$+{~%2SEdfz&*m%43J4D!E7TGWRE?z6DtIT zb?TuK$1i{0*m6`0GLma#^+d!j|&wMfehx2}_TG+2A zF)`vc8-~HaYqob1$xe7a@WY|0cWuHDmPF_7($l28v|91>pA|@PRQcF1yE&oLa5UQH z^_v)pmEz)HiuNV$MdN=Pqfi}XupI^&86#wM3JbeC$LqLgt zCM&I8yvE4nYu5FHMIwmB5L7Gyx`B_yZ5g-8V`t>C0GC}tt%f(( z7#+cK$f9Gri4A)6#r)>5nb`9ZLQVT=rie4;md;LSv4`*WU3{=hHw03Q$hsjKKXPuP z9Y5LGJIPfA44-v-naeoj7B2sJy7;L63` z+zUa2fg5e_K$0Is#t=v^i4rZUWWr2pXv>K*DDyAE#DeEeQ)<*~W&q2&OQPG22$r&? z2m^VH%x!=s-xpG9`;xW|V&JeZSoitnsB%5YKp;bS|CY&Cti6`+kks)}6Dwo=v%Oi% zXOJV^f(xjMrFDmV0B5c;4rx<%KV?mgyfIfMuOTq<^ZTu|!Kho90aXsh62w=r2v#;z z%a#9e*ykSnlYz$Y8GNTZW7~7kn-LZq0D(SG3ZU?Y!Ra=O7Bo>{7UXV`hv+d)L-0$R zE<5rZ&XBuvSCug==i~>wl7UJwKRE0|=o4id9 zoxgXzB;if`5{O)3u3PkayjT9~;FEWR2GnC_57uzX4m{(}C}|pf!xbj;hASUTc1j38kis1LScT`W+!z}6mcAEnqcAz()ezfeVG!3Kv@P(i%75+c|CH^}e(oBWfPs{>N z2V6T=#0|%%yoT!q_#c;AVf*A{LIPKb`$HL?{`Y#;uhi8d*f-zoLwrOqFLwtP08{k= zOoE(&KclHKAYl6;_}k^bJPH=s=AT7w!M4kQ07Z~cjix6O!BQ_5Pqsb9gRMtO#BN3t zC-88eR|Ixz%oDADiTOFU-30u#1Odn?zJVe^(Z7;`U{5Z)XTk9~O|zgZb&I0b z%4SLcdfH&LCQq2P#K2E7t9U~CwM_Kb&&A~E5(QJG2ZG;(G_Q9YyZjB@#vw?Xn+j?G zp?-7e*@AQ1>R~R^Lp{%eH#`?loI!e6qSFZ+Skfx? z>&9$6fE9LDU`o1AUDz2Lw(IpT{$XW5;ZXhNK#En@;Uy90cN^%OFR2gi2qj!|0J9z! zLQUpl!(^h34EAYD#XXeUcV4a(`wYe7g9?{q84EOqiNkBt_*9!I(vvMXTj5y`DwE%{{+l!qqqy=E*F2EUEvHm1or+0xST=S zUTy2`ouGdgU-%HKvPRP}#X(Wu@Z~iJYS=cyPe8A!z3X;#*7e&@G_~hI$f-jBkESQ)u%HlitIe1lm-*FzX`q)?By92=qI+g5TPZjRC*Vd3+RMP1zT$VI~1E0m;S#a&Qu!Cdk}*>Pr?G*X4~Kn)&n8J|ex) zK}=WwyDewcE8E1IH%++I)L69VP_IBzyS*2%GsU;jD_ZFiS#(K0R`4zFhXwf4#QqBGB|neo|%1RnuQD2rW?CsbR|lhGhlSf2husPCuU1^eQB z@Gz(k1mLF7om4DC;6X6>T2?qUgD(GO31Q(#oe$SxDR(OEOZ#xKg2cK-AXQdi7K$m@QdTMxZ z>Z7!FFVr7DyKkRFxmj_Ik(z`AW62@&6$YBlewU@^KVDDX91@i`z4hpT+-K|1 z?n;DSdp1E6K8kH0M^+L$xdC`B(iJF_*1Fd?6jUevGW&F^sjA_Ywz5E0h|iD__HZp*9xTW>qdINE&L{EiPjuXUR+09%+_iG} zIG}cfCOZ-*hG5nr%K6@Q6rYY|Au1ZLtT&~Sp8Zuq-}97oywbP*(*0_TjaO*hMr(*G zI3aR8O(Jhs465&bd^;KRqO@Xai0g_CWx?Y?GZQyzYo})d`_K=jM3{9pQE&sc5)}V*58DLQrbJ+&=*W3#HrZgMSfD!Sa9` z>9#2C5dP=fjBGn9f&e6!e*>9C#T++1M#y_cOg>+8rlt9`9K!u*io1+VJ>^87G@rkc zg+1*!1wg1kBbLb#*ZX)h70LSBV<3`{k%zlb_2p1y#)sI{C(_O5Y@-&ed9TPL6k-qz zlx>ZH2m7SHG&b!z<=L>beCWw`~(he z2PPu2z-EQmj)SCB*RZ=NAU8#GpRou%#f6)-;B+jbru*vU6griau5KUUwYX17x%^As z<%q2!hgzs^FSVW?ttSYu5FB%4Q0Lgc9IQd!5_S*TUpr;WKXtKjpUFa0qRA9f(brA| zlaoxKO(E#u_fScIR;q@U^xjl&=a-R*a|%f+HuG)@A^aK zHJ=ib$VF3>ji1--?Rmyeh4=wboC|9qIN!(4Yfu=;U)*k%H52lYTjFdXuos|76+I>=&nCHd-@Ig~q z@agg2zd?vYf{43(5M%rF{$Z>~nNV+clBvJ)Qlw9_t!t`njaB&34D*_VL5XTX!|MRi zoj@!7ybqVoNPP{iQ*j9fAxtzVdmurLZGIA=hr;&1M6!X&=u?E97b}909eOIY2OWp4 z4n&bJvJ(U`al7xfan-R)5s55`EIDvr(@0>Y_|-b{TLh1wv&joZX{k#B1<2>%CI4l+ z-1Y(e4-hvV@uGq=S78(doB>{O=MeC;q-V&3*jv~rWOYxi>+X@|CTwKj=^I5gO+AU9 z&4dahI%AZ^_xbt1wzFaR(0xD8=y>$5(_SGX5hw?t!NBvxHfTo7=}LEE`CH5b+YWz` z^2i%e58%godvJ?6xmf8`{5kF2H+k`R=o!6wd@l;p4S+4n$Ww5?3m#7Ued@O-H_Sz% z$JVyq`}hHR-t*<=rcdO$u%1}5RG9bv>Gg&gqlXu{5E|PB-Cj=#oFHIAO>8c8w#fr- zjAux$7Ec%PL4H(*d*PX>ThX@j`Yb5lVUw?FFGNhFL=iCW9-|>=00>1*;O|K4z6Rhe zG=2vIf8M2ztN4qc^90EPf3iaWYcoEIs;*xfA@@@GAu;6BU`fc&kj(vL3#IpGCR)Fu z>!Oq0V%k-eCC1k>5A_R0fGhy_j9S}X5yKI?aEi#N*R&1i>?8+mQwsKUZbH?j1JRUo zlLZ~?23JPs$L`tx{AE-%mOWCQr&jkzd04D)&h~ZHmBR|*(rnCvzk!C;92NNPP+tf~ zeM%9E-&KNcBKOmbMr+G+oM zx!Z@DWccf$osfCAb<{!5p+QtNZrvE1Nv}$XG9%KWOCJ7ZFrG7Z-4k_bTCajRL zPJ)z_WK<`i%GA%l=vu!N+_#zP^Vv~Ryi*C=6+0vq6JZrBHwt{oPp_|2=A-4xAcCR! zKA`%GqBN15=RXgG8ZDfTxbmBMI0SXk4G&6CnJBV7It1v)e8<{1Vh_-8*mC6T^V8pyU#fIJ+1~o zowy;{3ly(s>Yip@srvn2Hm-e6>uEo7{gG02wP3tABBWmI^}{oL5*9RJtvRhVKk_Q6 zH^K)DdQF*DQZ!%t6~5lmRglY1y~4Nu(Bkv8O=Gol2L#{WFMSkx@JR^wTsfhf8`2Pl zFU6CG;Eu;jaE6M3i5uwW>u#36XrLcnDaoo{J&D>IuP4T6Ou^_x-h>nF4suzUAEbH`SBi0@&>J>#AfT z6+wTa0l49(tVRM4k+3Y2GRKWM%no*xfAv((x_0lOfPzuj+qJg|WGu~aX;%zPs=Cf{ z!~A@Ea!ET+3^5jKjKIVnhRneCfp4Q$1{5o{R%*$d7Wim(t~UaC_8Qm z&QB|j63@}I)u(g?%NI0x&=~9drU;_{=CF-JuFZ?1jB3T9XAhkrAp)iUT>zNJl4j9= z%u&67NPyb_5HVaG+XW#k6UcGG9BaUAsNcGW7bUEV$WNHLu?kDnq#?PczPfwnJA55@ zz*P5Z@R-_rm!qtXybp`Ds(HFpzJcj3VVY=c0l02`f(PAemiV+y0o%dsKih3u+bl+x zs5Mz^I%nHeQ8g$Md;V}}fwZ)=NH+p|Y>d7bJcQglqXT^8O?o7pWQ>kLgDSy$+2OWB zZre^O&~J}Y%A_=T#xqR9Y5+8V;<)MK8LDM36VNfv}W|-TqY0PA#UV z(@?uN;E#Y8*LgOj1omn-iop zL2z5=Ri^c9rl47xO^PxZc^@%oubcscdmx0|lR$tu-+$gK4iOI0&QgT0LeGdiEsllD zgE?PKJqj{T4R;J18NXD&ZOV~&!U%W#I>PCIsXD`suc5|LoOFcmvuNXK!#%+#t0 zvE9=p!na@ae2R>TVv0P}=_A8~>0QFn&MeHLceDpVgv>I5{i7BlodjR;iHAFgyNzie zlUZvdkfLy$Q!p#|3yv^)2_D-aqom2{Fo?Xn>o*C1r?v!r3&L+_qQ6gVZs7+|(362K z!q$T{B+Od{f^8SX0}$Zda;0Ue_F#AUJwXQ6bc*H))q()Vri(}wY67A85c&(2#xnx5 zNbG%hi9AX~qu-o+i~pR3@VOK|r|w&KmDWoWQVe5ns2)wpmHfucReMc*p0jJ+E-RR*O-qg($iT3o=473)$m<3MbLm<2?coE z{>#|gxLt)dVQ_F;%dlR_5NPY{#r0ulwl*eIsU2WKV?#9XW_+(uFvIBE@tSqn3!fd# z-F?~JtsGoaF6%t)N|bsSvakCJe?1_)P)FFzN7D9WaOASc1p3QXEl}j!$q23={{Sxk zDjqr-8Hpv85Gtqh$n{QNR0DN=CFUQ0j4GbOzfz|g9b7QUO;mcKWOzHlQH|9`M(}zm z0u|M?M=%_wi=Cf9dB*QDxFdH=%RuJ&pDP>;C}~9rq=-y#Euoj-k8P*P8NE>(IuB-5 zqxU;gR0=E7#r&0L?dIwuSp1I1DEaZF-2CuilmErOLICaZSVu!zmjO8tXqt8=fz8A@ z82plU0lcx)t+8OKRbf4TX#4-7?#-j2{{Q~rk;=Z7eHmFIYh^3RkZehc>}e`lCrt|> z%vhrAAw+5NK|)NjW*fU~2_b|rmTbd}YP>Vc=YIKquj~3<_x1grbD#T<-|s&69}dSc z-plK`FOLT*AMzuS7e;Oc3h~-P*987J$_AWVu_TspV30-_ z!e|E9l8PA0Su+QmUYJ(AADXgGmMSt&?9z9$J#yPg1pxc>+z{sGbSH|#t6 ze_%A>SAj8>*|l9OSc*M*a&iqYSg{W4x95znjN%vNho$fe# zUgO#ubpovY*h0zxR(+6&V5Sp*l&{pxH068*4R;*B4lZhzCrp2ikN>2`F9g@hgLr)cy0ha4k?Qf=>bVgk+e#7TxJ zQ3L8@2#*=48ycG473qC-wVtU;@s|~ao@MIX5;U2L;yfIIh;e0NfJmu=>RPiA@YoI@ zw_TaFz~2OCl@+JmO_s=;XLalvpyppdm?XiflmWPXI!P=cEU*~bLO+o`L6v3!XZw%YkT)!)O z@l*e0kx!qG2fxLFktj&pwVh}|)|~_72($B-jgEd{sJ%5S0crM1&H>*__dy=6L;F4M z9{wt%`(`Dqm|$dmX{)ve>45ThzjWIC&5`f9P?OCXxxVS zeO7>RckYGIU$%j#xnoByv-9FA0q^qwbT)ClP}!eu2xW*qrHl(#}uDg18TYkQtUcgpeo9dbhhI8gwvsqhw4snC`MpatC`ix+38U8@g^Jji+>WxaAUS@xfS6~td3!v+Z{T5RWB)IwE@09C1o)=yt%1g?LjbE3~1_!AR$O1*Pxz-UHzT98Vw+$B}<7 z)V2(kH~kHJ%PK@Z44z`~KKl5Pnp}K1l5XltGmy@@SA+8ppYy7#uZ^hK8g>y~5&5pk zcCCbKFEIcQ66R!pP?4Dh&S8jIzRo%dy&ye-V#_9P#qy zC49^nby|dd^hgtf!Q&DCg7iH8n+?4E+dxY%R{>#HA&|pM!`vyjPA$kdP)o!@YOTg; z$J8(N%9U=(vuSOPKK<9a;<39j;yu+R%4KLni(3Et%X6l6@sZh^J+27(ydyc;dQ z#q?wLQEZCh@P-2A49$+K5~Gp_!qFlP1wyvGlnqlAA7Toxgo`^{&%S3YUu*~eRV7d+ zN(6^VNdoUAU1blTh7({ps59Rq)=%NVMe;52ZN01>#0>%@FV+qaHczp*?vrdMGB=&a z49yA|#_=0yQ9qzXUYR?zRnELj=L3po;-8**n|+J3^+vxluK%<(i#K2#eyL6mTjiK8 zlbxE#qyI=TL_~Wl_WA?1TONuqu5}gmB^;#P&awlFdCwfEf8LZ6@DCJ6?4NU3vIel4 z_8q_$J*dYoudz6R$nT$8K~k)6$&%w8N%u+8V9VGgu&6oxjc$3)q<>*jW1*}6Q1;E^ zrw*Q%j)1MC=TJdfI?j+3c!L>-*T9pEl!KM$_hR-{QaAGS8j5`lUd2zouz36Z+}8J= zgMS5oid(obCoWw0oW0iPqQDP$3wFzJnedwlkmg;Pzzj#vU>8;~KykhmgYcy=s>z-P z!tLlI?;3@u?d=AG$;~I-prW>i?yJ+QlIB@++anv*`_6ypsC*G;`2LC}>0oX{gsSEB zZ8$?rjm4?88^PTmGvK9e!Ncf1!6U6}N-xE>Twc&mC7hHB?Vxa&HX3o>X|I!`MhGmZ z{+k@12>+ip4l@RzfDzYsi)p)nj5oEo!VXcY7iooWAg?2m)_q<5wOIfT)AoqWL7n6U zl{pjkYFKFCe}KR3J2;+k1Fo&gyl{koN?<*-_!}0^4P$_US9C98QP|a%`Ah@24J-88 zA5T2Fl*Qcu)ag4QJrS{{of-Wv4vxZX|pb|)6KBoQn3Jc+G+G(e|M3CiUib895gC8NuXZO99aWc zn)Plt^QP2qN4Ob?`M=BEjb+{FcZ0W`AP2fg*fpxf<4q}>@kV`OF6#2g?NO1*u_Yd^ z!g%9%+-c<{cUq#8O+*;SA$;e2xVp(DM4b^b(K}4hQ4z)_9m6DmOo6j>y8$@SNL>UdoQ1y zGVCRLgsli8460;tn$nRFz#Pp?S6-)6Rd;MOSKTb?Du%DFj@_4frggCRS(;v#!`G)R zQC#~YPLHtoPC?Tou;)Q>LBfJ<&Bm+w0I-*!zpfiX$TyX+qgf_XlMbYA>CmTXX{wr9 zqCwOu|AdU&&M~K~+5aRqH@~oN%2`_Sy{)M!Q<&*;rJ;BIo=e<9+6AlAC>mJS z?uLd)Wl}8eBcv3zDOCL<4Bxonp^#FfudJ89cS98Fa(sBv#-~!X^2JuUhtaB+WMHzj+(9gvGT_~%pPsrkIaPK`~ib!A}DUQNwTWYGzbLT zA!Z`q3q;Hy$V;9GvAmrYX6>+OpZQuITs4e|`C^gP9qdg989> z=VpL zHTFE|$QB7>AH)*l5Rzd|6jK&=6J`He+e2Djo(%~uvkWLxSp$c~E{@;>d4e_lJ?%4s zXY>sP@3gw#S$>D8KF6}f5}zU{XJ#jpZ2_5YYg1?kmMhGGh~N(v>&)>^&~TgfBG{+H z)W0A*U9;NBxBbs-TU7V?iQO3#?vR`@#8C2KE+%}@Px<)~IQ+k;77PPU9)>_^J(&bM z#0hj;)Fr6la=6nRU&@p8)^Y9Q1BauY`rMsvOGR_2NpJL*Ohx_hkuqlgvpil+-xNV* z)xRkY4J!6?eYP1G6cng!bNcJ&&R144ab6E@SlzHvOhfJ50XGUV*H%T0Rv5?*qC~=i zAp3b6qwcn=dJx&WELZu$*P8+9=frZ(Ih)eHl$<~BuAgzS|82XR-ertba5Rh4fsPCg zy4pOlzK@vD=Wz42rxt^Z8>y_@4x<1o1N+C^fLo|t{L2<3`nhP zBY<*WD#gZQg^XlTSb;DtB3XWG#+1YpR9W}ct08hONh4oDt5zwu&Vjw$#e|jQ^WUjN z*`Y}Al!d!D0b&sqWEUT*OL8NR#Ai3#nUk5#yM6EWQP-(yYhJHNJl@UG%Adjw~oKy^mMp-6HR3OuGE zmKX&VpS%+g{^(f3}73hRtIcMh+9e8nVE3E)bAV;D8 zhyW&T7QX~DzVR+S6zr@Rx)aflA7oBu#rtpG`du0-adT^isd4kT;W6oJy?xu4i3hsw znPHEDb67x3{serCB*2&enn2tz<`obu%CmHWOWhIDh8&^gt72CLD9q;U65smEAJ9>4 zOC;-zY!XMtjQshC2Wi@gX||iA7!Xvpf?$|(9+`-(ki}8GyJkhg)Ix|)8vSLGv)`<^ z=I9uthiCi!CH3XfjoHbU<>1s zrJ-MM(nV{8ue5=cdJd`Sh(75GZh%Q4uo_FqML-7HEY3f46U-MlVya#c<_B7asz(sDzqDq^FwgTtx9{~PooH)W8G)LfM zwz1X=?q+Iddm*&p#2T+}ne3NeE&zu6jEv3|Zo{ z6$XP+-Zp0iY2>^6PTEhqUU_z-fs}FlgZa1K4=)8xM7C`iQcmE7-el$*2?7DNH-uBI zyJz{THy<|(L-$c{zhAW_sd;}$zZeSMRj-Qw)5K46Mf}|{mrX&0_l}LJ@t+=a_q!9- z9UN5iZn&8Y5H1B$vaV`^aPuLW51MzO%w*8AdVMGTGngHsV-+fG1wIKw#^3ulCW5>W zlSx<=Kr@Pj4mpv{048?PylB%L;So!Px*f|10gLUO^^RaF^@vuP>m+Ln^(g+#LGd3c zAr|KB`|dTQrha4p2?+b#{}FS{(;icTk_-=km_w3k9$jL3Iyv9-aL+Zx_3O)@mirSE zx$947 zc;TSxOBm0wwzpo_1ZXG!SrPqDJn{Y_hHL^gztgpdaZs!k1H-w0g||LLL~H1l+Q9p# zP$KKwP8HU7E{Cu1-syT>^B}^wYiV#kqwsr;*rma}X{I`wipiGWO3hmq1TEyEKbv6~ z;3&Sq0c1c5yjqK#ItP!R`wW+;=zi z*|i6YK^R9@Y%4zj)z|lS4mmMJD$ywy)X6#cwlqC4%^bE3iY+^!QWd2P=z8_LQy?Xa z762g$wL3RbIpP9S_!VaT^lwRq)>z$4)PB!?-`{Lsieh4#)#2s>JPrW{>noL?5IjJE z+`A1KYov+VPt|&)qgeIT?AaS1;TxNW%)_1O)X(e?`))^e z2TYSjn-dm!UrB&YUzQ6X=8I?Q;rnrr4hqmK&Ee=0fjFiLe#8&P2n2sRj#!k0^CQ;Z zOk!6`FXNZ=2*sO_p*=Ax4g*Rwt!h=d_U${jb~5l03}0XTn$fweP?Xk(t8O||#)>69 z_a!`robC(5bQ$lz{1@dAvI$P{z%o+cRawkn2xOV&eh!cp@=FmzLWpAB(JV_z>9&Q> ztW-_~BRhtDT`y#Z$Y;3qkol$tD;1U;=x06de44Hcv^Rz~iTv-)bazql|LZGw6YxIK z-el1fnei+|hI6!4%GElGW`brB-=lPKyTx(tl76;BK}VIPpK{KZg^Gl|07^gh1>I;1 z3`sQewR*smTG<&In(MnbnVjeA8U8NuP6^kEi&ovnT(1RwyrQ%CzcVX0&|Yb9QOIcK zh)lLTlE?LEoYkC(s#W*-SEyeo79S`%Mn^Os$F`aQL|CM7m@?EP`ndf3r1Xs3h;9PU ziC^25f>yj}_D)H1D)O@DhLh2TqxTuST6snaD6wb zO=1pt=_B@5R)d_?RiRr?PfQ+fskoos7Qvy&vd4C~&emix`7uJEJZ?AmHz2&q2NZj! z_y0kl1fq8GUEGuOY`Lzu7EZ)8#%7~sa*je%jcP~KUPiQRi$Dp+@UZ|KAX^j`fb_Gm zs3GFCHtWFtD%K#Ymu?;Y1PPuU>lvQ^W>al|xt;Nn{73(9mzD*K-n#> z2%hQoW+s$|;|Ez^w~Ts}vuU=FHDC9n-n@}hDt-EISdjNXZsR!79frLGr`#ei#NdF6 z@!;REPQc~PL-m@n_*y;4Jz-JqvX&{@;)+iyZ694B>Vb|6k04W;NCP)N`Hh+wtQLD<-?Kx%`8GrQ9_?whcgeS9*qHHu6xr9^)*C)_LF7C76 zEvTp3Yi&iUG}PW5U3L6vbBLRZOWJZ{V8}rt8FrXZtNM>=5;ZtUIU%bBF7SIGw@W9=ANn6pNzYl zf1g|#ZvG3HZjivV_T#EJ$o4?*xB<;8E`e`uR>TJ(DC=%CgF3AzqZfWSAl~XxOBbQh zPC@ljlgVBWTx+hyg}Px~F0v*uXKZ!MJl4nSOxl^%JH5Rjyqt2;pX8e-fy&*{ks{w` z@z_!{-o4wvi?%y_tp9lQ%rH-|Sw>qn97Z;FG4Y7t5dgS8g#hwNj3J#$ouUEz7!`w> z-QO%&9yu4TJ;yA}S=HaJ5%RM~1&L)kj_0+c-sQbpuqN)%qb`#m=wLsu#R;Ba8d0Me z2fDI9?Wr8mfn|Eb8%N5k{ut`}b0g1|^R!((D#QL(?W5)Omr;Coag;!g$G`c+SUefj z_@%9`t>Ib6+Arnno;jjPTGu6&9v!xf6AuvXR1eKk2)od9Ir0x#Qm4G6q`&_TOhWMs z$94wYFzkwL3fKTQ-O3p65ZbXYHEvJGINbrq7duN~g9*__b0S+`Y7Bu_7rr=t^p>?dSRV_EzQ%n#Kz z4`{1{9W#6*>Kn>w&g%1&&UbM>i5>6!KJy&CXuK~@Swb;~1?E*t>iM?g0WJ)o$~?Pp z!hO{>GmkRfp*L2O`Mzr8d+&}v2&v+Ew6z8IL|7?{Fc4oje%XMWas%7iiI7HsPR`V0 zmWV*-;+HTh6v_Yi(odVO?{0aF6b^e<$1PrS)SXhb%ojhQdgkffr`W*KU0V}v86ZWr zXYuIHSvR268WP;~jmMN4b5vaBtIewd98)9O_NHe`SAE;Y7x z0rihV5aSWD1X8Oe|N zSA~)dWMugaZe8hp|E^dEm2uOXFNzZ^J)r>f3Afnm*H^+8oz{#vn(m|^k&I0FD1N1p+4M{ma2r1XbH>(7rtt!V1P5C-O52|d0#N)DYrK_@j z{)Z5g{POmRp04|ht$*W@CgIR2Ta)H-0x^-LPUp**{_*WM^YH5Yr*sDqsj~^kRb0Xfj|{x--RzM1qekde3KpomU(%sm#Z5%MTgLG^}Y_*g9ZJ0 zV-Nf+mf&TRkk}<){sqJ)MUrm|OyJQTjY{^C115BslB1ga8ViPfN>e)5sa183c1|bR zA~)1_ME?oW)+_*II~G(L-^!|khjn1J)ifUi@HN*%NVB3S3gm6m&?I|P(z*wI_s;F} zXK&bvPCq*&cl*u-V^s_V-#rVRK4sE8N+4z%?X9o=(D1Qln5E$I{i^jrb?T(~BqmR| zF>$}p`84@G-}dEN*zJR1T>i}q`i2E0cMLK1uGw0BFeWoEmePFK=OHGKgpZ_+i`53X zhPa(zmT)W;s) z3%)lwck70i(jj3>p61a*LB|y5m2>C{EKYV17%7GYGpn$~yct!pSr1%f-+mFbY$X4(0 zfxG_Tqd(aVnuEmY&cfW$(dI1vhQLtuqg8RT zOFn>Ad2tg`R|XawWFdPR+59P8-PX2ZVOE5OWDPHt(#I6bsJ`_<6=!k()Wx5prpMY; zu+J(SkzHC)f*WlFbX^DA5YRc02T*mlu?#rSpDT>x%-p~e>*ZjO>G$b}lO*zW53mL= z5^W*XtXFE#dBj0kUE!W~)%>4%z*&8t#kUP)B?o%J;x3}Yy5|X!zIp~h31QlJ+c|ef z(kF85=u*_o*xd&e9@|R8E9`Pgruib1?}RW!d=DJb<}oG=fz@PLmAb}?#_HP9W+gu- z$KwM$=2=gZ5a%EDNGR|}ITxnfTRzX*9s+ms_}z^*8v!we+rr{J`%LZhLQEsyn^_A0|L2^I1Gt zsws;rq=a#bf@}+QSFjy#M56R60?Uz8g+J2w+Jzv!Ev5bYj(xwfK!m{-a7(6hnji|p z#bYE#0vb*=S=)WpwMK=yxjDP~Zg~1SI0zjS^}1mU+he6mzXU|Faaq5emIT2|!s9pT zx{dnbewvKTE#9`}V<{ND%}JA~m&BzAqnp5PR}=gZJJJjV5J(l`fW#t)uneJg8f6qg zEFb^^t33?Yah*_i0k^TSrNN7}HBTG0PM39LZk@}Pa7fM9Jow}G`ptP=#-;J!6F_Fm zee`EQu`%i}LnGy4Wqm~&>8DeD-0ed?InFn~mdlF4%%=a~dPkBU? zRk>@40t~YM9{jq8Pi`4(8*N|rM0o5ss7iu zv=n|S%7qm-TaMpt`Mzq9itHVLVV!1oLw!+ek{Cne{v4vCOq&#woTZ##BtE^eJ~{fx z+F!O5jipse>eeT{rAk@sdp0*ZvdH^xU?&_INuYRKh{XE8W{tQOWa_7`1ZBxeRM$58 zIJ*t%SeNQlz$$3i+@W>`o;KlbbKOr-Fd zRD@y5ujJo5ca4RR&h5n$edznVPNF}Ht>jZp72|{DqFXt>j1P{3O05-q!rOt6v${sR z2f^>&zpr4e{vlZTf_Lh-#pmfDh)=V^=K&gnT*TZLSVo1M+7C667w_a=i-s4_hceY#p>z&8H~9-EoHN{|!qV0<3>g%x8qh9;S9Ua0i#iTff2gRn z4h;>E-gow*i1DvIRc~MKX$I**)DpH=4(o%b3xliu;@Qd9 zXAV6E8~ERMFkiNjynzE+T{F5=y$tT=cg<3@HEA-3&Dk=!Ttoh1&e=)vG{6#{UQqjW zgOG%TPGqs{ZK+@&y8R27$PuZAVJ5?l0U7THQ#V+6n90yDchfdFsFZt_xEEH{Hy3%{ zxb=(JEqklkp|9tCkZM8X=VQpS$qQxvcwfibvHKSHy+ed;_YURusNVv~;&(w41`6WA zZ)(#(fRM&UMAOCA5x;T_VC!m`&@}aNJG^qc4Ivaf-j$PElQM_o{nis+a`g&t_g-$J z8^%2+mqGpzIa8nW2hgi6g6Rb`GY`9N4^GtF!Wf<*`GX-=$>K^6ESs@XUbjcTwNl&| z2*Ah&-)AajD)#(5FSWGx6BhC#hYq&vfhO-e&rCzy(!h4U;%4Gl5)%LkkUSLVHo(CW zqVvhy+z^?JT{L@f+)_ckBuJ+3V{q$!LYgEK)B&RPz)+jUP8Ktr3C{K0Eg$xs+}$!) zA~Z)g2tzR6G%E8)17t#@gVd!s^qtzf&ZwQsl1WwfRa6y)`rqY6cuBCIINXT)Hxl*B z{UF<`FttJ)FBrc9WP8p4!;bl49eaICY@KC3-RzW=)SySL74H@y;wZ~n6V|MxS|XAq zxPu05TKr$juley^umf|*JwErp1F(6?{Tcgn{T2IRzPxuzif@RR zD0;)3=55KX;7u$5X}NDozZ~2)+li>WjiY9Di0VgurgsTKYG%zTPHBqHJuVvsRt1-* z7GtkRDg3dza)6$SY=w8?|EM+}2IGy@;qp#tB8P{YrW>V|av5Ro@`@=Rfgy`{OGCd_AeurkAS3ItR(%pB}`2?XBs z1zI(;`JdS44!@JGjvEi(#&3Ejk9`Jk46l6ozsBqA%Z^)z?Zp*3muu6K!C+Yk^#?Pd z9&9XF8N;8cSFVVz*Kd zpy5foz(9bRoT-O2KN#mHJdR-klJGp{b3-RQ(l4iIol3F4__bOsPP2TY+)%8@O1B^M zc`ZCd0nm}+D=WqEFJQN;u~^8V&~7xS_dKhK4Iny00bKKa0y304H*s2W#D7b`l}d08`<%5nVUdHDi?ax*Ae#A-u>RT?v16CP zEdRm=45%PR79gXBi4_NF?JTJx&1G_r;@4{TD8l&6m3I6k_sU$m+6~h#Cvj2veq|^QT z6J0k0w%!day==aD@kfxH^!O;iju3#<#@8L_=P^UBMeVbZvmEP~{lW7cqDBEnYh!h$ z4>exMTvNmgTix7jnoBf%p(sb5>B}ukS#S= zomdUJJ7rVqr{}2QTU07&IC0j=_HgIt3!^tU{1~A97@dk;)SCssI%=NaCG;6Um^BFPOT=JfF%YQh^EA%zB6&Ans@4~e4 zH>_Ul7K*?}Sa$@Qp=qj{Z`rk4So?=C62YgeDVvXVa`jEV1i1OES20i7)>xz5#D&fj zTr!VPm46FUjev<1?{-Mj7-(iBivu&vI+oIb6fzQEXj8`sG9O&tbIlV3cHij2VlwB?C4B041ibO3_2h(T>gF6f1oH%i9Il2VwIOU?`M z`PQv|m&5MT#0AbzHw+~IY_oV`fNIGxv#~c2u~(&r5I)Gg6AT9sk)~~uAEF<!3 zq-=<&y6}R;b6TJHIjcaUcTl7@i*__wW31meU;g?r=e_5Q{&1giqADXBW+VVtjf|lg za9bc#`}%n%!YCft+izHJ#Li)aB!P0^JlMwWr6X8eAmrUw4*{>R?|wBPhmY`}?vEsC z@rY#pJTpEQuQpcgpLaXGH|?xXZB2XZ_p-zY2`kodHCh&n_ZiTJ73d?-26+%ywoOJD zv6zJf;wEkaPkR8(leG3u1&5L>k~v~Ga(dl8OsXsWjg~HxOb*!Gm3r^=y`d)JTpE|S z@_ibZ8;Ztby^&KEw5OEK_EX^kB#y{xThR+O2JYKW%F2X$m#E*=^`>6GGIvq(sOC|z z`*vk90?`*lW&=BI^IfwvIGE}gj#C5Ke->dY4MfnhtX?#zKK)wrH->*V1IeN+Fh=_h zGs~I}vN{PBn+ai_)K^mO7A;$DMp<|BZG5tlNJo?{c&EkCpJMbIu*5tvE@{9dYUAFJz(RJL-2*Ni^lTj0>K6xQto z)}k$Qz-eMnt#zZnPrbtVjWn&>G!OCGS){(t#D)#e8w zH7AC&?QTiDmUs%jZ`IX2iHKl3Gj!#uPyH4AoxoGir?YS7c*8P~>9?-IB=m3oVE?%l z4^lil(dJh@o)A)X_U?~~n78rzTkO=lvww{Suj z(5hr-z)3u0Wk;xwQe-_E&=9Agbqi1sT;JDOFjQeBpQiWhxWRi>hX}bS&H*zz26|yf zaavcP;$4O|H+9O8Rus{OP>)bwZPUBnh7ZFwr7P;6ep>AnKPnf+COuB+3=Gux-Jfow zm~NvjmS%CIfBk}$;+7TI>Y;Z*Hgq$_X^1wyMnuX$`dvwdPX?C11eW?+uVcK(ue{U9 zl@GeR)Df+Ma>(_dtaTrC39}>vixA2xcJnK zlBbO_uZj1OyNc@v0IgF%%GHW)KCrTLyY%fh=ivB}D<4Mv%Qj?B1gab`Z>qC=;~;05 zpd;8T@`1f0Y19Nh4rm)y|3PAI1G=Vef|=rcnaVcLJWtkfJWYf8rLadYb^NImG2=U`$;fk}^+sKClk(V`)~ zIoxD45f|4i3-xwZYwmx}A*G-Y^OU+?=9ayyaZg%5Nta?Fvp84`SyuF#2B z0?AiTMsYmwAsJ|tKL7ek>;d&9zjiWeT>w(lVn_`D{#dR`fH-K*`wEOC@F%A9~ z=wiQ%g(WLvi4FyE@=5U(56&vRP+o+2yktcJYlx*HXcdqP#iGM19~vP!weK=23>()T z@hl6D9kH$RdH&&D)XAFw?Cd>XqHxidQS9T{8rwE_0!F>&OPCtfX^s$AJ=L-_dto@t z#xpeNgqbE?KH2=iWJ&*>6WmrrCn`t>_C28{m(61#)mgGkV3}V|EEuh#DtAbR9iO>k zxG7TIoNp*fDZXKutm2WZbE)kw;KKJbrqj^iQxuS^L{ctP=YdZpyz{MJd^wgirZ0fL zdy90lUGa&jn!rU9KpCDz3hD$o6+@UpT9efV3j|>b1%`dMYA44E`C=;1wPYk(si|%p z+U*GBDk}$9-o)~zYejacQ0Ip1dbv;UF^`DVuP%84&-$TA=v-|8q*+CvB z6utKoq_93X@zq&Dqtnoc>v*+|sq2`o#_g@phPAcFb#?Vm`~L(1USy9|uu0DZAjnz> z;xu9cLHq#*tN>GM5ben@<^+nhKQ)l7-A*$oe>Lk}ALTqLH6c3Yoti13cXaQ*8T-2m zi|BE{m@3mKL{YY>Zlrc6HmVV)POY5R8(GV#ZqAkIh`9Ne&gZ=6{JY$b;QSdv<6Q=| zk-(Rf;wodt5=#6eZAQmCa@lQj+)h1?T&$`4dM%11s%mnqk2lSHHTdQOJ5EKXDI6sP zofznE&En3^^U0EYlV_At2+6jQK6hC^Mxm?#6V}hGCweB|v0QJ}3ysV`1@&;>Xk870D%!?rwaw z3NYMDpiLH+5{1KY?6(@z<8mVm8-XH5pp{4_|uxc{EHl1L)> zzkSVbCjgleyb#e3^oD1ASbP*jw9Z<%5f^&Xac1y*Rg9hZr8n8{OAg9Mk<;$>9Ii`& zf>YLav6sz8h?^|Y9Q5BXS$rShmCx@os7av9hsg}z`fYP&+tSEE9~SHnAPi<XEVyOhb6Xgi5g)uS~ zfu)2Y5(7~-_Oy^MH%4_if1O$mY?w+qeOuaf(9q^)YA^K%Xia=jt&rGb4%_l!rzNmt z73d(D7z%(u%`fK73s9FRMbi?QI0hd*0MpNsqX;A@q;$ycqtbFJ`e{y~XDJmf*I!9o zOwI62K6^bw>GIQdIbhNqn4Z|yLRs*TNbr#NDVwyK;gue~J3qrSr)9WJ zWl0aZ3SV#^-ES-Xg=4^*#Ty0PvSnnUrLn|#AU6s4Y6LL_NinvK#ZoNc2LKiPjwqoF zn(AskT+QlZO?>^!4gG#)k(5!EAA8*XM~=#L?g!h<2!FQP@I^CN$ZQ90^5-0)6@Q>6 zZeAcB&d)efrijvZM#(BvR@7fc*LMLQITdLlOqh(0kN1D-619?tOj8rWoiaZGnVwH(j&v8gU(>4?li zMOVY~^y*6v#>;ymPKemk_5hDaL`*%N6D`VH&52EU+dj{CM0mA#j+ia7L9$@`8pW=5 zJoS}{LLTsm7{FR405l~khSh~(V>V%L4vK`GgfbLdAmqhpGWxAX_tVMq@1v;Gv5G2N zPre5+3fOU%plDE1>A`f({tCakt7grQ0pgSH4hYaWZb0phlpuysY3D!zxO(Pf?8m}y z(^A%tAAO(dL2MohFSwob{&-PzvKdw>fQ=akKFz<25WN>X=Su{ap-FMrVd^NM(;6@F z{mE~*s7dI)4*g>-ZhbW-x5>g!`aYgeit~7RhvUaK8Mxy8ZLF>ucojVU8RqLGyv-r= z=OB{vDoWBN%S2fI(;={Jp6fLc`#rMj%<-++;GY1^;0KV<6fxEpN#2|{d50gPJPVMf zMb+&uoXcgc3ql+X$^5-f3L|xEcRP>!4xU!CCrhtY@`O!wJ&lULXmK5v7t z{TOQfIHp|VO}0mln5y~mg1M)sAwXT18>eOcb4=6|nr+QNc3}6d$;yn84_q7yTsu;= zMtAaZ_*k1Qc~U{iQtr4mZXKiPJ#0B`=^1#9^P0)tJu2)}IBA?A#ZH%6{u_2M%*+pa%U__iYu2tVh0eY0JT&>k zyd%-5?@(G`IF;GnWa=3L$|C?r@gRW+grpL&9z?fs53{3ibPb%EFG^_C^7^u})Mxi0 z6TvU_d@K$8F zhF}n14e*)*0SEiKnjCGX_z;)Bs^1DTL6=SVJ)mt1NV_KT{55-_;{0pU*#oA$k#2H^ zmeZF7*>?T_x3HUg?q~HQs?ocm!=v2|m;N;yiBY9`z7IE_*|x7p(jAX;RQhF6o=FOT5fzOI*P6+bwDC)d_XR9y~K0zA}xwKzV9-TVnNA&#m z`*1T@Cnz1{Tcy><;sWn_U_Dd^(f84ZlARUpGXgvT^XVSW1tDYjj8`=Y&V~65?mfS5 z^tZoQD@%|vT_t*B)4ti-T8%T(^15S=Blkxo!P~deF5k1cgZ(B*Ybh6{*7qCmy@|EBA7-K4@JR&D zvhuP5pEt2c@s%FM3q`!|-rS1%$A}|Agg{K56%TVM!nPL8R{dTi#5J1i`|&97G^Wq= z)#)=1=@t7|sNY^9tA65T<_o;u$GCe7A0 zHyQ=z5nNF3xb0K3K)smKsnugsCndPaJntP9p7L{vj5%!$0IP8a>UJF7jQ|Z|6Z}j3 zM0bqJhy!Qbro=a7{5NKf$Q{#3jjry!?Oza`)GdSB$p+oq{Xlq@fU#0C0zb8yF}UVZ z1jC%5NHitgj`DJ;x4TncY>Y;!I=+~cpYYi7d8Ov!NnX8I`9-t346v)ZDSLGEarbPL zfdK7I9$=S2gkk~!vK(TC$P5VNUzksS`+VZF${t%N)|Lj^i}|aRA{O5-1j`El z4Z-)1OY>+k4hUlUYH0dS1n%E_>Su-v*K`N^2aZoCEIyMC!TIu>=MioC*!e7qw00dR+dhi%M|Ib?lNrYsz}Za|TL1A;nn>pOIHBc{}G|EWRhg$P^S% zLBCfr;D*j~HCd8lg_2WMqURhs zKb(2=dcWDs=X)0p^1|+K%K(;jH~JlHYZan`RY|`TTx-|JEYd+n*Sa>1(;m$nXd3-B zt9?G?ndLVz?RR<76@k1boTO8u-NEdY7K^W+nQnA|^$jCNeHBeN_eDqr4({u%U&@`E zl5x7=A=59cVESMUrTwXo_l27{PlJ~g{0Iv~uztH16m(x%awy3SV6Y9Ux1XR3J~lPi zg(k#Jk@=5rbW}>S_2@3oQmkyn{bo48slH|n2IXiLzxCC>A9nW3yMmn0bHr34?qZY# zNxeWg?d(QklZ1QZCq=NaPTt*He*})g90(5%!wEi}gXuFEY#Wqeewb<7F>4z6&0{fj zw_w)1{Ln{}TzqZlo}_lQqr|7}TLhL|~0o}FtJ`@68 zF9xzD@d>o<#1@0xyi4rvm9=`->%I4p?__;yo}4Wx@D(1=JeD?~kapp+*7)WBwj}W1 zCS3OYZ~Phla+Y#M3(_aF5-5%ungSxhpi z`zP(K_kU27?rtQ-4@+*H(-QNOyp5yat*BY|1k#=fs}4xZy%fEc_3oR@g;C{Q^Z`1^ zR`vpojYb9_A!S^O9Q1-3?F5P;9RhKV3I=9xaJ(r6!694sbr`*BD<+5%9iZwv{90XLe}0%mI>Jh!-(u#6v|#g_GHb@Sd);g#yT@1YcrPd z!kEt8=X<{AI=^#$uXCM0&N+X*U6;#p4cBX4ujhT=&*iZ=Rf@An%$JWJk^BqRVUF6* zYv!t5h3}~WCesmTws$S@nkw&qvH1_~w11f*^up6jJbq8iyH3H?idic}lhEC4AOLW#3r7qbx`(Y)tS zXo)7cz8VW>ELo;jrU>=Dn)wozQMSn*-@wl+Jyehx)Y}3Yf2&a>Bmvbnh#77LjzKL) zx(ho!3!#vW4hz#yYwiAhC{_JjRx1|0REs%qnDRs;hz zO?%1$3EozUGNeArC1z?H900bw-)t=x$4{vwFIc{x|CFM7@6u@jHYxT(p*EM6BwH$5 z!YXL}Tvho4lS3hZ8lN}Lw#o?pb26K9Wk56=QenY8%!c$DQQ*YqcEMFK#w$Qd)&HAT zODq4rz2MV7Q@2sUnEryThCpwftBx_1y?F(M@wQApLu!YLAiE%IZZ6Kr z!YS3lsn+k^z3r)fVUI3`yOoryaqiC(b3qFr@OrefLD*FVP4T5t@Dgp;UBps!bhB&) z9M7m;Qk89!yZ1pov*;EamL;_dIvpY2jK12L3d0pGN;mOVL6OC>f9lgJNw2=*WydsC z*V9%hk|my1$4?v!x%id!i53t_mEt2AJ~Qk9&3(1^6}~DR&S*Q0+co*4&j{h^Ztp%FpD9 zMtExal$LZWqK}oBd|u{peKi$zq!@;Oa)R}*19YSn&UIk-WXp~pgab3Y;@y2PQPBbq zb^-^TQ329+slZtxk2nyOW(v&t zAe8V6oT2}r{tUz(6}M>coSu-{PCKmhSJ<|1=}l_vAKSp z?W5u+x?H+$|L~PcI(&||<=+2%LX`gkoNy(uE7GZE^kasM!~MkU2vOCh;kCxQN~>RM z-?u&V#pm~V+>I-Lb@*siz(oh!mo+|fu^nN^OLLJu%rk&yQpr;GcVHNdwC+)UpNr9`)EQ zqCPc&jwqn3QL-bpime@t#>(iI@|G?u)+C=l13vE%|!5Uu*yJ`5YQM`!HeQ z*#CSk(i8q(4R}y-GkgFVUoY1&a2)%fa+t1ChN0YBq&16#ue0r}o~(9CA4m!C6OY&q ze{eubp2_37gd;uU@8-su#rT2aY?eV20gA=S>kO87hRxq6O(TJKUm2^nacStn3OHg;>l8Hy^-|+oZ7d zXbc`ew^@iWUyp~+PhZFOzqX$6U7JYEuoElw-ZX-*%Yckm$&5Bl>5lcbRxyAK>_`UM zU2B#BE~b5*1Vv@SOXm)0>N{{wy`N=0aDAv%)gFGTc^KR+47gj<2*Nd*0SJTV5Sk3{ zLUMrjjnG6|KnMBw#p zSlsO748e}}x8#h6@TA&y6dU+Gx26ge(bJOnVC=DCL)>x4Qc|tsY0JLI#{TvbOVlLo z2ZxMVe<%ITzuO91YPx~?RPLhxc>6NquoK_5JKnF<1}eYysdo63IQ%uawjs&k$OWIS z*zO#W@C4ame?0q4%R$PIiN?Wx z#Ta%sfoe{>@ zqL^*9CFT*jLBkTWdOAs2O5;SeiCC8pYhp5cbci16cr^-AN~it$Fj96RVz!EJ=!%qX zDi_punR`-F9B#B-M6`TAS@c~q7*oDXh=$S6f`{;^Wj166+47Iv1i})SU#4Fjrn8~{ zc{A|VY4;i}GpoS6NC~goue_)CE?$z%{d7ZlUZy?x0?zwAFhu+-rx=Db!kMatZg2ab9M`>?qC-$?Z(Z#Q2Em-WfJ02ONQ1i4Agm14wj-Mb{4J*@a@LmH9D2Kb=hJvqzd@OHrzdYa zp_3!;OOn9phisoh-pbC4KL?Y9;h+CS0{cIU0{r{^zxU+d82C2^{*8fuW8mKy`2Rl! z%v4J0_Xe?QpAau}$?;F?b zcE&#%E^A(lHe5?mBYcPLKJkx+{g#S_$W#oLmy)4gm1f_S$DVrDO(OM#p@dJ6T^l(z z@QG@_-7?7H>3ZaZgOJa-5j7?6-4ek{x;9nNQWNdpU(52LFX@f{ws;0RsZo7_=u=IV zA}ftQU7@>>M9{hwvD(Gn-kM1|#|nFm!ai4UQ~`JtmmOV^X{#x2#UlBo3YU5>N4s=l z-LKEOm*|%Y(F+snE6H7qy$$ZT2sfa;otDDUBS*!Jhe9v#eS{8K;aA%hkce@<}b=(2X&LHv=35n;ni{a`t{)mF8zbXPfw8zH+WRO2FyC-01`)| zEiIjzRZ5OC1Ra^W`bq{NOW$`?gpaG+`|M%C#DOFHN#{N5gnbHvzG-tq$#1cebuT)C zB!{-FMIJ@j5ZDysHtJl?UA>{psm|RM^pa`ogj$L1F|6*OCZD&kdUJKY(~Fx8L$a>N z&v<{mz@(*M;|tv4k-Ul%tD@E>U4jPl35Q|_oKuSD^+e&Hk)Pq(I|s*VS0!eDvs+_P zqt*Mq`nAoUhNZ>z{ltBYOFXzQ9rX4Qr4O<$(B-nhpm*>G=;zWA33O)^FE)sR3-R{( zFsjrSpkJ-?ysKw*{ln*sFtO*JHUWEDL0@e68lfw>Ht4j|O|-oCIus&fIeoTMJ-kZn&i_7wy8-+0wlSeIyJ9Sq z^q1ng1WZaL!d))=)?sGrx1(XO&s6HTR4T z+ENCa!R+HrF=!vR=H)U3`C0+~S;-Lx@l>}nhEPm&B(Xe%uusb-?rVQrA7jDf$KCiS zvOG9nLg#>(X8S#f@CY|KO^t#nYwLU0q`&p#rRn~Yq}LDgtw(m@02Kaic();yHQCSa z$@OUdGTzLdgBge%2})VoClrFtPDuz;iuC4C!S4oUu2!5mn|t%Z(L<^WgNe|%PU)mp zR-B)rOU?B+z9zA=Vd6$CMdTnk`acmw37TA{-Ra-7(WM`Wl^&yxgIM?neLkpd|qo|6U$A-=|d4D7fXY051Q%{baCebR-m zx<&j17EoxzC=&p$(D9_HJRGTxVhO9TZA^bEi=UJ{A=(i=1D&NTA)aNkD^8ugYwY-7 zqa3-nalY@xrY6FkhiR;##Q9$2UEw6d8`mB5pG^ss#4Ht@YIdMxMZFIqWXX+gaai7I z{y+&&RSp04`A5iW!`m4lDlTo}Awl<9;n(PQAs|yAvseM)+tj<>h%{vkm2M?}c}h+e z2#KAA;oV^&3L6i(bYOvUR^=TDaoR$V-s!aE-B9ytiYM&&x_+aqN?@mI7jb;7TJtce zr_p=1U(P&?-F?Zl_X$_Vt3KF`>^D7hscxJz?jrz=twGkb-i+F-CItE2)ZrUVorABaKwf+HG*i--uJryi_`;tqp zM6i0=&R(0ZLeMv4RoGbV{Pk&Gir=~mkILFw{y}0Br={HH9e61$kolgrEOk0zR_c10 zQl{>>LAb*@*>A?V!A=z!hBgZenldUK2-lPnjuuo?dTDAh$#Y)^<`1VUx;@4HavJ&r z`wfA0s?yrBsrc&;Y@2?^mhxp9j2Pb7-UP-bh7foKd;hH5F7M{R9ys{oV9t;O85u@* zD3*y&E_2%6j?YWsx1a3b)a$PHs@|A>Iy)uTnq4!R8J}R!neDJ23qT2oRCm+feU<>T z4%A@-a>8M&#gyGVT8vWLWjX1|cSA1sC%LX!-QO5YgT3g^l5}Iipo~USKD%7m(>|tQ z*?!UYcAuSgs05{GhF$yQJ-RVPIlJs0N8)sMN@UG^SU3B#JuckQJx`Z3|e-qM_N*I2G!U>p=IpPl)!ed$A@*9?Mg(->CvPPV4rCr|LTt?=_G35g?W zT~D;$GmmMX!=kGp`%*{$2j3^1>$94koD3hE3Qjb8HzYH_BD z`>tazQpM9c>mO*Oq-{$BvLB`c@QsMR2S1I0dmjOkzn=KTAT)$YXPjvTiu}H_qH~=1L93b z+OIPo3|0u;=HKXq1$q;;9=UO291L6PD$$3_@9Pn!vOTZVF;``9>iUtYq}t#sc%u9t zn9?y@(7$R#(G9d`X!j^hbIUv^E=?WjB$1-1;;Ub}2RW*ke)E2&KlcyJksNQ9XN;CjBLcOZ=wgqT(X^WV zzEzFWPc$;uX9G@v^jC>@BdDBujeI@)I>BEqP#WrluDEov~KvEwNW~;ScOC@1gx_x`;b~oqkx} zo*pzmJ}|H2m)1uvb|aK2h=vlq01~oNwbty^;FZvg1+}RxRK8ldx!^)41vSSq+d@B* zN_S*#q{zkIr*4M_bB5t)IQ!H@AN2IO%&-6k~_hn8RdZ)5t-{ zr65$pcr=uLld_th{!<=dry(o0tPH^&kv=Jq!`B2J%xs2aysi4%y%p$WX zR`l^W7anP_Ml{<{+K>aoZF!7Y+Xswm_mvQujG67|+se!1tQSw6`284p*(__cVowh% zJw#`tY_|o8R7C%i*GrNdP14hsKgv3_;k|GsSj(^j++y1h%~J`qmZMH3REk~Ai`AnJ z!;b?J7A%{YsMFb0iT0ZMuPN(#_P!1OcoSx%c-2Hcv3{T{?Wsdz?Kk5bh)pVA5|aG} zCsnm`dd$b|h0*NJ?~96ZnJ%^-2v(b6#oili3@mddxW4vVuyOqO%GJ(ZvNis&k^lS0 zWr5e~&TPjwnjxPy3KgP{wr_JU*8F}lRESr-oxyD5OVyvfm51uIMs{fsqKKZ^_J*=6SCy9ZQq`sOqxv<`yx2 zMhOf|+5!+U>sXl~OPyv)yOWGt`rUey1!+Iksb}4PV3q>&pJsk!4Uir)M90$cd=~5g z*hKJPt8PZB)e#yZ3GeLqX*q}z^b07PT8W>5_b1S{M#g3cq7I24;V*pp8=8Nq?k$hi zzpwSvZxQ;SVmDT}tSvB%)%gi3WG9-+^et>l)y0o1Ulrn7$*LkiPBn(&wyHmyAhP^n zu|~a1&Ncn9e7~FcWz&^6uk0i~d^L35Jhy|RO;inOGL>U=$$G=mkCh}JB>^vao@XD$ zbj^GQ6*ibdJY$g-yB?p{wJv`ZinR1H$w}n*HNL8d7a`0&pX257w&9mRjreVvKE&%s!QoTK)3UyM%RK2? zw>UYO;+S5O@1f_GvGC%%XUKKdh1eFS-z*SwN3Df!VUeC2?$I56{;K5mH6whoF@jxP z@C#rW3_!ks_FULPVZ|{{m4MQNOPX=iaZEjxF6JZnbF9F>Cg9(!2yi&e%u)JHKMo$%tO z3DXuYiuTXiiaOzrYN{PC&7D;0kNxsFNzq00MiPRZY3W>fCqkqbfG=A8Q>tAa)xS&g zvT?6H7TG7c*R86>sz=UIs^&~~bGbkE0ngjWe(1;1gSWF9kLw@q0L z>xv2GQHraqTI1t1Q)gDv@Yjl1e|m>mzX#)s4ev;clOv#)a{T(HxG%Aq7P6mHaY2Er zxZyqVIa!muY}u>`IlP57XpNOa)o9(u;ZnM0Y=waJh1;)WZ#>6YzUm|qNqQ2a2A_)yc2m}0-5rIzPw){Ga@34espQqCN%HI=cd!wdHR5tPe$#s zwaEXc>{Q;4B#z*&_l!;h`sRrpXmaIDu)oOjUqXZlv2ZQg)(M z>q}I^J|4ad@2@>b4rz67W3B761!%S(Cb;%MYN%-lI>s$+_-2ozjG@s}IOk;GXb)Mu z_WmJlzL>b<1Tty7qalCKxjiV)Hb zva{8dB%D`e4u(7bM>XiLv+#V-S?n2L4f%^U=09iL>*jzWnDxEv^59K981Ejw&w_CP z`V`GV)QE1&)c~Dg;zvrtv+E%Q-Q-NSvQ&vL(sPtgQ{y=aV@=&;WJon45{GcBic>Il zIk&Xoyfd5ecv7XXTUDygZo0AH)_!~=iiqUI^F=rQ^L`+IB8hRusx&$6MXv@h@dA=( z3D5-P2Duf-MOS%$Zz)bCQqzFskTOkk%XDc=665@-#!;Ov5i9zt(|SK2*pi4{hwi$o zL^j^dTO!;T;7XYmQVi>rS$}@@KIJ=EY}XV7hzX`yvI(+bh>n~ztsiY{U>ohI$o(C4 zdoP^nLqKTs0w^XE4g&lqgV)Lm6WVp$e=Mbz`>;z=@5u-H*mt}aF^}C7_a&J}f_MSV z(=@Yq(N}#tMzQJ2HSZa_F1}}j5n3p z=T7)^2@T77LUN_P*+=LcRK6QyUJXs2U8epi@rjSnk}J0b+G4|)+L_lY!2gKu-!I1^ z@(F+p;ShOr3Rf*BQ&K9P%6OeO`Ot1TG21$3@uG0x)=Y<7xLs?lS=I&1RTZ0-s_H#y zpTzEuDg>VwwfWPkrEUZ^^HCt0dIPewsHHF}OU)hY0~FXa9G-mBTaZ|^p178x89PzJFjwMw9V%bB+AQVNt za%Zp4yNurmI8!E^MdpHrOVQ3YsR7r%zF+7icevoUVm4d(gf~ok8TI)KqMTl7=(1|_ z0~pmzK!`R%N`*Hg_g1Y_Bx8p;OmuD*;b~*A#{(IYaGcAc7e<-`^*dEui9MiLbM6-R zclDSFDTM_*F6uwVqG2Aww5deKQ3{%+a6>m@K(V;U;9VJIEhqv4rr2@c7kD(5;_uOL`#tm`@-`wAp9_72vqbWGgDkISpIRBlD;xhP&=`7wfD-7h6EsseFhTwB~?i zCS;aB?Ks0FHoRhoa7_vCe~5O8y=dIwe}A9f^jZD0nHofA@%2?wGM0tt*SL7vo1kZ7$ZOqxC=7&pn)Sb9&QMBX{9(L1QEn)^(Y2%=~z z1>Jso(cgLPk;a;s%^hxuqJ&?|9!CzVT$?+NcVGT1W^9fK@#g!wD)=~j?!ehWQdpyLxnD=o39Z-iRd4sCSK#X00r2{caBR$XBVrGE3h0i5xmXt{x!d z(w+~}17@?Kt8@hgG+$Z=RXl;`F*nTH-srl}>sOMkWPye1fZ7fYIDWV0sZw_o5O!~1 zJ4Emes58|GyV!gGI{7J~QG#|K@&poBB-dPW#XG{HzesEshu}BKtLQrSc`o&2k)54a ze_#Nj^$9CQc>y7aT~WC`ZTZ71ZtPVQe&FV(cqE1{9~qcpOamF-eA3VxvMC-*Y&J z8T4pT0oR!5PLou%!rh^EsY*i!)C ziF&Vu`c4|tQYN+#<*6@LqvFRKSu8`_g#3-b^Zt-Ze?G7m*-uYudKLS>bo+o%6=h}u zV8yS01p}dfUbSrA+*)Y{P2EIau(g#tf*v=wjnWaM2|%{@mn^d+$>9y7__C{no27@n0ws(i4EMjjs!e74+41y= z(0y->kLDYb3&Z@`Rkxa`rh7!)H^^v3Th1T9N~nI{2zsyY?8MX}n_Nqu4udfOA7!J% z6N7J6mo3$F?^Z6b;~0q&^(j_oJ>ihyIo%jQqh;y4u|lrbO}?aQnO~4SJlV$o^yz@b z#D1!`t%%hxK~2Hw^mmq!jh3j~2dj%e>QdSG6EM^M`J|xR0|htw_Uk~O+lG?8AO3e+ zP-K#z%OYc1HgCD$D3Um;7p|T7*W~5uJZKCYi1-8JvzY=aNkf$34x1A|$oPeuJIUaT z7N3!zWlqhNx4${wCchK2>syEcv5+YH0@0YBRhJz9qcDX&W+{!N$7Q)I`%2}*6qq+O;Ls4X7<$h3Q!~vvRmCsFs8O_ z21mY##r-_G|K!mEs&}8HovC%ouJph@Gq9PbgG9XG06?Mf4 z;qY$G{na&~;!w_2C2K#tYfA(^=@K2N-Ayo-96XhSa4u)wf1&S$6KIxPECwAF>XJ_4 zHJy&(gt==!ZMMW?*dKpsyQ{0G1uKW|$K-%EcO*8D6gEW*w7V;HSggEXJPy!&jDy9X z1?spl6k2GDk!nSs?(A^y1;PS_6CzdU%%#$aIvGe-uV*9sX$=lurj5KfE$=vPwej$1 zc+zT+^^DwZOKQHl*@*8n_pqis$ZWX`3cgPe`U}O!RSX4*1e+hD6M5RWxP;~g9aUl{C=@g{CpVKT- zlTOUMhR-K#MzKsqTU2-5laG71#tbhvnBHXC?2e*e;tSK%2rM)dUdP+9-CAUSq{lD$ z6((YU8Cd!G6N4u`Dj>CRO<7(eW!fMV42RiYCv|mD3fvyxMfYy~CYO7abL>}off*8k zBO#__4de#cd-(&xrbqD|sdi6T4*jmy7tWGT)>M{cRv~%sPSG8+J4E5NjU!6^te8yz zU%w*27Qmh@c``!aAg@HokUHvvNiZQ z!}evvV9f4Nj%RD<&Dus<^lwCLxGDb13%D1V82tUp(^i^lJzdT zD_C=@Q3H@Lt^3@77>I9t_zN29{9+KeA&z`)LevJ|REUvbTHeXQ+)i+Ns&u$&hKYXl zn_QyV=f;r-$uc>8?@#x7s2_z3C!OOPTb-aqSSCOgd_u?j+-TcB)I|!mg*c#`3XXM&j$qUnJl!Y*FTLPN_Wt)D=;#Q$wQJd_WqN%=yzY9|EMVNs>8PKFTE zA0H>_n$#kv4kB}wJAG`AmnvE|p^36l#B_8(pvLWyVbkkGqN#LOXT=&@Hgq_L9T_*n zQr^t(tj!lc(3N$zF2Q;vsohO|wUYlquE%WrYfa&W_{{yvwi$RC7Cvk3pjSy1E%){- z6wlJxTFH8x;%Y}IX;L)_?=JN~R(I3GwZTC2dnbbNta@dWX=743ND;E6j<=csi0Y~%GPTH5pVc1|+0MS* z*+PrzzN=c4okuOTv5PKpg!M_wot%S{0S+@cbF-cn*)@C#HJBR+)*T(oxd$5I>VTGmSl^pl?1n z4mzszlKY$Gw9qBN^y%=AQ4h`WMRD(v)OHL40pFiPE)*-2=OkbAbJ*on&D8#CG*5Q7 zK>6AIDBPS0)oefZ7?|~O55mhZ4pz3uAWy@Q+;Ftlhz$9YdA`uYSoLrMzLI@*Z*=~C z^u2f<54}|#OvEaJ>ao$8_8G`-dj=L*j10VZZ>1Zg&)*WK&l!Ql(1ZiYt#CCrq8#Kr zBML2BwD;zH|6UXKw!x$H{LZsb>P~Y~^rqYoF%q19#`dsFi7T6IA#oUpqaSv8d#D7PXf-`msotYzmC5%Y&uu3bi)7|#GRD>UhM%(O!7kcp9jRGaqY zK0j&rvB~D^@<^ND1foE1FU$xoL@~6i9K|`c9M{WRs+h4PKtUSjTLR}zO`}`}d;AZm zoE4AfjZ5d}IjaG6Yb*WwJ7&ML>-E}a-80OJv6|HW zf*87dp27YbJ9nJ7J#WVl{z4_IWL~q82dg2`Mq*Eo(La5d!q;YCc5lclyK=_z6}tjk z?>m@cPeHgx!MW~Q|4meIHglu=Y#id+c-53z=@xoQ_N3xG$!w;HJN@UQa~JTMOwq8P zC^yxIbrg6+Zn)t&C#6@TNO$X4yBBb=#L($-Tzmar6GD3TO+0r1co~ zBml;|-Y-4<8l7%w=QY+R-exW-+oX`BM*cgO-sNjo+YG!c$Kz-Mp*sXMACH^zsSFpS z-Er8y+R~rZH=cDtE$0x*0STx@_g-{>_LR0(?I1n#YF?Z2&;YG_r^R1Il!>+s-32q{ z=Hpc5uu_b&foInT=q-GVXin6>6VdC^y0!zRe8QvUKRjtxl#4w8z8!fgV5@ zJyG29KL73Z#EcKFf$}z)`9VFROy87*0vZBEAJ$bTK8kCekSBb;1ZSxjF@pP&_C?1# zw0TD!yXB(;q>6o`{Y)=IyHAP6zPq0Q`k}k!VT}@`84m0%7mapx zgVNNn`zdTPmu03OR=#^|$z1TcFC)Iwkr_!uwIjlI$Ahc-ub5Mv&6ic4{BDQ^qmiD@ z_)8`|t8Jxj{B((0YQn0b$)A_ybN1hLr~>|LFN)Eb>g*&jOdW42yd*JftFbkQUiN~ms%Zf#~i`^ynLR1p5y zYo^-DWQ5zdW=IJbbFJl+8%=b7sUbqFlZ{aSz5#8Nmyk zX`dyTyS3$l22@N1i=ay4URfMTU;eesKL=Fh2EC12yyUyDtk+^i(bbf}k>;~>(RHkj z-I<7hQuAJ*H1D$4bEVJ#4{WI58n>>0Da?g+j8NT2V^NRYveg#y8aAEiGL`m^q}|yF&iveZbi9TWv*I7xvS&Phu|(P8f`7pRk<%-m|oN%$((5xma==scAkMwvOVUpMg@}IoAhR zV=sp^DrejilMgyfHz4AXh%-F9fA4+~HB z2(R-Yu6Y#z1CgJXYFUL<7J5-1>?e%{u++!zYJ~f_`?lqS&$+|#>?D0VFhF3o#MMEm%VUq?pfDX>s08>TJ$1aj<}3F zjPdB!ue7eI#~ya$2Ythf!h0Ey%5MWt0mVLh`C9YoSuy_PucME;VOe;k+2?dV>i8;i zQR#TX@0JQL)rV7}2@lB^?r*o9z%_QJmmT=4UCoLM*OV)O6t2_)D-83;D=ub_-F45y zKW#}AN+dZ=i&3T$VyZ0TG>ogm*FP`%^zwwB!|!sBaQ7$yaPAIJ|DBD9*Jg6QV2jll z5i%W`3tV>`j}>yoH6$?yUNW5xT-s*qDJ!F&RrI`uoUpHwI(5~f`JY>t7i#<;Re=?v zu%FZ@QiM4+yV{&x{G-nflWL+CebW_|5!&|K<~1(F8(#>!N-xoG=I8o$?_P&-_Z7xA+^7N^3L5kEk8q9%t6f`qWDIM;yYK4zw}^D!zF^^kR)e&u*7{z>Nd-)Wo;N$C7!N zEX^Cz@Ot%Wgwb-@lTlex-uLq#F`vAo$)0rq-HUpJ4Bu@>r)|g4@4C`Asy*1Y;~AENC|*VdcD3zWAUEh9X;WEg@bh8!C8S!c>lN!3k9pIbqKbGy&*% zG1D9&+W<-e@orgap`z~2d=kk%u2$MUvVEZ$Ca{{ZwQRl5V+7~eyNCq=7PdIs7bv!; zzbCh{9pyW}pFJ>DE#k8}*)sFhb0W{j-zA$Clt{<{i<@{dlZ`IR5@ zD;FDT=E$mFFSsw;*zMWTmBzFsFostvhgNm8hkWwD0ja<}5)urmaSWUIa?3~X&Si)t)el%G03ID7^ z4O;B#5Z#4pNuCL(Un<4eGPs*m@IyEFo4$Rvmo8zI5A7AJfx)vvnx>;+#2_gQxl8ou zXlW|Gna7a(ULWV8mGvPDzFso?y%_xIvQ^8<#%ux)GAipFH6m*B`=(!Yb-bj@)uNRH zD=~wE8G-Z(CUiRpOqnKi4pATI@7_zJQ^Ig~lrWHLco^cb3 zC$SNAm?7d+FD(60?{Rv;OT^a4TG^SK8HoE*dFS7xe)-|;t+VF_f$;O^JJduV=L_4@c*0v_w3 zcP4|xctJ^XkFpoq#9y+9`N(8e6fF(%6 z2{#`>3&LX}60odTP!Fc<_QLyuntE(W;+$qso!Gqa?sj;&oHqh-&+t7VFPT1=J87A_ z6T_C6)bK$elwn{4&P$VaQPB}a)CAnz(fiYuhQ;9Xi=l(OUBr>G*$O5iJFw%-4PFdc< zgx*q8JJj(s9wyMn-?lrB+)t$PF`n5*;4%ZiPPyR(DNtebyP|;3{>-eEaldvC^?e(c zyV*s@;c5a^%Zm5JNOJp!S~&l})M!cb7sN6|6+rKimBHq(s5voQ%-jvS>{D4aXfq5!T-SUY?Ox2i}4r1 z?qAP!w`Kk0bw0z2+65d~E~KlzyVf4SDeCvQ0C)X}i*jAPVdx*x&b1oqDcD zD*{(^mXi8RNgqX+>9qJ#qr$!Dn|8+t?`5*>p~_`{X=}%FUFlbl3?)3B41Cr7^WtjW z9O$VO_bmI#S#S=A`f^nLsMIZR@Gp>(XFwSh8^n>0=a|(9JR{UAc62fw`oVo0&+c*l zLT>&o4YNj!-8B%40eFj`>^jtUm)zP(RLq00fkm_TIf&ELh0foRY^3FVKD;PPK*Gm5 zoe|fcrwA7>o=nx?dxPwHf+=$cW+1BpBu1H%TV(TzDfRA;Le?_en#oJ7G%NQ-aBMW$ z0QC<{UU?fLvW0#18{!Z&CbX&yJ1!++Bhs!Y!lTOmM1!fyxHch3(UGHOxvieF1f;q$ z(rcF;C>KD>BYwR4_C;a0kgWk8%rdr*$B~r|b03UJd zkGh{u-;J=_P~e{M(J(6lact%=J9Qg{ZYXsxaS+u174pNiVz2VdPmQ2P!L1vT|K%+R y1P%}U`}yxZ_%{arje&n-;NKYdHwONVfq!G*-x&Be2L6qKe`DbPHU Date: Mon, 14 Sep 2026 11:22:05 +0800 Subject: [PATCH 052/209] [Frontend] Add ComfyUI FastH3 node, fix t2va aspect ratio and dropped audio (#7456) Signed-off-by: princepride Co-authored-by: Claude Opus 5 --- apps/ComfyUI-vLLM-Omni/README.md | 66 +++- apps/ComfyUI-vLLM-Omni/__init__.py | 9 +- .../comfyui_vllm_omni/nodes.py | 134 +++++++- .../comfyui_vllm_omni/utils/api_client.py | 14 +- .../comfyui_vllm_omni/utils/format.py | 7 + .../comfyui_vllm_omni/utils/models.py | 27 ++ .../comfyui_vllm_omni/utils/types.py | 7 + .../vLLM-Omni FastH3 Text to Video.json | 322 ++++++++++++++++++ .../vLLM-Omni Video Generation.json | 2 +- apps/ComfyUI-vLLM-Omni/web/main.js | 146 +++++++- apps/ComfyUI-vLLM-Omni/web/vllm-omni-mark.svg | 5 + docs/features/comfyui.md | 19 +- .../comfyui/test_comfyui_integration.py | 88 ++++- 13 files changed, 823 insertions(+), 23 deletions(-) create mode 100644 apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni FastH3 Text to Video.json create mode 100644 apps/ComfyUI-vLLM-Omni/web/vllm-omni-mark.svg diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md index cbd53a9a0..0d607b30a 100644 --- a/apps/ComfyUI-vLLM-Omni/README.md +++ b/apps/ComfyUI-vLLM-Omni/README.md @@ -13,7 +13,6 @@ It can send model inference requests to either a locally running vLLM-Omni servi > [!TIP] > If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them. - ## Installation Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`. @@ -21,9 +20,10 @@ Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Y If you are running ComfyUI during copying, you should restart ComfyUI to load this extension. > [!TIP] -> You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info. +> You can use utility websites such as to download a subdirectory of a repo. Also checkout community discussions (e.g., ) for more info. On the device and virtual environment you run ComfyUI, launch ComfyUI with + ```bash cd ComfyUI @@ -35,6 +35,7 @@ python main.py --cpu ``` On the device and virtual environment you run vLLM-Omni, start a model service with + ```bash vllm serve The_Model_ID_to_Serve --omni --port 8000 ``` @@ -48,6 +49,7 @@ This extension offers the following nodes based on the output modalities (at **C - **Generate Image** for text-to-image and image-to-image tasks - **Generate Video** for text-to-video, first-frame/image-to-video, and reference-conditioned video +- **FastH3 Deployment** for routing text-to-video requests to a MiniMax-H3 server with FastH3 fused at startup - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks - **TTS** and **TTS Voice Clone** for TTS tasks @@ -56,6 +58,20 @@ This extension also offers example workflows (at **ComfyUI sidebar -> Templates > [!NOTE] > The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support. +Every node carries the vLLM-Omni mark in its title bar and is tinted by what it outputs, so a graph is readable at a glance: + +| Colour | Nodes | What they produce | +| --- | --- | --- | +| Blue | Generate Image, Generate Video, Multimodality Understanding, TTS, TTS Voice Clone | A generated image, video, audio, or text. These are the only nodes that reach a server. | +| Amber | AR / Diffusion / Multi-Stage Sampling Params | Sampling parameters that apply to any model | +| Purple | Qwen TTS Params, Wan Video Params, MiniMax-H3 Video Params | Parameters that only one model family accepts | +| Red | LoRA, FastH3 Deployment | Which weights the server is expected to have loaded | +| Teal | Video References | Reference media | + +Recolouring a node by hand (right click -> Colors) overrides its tint, and the choice is kept. + +**Generate Video** takes a clip length in seconds (`duration`), not a frame count. Frames stay the wire unit and are derived with the node's `fps`, so the length is always measured against the rate that is actually served; models that accept only certain frame counts still round to their own lattice server-side. Graphs saved before this widget existed stored `num_frames` in its place and are converted on load, using the fps recorded alongside it -- the browser console names every node it rewrites. + To build a simple workflow yourself, - Drag a generation node onto the canvas. @@ -123,6 +139,40 @@ You can configure per-stage sampling parameters for multi-stage models. > > Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect. +#### FastH3 text-to-video + +FastH3 is fused into MiniMax-H3 when the vLLM-Omni server starts; it is not a request-switchable LoRA. Download one adapter variant and start a non-offloaded FL2VA server. For the Dense / Data-Free profile: + +```bash +export FASTH3_DIR=/path/to/fasth3 +hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \ + dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}" + +vllm serve /path/to/MiniMax-H3 \ + --omni \ + --task-type fl2va \ + --lora-path "${FASTH3_DIR}/dense-datafree/adapter_model.safetensors" \ + --port 8000 +``` + +For the VSA / Data-Free profile, download `vsa-datafree/adapter_model.safetensors`, use that file as `--lora-path`, and add: + +```bash +--diffusion-attention-backend FASTVIDEO_VSA \ +--fastvideo-vsa-topk 64 +``` + +The VSA profile also requires a compatible `fastvideo-kernel` installation. See the [MiniMax-H3 FastH3 recipe](../../recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter) for the full serving contract, supported parallel layouts, and kernel requirements. + +Open the **vLLM-Omni FastH3 Text to Video** template, then: + +- Set the server URL and served model name on **FastH3 Deployment**. +- Connect its output to **Generate Video → fast_h3**. When connected, the deployment node's URL and model take precedence over the corresponding Generate Video widgets. +- Keep `frame`, `references`, **LoRA**, and **MiniMax-H3 Video Params** disconnected. FastH3 Preview v1 supports T2VA only, is already fused, and owns both flow shifts. +- A connected **Diffusion Sampling Params** node may set seed and other ordinary sampling options. The integration always enforces four denoising steps and 24 FPS for FastH3. + +The node records which server the workflow targets; it does not start one, nor switch adapters or attention backends on a running server. + ### TTS (e.g., Qwen TTS series) (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**) @@ -152,6 +202,8 @@ You can configure per-stage sampling parameters for multi-stage models. Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/). +Node tints and the title-bar mark are applied in `web/main.js`, keyed off each node's declared output types rather than a list of node names. A new node that returns an existing type is themed with no front-end change; a new output type needs one entry in `FAMILY_BY_OUTPUT` there. + ## Limitation and Non-Goals - Single server mode only. No automatic load balancing or failover. @@ -178,11 +230,11 @@ Whenever you find an issue or problem, please Features -- https://github.com/dougbtv/comfyui-vllm-omni/ The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API. -- https://github.com/Comfy-Org/ComfyUI/tree/master/comfy_extras ComfyUI's built-in node implementations. +- The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API. +- ComfyUI's built-in node implementations. UI/UX design references -- https://github.com/sgl-project/sglang/pull/15271 SGLang Diffusion's official ComfyUI integration for image and video generation. -- https://github.com/SXQBW/ComfyUI-Qwen-Omni A third party ComfyUI integration for Qwen Omni series. -- https://github.com/flybirdxx/ComfyUI-Qwen-TTS https://github.com/DarioFT/ComfyUI-Qwen3-TTS Third party ComfyUI integrations for Qwen TTS series. +- SGLang Diffusion's official ComfyUI integration for image and video generation. +- A third party ComfyUI integration for Qwen Omni series. +- Third party ComfyUI integrations for Qwen TTS series. diff --git a/apps/ComfyUI-vLLM-Omni/__init__.py b/apps/ComfyUI-vLLM-Omni/__init__.py index 268107899..55b003ef1 100644 --- a/apps/ComfyUI-vLLM-Omni/__init__.py +++ b/apps/ComfyUI-vLLM-Omni/__init__.py @@ -1,4 +1,8 @@ -"""Top-level package for comfyui_vllm_omni.""" # noqa: N999 # This is not a python library intended to be imported +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N999 + +"""Top-level package for comfyui_vllm_omni.""" # This is not a Python library intended to be imported. __all__ = [ "NODE_CLASS_MAPPINGS", @@ -13,6 +17,7 @@ __version__ = "0.0.1" from .comfyui_vllm_omni.nodes import ( VLLMOmniARSampling, VLLMOmniDiffusionSampling, + VLLMOmniFastH3Deployment, VLLMOmniGenerateImage, VLLMOmniGenerateVideo, VLLMOmniMiniMaxH3Params, @@ -40,6 +45,7 @@ NODE_CLASS_MAPPINGS = { "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling, "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList, "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA, + "VLLMOmniFastH3Deployment": VLLMOmniFastH3Deployment, "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams, "VLLMOmniWanParams": VLLMOmniWanParams, "VLLMOmniMiniMaxH3Params": VLLMOmniMiniMaxH3Params, @@ -59,6 +65,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VLLMOmniDiffusionSampling": "Diffusion Sampling Params", "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List", "VLLMOmniRemoteLoRA": "LoRA", + "VLLMOmniFastH3Deployment": "FastH3 Deployment", "VLLMOmniQwenTTSParams": "Qwen TTS Params", "VLLMOmniWanParams": "Wan Video Params", "VLLMOmniMiniMaxH3Params": "MiniMax-H3 Video Params", diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py index 8fd00cd9c..df91ece60 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from typing import Literal import torch @@ -10,6 +13,7 @@ from .utils.types import ( AudioFormat, AutoregressionSamplingParams, DiffusionSamplingParams, + FastH3Deployment, MiniMaxH3ModelSpecificParams, QwenTTSModelSpecificParams, VideoReferences, @@ -22,6 +26,23 @@ from .utils.validators import ( logger = get_logger(__name__) +FASTH3_INFERENCE_STEPS = 4 +FASTH3_FPS = 24 +# A FastH3 deployment may expose H3 under any --served-model-name, so the payload +# spec has to be named outright rather than recovered from that alias. +FASTH3_SPEC_MODEL = "MiniMax-H3" + + +def _resolve_fast_h3_deployment(deployment: dict) -> tuple[str, str]: + if not isinstance(deployment, dict): + raise ValueError("FastH3 deployment must be provided by a FastH3 Deployment node.") + + url = str(deployment.get("url") or "").strip().rstrip("/") + model = str(deployment.get("model") or "").strip() + if not url or not model: + raise ValueError("FastH3 deployment requires both URL and model.") + return url, model + class _VLLMOmniGenerateBase: """Base class for vLLM-Omni generation nodes with shared functionality.""" @@ -162,7 +183,20 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase): "width": ("INT", {"default": 832, "min": 1}), "height": ("INT", {"default": 480, "min": 1}), "fps": ("INT", {"default": 16, "min": 1}), - "num_frames": ("INT", {"default": 41, "min": 1}), + "duration": ( + "FLOAT", + { + "default": 4.0, + "min": 0.1, + "step": 0.1, + "round": 0.001, + "tooltip": ( + "Clip length in seconds, converted to frames with the fps above. " + "Models that only accept certain frame counts (e.g. MiniMax-H3) round to " + "their own lattice, so the served clip can be slightly longer than requested." + ), + }, + ), }, "optional": { "frame": ("IMAGE",), @@ -170,6 +204,7 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase): "sampling_params": ("SAMPLING_PARAMS",), "lora": ("REMOTE_LORA",), "model_params": ("VIDEO_PARAMS",), + "fast_h3": ("FASTH3_DEPLOYMENT",), }, } @@ -194,20 +229,66 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase): width: int, height: int, fps: int, - num_frames: int, + duration: float, negative_prompt: str | None = None, frame: torch.Tensor | None = None, references: dict | None = None, sampling_params: dict | list[dict] | None = None, model_params: dict | None = None, lora: dict | None = None, + fast_h3: dict | None = None, **kwargs, ): if kwargs: logger.info("Uncaught kwargs: %s", kwargs) logger.debug("Got sampling params: %s", sampling_params) logger.debug("Got model params: %s", model_params) - validate_model_and_sampling_params_types(model, sampling_params) + + # Which spec builds the payload. Only a FastH3 deployment separates this + # from the served name; every other path keeps them equal. + spec_model = model + + if fast_h3 is not None: + if frame is not None or references is not None: + raise ValueError("FastH3 Preview supports T2VA only; disconnect frame and references inputs.") + if lora is not None: + raise ValueError( + "FastH3 is already fused into the selected server; disconnect the request-level LoRA input." + ) + + url, model = _resolve_fast_h3_deployment(fast_h3) + logger.info("Using FastH3 deployment at %s", url) + fps = FASTH3_FPS + # The served name is whatever the operator passed to --served-model-name. + # Left alone, lookup_model_spec would miss H3 for an alias such as + # "fasth3", drop the params builder, and send a t2va request carrying no + # aspect_ratio -- which the server refuses. + spec_model = FASTH3_SPEC_MODEL + + if sampling_params is None: + sampling_params = DiffusionSamplingParams() + elif isinstance(sampling_params, list): + if len(sampling_params) != 1: + raise ValueError("FastH3 expects a single diffusion sampling params group.") + sampling_params = sampling_params[0].__class__(sampling_params[0]) + else: + sampling_params = sampling_params.__class__(sampling_params) + sampling_params["num_inference_steps"] = FASTH3_INFERENCE_STEPS + + # FastH3 owns both modality shifts. Sending the ordinary H3 values + # from a connected H3 Params node would turn a deployment choice + # into a request-level override, which the server intentionally + # rejects when it differs from the fused adapter contract. + if model_params is not None: + model_params = model_params.__class__(model_params) + model_params.pop("flow_shift", None) + model_params.pop("audio_flow_shift", None) + + # Frames stay the wire unit. Convert after the FastH3 branch above, so the + # duration is measured against the fps the server will actually apply. + num_frames = max(1, round(duration * fps)) + + validate_model_and_sampling_params_types(spec_model, sampling_params) # Currently, all video generation models are single-stage diffusion models if isinstance(sampling_params, list): @@ -227,6 +308,7 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase): client = VLLMOmniClient(url) output = await client.generate_video( model=model, + spec_model=spec_model, prompt=prompt, frame=frame, # frame present => fl2va / Wan I2V references=references, @@ -695,6 +777,52 @@ class VLLMOmniRemoteLoRA: return (lora,) +class VLLMOmniFastH3Deployment: + """Select a vLLM-Omni service that fused FastH3 at startup.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ( + "STRING", + { + "default": "http://localhost:8000/v1", + "tooltip": "URL of a vLLM-Omni server started with a FastH3 --lora-path.", + }, + ), + "model": ( + "STRING", + { + "default": "MiniMaxAI/MiniMax-H3", + "tooltip": "Model name exposed by the FastH3 deployment.", + }, + ), + } + } + + RETURN_TYPES = ("FASTH3_DEPLOYMENT",) + RETURN_NAMES = ("deployment",) + FUNCTION = "get_deployment" + CATEGORY = "vLLM-Omni" + DESCRIPTION = ( + "Selects a server with FastH3 fused at startup. The connected Generate Video node uses T2VA, " + "four inference steps, and 24 FPS without sending a request-level LoRA." + ) + + @classmethod + def VALIDATE_INPUTS(cls, url, model) -> str | Literal[True]: + try: + _resolve_fast_h3_deployment({"url": url, "model": model}) + except ValueError as exc: + return str(exc) + return True + + def get_deployment(self, url: str, model: str): + url, model = _resolve_fast_h3_deployment({"url": url, "model": model}) + return (FastH3Deployment({"url": url, "model": model}),) + + class VLLMOmniQwenTTSParams: @classmethod def INPUT_TYPES(cls): diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py index cd509244a..e051e9a09 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API, and transforms the API responses back to ComfyUI formats. @@ -271,8 +274,15 @@ class VLLMOmniClient: sampling_params: dict | None = None, model_params: dict | None = None, lora: dict | None = None, + spec_model: str | None = None, **extra_params, ) -> VideoInput: + """Post a video job and return the decoded result. + + ``spec_model`` names the model whose payload spec builds the request, for + deployments that serve a known model under a different ``model`` alias. It + never reaches the wire; defaults to ``model``. + """ if frame is not None and references is not None: raise ValueError("Provide only one of frame or references, not both.") @@ -358,12 +368,14 @@ class VLLMOmniClient: model_params = dict(model_params) model_params.pop("type", None) - spec, _ = lookup_model_spec(model) + spec, _ = lookup_model_spec(spec_model or model) params_builder = spec.get("params_builder") if spec else None if params_builder is not None: form_fields = params_builder( model_params or {}, extra_params={**extra_params, "task": video_task}, + width=width, + height=height, ) for k, v in form_fields.items(): form.add_field(k, v if isinstance(v, str) else str(v)) diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py index 42d396f06..c2aba4cff 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Image/tensor format helpers. The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv). @@ -170,6 +173,10 @@ def bytes_to_video(video_bytes: bytes) -> VideoInput: audio_stream = container.streams.audio[-1] audio_frames = [] resampler = AudioResampler(format="fltp") + # Decoding the video stream above left the container at EOF. + # Rewind before demuxing again, or this pass yields no frames + # and the audio of a T2VA result is silently dropped. + container.seek(0) for frame in container.decode(audio_stream): if not isinstance(frame, AudioFrame): continue diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py index 3a37db16d..82c7e4a51 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import json import re from typing import Any @@ -34,10 +37,28 @@ def _qwen25_payload_preprocessor(payload: dict) -> dict: return payload +MINIMAX_H3_ASPECT_RATIOS = { + "21:9": 21.0 / 9.0, + "16:9": 16.0 / 9.0, + "4:3": 4.0 / 3.0, + "1:1": 1.0, + "3:4": 3.0 / 4.0, + "9:16": 9.0 / 16.0, +} + + +def _nearest_minimaxh3_aspect_ratio(width: int, height: int) -> str: + """Pick the supported named ratio closest to the requested frame size.""" + target = float(width) / float(height) + return min(MINIMAX_H3_ASPECT_RATIOS, key=lambda name: abs(MINIMAX_H3_ASPECT_RATIOS[name] - target)) + + def _minimaxh3_params_builder( model_params: dict[str, Any], *, extra_params: dict[str, Any], + width: int | None = None, + height: int | None = None, ) -> dict[str, Any]: """Build multipart form fields for MiniMax-H3 from model_params + routed task.""" params = dict(model_params) @@ -53,6 +74,12 @@ def _minimaxh3_params_builder( if params: logger.warning("Unused MiniMax-H3 model params ignored: %s", sorted(params)) + # H3 refuses t2va without a named aspect ratio, and the Generate Video node + # carries width/height rather than a ratio. fl2va takes its ratio from the + # input image and ref2va defaults server-side, so only t2va needs this. + if merged_extra_params.get("task") == "t2va" and "aspect_ratio" not in merged_extra_params and width and height: + merged_extra_params["aspect_ratio"] = _nearest_minimaxh3_aspect_ratio(width, height) + if merged_extra_params: form_fields["extra_params"] = json.dumps(merged_extra_params, ensure_ascii=False) return form_fields diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py index e1ae43916..bfba44142 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Callable from enum import Enum, auto from typing import Any, Literal, TypeAlias @@ -31,6 +34,10 @@ class VideoReferences(dict): pass +class FastH3Deployment(dict): + """Descriptor for a server that fused FastH3 at startup.""" + + class ModelMode(Enum): IMAGE_GENERATION = auto() VIDEO_GENERATION = auto() diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni FastH3 Text to Video.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni FastH3 Text to Video.json new file mode 100644 index 000000000..5b3d54f39 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni FastH3 Text to Video.json @@ -0,0 +1,322 @@ +{ + "id": "c6db00da-2ccb-4ff8-8334-d67e1a267781", + "revision": 0, + "last_node_id": 5, + "last_link_id": 3, + "nodes": [ + { + "id": 1, + "type": "MarkdownNote", + "pos": [40, 20], + "size": [420, 250], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "FastH3 server prerequisite", + "properties": {}, + "widgets_values": [ + "FastH3 must be fused when the vLLM-Omni server starts. Use --task-type fl2va and --lora-path pointing to exactly one FastH3 adapter file. Do not use CPU or layerwise offload. For VSA / Data-Free, also install fastvideo-kernel and use --diffusion-attention-backend FASTVIDEO_VSA --fastvideo-vsa-topk 64." + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 2, + "type": "VLLMOmniFastH3Deployment", + "pos": [40, 320], + "size": [420, 160], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "localized_name": "url", + "name": "url", + "type": "STRING", + "widget": {"name": "url"}, + "link": null + }, + { + "localized_name": "model", + "name": "model", + "type": "STRING", + "widget": {"name": "model"}, + "link": null + } + ], + "outputs": [ + { + "localized_name": "deployment", + "name": "deployment", + "type": "FASTH3_DEPLOYMENT", + "links": [1] + } + ], + "properties": {"Node name for S&R": "VLLMOmniFastH3Deployment"}, + "widgets_values": [ + "http://localhost:8000/v1", + "MiniMaxAI/MiniMax-H3" + ] + }, + { + "id": 3, + "type": "VLLMOmniDiffusionSampling", + "pos": [40, 540], + "size": [420, 230], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "localized_name": "n", + "name": "n", + "type": "INT", + "widget": {"name": "n"}, + "link": null + }, + { + "localized_name": "num_inference_steps", + "name": "num_inference_steps", + "type": "INT", + "widget": {"name": "num_inference_steps"}, + "link": null + }, + { + "localized_name": "guidance_scale", + "name": "guidance_scale", + "type": "FLOAT", + "widget": {"name": "guidance_scale"}, + "link": null + }, + { + "localized_name": "true_cfg_scale", + "name": "true_cfg_scale", + "type": "FLOAT", + "widget": {"name": "true_cfg_scale"}, + "link": null + }, + { + "localized_name": "vae_use_slicing", + "name": "vae_use_slicing", + "type": "BOOLEAN", + "widget": {"name": "vae_use_slicing"}, + "link": null + }, + { + "localized_name": "vae_use_tiling", + "name": "vae_use_tiling", + "type": "BOOLEAN", + "widget": {"name": "vae_use_tiling"}, + "link": null + }, + { + "localized_name": "seed", + "name": "seed", + "type": "INT", + "widget": {"name": "seed"}, + "link": null + } + ], + "outputs": [ + { + "localized_name": "diffusion sampling params", + "name": "diffusion sampling params", + "type": "SAMPLING_PARAMS", + "links": [2] + } + ], + "properties": {"Node name for S&R": "VLLMOmniDiffusionSampling"}, + "widgets_values": [1, 4, 1, 1, false, false, 1101, "fixed"] + }, + { + "id": 4, + "type": "VLLMOmniGenerateVideo", + "pos": [560, 210], + "size": [540, 610], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "localized_name": "frame", + "name": "frame", + "shape": 7, + "type": "IMAGE", + "link": null + }, + { + "localized_name": "references", + "name": "references", + "shape": 7, + "type": "VIDEO_REFERENCES", + "link": null + }, + { + "localized_name": "sampling_params", + "name": "sampling_params", + "shape": 7, + "type": "SAMPLING_PARAMS", + "link": 2 + }, + { + "localized_name": "lora", + "name": "lora", + "shape": 7, + "type": "REMOTE_LORA", + "link": null + }, + { + "localized_name": "model_params", + "name": "model_params", + "shape": 7, + "type": "VIDEO_PARAMS", + "link": null + }, + { + "localized_name": "fast_h3", + "name": "fast_h3", + "shape": 7, + "type": "FASTH3_DEPLOYMENT", + "link": 1 + }, + { + "localized_name": "url", + "name": "url", + "type": "STRING", + "widget": {"name": "url"}, + "link": null + }, + { + "localized_name": "model", + "name": "model", + "type": "STRING", + "widget": {"name": "model"}, + "link": null + }, + { + "localized_name": "prompt", + "name": "prompt", + "type": "STRING", + "widget": {"name": "prompt"}, + "link": null + }, + { + "localized_name": "negative_prompt", + "name": "negative_prompt", + "type": "STRING", + "widget": {"name": "negative_prompt"}, + "link": null + }, + { + "localized_name": "width", + "name": "width", + "type": "INT", + "widget": {"name": "width"}, + "link": null + }, + { + "localized_name": "height", + "name": "height", + "type": "INT", + "widget": {"name": "height"}, + "link": null + }, + { + "localized_name": "fps", + "name": "fps", + "type": "INT", + "widget": {"name": "fps"}, + "link": null + }, + { + "localized_name": "duration", + "name": "duration", + "type": "FLOAT", + "widget": {"name": "duration"}, + "link": null + } + ], + "outputs": [ + { + "localized_name": "video", + "name": "video", + "type": "VIDEO", + "links": [3] + } + ], + "properties": {"Node name for S&R": "VLLMOmniGenerateVideo"}, + "widgets_values": [ + "http://localhost:8000/v1", + "MiniMaxAI/MiniMax-H3", + "A cinematic live-action singer performs naturally on a neon-lit stage, with synchronized vocals and environmental sound.", + "", + 1344, + 768, + 24, + 4.458 + ] + }, + { + "id": 5, + "type": "SaveVideo", + "pos": [1200, 280], + "size": [380, 310], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "localized_name": "video", + "name": "video", + "type": "VIDEO", + "link": 3 + }, + { + "localized_name": "filename_prefix", + "name": "filename_prefix", + "type": "STRING", + "widget": {"name": "filename_prefix"}, + "link": null + }, + { + "localized_name": "format", + "name": "format", + "type": "COMBO", + "widget": {"name": "format"}, + "link": null + }, + { + "localized_name": "codec", + "name": "codec", + "type": "COMFY_DYNAMICCOMBO_V3", + "widget": {"name": "codec"}, + "link": null + } + ], + "outputs": [ + { + "localized_name": "video", + "name": "video", + "type": "VIDEO", + "links": null + } + ], + "properties": {}, + "widgets_values": ["video/ComfyUI-FastH3", "auto", "auto"] + } + ], + "links": [ + [1, 2, 0, 4, 5, "FASTH3_DEPLOYMENT"], + [2, 3, 0, 4, 2, "SAMPLING_PARAMS"], + [3, 4, 0, 5, 0, "VIDEO"] + ], + "groups": [], + "config": {}, + "extra": { + "workflowRendererVersion": "LG", + "ds": {"scale": 0.9, "offset": [100, 100]} + }, + "version": 0.4 +} diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json index 9785845a5..b03b463d8 100644 --- a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json @@ -1 +1 @@ -{"id":"3aea1d08-7de5-4997-ba71-b2f66ad72a21","revision":0,"last_node_id":24,"last_link_id":35,"nodes":[{"id":13,"type":"SaveVideo","pos":[1315.729153459605,-137.92038948394355],"size":[374,313.7692307692308],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":32},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"codec"},"link":null}],"outputs":[{"localized_name":"video","name":"video","type":"VIDEO","links":null}],"properties":{},"widgets_values":["video/ComfyUI","auto","auto"]},{"id":10,"type":"MarkdownNote","pos":[183.74636797766794,-571.2529158046766],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both *text-to-video*, *frame-to-video*, and *reference-to-video* generation. The node will create corresponding payload depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[452.07189605801125,-569.7242536425791],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[29]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["rap.png","image"]},{"id":23,"type":"LoadImage","pos":[9.206110246057657,-7.6955291509170785],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[34]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["white-cat.jpg","image"]},{"id":22,"type":"MarkdownNote","pos":[-9.783113411416666,-148.8790300038605],"size":[297.3012007349086,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Reference to video (r2v)","properties":{},"widgets_values":["For r2v tasks, connect this Video Reference block to the main Generate Video block. Then disconnect the \"frame\" input above."],"color":"#432","bgcolor":"#000"},{"id":24,"type":"LoadAudio","pos":[317.9929742385805,-141.44821500065095],"size":[282.798828125,136],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[35]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["nothing-beats-a-jet2-holiday.mp3",null,""]},{"id":21,"type":"VLLMOmniVideoReferences","pos":[376.6129172897595,105.41557277556292],"size":[156.9380859375,126],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image_1","name":"image_1","shape":7,"type":"IMAGE","link":34},{"localized_name":"image_2","name":"image_2","shape":7,"type":"IMAGE","link":null},{"localized_name":"audio_1","name":"audio_1","shape":7,"type":"AUDIO","link":35},{"localized_name":"audio_2","name":"audio_2","shape":7,"type":"AUDIO","link":null},{"localized_name":"video_1","name":"video_1","shape":7,"type":"VIDEO","link":null},{"localized_name":"video_2","name":"video_2","shape":7,"type":"VIDEO","link":null}],"outputs":[{"localized_name":"references","name":"references","type":"VIDEO_REFERENCES","links":null}],"properties":{"Node name for S&R":"VLLMOmniVideoReferences"},"widgets_values":[]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[83.24620208113915,422.08543477542605],"size":[284.205078125,226],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[30]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[1,50,1,1,false,false,213,"randomize"]},{"id":18,"type":"VLLMOmniMiniMaxH3Params","pos":[483.0944676363692,426.7197881452892],"size":[276.56796875,82],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"audio_flow_shift","name":"audio_flow_shift","type":"FLOAT","widget":{"name":"audio_flow_shift"},"link":null},{"localized_name":"flow_shift","name":"flow_shift","type":"FLOAT","widget":{"name":"flow_shift"},"link":null}],"outputs":[{"localized_name":"MiniMax-H3 video params","name":"MiniMax-H3 video params","type":"VIDEO_PARAMS","links":[33]}],"properties":{"Node name for S&R":"VLLMOmniMiniMaxH3Params"},"widgets_values":[3,12]},{"id":16,"type":"VLLMOmniGenerateVideo","pos":[826.1428043381621,-148.53719710697413],"size":[425.5948726785682,551.767376721426],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"frame","name":"frame","shape":7,"type":"IMAGE","link":29},{"localized_name":"references","name":"references","shape":7,"type":"VIDEO_REFERENCES","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":30},{"localized_name":"lora","name":"lora","shape":7,"type":"REMOTE_LORA","link":null},{"localized_name":"model_params","name":"model_params","shape":7,"type":"VIDEO_PARAMS","link":33},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null},{"localized_name":"fps","name":"fps","type":"INT","widget":{"name":"fps"},"link":null},{"localized_name":"num_frames","name":"num_frames","type":"INT","widget":{"name":"num_frames"},"link":null}],"outputs":[{"localized_name":"video","name":"video","type":"VIDEO","links":[32]}],"properties":{"Node name for S&R":"VLLMOmniGenerateVideo"},"widgets_values":["http://localhost:8000/v1","/workspace/models/hub/models--MiniMaxAI--MiniMax-H3/snapshots/428986c4aa5280855b9ac26eff8d1a46b5999cd0/FL2VA","A scene of urban fantasy art. A dynamic graffiti art character. A boy made of spray paint comes to life on a concrete wall. He sings an English rap song at high speed while striking a classic, energetic rapper pose. The scene is set under an urban railway bridge at night. The light comes from a single street lamp, creating a cinematic atmosphere full of high energy and amazing detail. The audio of the video consists entirely of the rap, with no other dialogue or noise.","",832,480,24,120]}],"links":[[29,4,0,16,0,"IMAGE"],[30,8,0,16,2,"SAMPLING_PARAMS"],[32,16,0,13,0,"VIDEO"],[33,18,0,16,4,"VIDEO_PARAMS"],[34,23,0,21,0,"IMAGE"],[35,24,0,21,2,"AUDIO"]],"groups":[{"id":1,"title":"Image-to-video input (first-frame-to-video)","bounding":[168.92390079062986,-653.2124402506535,560.8246471437027,407.4246532472182],"color":"#3f789e","flags":{}},{"id":2,"title":"Reference-to-video input","bounding":[-66.75175444110151,-225.24978633666066,805.4597228678848,554.7384400819777],"color":"#3f789e","flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":0.8,"offset":[1240.4751945653188,770.6151743053379]}},"version":0.4} +{"id":"3aea1d08-7de5-4997-ba71-b2f66ad72a21","revision":0,"last_node_id":24,"last_link_id":35,"nodes":[{"id":13,"type":"SaveVideo","pos":[1315.729153459605,-137.92038948394355],"size":[374,313.7692307692308],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":32},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"codec"},"link":null}],"outputs":[{"localized_name":"video","name":"video","type":"VIDEO","links":null}],"properties":{},"widgets_values":["video/ComfyUI","auto","auto"]},{"id":10,"type":"MarkdownNote","pos":[183.74636797766794,-571.2529158046766],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both *text-to-video*, *frame-to-video*, and *reference-to-video* generation. The node will create corresponding payload depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[452.07189605801125,-569.7242536425791],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[29]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["rap.png","image"]},{"id":23,"type":"LoadImage","pos":[9.206110246057657,-7.6955291509170785],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[34]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["white-cat.jpg","image"]},{"id":22,"type":"MarkdownNote","pos":[-9.783113411416666,-148.8790300038605],"size":[297.3012007349086,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Reference to video (r2v)","properties":{},"widgets_values":["For r2v tasks, connect this Video Reference block to the main Generate Video block. Then disconnect the \"frame\" input above."],"color":"#432","bgcolor":"#000"},{"id":24,"type":"LoadAudio","pos":[317.9929742385805,-141.44821500065095],"size":[282.798828125,136],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[35]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["nothing-beats-a-jet2-holiday.mp3",null,""]},{"id":21,"type":"VLLMOmniVideoReferences","pos":[376.6129172897595,105.41557277556292],"size":[156.9380859375,126],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image_1","name":"image_1","shape":7,"type":"IMAGE","link":34},{"localized_name":"image_2","name":"image_2","shape":7,"type":"IMAGE","link":null},{"localized_name":"audio_1","name":"audio_1","shape":7,"type":"AUDIO","link":35},{"localized_name":"audio_2","name":"audio_2","shape":7,"type":"AUDIO","link":null},{"localized_name":"video_1","name":"video_1","shape":7,"type":"VIDEO","link":null},{"localized_name":"video_2","name":"video_2","shape":7,"type":"VIDEO","link":null}],"outputs":[{"localized_name":"references","name":"references","type":"VIDEO_REFERENCES","links":null}],"properties":{"Node name for S&R":"VLLMOmniVideoReferences"},"widgets_values":[]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[83.24620208113915,422.08543477542605],"size":[284.205078125,226],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[30]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[1,50,1,1,false,false,213,"randomize"]},{"id":18,"type":"VLLMOmniMiniMaxH3Params","pos":[483.0944676363692,426.7197881452892],"size":[276.56796875,82],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"audio_flow_shift","name":"audio_flow_shift","type":"FLOAT","widget":{"name":"audio_flow_shift"},"link":null},{"localized_name":"flow_shift","name":"flow_shift","type":"FLOAT","widget":{"name":"flow_shift"},"link":null}],"outputs":[{"localized_name":"MiniMax-H3 video params","name":"MiniMax-H3 video params","type":"VIDEO_PARAMS","links":[33]}],"properties":{"Node name for S&R":"VLLMOmniMiniMaxH3Params"},"widgets_values":[3,12]},{"id":16,"type":"VLLMOmniGenerateVideo","pos":[826.1428043381621,-148.53719710697413],"size":[425.5948726785682,551.767376721426],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"frame","name":"frame","shape":7,"type":"IMAGE","link":29},{"localized_name":"references","name":"references","shape":7,"type":"VIDEO_REFERENCES","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":30},{"localized_name":"lora","name":"lora","shape":7,"type":"REMOTE_LORA","link":null},{"localized_name":"model_params","name":"model_params","shape":7,"type":"VIDEO_PARAMS","link":33},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null},{"localized_name":"fps","name":"fps","type":"INT","widget":{"name":"fps"},"link":null},{"localized_name":"duration","name":"duration","type":"FLOAT","widget":{"name":"duration"},"link":null}],"outputs":[{"localized_name":"video","name":"video","type":"VIDEO","links":[32]}],"properties":{"Node name for S&R":"VLLMOmniGenerateVideo"},"widgets_values":["http://localhost:8000/v1","/workspace/models/hub/models--MiniMaxAI--MiniMax-H3/snapshots/428986c4aa5280855b9ac26eff8d1a46b5999cd0/FL2VA","A scene of urban fantasy art. A dynamic graffiti art character. A boy made of spray paint comes to life on a concrete wall. He sings an English rap song at high speed while striking a classic, energetic rapper pose. The scene is set under an urban railway bridge at night. The light comes from a single street lamp, creating a cinematic atmosphere full of high energy and amazing detail. The audio of the video consists entirely of the rap, with no other dialogue or noise.","",832,480,24,5.0]}],"links":[[29,4,0,16,0,"IMAGE"],[30,8,0,16,2,"SAMPLING_PARAMS"],[32,16,0,13,0,"VIDEO"],[33,18,0,16,4,"VIDEO_PARAMS"],[34,23,0,21,0,"IMAGE"],[35,24,0,21,2,"AUDIO"]],"groups":[{"id":1,"title":"Image-to-video input (first-frame-to-video)","bounding":[168.92390079062986,-653.2124402506535,560.8246471437027,407.4246532472182],"color":"#3f789e","flags":{}},{"id":2,"title":"Reference-to-video input","bounding":[-66.75175444110151,-225.24978633666066,805.4597228678848,554.7384400819777],"color":"#3f789e","flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":0.8,"offset":[1240.4751945653188,770.6151743053379]}},"version":0.4} diff --git a/apps/ComfyUI-vLLM-Omni/web/main.js b/apps/ComfyUI-vLLM-Omni/web/main.js index da99b2539..9aba6260e 100644 --- a/apps/ComfyUI-vLLM-Omni/web/main.js +++ b/apps/ComfyUI-vLLM-Omni/web/main.js @@ -1,19 +1,153 @@ /** - * @file This file is intended to add dynamic fields to vLLM-Omni nodes - * based on widget (in-node form fields) and input (connection link) values and changes. - * However, this functionality is currently disabled/commented out. - * Because it introduces too much complexity, - * and it may even conflict with the current backend (Python) validation for unknown reasons (pending ComfyUI upstream fixes). + * @file Front-end extension for the vLLM-Omni nodes. + * + * Branding. Every vLLM-Omni node draws the vLLM-Omni "v" mark where LiteGraph + * would otherwise draw its default title dot, and is tinted by the kind of value + * it produces: blue nodes reach a server, amber and purple ones only assemble + * parameters, red ones pin what the server has loaded, teal ones carry media. + * A graph is then readable at a glance, and a node that turns out to be the + * wrong colour for where it sits is usually wired wrong. + * + * The tint is keyed off each node's declared output types rather than a list of + * node names, so a node added later is themed as soon as it returns one of the + * existing types. + * + * It also migrates Generate Video graphs saved before `duration` replaced the + * `num_frames` widget. See migrateNumFramesToDuration below. + * + * This file also used to host dynamic fields on vLLM-Omni nodes, driven by + * widget (in-node form fields) and input (connection link) values and changes. + * That is deliberately absent: it introduced too much complexity, and it even + * conflicted with the current backend (Python) validation for unknown reasons + * (pending ComfyUI upstream fixes). */ import { app } from "../../scripts/app.js"; + +const MARK_URL = new URL("./vllm-omni-mark.svg", import.meta.url).href; +const MARK_ASPECT = 157 / 165; + +/** + * Title bar and body tints, one per node family. + * + * The hues are the logo's: the amber and blue of the "v" mark, and three stops + * of the gradient ring in the "O" (#9128A3, #CC5260, and the teal neighbouring + * its blue). All of them are dark because LiteGraph paints the title bar darker + * than the body and expects both to sit under light text; the ComfyUI light + * theme lightens custom node colours on its own. + */ +const FAMILIES = { + generate: { color: "#152A43", bgcolor: "#1F4269" }, + sampling: { color: "#3A2C12", bgcolor: "#5C461C" }, + model: { color: "#2E1233", bgcolor: "#4A1C52" }, + deployment: { color: "#3A1A1F", bgcolor: "#5C2931" }, + media: { color: "#0F2E2A", bgcolor: "#194A44" }, +}; + +/** Output type -> family. Types absent here fall through to "generate". */ +const FAMILY_BY_OUTPUT = { + SAMPLING_PARAMS: "sampling", + TTS_PARAMS: "model", + VIDEO_PARAMS: "model", + REMOTE_LORA: "deployment", + FASTH3_DEPLOYMENT: "deployment", + VIDEO_REFERENCES: "media", +}; + +function familyOf(nodeData) { + for (const type of nodeData.output ?? []) { + const family = FAMILY_BY_OUTPUT[type]; + if (family) { + return family; + } + } + // What is left returns a generated artifact (IMAGE, VIDEO, AUDIO, STRING), + // which is only ever produced by a node that talks to a vLLM-Omni server. + return "generate"; +} + +/** + * Rewrite a pre-rename Generate Video graph in place. + * + * `duration` (seconds) took over the slot `num_frames` used to occupy, and + * ComfyUI restores widget values positionally, so a graph saved before the + * rename reads its frame count as a duration: the workflow this repository + * shipped stored 120 frames at 24 fps, which would come back as a 120-second + * request. Saved files still name the old widget under `inputs`, and that is + * what makes the rewrite both detectable and safe to apply exactly once -- once + * rewritten the node serializes `duration`, so a later load finds nothing to do. + */ +function migrateNumFramesToDuration(node, info) { + const savedBeforeRename = (info?.inputs ?? []).some((input) => input?.widget?.name === "num_frames"); + if (!savedBeforeRename) { + return + } + const duration = node.widgets?.find((widget) => widget.name === "duration"); + const fps = node.widgets?.find((widget) => widget.name === "fps"); + if (!duration || !fps) { + return + } + // configure() has already placed the stale frame count in the duration widget. + const frames = Number(duration.value); + const rate = Number(fps.value); + if (!Number.isFinite(frames) || !Number.isFinite(rate) || rate <= 0) { + return + } + const seconds = Math.max(duration.options?.min ?? 0.1, Math.round((frames / rate) * 1000) / 1000); + console.info( + `vLLM-Omni: migrated Generate Video "${node.title}" from ${frames} frames to ${seconds}s at ${rate} fps.`, + ); + duration.value = seconds; +} + +const mark = new Image(); +let markReady = false; +mark.addEventListener("load", () => { + markReady = true; + // The image resolves after the first frames are already on screen. + app.graph?.setDirtyCanvas?.(true, true); +}); +mark.src = MARK_URL; + app.registerExtension({ name: "vllm.vllm_omni", async beforeRegisterNodeDef(nodeType, nodeData, app) { if (!nodeData.name.startsWith("VLLMOmni")) { return } - // Stub frontend plugin for now + + const { color, bgcolor } = FAMILIES[familyOf(nodeData)]; + // On the prototype, not the instance: a colour the user picks by hand is + // written to the instance and shadows this, so manual overrides survive. + nodeType.prototype.color = color; + nodeType.prototype.bgcolor = bgcolor; + + if (nodeData.name === "VLLMOmniGenerateVideo") { + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function (info) { + onConfigure?.apply(this, arguments); + migrateNumFramesToDuration(this, info); + }; + } + + // Defining this replaces the default dot outright — LiteGraph draws its + // own square or circle only for nodes that leave onDrawTitleBox unset. + nodeType.prototype.onDrawTitleBox = function (ctx, titleHeight, _size, scale) { + // Zoomed out far enough that LiteGraph itself drops to flat shapes, + // the mark is a few muddy pixels. Leave the title bar clean instead. + if (!markReady || (scale !== undefined && scale < 0.5)) { + return + } + const height = titleHeight * 0.62; + const width = height * MARK_ASPECT; + ctx.drawImage( + mark, + titleHeight * 0.5 - width * 0.5, + titleHeight * -0.5 - height * 0.5, + width, + height, + ); + }; }, async setup() { console.info("vLLM-Omni Setup complete!") diff --git a/apps/ComfyUI-vLLM-Omni/web/vllm-omni-mark.svg b/apps/ComfyUI-vLLM-Omni/web/vllm-omni-mark.svg new file mode 100644 index 000000000..5ce519c1b --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/web/vllm-omni-mark.svg @@ -0,0 +1,5 @@ + + vLLM-Omni + + + diff --git a/docs/features/comfyui.md b/docs/features/comfyui.md index ff766cbb3..b4d6fa101 100644 --- a/docs/features/comfyui.md +++ b/docs/features/comfyui.md @@ -13,7 +13,6 @@ It can send model inference requests to either a locally running vLLM-Omni servi !!! tip If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them. - ## Installation Copy the `apps/ComfyUI-vLLM-Omni` folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`. @@ -21,9 +20,10 @@ Copy the `apps/ComfyUI-vLLM-Omni` folder to the `custom_nodes` subfolder of your If you are running ComfyUI during copying, you should restart ComfyUI to load this extension. !!! tip - You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info. + You can use utility websites such as to download a subdirectory of a repo. Also checkout community discussions (e.g., ) for more info. On the device and virtual environment you run ComfyUI, launch ComfyUI with + ```bash cd ComfyUI @@ -35,6 +35,7 @@ python main.py --cpu ``` On the device and virtual environment you run vLLM-Omni, start a model service with + ```bash vllm serve The_Model_ID_to_Serve --omni --port 8000 ``` @@ -56,6 +57,20 @@ This extension also offers example workflows (at **ComfyUI sidebar -> Templates !!! info The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support. +Every node carries the vLLM-Omni mark in its title bar and is tinted by what it outputs, so a graph is readable at a glance: + +| Colour | Nodes | What they produce | +| --- | --- | --- | +| Blue | Generate Image, Generate Video, Multimodality Understanding, TTS, TTS Voice Clone | A generated image, video, audio, or text. These are the only nodes that reach a server. | +| Amber | AR / Diffusion / Multi-Stage Sampling Params | Sampling parameters that apply to any model | +| Purple | Qwen TTS Params, Wan Video Params, MiniMax-H3 Video Params | Parameters that only one model family accepts | +| Red | LoRA, FastH3 Deployment | Which weights the server is expected to have loaded | +| Teal | Video References | Reference media | + +Recolouring a node by hand (right click -> Colors) overrides its tint, and the choice is kept. + +**Generate Video** takes a clip length in seconds (`duration`), not a frame count. Frames stay the wire unit and are derived with the node's `fps`, so the length is always measured against the rate that is actually served; models that accept only certain frame counts still round to their own lattice server-side. Graphs saved before this widget existed stored `num_frames` in its place and are converted on load, using the fps recorded alongside it -- the browser console names every node it rewrites. + To build a simple workflow yourself, - Drag a generation node onto the canvas. diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index 61cf76dcd..0b0f99852 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -24,6 +24,7 @@ import requests import torch from comfy_api.input import AudioInput, VideoInput from comfyui_vllm_omni.nodes import ( + VLLMOmniFastH3Deployment, VLLMOmniGenerateImage, VLLMOmniGenerateVideo, VLLMOmniTTS, @@ -78,6 +79,7 @@ class SamplingKind(str, Enum): TTS_DIFFUSION_SINGLE = "tts_diffusion_single" VIDEO_NONE = "video_none" VIDEO_DIFFUSION_SINGLE = "video_diffusion_single" + VIDEO_FASTH3 = "video_fasth3" VIDEO_REF2VA_IMAGE_AUDIO = "video_ref2va_image_audio" VIDEO_REF2VA_MULTI_VIDEO = "video_ref2va_multi_video" @@ -89,6 +91,9 @@ VIDEO_WIDTH = 32 VIDEO_HEIGHT = 32 VIDEO_FPS = 8 VIDEO_NUM_FRAMES = 5 +VIDEO_DURATION = VIDEO_NUM_FRAMES / VIDEO_FPS # 0.625 s at the generic 8 fps +# FastH3 pins 24 fps, so the same duration is a different frame count there. +VIDEO_FASTH3_NUM_FRAMES = round(VIDEO_DURATION * 24) DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams( { "n": 2, @@ -394,6 +399,24 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam LORA_PARAMS, ) _assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS) + elif sampling_case.kind is SamplingKind.VIDEO_FASTH3: + assert len(received_sampling_params_list) == 1 + received = received_sampling_params_list[0] + _assert_sampling_param_values( + received, + { + "width": VIDEO_WIDTH, + "height": VIDEO_HEIGHT, + "num_frames": VIDEO_FASTH3_NUM_FRAMES, + "fps": 24, + "num_inference_steps": 4, + }, + ) + assert received.lora_request is None + # t2va is refused without an explicit ratio; 32x32 derives "1:1". + _assert_model_param_values(received, {"task": "t2va", "aspect_ratio": "1:1"}) + assert "flow_shift" not in received.extra_args + assert "audio_flow_shift" not in received.extra_args elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO: assert len(received_sampling_params_list) == 1 assert isinstance(prompt, dict) @@ -861,7 +884,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b "width": VIDEO_WIDTH, "height": VIDEO_HEIGHT, "fps": VIDEO_FPS, - "num_frames": VIDEO_NUM_FRAMES, + "duration": VIDEO_DURATION, "model_params": VIDEO_MODEL_PARAMS, } if image_input: @@ -878,6 +901,67 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b assert isinstance(result[0], VideoInput) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_case,deployment_model", + [ + pytest.param( + ServerCase( + served_model="MiniMaxAI/MiniMax-H3", + stage_list=["diffusion"], + stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}], + outputs=[_build_diffusion_video_output()], + ), + "MiniMaxAI/MiniMax-H3", + id="canonical_model_name", + ), + pytest.param( + ServerCase( + served_model="fasth3", + stage_list=["diffusion"], + stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}], + outputs=[_build_diffusion_video_output()], + ), + "fasth3", + # A --served-model-name alias lookup_model_spec cannot resolve to H3. The + # request must still be built by the H3 params builder, or it goes out + # without the aspect_ratio (and task) a t2va request is refused without. + id="served_model_alias", + ), + ], + indirect=["server_case"], +) +@pytest.mark.parametrize( + "sampling_case", + [SamplingCase(kind=SamplingKind.VIDEO_FASTH3, sampling_params=DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS)], + indirect=True, +) +async def test_fast_h3_deployment_node(api_server: str, sampling_case: SamplingCase, deployment_model: str): + deployment_node = VLLMOmniFastH3Deployment() + (deployment,) = deployment_node.get_deployment( + url=api_server, + model=deployment_model, + ) + + result = await VLLMOmniGenerateVideo().generate( + url="http://ignored.invalid/v1", + model="ignored-model", + prompt="A singer performs on a neon-lit stage.", + negative_prompt="", + width=VIDEO_WIDTH, + height=VIDEO_HEIGHT, + fps=VIDEO_FPS, + duration=VIDEO_DURATION, + sampling_params=sampling_case.sampling_params, + model_params=H3_MODEL_PARAMS, + fast_h3=deployment, + ) + + assert isinstance(result, tuple) + assert len(result) == 1 + assert isinstance(result[0], VideoInput) + + @pytest.mark.asyncio @pytest.mark.parametrize( "server_case,sampling_case,ref_mode", @@ -944,7 +1028,7 @@ async def test_video_generation_node_minimax_h3_ref2va( width=VIDEO_WIDTH, height=VIDEO_HEIGHT, fps=VIDEO_FPS, - num_frames=VIDEO_NUM_FRAMES, + duration=VIDEO_DURATION, references=references, model_params=H3_MODEL_PARAMS, ) From 58adeec05f151e18542323cb4644b009b83cfef3 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Sun, 13 Sep 2026 21:02:28 -0700 Subject: [PATCH 053/209] [CI/Build] Add experimental AMD MI300 nightly lane (#6978) Signed-off-by: andyluo7 --- .buildkite/amd/scripts/bootstrap-amd-omni.sh | 15 +++-- .buildkite/amd/scripts/select_test_suites.py | 33 ++++++---- .buildkite/amd/test-amd-nightly.yml | 66 ++++++++++++++++++++ .buildkite/common/scripts/skip_ci.py | 1 + docs/contributing/ci/ci_settings.md | 14 +++-- docs/contributing/ci/test_system_overview.md | 2 + tests/buildkite/test_amd_bootstrap.py | 2 +- tests/buildkite/test_amd_pipeline.py | 17 ++++- tests/buildkite/test_amd_suite_selection.py | 19 ++++-- tests/buildkite/test_skip_ci.py | 8 +++ 10 files changed, 147 insertions(+), 30 deletions(-) create mode 100644 .buildkite/amd/test-amd-nightly.yml diff --git a/.buildkite/amd/scripts/bootstrap-amd-omni.sh b/.buildkite/amd/scripts/bootstrap-amd-omni.sh index 9eff6b6c7..60732d801 100755 --- a/.buildkite/amd/scripts/bootstrap-amd-omni.sh +++ b/.buildkite/amd/scripts/bootstrap-amd-omni.sh @@ -60,7 +60,8 @@ resolve_test_specs() { if ! selected_specs=$(python3 "$AMD_SUITE_SELECTOR" \ --branch "$BUILDKITE_BRANCH" \ --labels "$PR_LABELS" \ - --debug-test-yaml "${DEBUG_TEST_YAML:-}"); then + --debug-test-yaml "${DEBUG_TEST_YAML:-}" \ + --nightly "$NIGHTLY"); then exit 1 fi @@ -80,6 +81,12 @@ filter_test_specs_by_skip_ci() { case "$suite_spec" in READY_TESTS:*) level="l2" ;; MERGE_TESTS:*) level="l3" ;; + NIGHTLY_TESTS:*) + # Explicit PR and scheduled nightly selections must survive + # L2/L3 diff gating, including docs-only main commits. + runnable_specs+=("$suite_spec") + continue + ;; *) echo "ERROR: unknown AMD test suite spec '$suite_spec'" >&2 exit 1 @@ -101,7 +108,7 @@ filter_test_specs_by_skip_ci() { ":memo: CI skipped — docs or pytest skip-mark changes only" \ --style "info" 2>/dev/null || true fi - echo "No AMD L2/L3 suites remain after skip-ci filtering." + echo "No AMD suites remain after skip-ci filtering." exit 0 fi } @@ -219,8 +226,8 @@ if [[ $BUILDKITE_BRANCH == "main" ]]; then file_diff=$(get_diff_main) fi -# Resolve PR tier labels before skip-ci so ready and merge suites can be -# filtered independently for CI-YAML-only changes. +# Resolve PR tier labels before skip-ci so ready/merge suites can be filtered +# independently while an explicitly selected nightly suite passes through. resolve_test_specs filter_test_specs_by_skip_ci diff --git a/.buildkite/amd/scripts/select_test_suites.py b/.buildkite/amd/scripts/select_test_suites.py index b7c72f6e4..daed05bd3 100644 --- a/.buildkite/amd/scripts/select_test_suites.py +++ b/.buildkite/amd/scripts/select_test_suites.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Select AMD L2/L3 suites after Buildkite has created an AMD build. +"""Select AMD ready/merge/nightly suites after Buildkite creates a build. This repository-side selector does not decide which GitHub label events start the external ``vllm-omni-amd-ci`` pipeline. Its Buildkite PR build condition @@ -17,6 +17,7 @@ from collections.abc import Iterable SUITE_SPECS = { "ready": "READY_TESTS:test-amd-ready.yml", "merge": "MERGE_TESTS:test-amd-merge.yml", + "nightly": "NIGHTLY_TESTS:test-amd-nightly.yml", } @@ -28,7 +29,7 @@ def _parse_debug_suites(value: str) -> tuple[str, ...]: seen: set[str] = set() for suite in suites: if suite not in SUITE_SPECS: - raise ValueError(f"DEBUG_TEST_YAML entries must be 'merge' or 'ready', got {suite!r}") + raise ValueError(f"DEBUG_TEST_YAML entries must be 'merge', 'nightly', or 'ready', got {suite!r}") if suite in seen: raise ValueError(f"duplicate DEBUG_TEST_YAML suite {suite!r}") seen.add(suite) @@ -40,18 +41,27 @@ def select_amd_test_suites( branch: str, labels: Iterable[str], debug_test_yaml: str = "", + nightly: bool = False, ) -> tuple[str, ...]: """Return ordered suite names while preserving legacy unlabeled PRs.""" if debug_test_yaml: return _parse_debug_suites(debug_test_yaml) if branch == "main": - return ("merge",) + return ("nightly",) if nightly else ("merge",) label_set = {label.strip() for label in labels if label.strip()} - selected = tuple(suite for label, suite in (("ready", "ready"), ("merge-test", "merge")) if label in label_set) - # AMD historically uploaded L2 for every PR build. Keep that behavior for - # triggers such as amd-test while adding tier selection for ready/merge-test. + selected = tuple( + suite + for label, suite in ( + ("ready", "ready"), + ("merge-test", "merge"), + ("nightly-test", "nightly"), + ) + if label in label_set + ) + # AMD historically uploaded the ready suite for every admitted PR build. + # Preserve that fallback when no explicit tier label is present. return selected or ("ready",) @@ -60,6 +70,7 @@ def main() -> int: parser.add_argument("--branch", required=True) parser.add_argument("--labels", default="") parser.add_argument("--debug-test-yaml", default="") + parser.add_argument("--nightly", type=int, choices=(0, 1), default=0) args = parser.parse_args() labels = tuple(args.labels.splitlines()) @@ -68,19 +79,15 @@ def main() -> int: branch=args.branch, labels=labels, debug_test_yaml=args.debug_test_yaml, + nightly=bool(args.nightly), ) except ValueError as exc: parser.error(str(exc)) label_set = {label.strip() for label in labels if label.strip()} - if not args.debug_test_yaml and args.branch != "main" and not ({"ready", "merge-test"} & label_set): + if args.branch != "main" and not args.debug_test_yaml and not ({"ready", "merge-test", "nightly-test"} & label_set): print( - "No AMD L2/L3 tier label found; preserving the legacy ready-suite fallback.", - file=sys.stderr, - ) - if "nightly-test" in label_set: - print( - "AMD nightly-test has no suite yet and does not affect L2/L3 selection.", + "No AMD ready/merge/nightly label found; preserving the legacy ready-suite fallback.", file=sys.stderr, ) diff --git a/.buildkite/amd/test-amd-nightly.yml b/.buildkite/amd/test-amd-nightly.yml new file mode 100644 index 000000000..0da365d0c --- /dev/null +++ b/.buildkite/amd/test-amd-nightly.yml @@ -0,0 +1,66 @@ +# AMD MI300 `full_model` nightly coverage during burn-in. +# All jobs are non-blocking until the lane has stable runtime evidence. +# Historical build #11646 was red in Function Expansion and Accuracy (one +# async_chunk similarity mismatch; Seed-TTS WER 0.376480 > 0.35). + +env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + HF_HUB_DOWNLOAD_TIMEOUT: 300 + HF_HUB_ETAG_TIMEOUT: 60 + +steps: + # MI325 is the repository's legacy ROCm pytest marker. The AMD template maps + # agent_pool to the current MI300 Buildkite queues independently of that mark. + - group: ":full_moon: Qwen3-Omni MI300 Nightly" + steps: + - label: "Qwen3-Omni Function Expansion" + agent_pool: mi300_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 180 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "full_model and rocm and MI325 and cards_2" --run-level "full_model" + + - label: "Qwen3-Omni Accuracy" + agent_pool: mi300_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 180 + artifact_paths: + - tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json + commands: + - export SEED_TTS_WER_EVAL=1 + # PyTorch ROCm uses the torch.cuda device namespace. + - export SEED_TTS_EVAL_DEVICE=cuda:1 + - pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model and rocm and MI325 and cards_2" --run-level "full_model" + # Native tests run from the image copy under /app/vllm-omni, while + # Buildkite's artifact hook searches the checkout workspace. + - | + if [[ "$$PWD" != "$${BUILDKITE_BUILD_CHECKOUT_PATH:?}" ]]; then + artifact_dir="$${BUILDKITE_BUILD_CHECKOUT_PATH}/tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc" + mkdir -p "$$artifact_dir" + cp -v tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json "$$artifact_dir/" + fi + + - label: "Qwen3-Omni Documentation Examples" + agent_pool: mi300_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 120 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "full_model and rocm and MI325 and cards_2" --run-level "full_model" + + - label: "Qwen3-Omni AITER-on Smoke" + agent_pool: mi300_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 90 + commands: + - export VLLM_ROCM_USE_AITER=1 + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py::test_text_to_text_audio_001 -m "full_model and rocm and MI325 and cards_2" --run-level "full_model" diff --git a/.buildkite/common/scripts/skip_ci.py b/.buildkite/common/scripts/skip_ci.py index 8d5fdd5e0..3b4168c01 100644 --- a/.buildkite/common/scripts/skip_ci.py +++ b/.buildkite/common/scripts/skip_ci.py @@ -44,6 +44,7 @@ L3_YAML_FILES: dict[str, str] = { L45_YAML_FILES: dict[str, str] = { ".buildkite/cuda/test-nightly.yml": "cuda", ".buildkite/cuda/test-weekly.yml": "cuda", + ".buildkite/amd/test-amd-nightly.yml": "amd", ".buildkite/npu/test-npu-nightly.yml": "npu", } diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index 9931a418d..4032891c0 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -35,6 +35,7 @@ Canonical layout (prefer these paths for new changes): ├── amd/ │ ├── test-amd-ready.yml # L2 job definitions (template input) │ ├── test-amd-merge.yml # L3 job definitions +│ ├── test-amd-nightly.yml # Nightly full-model jobs (experimental) │ ├── test-template-amd-omni.j2 # Renders final pipeline.yaml │ └── scripts/ │ ├── bootstrap-amd-omni.sh # Entry: skip-ci → Jinja → upload @@ -64,7 +65,7 @@ Canonical layout (prefer these paths for new changes): | -------- | ---------------- | -------------- | ---------------- | -------------------- | | **CUDA** | `cuda/pipeline.yml` | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml` | `upload_pipeline.py --upload` (expands uploader-only keys) | omit `mirror_hardwares` (from `-m`) / preset string + `MIRROR_HW` | | **NPU** | `npu/pipeline-npu.yml` | `test-npu-ready.yml`, `test-npu-nightly.yml` | `upload_pipeline.py --upload` | `mirror_hardwares: a2b3_npu_1` / `a2b3_npu_4` / `a3_npu_2` | -| **AMD** | `amd/scripts/bootstrap-amd-omni.sh` | `test-amd-ready.yml`, `test-amd-merge.yml` | Jinja (`test-template-amd-omni.j2`) → `pipeline upload` | `agent_pool` + `mirror_hardwares: [amdproduction]` (array, template filter) | +| **AMD** | `amd/scripts/bootstrap-amd-omni.sh` | `test-amd-ready.yml`, `test-amd-merge.yml`, `test-amd-nightly.yml` | Jinja (`test-template-amd-omni.j2`) → `pipeline upload` | `agent_pool` + `mirror_hardwares: [amdproduction]` (array, template filter) | | **Intel** | `intel/scripts/bootstrap-intel-omni.sh` | `intel/pipeline-intel.yml` (steps inline) | Direct `pipeline upload` | Inline `agents.queue` on each step | ## Platform configuration style @@ -171,9 +172,9 @@ Canonical layout (prefer these paths for new changes): **Bootstrap:** [`amd/scripts/bootstrap-amd-omni.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/amd/scripts/bootstrap-amd-omni.sh)—skip-ci, diff filtering, Jinja render, then `buildkite-agent pipeline upload`. - **Test YAML (data):** `amd/test-amd-ready.yml` (L2 / `ready`), `test-amd-merge.yml` (L3 / `merge-test` and `main`). On a PR carrying both labels, AMD combines both suites behind one image build. PR builds without either tier label retain the legacy L2 fallback for compatibility with the `amd-test` trigger. `DEBUG_TEST_YAML` remains an explicit override. AMD has no L4 file yet, so `nightly-test` does not select an additional suite. + **Test YAML (data):** `amd/test-amd-ready.yml` (L2 / `ready`), `test-amd-merge.yml` (L3 / `merge-test` and ordinary `main`), and experimental `test-amd-nightly.yml` (`full_model` / `nightly-test` or `main` with `NIGHTLY=1`). Multiple PR tier labels combine their suites behind one image build. PR builds that reach the bootstrap without a tier label retain the legacy ready-suite fallback. `DEBUG_TEST_YAML` accepts `ready`, `merge`, and `nightly` as an explicit override. During nightly burn-in every AMD nightly leaf is non-blocking. - **Trigger boundary:** AMD label handling has two layers. First, the external `vllm-omni-amd-ci` Buildkite pipeline condition decides whether a GitHub label event creates a build. Only after that build starts does this repository's bootstrap inspect all current PR labels and select L2, L3, or both. Repository-side selection therefore cannot make a `merge-test` event start AMD CI by itself. The external condition must admit both `ready` and `merge-test`, while preserving the legacy `amd-test` trigger and non-PR `main` / scheduled builds. For PR builds without a `DEBUG_TEST_YAML` override, an unreadable label set fails the bootstrap instead of silently selecting the wrong tier. Debug overrides skip label lookup and are validated directly; invalid values, including whitespace-only values, fail the bootstrap. Keep `nightly-test` out of the AMD condition until an AMD L4 suite exists; otherwise it would start a build without providing nightly coverage. + **Trigger boundary:** AMD label handling has two layers. First, the external `vllm-omni-amd-ci` Buildkite pipeline condition decides whether a GitHub label event creates a build. Only after that build starts does this repository's bootstrap inspect all current PR labels and select the ready, merge, nightly, or combined suites. Repository-side selection therefore cannot start AMD CI by itself. AMD-only nightly validation uses both `amd-test` and `nightly-test`: the external condition should accept an event for either label only when the complete PR label set contains both. `amd-test` admits the AMD pipeline, while `nightly-test` makes the bootstrap select `NIGHTLY_TESTS:test-amd-nightly.yml`. Existing `ready` / `merge-test` behavior and non-PR `main` / scheduled builds remain separate. For PR builds without a `DEBUG_TEST_YAML` override, an unreadable label set fails the bootstrap instead of silently selecting the wrong tier. Debug overrides skip label lookup and are validated directly; invalid values, including whitespace-only values, fail the bootstrap. Until the paired-label condition is configured, use an authorized manual Buildkite build with `DEBUG_TEST_YAML=nightly` and verify that the bootstrap reports only the nightly suite. **Rendering:** [`test-template-amd-omni.j2`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/amd/test-template-amd-omni.j2) wraps data steps with `amd-build` image build and `amd_` queues. Do **not** hand-edit generated `pipeline.yaml`. @@ -189,7 +190,7 @@ Canonical layout (prefer these paths for new changes): **Adding a job** - 1. Edit `test-amd-ready.yml` or `test-amd-merge.yml`. + 1. Edit `test-amd-ready.yml`, `test-amd-merge.yml`, or `test-amd-nightly.yml`. 2. Copy a neighboring block: `label`, `agent_pool`, `mirror_hardwares`, `commands`, optional `grade`. 3. Regenerate via bootstrap / Jinja; update `skip_ci.py` if you add a new YAML path. @@ -245,7 +246,7 @@ Label triggers (`ready`, `merge-test`) are unchanged—diff-aware logic only red | Non-qualifying skip-mark and no CI YAML | normal CI | all on | | Diff unavailable | normal CI | all on | -**`skip_all` exceptions (CUDA/NPU bootstrap only):** PR labels (`nightly-test`, `merge-test`, `weekly-test`, `ready`, …) do **not** revive jobs under `skip_all`. On `main`, scheduled `NIGHTLY=1` still builds the image and uploads L4; `WEEKLY=1` or `NON_CRITICAL=1` still builds the image and uploads L5 (CUDA). `WEEKLY=1` on `main` also uploads L2/L3 with `--e2e`. +**`skip_all` exceptions:** For CUDA/NPU, PR labels (`nightly-test`, `merge-test`, `weekly-test`, `ready`, …) do **not** revive jobs under `skip_all`; on `main`, scheduled `NIGHTLY=1` still builds the image and uploads L4, while `WEEKLY=1` or `NON_CRITICAL=1` still builds the image and uploads L5 (CUDA). `WEEKLY=1` on `main` also uploads L2/L3 with `--e2e`. AMD applies `skip_ci.py gate` only to selected ready/merge specs, so an explicitly selected AMD nightly spec (`nightly-test`, `NIGHTLY=1`, or a nightly debug override) remains runnable even for a docs-only diff. **Yaml-gated nightly/weekly:** L4/L5 upload steps keep their normal label / `NIGHTLY` / `WEEKLY` / `NON_CRITICAL` conditions (for example PR `nightly-test` still uploads nightly). Only L2/L3 upload steps are matrix-gated. @@ -262,6 +263,7 @@ Register new files in `L2_YAML_FILES`, `L3_YAML_FILES`, or `L45_YAML_FILES` in ` | L3 | `.buildkite/cuda/test-merge.yml` | cuda | | L3 | `.buildkite/amd/test-amd-merge.yml` | amd | | L4/L5 | `.buildkite/cuda/test-nightly.yml`, `test-weekly.yml` | cuda | +| L4/L5 | `.buildkite/amd/test-amd-nightly.yml` | amd | | L4/L5 | `.buildkite/npu/test-npu-nightly.yml` | npu | ##### Yaml-gated category branches @@ -294,7 +296,7 @@ Docs/skip-mark mixed with any row below follows the same matrix. | CUDA ready + AMD merge | G4 | `cuda/l2` + `amd/l3` | | CUDA ready + AMD ready | G2 | `cuda/l2` + `amd/l2` | | CUDA merge + AMD merge | G3 | `cuda/l3` + `amd/l3` | -| Only nightly / weekly / npu-nightly YAML | G1 | (none) | +| Only nightly / weekly / amd-nightly / npu-nightly YAML | G1 | (none) | | CUDA ready + nightly | G5 | `cuda/l2` | | CUDA merge + nightly | G6 | `cuda/l3` | | CUDA ready + merge + nightly | G7 | `cuda/l2` + `cuda/l3` | diff --git a/docs/contributing/ci/test_system_overview.md b/docs/contributing/ci/test_system_overview.md index f432678bc..dde8f2a58 100644 --- a/docs/contributing/ci/test_system_overview.md +++ b/docs/contributing/ci/test_system_overview.md @@ -182,6 +182,8 @@ CUDA **L4** (`full_model`) jobs are defined once in [`.buildkite/cuda/test-night - **Nightly:** scheduled `main` builds with `NIGHTLY=1` (and PR labels such as `nightly-test` / `omni-test`). Default fleet is typically H100. - **Pre-release:** the same L4 suite runs again on extra GPU SKUs before a release (for example B200 via `MIRROR_HW=b200`), so `full_model` coverage is not limited to the nightly machine type. +AMD has an experimental, non-blocking `full_model` nightly vertical slice in [`.buildkite/amd/test-amd-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/amd/test-amd-nightly.yml). It selects ROCm Qwen3-Omni function-expansion, accuracy, documentation-example, and AITER-on smoke coverage on the MI300 pool. AMD-only PR validation uses `amd-test` to admit the external AMD pipeline and `nightly-test` to select this suite; scheduled `main` builds select it with `NIGHTLY=1`. Keep it non-blocking while the lane accumulates stability and completion-time evidence. + How count-form `mirror_hardwares` and `MIRROR_HW` pick a SKU: [CI Settings](./ci_settings.md). ### L2 / L3 diff-aware CI (CUDA) diff --git a/tests/buildkite/test_amd_bootstrap.py b/tests/buildkite/test_amd_bootstrap.py index 593ba8b26..ac0c8dba4 100644 --- a/tests/buildkite/test_amd_bootstrap.py +++ b/tests/buildkite/test_amd_bootstrap.py @@ -136,7 +136,7 @@ def test_label_lookup_selects_tiers(tmp_path, labels, expected_levels): ] == expected_levels -@pytest.mark.parametrize("debug_test_yaml", ["ready,ready", "nightly", ", ,", " \t"]) +@pytest.mark.parametrize("debug_test_yaml", ["ready,ready", "nightly,nightly", ", ,", " \t"]) def test_invalid_debug_override_fails_without_label_lookup(tmp_path, debug_test_yaml): result = _run_bootstrap(tmp_path, debug_test_yaml=debug_test_yaml, curl_status=22) diff --git a/tests/buildkite/test_amd_pipeline.py b/tests/buildkite/test_amd_pipeline.py index 1e74b50a1..c2c45f5ef 100644 --- a/tests/buildkite/test_amd_pipeline.py +++ b/tests/buildkite/test_amd_pipeline.py @@ -10,10 +10,11 @@ import yaml pytestmark = [pytest.mark.core_model, pytest.mark.cpu] AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml") +AMD_NIGHTLY_PIPELINE = Path(".buildkite/amd/test-amd-nightly.yml") -def _find_step(label: str) -> dict: - pipeline = yaml.safe_load(AMD_MERGE_PIPELINE.read_text(encoding="utf-8")) +def _find_step(label: str, pipeline_path: Path = AMD_MERGE_PIPELINE) -> dict: + pipeline = yaml.safe_load(pipeline_path.read_text(encoding="utf-8")) def walk(steps: list[dict]) -> dict | None: for step in steps: @@ -40,3 +41,15 @@ def test_qwen3_tts_base_preserves_advanced_model_arguments() -> None: run_level_index = argv.index("--run-level") assert argv[marker_index + 1] == "advanced_model and cuda" assert argv[run_level_index + 1] == "advanced_model" + + +def test_qwen3_accuracy_defers_artifact_path_expansion() -> None: + step = _find_step("Qwen3-Omni Accuracy", AMD_NIGHTLY_PIPELINE) + staging_command = next(command for command in step["commands"] if "artifact_dir=" in command) + + # Dynamic pipelines are interpolated once during upload. Double dollars + # preserve these variables for the GPU job's runtime shell. + assert '"$$PWD"' in staging_command + assert '"$${BUILDKITE_BUILD_CHECKOUT_PATH:?}"' in staging_command + assert '"$$artifact_dir"' in staging_command + assert step["artifact_paths"] == ["tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json"] diff --git a/tests/buildkite/test_amd_suite_selection.py b/tests/buildkite/test_amd_suite_selection.py index b1b25aa26..dc447f208 100644 --- a/tests/buildkite/test_amd_suite_selection.py +++ b/tests/buildkite/test_amd_suite_selection.py @@ -24,10 +24,12 @@ SPEC.loader.exec_module(SELECTOR) ("feature", ("ready",), ("ready",)), ("feature", ("merge-test",), ("merge",)), ("feature", ("merge-test", "ready"), ("ready", "merge")), + ("feature", ("nightly-test",), ("nightly",)), + ("feature", ("ready", "merge-test", "nightly-test"), ("ready", "merge", "nightly")), ("feature", ("amd-test",), ("ready",)), ("feature", ("amd-test", "merge-test"), ("merge",)), ("feature", ("amd-test", "ready", "merge-test"), ("ready", "merge")), - ("feature", ("nightly-test",), ("ready",)), + ("feature", ("amd-test", "nightly-test"), ("nightly",)), ("feature", ("not-ready", "merge-test-extra"), ("ready",)), ], ) @@ -39,8 +41,17 @@ def test_debug_override_takes_precedence_and_normalizes_input(): assert SELECTOR.select_amd_test_suites( branch="main", labels=("ready",), - debug_test_yaml=" MERGE, ready,", - ) == ("merge", "ready") + debug_test_yaml=" NIGHTLY, merge, ready,", + nightly=False, + ) == ("nightly", "merge", "ready") + + +def test_scheduled_main_selects_only_nightly(): + assert SELECTOR.select_amd_test_suites( + branch="main", + labels=(), + nightly=True, + ) == ("nightly",) def test_empty_debug_override_uses_normal_selection(): @@ -51,7 +62,7 @@ def test_empty_debug_override_uses_normal_selection(): ) == ("merge",) -@pytest.mark.parametrize("value", ["ready,ready", "nightly", ", ,", " \t"]) +@pytest.mark.parametrize("value", ["ready,ready", "nightly,nightly", "weekly", ", ,", " \t"]) def test_invalid_debug_override(value): with pytest.raises(ValueError): SELECTOR.select_amd_test_suites( diff --git a/tests/buildkite/test_skip_ci.py b/tests/buildkite/test_skip_ci.py index 4c90b3ff7..1a04ecab8 100644 --- a/tests/buildkite/test_skip_ci.py +++ b/tests/buildkite/test_skip_ci.py @@ -144,6 +144,14 @@ def test_npu_nightly_only_skips_all_l23() -> None: assert not decision.is_run("npu", "l3") +def test_amd_nightly_only_skips_all_l23() -> None: + decision = _decision([".buildkite/amd/test-amd-nightly.yml"]) + assert _is_yaml_gated(decision) + assert _all_l23_skipped(decision) + assert not decision.is_run("amd", "l2") + assert not decision.is_run("amd", "l3") + + def test_nightly_and_weekly_skips_all_l23() -> None: decision = _decision( [ From 64903b8e30b6595e4018c66232abbbd806aa9dfd Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:27:10 +0800 Subject: [PATCH 054/209] [CI/Build] Share nightly CUDA YAML between H100/L4 and B200 (#7028) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] [Your Email] Signed-off-by: [Your Name] Co-authored-by: Cursor Co-authored-by: amy-why-3459 --- .../common/ci_source_file_dependencies.yml | 6 + .buildkite/cuda/test-nightly.yml | 167 ++++++------------ .buildkite/cuda/test-weekly.yml | 14 +- tests/buildkite/test_upload_pipeline.py | 47 +++++ .../perf/tests/test_cosmos3_vllm_omni.json | 8 +- .../test_hunyuanvideo15_i2v_vllm_omni.json | 4 +- .../test_hunyuanvideo15_t2v_vllm_omni.json | 4 +- tests/dfx/perf/tests/test_minicpmo_4_5.json | 2 +- .../test_minicpmo_4_5_duplex_seed_tts.json | 2 +- .../perf/tests/test_minimax_h3_vllm_omni.json | 8 +- .../tests/test_qwen3_omni_async_chunk.json | 4 +- .../tests/test_qwen3_omni_no_async_chunk.json | 2 +- .../perf/tests/test_qwen_image_vllm_omni.json | 8 +- tests/dfx/perf/tests/test_tts.json | 8 +- .../perf/tests/test_wan22_i2v_vllm_omni.json | 4 +- .../distributed/test_cfg_parallel.py | 8 +- tests/diffusion/distributed/test_comm.py | 6 +- .../distributed/test_expert_parallel.py | 2 +- .../test_flux2_vae_patch_parallel.py | 2 +- .../distributed/test_pipeline_parallel.py | 6 +- .../distributed/test_sana_video_sp_numeric.py | 12 +- .../distributed/test_sequence_parallel.py | 2 +- .../test_vae_decode_parallelism.py | 4 +- .../distributed/test_wan_spatial_shard.py | 5 +- .../test_minimax_h3_quantization_quality.py | 2 +- .../offloader/test_diffusion_cpu_offload.py | 7 +- .../test_diffusion_layerwise_offload.py | 13 +- .../quantization/test_gguf_diffusion.py | 2 +- .../quantization/test_quantization_fp8.py | 24 +-- .../quantization/test_quantization_quality.py | 17 +- tests/e2e/accuracy/helpers.py | 26 +++ ...est_hunyuanvideo15_i2v_video_similarity.py | 6 +- ...est_hunyuanvideo15_t2v_video_similarity.py | 6 +- .../minicpmo_4_5/test_minicpmo_4_5.py | 6 +- .../test_minimax_h3_i2va_ref2va_similarity.py | 6 +- .../accuracy/qwen3_omni/test_qwen3_omni.py | 4 +- .../sana_video/test_sana_video_golden.py | 5 +- .../test_sana_video_i2v_alignment.py | 4 +- .../test_sana_video_pipeline_alignment.py | 24 ++- tests/e2e/accuracy/test_cosmos3_similarity.py | 7 +- .../test_diffusers_backend_similarity.py | 39 ++-- tests/e2e/accuracy/test_hunyuan_image3.py | 13 +- .../test_hunyuan_image3_pixel_accuracy.py | 4 +- tests/e2e/accuracy/test_qwen_image.py | 4 +- .../test_wan22_i2v_video_similarity.py | 6 +- .../test_async_omni_collective_rpc.py | 2 +- .../test_async_omni_qwen_image_generate.py | 2 +- .../custom_pipeline/test_worker_extension.py | 2 +- .../test_device_postprocess_equivalence.py | 2 +- .../test_dots_tts_expansion.py | 2 +- .../e2e/offline_inference/test_gepard_tts.py | 4 +- .../test_hunyuanimage3_i2t_expansion.py | 2 +- .../test_lingbot_world_v2.py | 2 +- .../test_longcat_video_avatar.py | 15 +- ...st_qwen_image_autoround_w4a16_expansion.py | 8 +- .../e2e/offline_inference/test_voxcpm2_tts.py | 3 + .../test_wan22_autoround_w4a16_expansion.py | 16 +- .../test_hunyuan_video_15_expansion.py | 26 +-- .../e2e/online_serving/test_lingbot_video.py | 2 +- .../test_lingbot_video_expansion.py | 4 +- .../online_serving/test_lingbot_video_moe.py | 2 +- .../test_minicpmo_4_5_duplex_expansion.py | 6 +- .../e2e/online_serving/test_pi0_expansion.py | 4 +- .../test_qwen3_omni_multi_replicas.py | 6 +- .../test_qwen3_tts_base_expansion.py | 6 +- .../test_qwen3_tts_customvoice_expansion.py | 10 +- .../test_qwen_image_expansion.py | 4 +- tests/e2e/online_serving/test_sana_wm.py | 2 +- .../online_serving/test_wan22_expansion.py | 4 +- .../online_serving/test_zimage_expansion.py | 2 +- .../offline_inference/test_image_to_video.py | 5 +- .../offline_inference/test_internvla_a1.py | 3 +- .../offline_inference/test_text_to_image.py | 2 +- .../test_vace_video_generation.py | 5 +- .../online_serving/test_qwen2_5_omni.py | 17 +- .../online_serving/test_qwen3_omni.py | 17 +- .../online_serving/test_text_to_image.py | 5 +- tests/model_tests/diffusion/case_filtering.py | 6 +- 78 files changed, 437 insertions(+), 321 deletions(-) diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 724794c95..33def4de4 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -329,6 +329,7 @@ source_file_dependencies: - *qwen_image - tests/e2e/online_serving/test_qwen_image.py - tests/e2e/online_serving/test_qwen_image_expansion.py + - tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py diffusion_qwen_image_accuracy: - *qwen_image @@ -417,6 +418,11 @@ source_file_dependencies: - tests/e2e/online_serving/test_z_image.py - tests/e2e/online_serving/test_zimage_expansion.py + diffusion_sana_video_function: + - *sana_video + - tests/e2e/online_serving/test_sana_video_expansion.py + - tests/e2e/accuracy/sana_video/ + diffusion_stable_audio_function: - *stable_audio - tests/e2e/offline_inference/test_stable_audio_expansion.py diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index cfaca90fe..1fdb1ef4a 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -35,23 +35,20 @@ steps: source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 50 commands: - - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model" - mirror_hardwares: h100_1 + - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and B200 and omni and cards_1" --run-level "full_model" - label: ":full_moon: Omni · Doc Test with L4 · 4-GPU" source_file_dependencies: omni_qwen2_5_omni_doc timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/examples/ -m "full_model and omni and L4 and cards_4" --run-level "full_model" - mirror_hardwares: l4_4 + - pytest -s -v tests/examples/ -m "full_model and omni and L4 and B200 and cards_4" --run-level "full_model" - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU" source_file_dependencies: omni_qwen3_omni_doc timeout_in_minutes: 90 commands: - - pytest -s -v tests/examples/ -m "full_model and omni and H100 and cards_2" --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -s -v tests/examples/ -m "full_model and omni and H100 and B200 and cards_2" --run-level "full_model" - label: ":full_moon: Omni · Accuracy Test" source_file_dependencies: omni_qwen3_omni_accuracy @@ -61,11 +58,10 @@ steps: - export SEED_TTS_EVAL_DEVICE=cuda:1 - | set +e - pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model" --run-level full_model + pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model and H100 and B200 and cards_2" --run-level full_model EXIT=$$? buildkite-agent artifact upload "tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json" exit $$EXIT - mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test" source_file_dependencies: omni_minicpmo_4_5_accuracy @@ -75,11 +71,10 @@ steps: - export SEED_TTS_EVAL_DEVICE=cuda:1 - | set +e - pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and cuda" --run-level full_model + pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and H100 and B200 and cuda and (cards_1 or cards_2)" --run-level full_model EXIT=$$? buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json" exit $$EXIT - mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · No Async Chunk" source_file_dependencies: omni_qwen3_omni_perf @@ -89,11 +84,10 @@ steps: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json -m "H100 and B200 and cards_2" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · Async Chunk" source_file_dependencies: omni_qwen3_omni_perf @@ -103,11 +97,10 @@ steps: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and full_model" + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and B200 and full_model and cards_2" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" source_file_dependencies: omni_minicpmo_4_5_perf @@ -117,11 +110,10 @@ steps: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json -m "H100 and cards_1" + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json -m "H100 and B200 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - mirror_hardwares: h100_1 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" source_file_dependencies: omni_minicpmo_4_5_duplex_perf @@ -131,18 +123,16 @@ steps: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json -m "H100 and B200 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - mirror_hardwares: h100_1 - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100" source_file_dependencies: omni_qwen3_omni_function timeout_in_minutes: 45 commands: - - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model" --run-level "core_model" - mirror_hardwares: h100_4 + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model and H100 and B200 and cards_4" --run-level "core_model" - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group @@ -157,8 +147,7 @@ steps: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - export VLLM_USE_DEEP_GEMM="0" - export VLLM_MOE_USE_DEEP_GEMM="0" - - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy - mirror_hardwares: l4_1 + - pytest -s -v tests/e2e/ -m "full_model and L4 and B200 and tts and cards_1" --run-level "full_model" --ignore=tests/e2e/accuracy - label: ":full_moon: TTS · Perf Test · Single-GPU" source_file_dependencies: tts_qwen3_tts_perf @@ -169,11 +158,10 @@ steps: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_1" + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and B200 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - mirror_hardwares: h100_1 - label: ":full_moon: TTS · Perf Test · 2-GPU" source_file_dependencies: tts_qwen3_tts_perf @@ -184,11 +172,10 @@ steps: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_2" + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and B200 and cards_2" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - mirror_hardwares: h100_2 - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]" key: nightly-tiny-model-test-group @@ -209,16 +196,14 @@ steps: # for the CI, since these tests should pass irrespective of whether or not we use # the full weights or not. commands: - - pytest -sv tests/model_tests/diffusion -m 'full_model and cuda and cards_2' --run-level "core_model" - mirror_hardwares: l4_2 + - pytest -sv tests/model_tests/diffusion -m 'full_model and L4 and B200 and cuda and cards_2' --run-level "core_model" - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 4-GPU" source_file_dependencies: diffusion_tiny_model timeout_in_minutes: 30 # CFG+TP extra_test_groups request 2×2=4 devices (see get_required_device_count). commands: - - pytest -sv tests/model_tests/diffusion -m 'full_model and cuda and cards_4' --run-level "core_model" - mirror_hardwares: l4_4 + - pytest -sv tests/model_tests/diffusion -m 'full_model and L4 and B200 and cuda and cards_4' --run-level "core_model" - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" key: nightly-diffusion-x2iat-group @@ -235,9 +220,8 @@ steps: - >- pytest -sv tests/e2e/online_serving/test_qwen_image_expansion.py - -m "full_model and diffusion and H100 and cards_1" + -m "full_model and diffusion and H100 and B200 and cards_1" --run-level "full_model" - mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU" source_file_dependencies: diffusion_qwen_image_function @@ -246,36 +230,29 @@ steps: - >- pytest -sv tests/e2e/online_serving/test_qwen_image_expansion.py - -m "full_model and diffusion and H100 and cards_2" + -m "full_model and diffusion and H100 and B200 and cards_2" --run-level "full_model" - mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2I(&A&T) · Z-Image Function Test · L4 · 4-GPU" - source_file_dependencies: diffusion_z_image_function + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · Single-GPU" + source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 commands: - - >- - pytest -sv - tests/e2e/online_serving/test_zimage_expansion.py - -m "full_model and diffusion and L4 and cards_4" - --run-level "full_model" - mirror_hardwares: l4_4 + - pytest -sv tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py -m "full_model and diffusion and L4 and B200 and cards_1" --run-level "full_model" + - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test" source_file_dependencies: diffusion_text_to_image_doc timeout_in_minutes: 60 commands: - - pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100 and B200 and (cards_1 or cards_2)" --run-level "full_model" - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test" source_file_dependencies: diffusion_qwen_image_accuracy timeout_in_minutes: 180 commands: - export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120 - - pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model - - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model - mirror_hardwares: h100_1 + - pytest -s -v tests/e2e/accuracy/test_qwen_image.py -m "H100 and B200 and cards_1" --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m "full_model and H100 and B200 and cards_1" --run-level full_model - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test" source_file_dependencies: diffusion_hunyuan_image3_accuracy @@ -283,8 +260,7 @@ steps: commands: - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct" - export HUNYUAN_IMAGE3_DEVICES="0,1,2,3" - - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model - mirror_hardwares: h100_4 + - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and H100 and B200 and cards_4" --run-level full_model - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU" source_file_dependencies: diffusion_qwen_image_perf @@ -295,12 +271,11 @@ steps: - export CACHE_DIT_VERSION=1.5.0 - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_1" + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and B200 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-GPU" source_file_dependencies: diffusion_qwen_image_perf @@ -308,16 +283,16 @@ steps: timeout_in_minutes: 180 commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - export CACHE_DIT_VERSION=1.5.0 + # Do not pin DIFFUSION_ATTENTION_BACKEND: H100 auto-selects FLASH_ATTN; + # B200 rejects explicit FLASH_ATTN without FA4 and uses CUDNN/TRTLLM. - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_4" + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and B200 and cards_4" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - mirror_hardwares: h100_4 - group: ":card_index_dividers: Diffusion X2V Model Test" key: nightly-diffusion-x2v-group @@ -331,36 +306,31 @@ steps: source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" - mirror_hardwares: h100_1 + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and B200 and cards_1" --run-level "full_model" - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · 2-GPU" source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and B200 and cards_2" --run-level "full_model" - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · Single-GPU" source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" - mirror_hardwares: h100_1 + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and B200 and cards_1" --run-level "full_model" - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · 2-GPU" source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and B200 and cards_2" --run-level "full_model" - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test" source_file_dependencies: diffusion_hunyuan_video_function timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and H100 and B200 and cuda and (cards_1 or cards_2)" --run-level "full_model" - label: ":full_moon: Diffusion X2V · LingBot Function Test" source_file_dependencies: diffusion_lingbot_function @@ -371,23 +341,20 @@ steps: tests/e2e/online_serving/test_lingbot_video.py tests/e2e/online_serving/test_lingbot_video_moe.py tests/e2e/offline_inference/test_lingbot_world_v2.py - -m "full_model and diffusion and H100 and cards_1" + -m "full_model and diffusion and H100 and B200 and cards_1" --run-level "full_model" - mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Doc Test" source_file_dependencies: diffusion_image_to_video_doc timeout_in_minutes: 60 commands: - - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model" - mirror_hardwares: h100_1 + - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100 and B200 and cards_1" --run-level "full_model" - label: ":full_moon: Diffusion X2V · Wan2.2 Autoround Function Test" source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and H100 and B200 and cuda and (cards_1 or cards_2)" --run-level "full_model" - label: ":full_moon: Diffusion X2V · Wan2.2 Device-Postprocess Equivalence Test" source_file_dependencies: diffusion_wan22_function @@ -395,31 +362,27 @@ steps: commands: - | VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - pytest -s -v tests/e2e/offline_inference/test_device_postprocess_equivalence.py -m "full_model and cuda" --run-level "full_model" - mirror_hardwares: h100_1 + pytest -s -v tests/e2e/offline_inference/test_device_postprocess_equivalence.py -m "full_model and H100 and B200 and cuda and cards_1" --run-level "full_model" - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" source_file_dependencies: diffusion_wan22_accuracy timeout_in_minutes: 180 commands: - - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model - - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model - mirror_hardwares: h100_2 + - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m "full_model and H100 and B200 and (cards_1 or cards_2)" --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m "full_model and H100 and B200 and cards_1" --run-level full_model - label: ":full_moon: Diffusion X2V · MiniMax H3 I2VA/Ref2VA Accuracy Test" source_file_dependencies: diffusion_minimax_h3_accuracy timeout_in_minutes: 180 commands: - - pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py -m full_model --run-level full_model - mirror_hardwares: h100_4 + - pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py -m "full_model and H100 and B200 and cards_4" --run-level full_model - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Accuracy Test" source_file_dependencies: diffusion_hunyuan_video_accuracy timeout_in_minutes: 180 commands: - - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model - - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model - mirror_hardwares: h100_2 + - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m "full_model and H100 and B200 and (cards_1 or cards_2)" --run-level full_model + - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m "full_model and H100 and B200 and (cards_1 or cards_2)" --run-level full_model - label: ":full_moon: Diffusion X2V · Perf Test · Single-GPU" source_file_dependencies: diffusion_wan22_perf @@ -429,12 +392,11 @@ steps: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_1" + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and B200 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Perf Test · 2-GPU" source_file_dependencies: @@ -444,18 +406,16 @@ steps: timeout_in_minutes: 180 commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_2" + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and B200 and cards_2" EXIT1=$$? - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json -m "H100 and cards_2" + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json -m "H100 and B200 and cards_2" EXIT2=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi exit $$EXIT2 - mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5" source_file_dependencies: diffusion_hunyuan_video_perf @@ -465,15 +425,14 @@ steps: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json -m "H100 and B200 and (cards_1 or cards_2)" EXIT1=$$? - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json -m "H100 and B200 and (cards_1 or cards_2)" EXIT2=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi exit $$EXIT2 - mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3" source_file_dependencies: diffusion_minimax_h3_perf @@ -483,12 +442,11 @@ steps: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json -m "H100 and B200 and cards_4" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - mirror_hardwares: h100_4 - group: ":card_index_dividers: Diffusion Test" key: nightly-diffusion-test-group @@ -499,65 +457,56 @@ steps: source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_2' --run-level "full_model" - mirror_hardwares: l4_2 + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and B200 and cards_2' --run-level "full_model" - label: ":full_moon: Diffusion · Distributed Test with L4 · 3-GPU" source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_3' --run-level "full_model" - mirror_hardwares: l4_3 + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and B200 and cards_3' --run-level "full_model" - label: ":full_moon: Diffusion · Distributed Test with L4 · 4-GPU" source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_4' --run-level "full_model" - mirror_hardwares: l4_4 + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and B200 and cards_4' --run-level "full_model" - label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU" source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_2' --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and B200 and cards_2' --run-level "full_model" - label: ":full_moon: Diffusion · Distributed Test with H100 · 4-GPU" source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_4' --run-level "full_model" - mirror_hardwares: h100_4 + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and B200 and cards_4' --run-level "full_model" - label: ":full_moon: Diffusion · Offloader Test" source_file_dependencies: diffusion_offloader timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/offloader -m "full_model and cuda" - mirror_hardwares: l4_1 + - pytest -sv tests/diffusion/offloader -m "full_model and L4 and B200 and cuda and cards_1" - label: ":full_moon: Diffusion · Quantization Test with H100 · Single-GPU" source_file_dependencies: diffusion_quantization timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model" - mirror_hardwares: h100_1 + - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100 and B200 and cards_1' --run-level "full_model" - label: ":full_moon: Diffusion · MiniMax H3 FP8 Quality Test with H100 · 2-GPU" source_file_dependencies: diffusion_minimax_h3_function timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100' --run-level "full_model" - mirror_hardwares: h100_2 + - pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100 and B200 and cards_2' --run-level "full_model" - label: ":full_moon: Diffusion · Quantization Test with L4" source_file_dependencies: diffusion_quantization timeout_in_minutes: 60 commands: - pip install "vllm-gguf-plugin==0.0.4" - - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model" - mirror_hardwares: l4_1 + - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4 and B200 and cards_1' --run-level "full_model" # NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon. # Do not add new depends_on entries or extend this step. diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index d53bc56e4..7809fcbc7 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.buildkite/cuda/test-weekly.yml @@ -339,11 +339,14 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · SANA-Video · H100 · Single-GPU" - timeout_in_minutes: 90 + source_file_dependencies: diffusion_sana_video_function + timeout_in_minutes: 180 commands: + - export SANA_VIDEO_GOLDEN_BASE_URL=https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/sana-video/v2 - >- pytest -sv tests/e2e/online_serving/test_sana_video_expansion.py + tests/e2e/accuracy/sana_video/ -m "slow and diffusion and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 @@ -488,6 +491,15 @@ steps: --run-level "full_model" mirror_hardwares: h100_2 + - label: "Diffusion · H100 · 3-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and diffusion and H100 and cards_3" + --run-level "full_model" + mirror_hardwares: h100_3 + - label: "Diffusion · H100 · 4-GPU" timeout_in_minutes: 90 commands: diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index a3192dada..b39dba2cd 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -27,6 +27,7 @@ from upload_pipeline import ( # noqa: E402 pytestmark = [pytest.mark.core_model, pytest.mark.cpu] CUDA_BOOTSTRAP_STEPS = Path(".buildkite/cuda/bootstrap-upload-steps.yml") +NIGHTLY_YAML = Path(".buildkite/cuda/test-nightly.yml") BOOTSTRAP_STEPS_TEMPLATE = """steps: - key: image-build - key: upload-ready-pipeline @@ -368,6 +369,52 @@ def test_cpu_step_without_mirror_hardwares_is_unchanged() -> None: assert _expand_mirror_hardwares(step) is step +def _leaf_steps(steps: list) -> list[dict]: + leaves: list[dict] = [] + for step in steps: + if not isinstance(step, dict): + continue + nested = step.get("steps") + if nested is not None: + leaves.extend(_leaf_steps(nested)) + else: + leaves.append(step) + return leaves + + +def test_nightly_yaml_infers_h100_l4_and_b200(monkeypatch: pytest.MonkeyPatch) -> None: + """Nightly inferred jobs share card counts across H100/L4 and B200 mirrors.""" + src = yaml.safe_load(NIGHTLY_YAML.read_text(encoding="utf-8")) + pytest_leaves = [step for step in _leaf_steps(src["steps"]) if "pytest" in str(step.get("commands"))] + inferred_leaves = [step for step in pytest_leaves if "mirror_hardwares" not in step] + pinned_leaves = [step for step in pytest_leaves if "mirror_hardwares" in step] + assert inferred_leaves + + monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "") + default_by_label = { + step["label"]: step for step in _leaf_steps(_render_test_pipeline(src, changed_files=None)["steps"]) + } + + monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "b200") + b200_src = yaml.safe_load(NIGHTLY_YAML.read_text(encoding="utf-8")) + b200_by_label = { + step["label"]: step for step in _leaf_steps(_render_test_pipeline(b200_src, changed_files=None)["steps"]) + } + + pytest_labels = {step["label"] for step in inferred_leaves} + pinned_labels = {step["label"] for step in pinned_leaves} + assert pinned_labels <= set(default_by_label) + assert pinned_labels.isdisjoint(b200_by_label) + assert pytest_labels <= set(default_by_label) + assert pytest_labels <= set(b200_by_label) + for label in pytest_labels: + default_step = default_by_label[label] + b200_step = b200_by_label[label] + assert default_step["agents"]["queue"] in {"mithril-h100-pool", "l4-k8s"} + assert b200_step["agents"]["queue"] == "b200-k8s" + assert _gpu_limit(default_step) == _gpu_limit(b200_step) + + @pytest.mark.parametrize(("raw", "expected"), [("", ""), (" ", ""), ("b200", "b200"), ("B200", "b200")]) def test_mirror_hw_selector_empty_or_b200(monkeypatch: pytest.MonkeyPatch, raw: str, expected: str) -> None: monkeypatch.setenv("MIRROR_HW", raw) diff --git a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json index bcd649de1..5e2b9d3dc 100644 --- a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json +++ b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } @@ -63,7 +63,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } @@ -132,7 +132,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } @@ -200,7 +200,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } diff --git a/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json index 8186edf2c..25d91c44f 100644 --- a/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json +++ b/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 1 } @@ -55,7 +55,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } diff --git a/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json b/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json index 3771c1c81..ffaaba436 100644 --- a/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json +++ b/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 1 } @@ -54,7 +54,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5.json b/tests/dfx/perf/tests/test_minicpmo_4_5.json index cac1880c5..7d7ca08ad 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100", + "cuda": ["H100", "B200"], "npu": "A3" }, "num_cards": 1 diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json index 801743134..bb0599419 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100", + "cuda": ["H100", "B200"], "npu": "A3" }, "num_cards": 1 diff --git a/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json b/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json index e16b40e91..7dfbe6f3d 100644 --- a/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json +++ b/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 4 } @@ -76,7 +76,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 4 } @@ -148,7 +148,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 4 } @@ -219,7 +219,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 4 } diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json index 6b0247aaf..1292a79e5 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100", + "cuda": ["H100", "B200"], "npu": "A3" }, "num_cards": 2 @@ -242,7 +242,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json index d73629666..186bac452 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100", + "cuda": ["H100", "B200"], "npu": "A3" }, "num_cards": 2 diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index 4f64d6ec1..61cb334eb 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 1 } @@ -66,7 +66,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 1 } @@ -150,7 +150,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 4 } @@ -197,7 +197,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 4 } diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index 592f61ca2..c627c01e8 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100", + "cuda": ["H100", "B200"], "npu": "A3" }, "num_cards": 2 @@ -121,7 +121,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } @@ -285,7 +285,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 2 } @@ -416,7 +416,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 1 } diff --git a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json index 234bbf1db..3244d1c41 100644 --- a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json +++ b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json @@ -5,7 +5,7 @@ { "hardware_marks": { "res": { - "cuda": "H100" + "cuda": ["H100", "B200"] }, "num_cards": 1 } @@ -57,7 +57,7 @@ { "hardware_marks": { "res": { - "cuda": "H100", + "cuda": ["H100", "B200"], "npu": "A3" }, "num_cards": 2 diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py index 2ade13549..20679cf76 100644 --- a/tests/diffusion/distributed/test_cfg_parallel.py +++ b/tests/diffusion/distributed/test_cfg_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for CFG (Classifier-Free Guidance) parallel functionality. CPU tests mock the CFG process group and validate mixin logic without GPUs. @@ -32,8 +32,8 @@ 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) +_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2) +_L4_THREE_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=3) def _set_random_seeds(seed: int) -> None: @@ -1009,7 +1009,7 @@ def _run_multi_branch_gpu_parity( @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "L4"}, num_cards=2) +@hardware_test(res={"cuda": ["L4", "B200"]}, 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}]) diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 647389048..effe2b541 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/tests/diffusion/distributed/test_comm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives. CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same @@ -24,8 +24,8 @@ 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) +_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2) +_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4) DeviceKind = Literal["cpu", "cuda"] diff --git a/tests/diffusion/distributed/test_expert_parallel.py b/tests/diffusion/distributed/test_expert_parallel.py index 8e9be1367..053fed9d1 100644 --- a/tests/diffusion/distributed/test_expert_parallel.py +++ b/tests/diffusion/distributed/test_expert_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ System test for Expert Parallel (EP). diff --git a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py index 8de3a42f8..204406ec1 100644 --- a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py +++ b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """FLUX.2-dev VAE patch-parallel correctness tests. diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index cccf116c3..08c3cc775 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -31,9 +31,9 @@ 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 -_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) +_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2) +_L4_THREE_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=3) +_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4) DeviceKind = Literal["cpu", "cuda"] _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] diff --git a/tests/diffusion/distributed/test_sana_video_sp_numeric.py b/tests/diffusion/distributed/test_sana_video_sp_numeric.py index b18cafe4c..872dcbcbb 100644 --- a/tests/diffusion/distributed/test_sana_video_sp_numeric.py +++ b/tests/diffusion/distributed/test_sana_video_sp_numeric.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Real multi-process numeric equivalence gate for SANA-Video sequence parallel. Runs the native transformer directly (no pipeline, no topology validation) with @@ -273,7 +273,7 @@ def _baseline_worker( @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "L4"}, num_cards=2) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2) @pytest.mark.parametrize("task", ["t2v", "i2v"]) @pytest.mark.parametrize("num_frames", [21, 5]) def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None: @@ -290,7 +290,7 @@ def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None: @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "L4"}, num_cards=4) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4) @pytest.mark.parametrize("task", ["t2v", "i2v"]) @pytest.mark.parametrize("num_frames", [21, 5]) def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None: @@ -307,7 +307,7 @@ def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None: @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "L4"}, num_cards=4) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4) def test_sp2_cfg2_matches_sp1_fp32() -> None: """Four processes: CFG branches with distinct inputs each run their own SP group; catches CFG/SP group crossover and deadlocks.""" @@ -324,7 +324,7 @@ def test_sp2_cfg2_matches_sp1_fp32() -> None: @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "L4"}, num_cards=4) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4) def test_tp2_sp2_matches_sp1_fp32() -> None: """Four processes: TP shards heads and channels while SP shards frames; the combined mesh must still reproduce the serial output.""" @@ -342,7 +342,7 @@ def test_tp2_sp2_matches_sp1_fp32() -> None: @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "L4"}, num_cards=2) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2) def test_sp2_bf16_drift_informational() -> None: """Records the bf16 reduction-reorder drift magnitude; the load-bearing numeric gate is the fp32 suite above.""" diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py index f083e1d26..ba68669e9 100644 --- a/tests/diffusion/distributed/test_sequence_parallel.py +++ b/tests/diffusion/distributed/test_sequence_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ System test for Sequence Parallel (SP) backends: Ulysses and Ring attention. diff --git a/tests/diffusion/distributed/test_vae_decode_parallelism.py b/tests/diffusion/distributed/test_vae_decode_parallelism.py index 28ded4639..acd916806 100644 --- a/tests/diffusion/distributed/test_vae_decode_parallelism.py +++ b/tests/diffusion/distributed/test_vae_decode_parallelism.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """VAE patch-parallel decode parity tests for diffusion models. @@ -145,7 +145,7 @@ def _run_generate( @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 4}) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards={"cuda": 4}) @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(): diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py index beb5ac2db..d8842ee3f 100644 --- a/tests/diffusion/distributed/test_wan_spatial_shard.py +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import os from types import SimpleNamespace @@ -538,7 +541,7 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_ @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parallel -@hardware_test(res={"cuda": "H100"}, num_cards=_SPATIAL_SHARD_WORLD_SIZE) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=_SPATIAL_SHARD_WORLD_SIZE) @pytest.mark.parametrize("split_dim", ["height", "width"]) def test_spatial_shard_decode_matches_reference(split_dim: str): manager = mp.get_context("spawn").Manager() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index a45631f7d..c7e53ca62 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -175,7 +175,7 @@ def _generate_joint_output(omni, config: QualityTestConfig): pytest.param( _QUALITY_CONFIG, id=_QUALITY_CONFIG.id, - marks=hardware_marks(res={"cuda": _QUALITY_CONFIG.gpu}, num_cards=2), + marks=hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2), ) ], ) diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 80a03180f..205fad8fd 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -9,7 +9,7 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.helpers.mark import hardware_test +from tests.helpers.mark import hardware_marks, hardware_test from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -27,7 +27,10 @@ MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_MODELS} MODEL_MARKS = { "riverclouds/qwen_image_random": pytest.mark.core_model, - "stabilityai/stable-audio-open-1.0": pytest.mark.full_model, + "stabilityai/stable-audio-open-1.0": [ + pytest.mark.full_model, + *hardware_marks(res={"cuda": ["L4", "B200"], "rocm": "MI325"}), + ], } _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"} diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index c386871d3..c579affe9 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -12,7 +12,7 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.helpers.mark import hardware_test +from tests.helpers.mark import hardware_marks from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -34,8 +34,14 @@ IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = { MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} MODEL_MARKS = { - "riverclouds/qwen_image_random": pytest.mark.core_model, - "stabilityai/stable-audio-open-1.0": pytest.mark.full_model, + "riverclouds/qwen_image_random": [ + pytest.mark.core_model, + *hardware_marks(res={"cuda": "L4", "rocm": "MI325"}), + ], + "stabilityai/stable-audio-open-1.0": [ + pytest.mark.full_model, + *hardware_marks(res={"cuda": ["L4", "B200"], "rocm": "MI325"}), + ], } AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = { @@ -309,7 +315,6 @@ def test_stable_audio_requires_both_outputs(audio_no_offload, audio_offload, mes @pytest.mark.diffusion -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) @pytest.mark.parametrize( "model_name", [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS], diff --git a/tests/diffusion/quantization/test_gguf_diffusion.py b/tests/diffusion/quantization/test_gguf_diffusion.py index ad2e72edc..dc91a1e40 100644 --- a/tests/diffusion/quantization/test_gguf_diffusion.py +++ b/tests/diffusion/quantization/test_gguf_diffusion.py @@ -109,7 +109,7 @@ def _generate_single_stage_image( return images, peak_mem -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) @pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"]) diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 5793ad95c..5aa7faaa6 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end tests for the unified quantization framework (PR #1764). @@ -270,7 +270,7 @@ def _generate_bagel_image( # ─── Single-stage diffusion model tests ────────────────────────────────────── -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_single_stage_zimage_fp8(): """Z-Image-Turbo with FP8 generates valid images.""" images, _ = _generate_single_stage_image( @@ -281,7 +281,7 @@ def test_single_stage_zimage_fp8(): images[0].save("test_zimage_fp8.png") -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_single_stage_zimage_fp8_uses_less_memory(): """FP8 should use less peak memory than BF16 for Z-Image-Turbo.""" _, mem_bf16 = _generate_single_stage_image( @@ -300,7 +300,7 @@ def test_single_stage_zimage_fp8_uses_less_memory(): assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)" -@hardware_test(res={"cuda": "L4", "xpu": "B60"}) +@hardware_test(res={"cuda": ["L4", "B200"], "xpu": "B60"}) @pytest.mark.advanced_model def test_single_stage_qwen_image_fp8(): """Qwen-Image (random weights) with FP8 generates valid images.""" @@ -316,7 +316,7 @@ def test_single_stage_qwen_image_fp8(): images[0].save("test_qwen_image_fp8.png") -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.skip(reason="This model is not authorized on Hugging Face Hub yet") def test_single_stage_flux_fp8(): """FLUX.1-dev with FP8 generates valid images.""" @@ -331,7 +331,7 @@ def test_single_stage_flux_fp8(): images[0].save("test_flux_fp8.png") -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.skip(reason="This model is not authorized on Hugging Face Hub yet") def test_single_stage_flux_fp8_uses_less_memory(): """FP8 should use less peak memory than BF16 for FLUX.1-dev.""" @@ -357,7 +357,7 @@ def test_single_stage_flux_fp8_uses_less_memory(): assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)" -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) def test_single_stage_ltx2_fp8_uses_less_memory(): """FP8 should use less peak memory than BF16 for LTX-2.""" _, mem_bf16 = _generate_single_stage_video( @@ -379,7 +379,7 @@ def test_single_stage_ltx2_fp8_uses_less_memory(): # ─── Multi-stage model tests (BAGEL) ───────────────────────────────────────── -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) def test_bagel_fp8_generates_image(): """BAGEL with diffusion-stage FP8 generates a valid image. @@ -392,7 +392,7 @@ def test_bagel_fp8_generates_image(): image.save("test_bagel_fp8.png") -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) def test_bagel_bf16_generates_image(): """BAGEL without quantization generates a valid image (baseline).""" image, _ = _generate_bagel_image(diffusion_quantization_config=None) @@ -402,7 +402,7 @@ def test_bagel_bf16_generates_image(): # ─── Quantization config routing tests ──────────────────────────────────────── -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_quantization_key_maps_to_quantization_config(): """The old 'quantization' kwarg should map to 'quantization_config' in OmniDiffusionConfig.from_kwargs for backwards compatibility.""" @@ -413,7 +413,7 @@ def test_quantization_key_maps_to_quantization_config(): assert config.quantization_config.get_name() == "fp8" -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_quantization_config_key_takes_priority(): """When both 'quantization' and 'quantization_config' are set, 'quantization_config' takes priority.""" @@ -428,7 +428,7 @@ def test_quantization_config_key_takes_priority(): assert config.quantization_config.activation_scheme == "static" -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_single_stage_quantization_config_key(): """Single-stage model using quantization_config (dict) key generates images.""" images, _ = _generate_single_stage_image( diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 47908daa0..4bace0da5 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -45,6 +45,7 @@ import pytest import torch from PIL import Image +from tests.e2e.accuracy.helpers import resolve_device_threshold from tests.helpers.mark import hardware_marks _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -70,7 +71,7 @@ class QualityTestConfig: id: str # pytest ID, e.g. "fp8_z_image" task: str # "t2i" or "t2v" prompt: str # generation prompt - max_lpips: float # fail threshold — higher = more lenient + max_lpips: float | dict[str, float] # threshold, or {"H100": 0.15, "B200": 0.17} model: str | None = None # HF model name quantization: str | dict[str, object] | None = None # quantization method/config, e.g. "fp8" baseline_model: str | None = None # explicit BF16/local baseline path @@ -80,7 +81,6 @@ class QualityTestConfig: num_inference_steps: int = 20 # keep low for CI speed num_frames: int = 5 # only for t2v seed: int = 42 - gpu: str = "H100" # minimum GPU requirement negative_prompt: str | None = "" guidance_scale: float | None = None sigmas: list[float] | None = None @@ -190,7 +190,7 @@ QUALITY_CONFIGS = [ quantization={"text_encoder": "fp8", "transformer": None, "vae": None}, task="t2i", prompt="a cup of coffee on a wooden table, morning light", - max_lpips=0.15, + max_lpips={"H100": 0.15, "B200": 0.17}, num_inference_steps=10, enable_cpu_offload=True, height=1024, @@ -528,7 +528,7 @@ def test_ltx_quality_gate_uses_official_eager_defaults(monkeypatch): assert captured.sampling.guidance_scale is None -_marks = hardware_marks(res={"cuda": "H100"}) +_marks = hardware_marks(res={"cuda": ["H100", "B200"]}) _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None @@ -588,8 +588,9 @@ def test_quantization_quality(config: QualityTestConfig): # --- Similarity metrics --- lpips_score = _compute_lpips(baseline_out, quant_out, config.task) psnr_score, mae_score = _compute_psnr_and_mae(baseline_out, quant_out, config.task) - assert lpips_score <= config.max_lpips, ( - f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} " + gpu_key, max_lpips = resolve_device_threshold(config.max_lpips, label=f"{config.id} max_lpips") + assert lpips_score <= max_lpips, ( + f"LPIPS {lpips_score:.4f} exceeds threshold {max_lpips} ({gpu_key}) " f"for {config.quantization_ref() or 'pre-quantized checkpoint'} on {config.quantized_ref()}" ) @@ -601,12 +602,12 @@ def test_quantization_quality(config: QualityTestConfig): print(f" Baseline: {config.baseline_ref()}") print(f" Quantized: {config.quantized_ref()}") print(f" Method: {config.quantization_ref() or 'pre-quantized checkpoint'}") - print(f" LPIPS: {lpips_score:.4f} (threshold: {config.max_lpips})") + print(f" LPIPS: {lpips_score:.4f} (threshold: {max_lpips}, gpu: {gpu_key})") print(f" PSNR: {psnr_score:.4f} dB (higher is better)") print(f" MAE: {mae_score:.6f} (lower is better)") print(f" BF16 memory: {bl_mem:.2f} GiB") print(f" Quant memory: {qt_mem:.2f} GiB ({mem_reduction:.0f}% reduction)") - print(f" Result: {'PASS' if lpips_score <= config.max_lpips else 'FAIL'}") + print(f" Result: {'PASS' if lpips_score <= max_lpips else 'FAIL'}") print(f"{'=' * 60}\n") assert np.isfinite(psnr_score) or np.isinf(psnr_score), f"PSNR is invalid for {config.id}: {psnr_score}" diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index cc1b7c589..0ea9074b8 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/tests/e2e/accuracy/helpers.py @@ -135,6 +135,32 @@ def resolve_similarity_thresholds( return table.get(profile, table[DEFAULT_DEVICE_PROFILE]) +def resolve_device_threshold( + thresholds: float | int | Mapping[str, float], + *, + device_name: str | None = None, + label: str = "threshold", +) -> tuple[str, float]: + """Resolve a scalar or per-GPU threshold for the current CUDA device. + + Returns ``(profile_key, threshold)``. A bare number uses profile ``default``. + A mapping matches keys as substrings of the device name via + ``resolve_device_profile`` (for example ``{"H100": 0.3, "B200": 0.36}``). + Prefer an explicit ``default`` entry when unmatched GPUs should fall back; + otherwise unmatched devices raise. + """ + if isinstance(thresholds, float | int): + return DEFAULT_DEVICE_PROFILE, float(thresholds) + + profile = resolve_device_profile(device_name, profiles=thresholds) + if profile in thresholds: + return profile, float(thresholds[profile]) + if DEFAULT_DEVICE_PROFILE in thresholds: + return DEFAULT_DEVICE_PROFILE, float(thresholds[DEFAULT_DEVICE_PROFILE]) + known = ", ".join(str(key) for key in thresholds) + raise RuntimeError(f"no {label} entry for device profile {profile!r}; known: {known}") + + _SSIM_RE = re.compile(r"All:(?P[0-9.]+)") _PSNR_RE = re.compile(r"average:(?P[0-9.]+)") diff --git a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py index eb8f2d769..f5df4004f 100644 --- a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py +++ b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py @@ -183,7 +183,7 @@ def _generate_online_video( @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_hunyuanvideo15_i2v_diffusers_offline_generates_video( hunyuanvideo15_i2v_image_source: str | None, ) -> None: @@ -203,7 +203,7 @@ def test_hunyuanvideo15_i2v_diffusers_offline_generates_video( @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_hunyuanvideo15_i2v_online_serving_generates_video( omni_server, @@ -229,7 +229,7 @@ def test_hunyuanvideo15_i2v_online_serving_generates_video( @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_hunyuanvideo15_i2v_serving_matches_offline_video_similarity( hunyuanvideo15_i2v_image_source: str | None, ) -> None: diff --git a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py index f5b90822d..f16813f12 100644 --- a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py +++ b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py @@ -144,7 +144,7 @@ def _generate_online_video(*, omni_server, online_client, timeout_seconds: int) @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None: if not torch.cuda.is_available(): pytest.skip("HunyuanVideo-1.5 T2V offline accuracy test requires CUDA.") @@ -160,7 +160,7 @@ def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None: @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_hunyuanvideo15_t2v_online_serving_generates_video( omni_server, @@ -182,7 +182,7 @@ def test_hunyuanvideo15_t2v_online_serving_generates_video( @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_hunyuanvideo15_t2v_serving_matches_offline_video_similarity() -> None: if not torch.cuda.is_available(): pytest.skip("HunyuanVideo-1.5 T2V video similarity test requires CUDA.") diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py index e688b9360..ad659541d 100644 --- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py +++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py @@ -141,7 +141,7 @@ def _inline_daily_omni_media(monkeypatch: pytest.MonkeyPatch) -> None: @requires_daily_omni_deps -@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", daily_test_params, indirect=True) def test_minicpmo_4_5_daily_omni_accuracy_bench(omni_server) -> None: _require_vllm_cli() @@ -174,7 +174,7 @@ def test_minicpmo_4_5_daily_omni_accuracy_bench(omni_server) -> None: assert _acc_bench.run_acc_benchmark(_acc_bench.parse_acc_benchmark_args(argv)) == 0 -@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", seed_test_params, indirect=True) def test_minicpmo_4_5_seed_tts_wer_bench(omni_server) -> None: _require_vllm_cli() @@ -200,7 +200,7 @@ def test_minicpmo_4_5_seed_tts_wer_bench(omni_server) -> None: assert _acc_bench.run_acc_benchmark(_acc_bench.parse_acc_benchmark_args(argv)) == 0 -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", duplex_accuracy_test_params, indirect=True) def test_minicpmo_4_5_duplex_seed_tts_wer_bench(omni_server) -> None: """Gate Seed-TTS WER through the explicit Realtime TTS contract.""" diff --git a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py index 60f829bee..4ae45e4e1 100644 --- a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py +++ b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py @@ -138,8 +138,6 @@ def _server_args() -> list[str]: "--vae-parallel-mode", "tile", "--vae-use-tiling", - "--diffusion-attention-backend", - "FLASH_ATTN", "--stage-init-timeout", "1800", "--init-timeout", @@ -228,7 +226,7 @@ def _assert_official_video( ) -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) def test_minimax_h3_i2va_matches_official_reference( accuracy_artifact_root: Path, ) -> None: @@ -304,7 +302,7 @@ def test_minimax_h3_i2va_matches_official_reference( ) -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) def test_minimax_h3_ref2va_matches_official_reference( accuracy_artifact_root: Path, ) -> None: diff --git a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py index 7c48053e1..0acf21692 100644 --- a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py +++ b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py @@ -110,7 +110,7 @@ def _daily_omni_hub_inline_media(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_qwen3_omni_daily_omni_accuracy_bench(omni_server) -> None: _require_vllm_cli() @@ -122,7 +122,7 @@ def test_qwen3_omni_daily_omni_accuracy_bench(omni_server) -> None: assert _acc_bench.run_acc_benchmark(ns) == 0 -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_qwen3_omni_seed_tts_wer_bench(omni_server) -> None: _require_vllm_cli() diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py index c0c468315..cb6ffe6c9 100644 --- a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py +++ b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -44,7 +44,7 @@ SSIM_THRESHOLDS = { } pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.skipif( not GOLDEN_BASE_URL, @@ -169,6 +169,7 @@ def _validate_metadata(payload: Any, manifest: dict[str, Any], variant: str, tas def _download_case(variant: str, task: str, output_dir: Path, filenames: set[str]) -> dict[str, Any]: output_dir.mkdir(parents=True, exist_ok=True) + assert GOLDEN_BASE_URL is not None base_url = f"{GOLDEN_BASE_URL.rstrip('/')}/{variant}/{task}" response = requests.get(f"{base_url}/manifest.json", timeout=60) response.raise_for_status() diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py index 8dd35eca2..ec7d67027 100644 --- a/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py +++ b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -11,7 +11,7 @@ import torch from tests.helpers.mark import hardware_marks pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.benchmark, ] diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py index a88aec7e2..aabe166d2 100644 --- a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py +++ b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -7,7 +7,7 @@ import gc import os from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, TypedDict import pytest import torch @@ -16,7 +16,7 @@ import torch.nn.functional as F from tests.helpers.mark import hardware_marks pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.benchmark, ] @@ -30,7 +30,23 @@ HEIGHT = 192 WIDTH = 320 NUM_FRAMES = 9 -VARIANTS = { + +class _VariantConfig(TypedDict): + model: str + revision: str + model_env: str + max_abs: float + relative_l2: float + cosine: float + cfg_max_abs: float + cfg_relative_l2: float + cfg_cosine: float + step_max_abs: float + step_relative_l2: float + step_cosine: float + + +VARIANTS: dict[str, _VariantConfig] = { "480p": { "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", diff --git a/tests/e2e/accuracy/test_cosmos3_similarity.py b/tests/e2e/accuracy/test_cosmos3_similarity.py index efa095fb1..a7c9a9ece 100644 --- a/tests/e2e/accuracy/test_cosmos3_similarity.py +++ b/tests/e2e/accuracy/test_cosmos3_similarity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -35,6 +35,7 @@ def _model_name() -> str: pytest.skip(f"Set {MODEL_ENV_VAR} to run Cosmos3 full-model smoke tests.") if not torch.cuda.is_available(): pytest.skip("Cosmos3 full-model smoke tests require CUDA.") + assert model return model @@ -56,7 +57,7 @@ def _image_data_url(image: Image.Image) -> str: @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_cosmos3_t2i_serving_smoke(accuracy_artifact_root: Path) -> None: output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) with OmniServer(_model_name(), _server_args(), use_omni=True) as server: @@ -97,7 +98,7 @@ def test_cosmos3_t2i_serving_smoke(accuracy_artifact_root: Path) -> None: ], ) @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_cosmos3_video_serving_smoke( accuracy_artifact_root: Path, name: str, diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index b12e9c1e7..a582666d4 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py @@ -29,6 +29,7 @@ from tests.e2e.accuracy.helpers import ( assert_similarity, env_to_apply_ftfy_mock_in_subproc, model_output_dir, + resolve_device_threshold, ) from tests.e2e.accuracy.helpers import ( parse_psnr_score as _parse_psnr_score, @@ -49,16 +50,25 @@ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter imp pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] -def _set_matched_attention_backend(pipe: DiffusionPipeline) -> None: - """Walk the same backend chain the omni server's diffusers adapter walks. +def _flash_attempt_backends() -> list[str]: + """Backends in ``CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS`` that can actually launch. - The server side (--diffusion-load-format diffusers) resolves its attention - backend through the adapter's preference chain, skipping backends that are - unavailable on the image (e.g. the FA3 hub kernel has no build variant for - the image's torch — build 2953). The reference run must resolve to the - same backend or the similarity comparison measures kernel differences. + Those names are Hopper FA2/FA3. ``set_attention_backend`` only checks import, + so a newer GPU still "selects" them and then dies with "no kernel image". + Do not probe by launching: a missing cubin can leave the CUDA context dead. + Capability >= 10 is outside that wheel, including future SKUs in that family. """ - for backend in CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS: + if not torch.cuda.is_available(): + return [] + major, _ = torch.cuda.get_device_capability() + if major >= 10: + return [] + return list(CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS) + + +def _set_matched_attention_backend(pipe: DiffusionPipeline) -> None: + """Pick a Diffusers attention backend the reference run can actually execute.""" + for backend in _flash_attempt_backends(): try: pipe.transformer.set_attention_backend(backend) return @@ -297,7 +307,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize("model_id", ["Qwen/Qwen-Image"]) def test_diffusers_backend_t2i_matches_diffusers(model_id: str, accuracy_artifact_root: Path) -> None: output_dir = model_output_dir(accuracy_artifact_root, model_id + "-diffusers-backend") @@ -336,7 +346,7 @@ def test_diffusers_backend_t2i_matches_diffusers(model_id: str, accuracy_artifac @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize( "model_id", [ @@ -363,7 +373,11 @@ def test_diffusers_backend_i2v_matches_diffusers( model=model_id, output_path=diffusers_path, conditioning_image=resized_image ) diffusers_latency = diffusers_latency * 1000 - latency_threshold_factor = 0.3 + # H100 keeps the historical 30% slack. B200 measured ~35.8% (6008 vs 4425 ms). + gpu_key, latency_threshold_factor = resolve_device_threshold( + {"H100": 0.3, "B200": 0.36}, + label="latency threshold factor", + ) latency_threshold = diffusers_latency * (1 + latency_threshold_factor) ssim_output = _run_ffmpeg_similarity("ssim", vllm_path, diffusers_path) @@ -372,7 +386,8 @@ def test_diffusers_backend_i2v_matches_diffusers( psnr_score = _parse_psnr_score(psnr_output) print(f"{model_id} latency metrics:") print( - f" Latency={vllm_latency:.2f}ms, threshold<={latency_threshold:.2f}ms, diffusers latency={diffusers_latency:.2f}ms, lower is better" + f" Latency={vllm_latency:.2f}ms, threshold<={latency_threshold:.2f}ms, " + f"diffusers latency={diffusers_latency:.2f}ms, slack={latency_threshold_factor:.0%} ({gpu_key}), lower is better" ) print(f"{model_id} similarity metrics:") print(f" SSIM: value={ssim_score:.6f}, threshold>={VIDEO_SSIM_THRESHOLD:.6f}, range=[-1, 1], higher is better") diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 3e37f22b6..4258a470e 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -12,6 +12,7 @@ import tempfile import time from dataclasses import dataclass from pathlib import Path +from typing import Any, cast import pytest import requests @@ -288,7 +289,7 @@ def _quant_accuracy_cases() -> list[pytest.ParameterSet]: def _make_config(enable_kv_reuse: bool, path: Path) -> None: - config = copy.deepcopy(_DEPLOY_CONFIG) + config = cast(dict[str, Any], copy.deepcopy(_DEPLOY_CONFIG)) config["stages"][0]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse} config["stages"][1]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse} path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False)) @@ -303,7 +304,7 @@ def _quant_tensor_parallel_size() -> int: def _make_quant_dit_config(path: Path) -> None: - config = copy.deepcopy(_QUANT_DIT_CONFIG) + config = cast(dict[str, Any], copy.deepcopy(_QUANT_DIT_CONFIG)) config["stages"][0]["devices"] = _quant_devices() config["stages"][0]["parallel_config"]["tensor_parallel_size"] = _quant_tensor_parallel_size() path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False)) @@ -448,7 +449,7 @@ 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") - from tabulate import tabulate + from tabulate import tabulate # type: ignore[import-untyped] output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-online") @@ -556,7 +557,7 @@ def _run_dit_model( torch.accelerator.empty_cache() -@hardware_test(res={"cuda": "H100"}, num_cards=8) +@hardware_test(res={"cuda": "H200"}, num_cards=8) @pytest.mark.skipif( 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)", @@ -564,7 +565,7 @@ def _run_dit_model( 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 + from tabulate import tabulate # type: ignore[import-untyped] """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 10 dB.""" output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-offline-kv-reuse") diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index 71dd6f308..bc656723b 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py @@ -299,7 +299,7 @@ def _assert_against_baseline(image: Image.Image, label: str) -> None: @pytest.mark.full_model -@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=4) def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None: model = _model_name() output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME) @@ -314,7 +314,7 @@ def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> N @pytest.mark.full_model -@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=4) def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None: model = _model_name() output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index 45fd20934..fb2f8cbc8 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -233,7 +233,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: model = _model_name() output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) @@ -253,7 +253,7 @@ def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: @pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_qwen_image_2512_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None: model = _model_2512_name() output_dir = model_output_dir(accuracy_artifact_root, MODEL_2512_ID) diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index 771ced5bc..419bf867d 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py @@ -195,7 +195,7 @@ def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]: @pytest.mark.benchmark @pytest.mark.full_model -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_wan22_i2v_diffusers_offline_generates_video( wan22_i2v_image_source: str | None, ) -> None: @@ -217,7 +217,7 @@ def test_wan22_i2v_diffusers_offline_generates_video( @pytest.mark.benchmark @pytest.mark.full_model -@hardware_test(res={"cuda": "H100"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2) @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_wan22_i2v_online_serving_generates_video( omni_server, @@ -244,7 +244,7 @@ def test_wan22_i2v_online_serving_generates_video( @pytest.mark.benchmark @pytest.mark.full_model -@hardware_test(res={"cuda": "H100"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2) def test_wan22_i2v_serving_matches_diffusers_video_similarity( wan22_i2v_image_source: str | None, ) -> None: diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py index a11c9b7e4..722d42a26 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Regression tests for AsyncOmni collective_rpc in inline diffusion mode. diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py index 12ad4f1b3..56a960121 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for AsyncOmni Qwen-Image generation with trajectory_* fields. diff --git a/tests/e2e/features/custom_pipeline/test_worker_extension.py b/tests/e2e/features/custom_pipeline/test_worker_extension.py index da7b0ac17..fff39cd1e 100644 --- a/tests/e2e/features/custom_pipeline/test_worker_extension.py +++ b/tests/e2e/features/custom_pipeline/test_worker_extension.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations diff --git a/tests/e2e/offline_inference/test_device_postprocess_equivalence.py b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py index a7fbcafd9..019d90a65 100644 --- a/tests/e2e/offline_inference/test_device_postprocess_equivalence.py +++ b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py @@ -83,7 +83,7 @@ def _as_uint8_frames(video: np.ndarray) -> np.ndarray: _MAX_DIFFERING_FRACTION = 0.35 -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_device_postprocess_produces_equivalent_frames() -> None: float_path = _generate(False) device_path = _generate(True) diff --git a/tests/e2e/offline_inference/test_dots_tts_expansion.py b/tests/e2e/offline_inference/test_dots_tts_expansion.py index 1c9362a23..4d13f03fd 100644 --- a/tests/e2e/offline_inference/test_dots_tts_expansion.py +++ b/tests/e2e/offline_inference/test_dots_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for dots.tts native AR offline inference.""" from collections.abc import Mapping diff --git a/tests/e2e/offline_inference/test_gepard_tts.py b/tests/e2e/offline_inference/test_gepard_tts.py index 2dec32ffc..e1b68eb5c 100644 --- a/tests/e2e/offline_inference/test_gepard_tts.py +++ b/tests/e2e/offline_inference/test_gepard_tts.py @@ -246,7 +246,7 @@ def _transcribe_for(wav: torch.Tensor, tmp_dir: str, name: str, accepts) -> str: pytest.param( "The quick brown fox jumps over the lazy dog.", None, - marks=pytest.mark.full_model, + marks=pytest.mark.slow, id="second_text", ), ], @@ -274,7 +274,7 @@ def test_gepard_offline_zero_shot(omni_runner, run_level: str, text: str, keywor assert keyword in transcript, f"expected {keyword!r} in transcript, got {transcript!r}" -@pytest.mark.full_model +@pytest.mark.slow @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_gepard_offline_long_text_skips_repetition(omni_runner, run_level: str) -> None: """The other branch of the prompt layout, end to end. diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py index 7e7c63900..925833216 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py @@ -57,7 +57,7 @@ def omni() -> Generator[Omni, None, None]: yield runner.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) def test_i2t_generates_text(omni: Omni) -> None: """Verify I2T output's first 20 token IDs match the HF greedy baseline.""" # Solid-color image keeps the input self-contained and reproducible. diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py index 88e5979dc..bf1d2ca23 100644 --- a/tests/e2e/offline_inference/test_lingbot_world_v2.py +++ b/tests/e2e/offline_inference/test_lingbot_world_v2.py @@ -47,7 +47,7 @@ pytestmark = [ ] -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None: """Load the real checkpoint and generate one causal block on one GPU.""" diff --git a/tests/e2e/offline_inference/test_longcat_video_avatar.py b/tests/e2e/offline_inference/test_longcat_video_avatar.py index 404ad9485..59f9ae4e4 100644 --- a/tests/e2e/offline_inference/test_longcat_video_avatar.py +++ b/tests/e2e/offline_inference/test_longcat_video_avatar.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections.abc import Generator +from typing import Any import numpy as np import pytest @@ -105,15 +106,17 @@ def _generate_avatar_frames( num_frames: int = 5, sampling_extra_args: dict | None = None, ): + extra_info: dict[str, Any] = dict(additional_information or {}) + extra_info.update({"stage": stage, "resolution": "480p"}) + multi_modal_data: dict[str, Any] = {"audio": audio_path} + if image_path is not None: + multi_modal_data["image"] = image_path prompt = { "prompt": PROMPT, "negative_prompt": NEGATIVE_PROMPT, - "multi_modal_data": {"audio": audio_path}, - "additional_information": additional_information or {}, + "multi_modal_data": multi_modal_data, + "additional_information": extra_info, } - prompt["additional_information"].update({"stage": stage, "resolution": "480p"}) - if image_path is not None: - prompt["multi_modal_data"]["image"] = image_path outputs = omni_runner.generate( [prompt], diff --git a/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py index e7b4d061e..3d370ec8b 100644 --- a/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E test for Qwen-Image AutoRound W4A16 quantized inference. Verifies that the W4A16 quantized Qwen-Image checkpoint loads end-to-end @@ -47,8 +47,9 @@ def _first_request_images(outputs) -> list: return req_out.images +@pytest.mark.full_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_qwen_image_autoround_w4a16_load(): """Load the W4A16 quantized Qwen-Image model and run a minimal generation. @@ -68,8 +69,9 @@ def test_qwen_image_autoround_w4a16_load(): assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)" +@pytest.mark.full_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": ["L4", "B200"]}) def test_qwen_image_autoround_w4a16_generate(): """Full generation: 512×512, 20 steps, CFG=5.0. diff --git a/tests/e2e/offline_inference/test_voxcpm2_tts.py b/tests/e2e/offline_inference/test_voxcpm2_tts.py index af55c51aa..0b6ff9940 100644 --- a/tests/e2e/offline_inference/test_voxcpm2_tts.py +++ b/tests/e2e/offline_inference/test_voxcpm2_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """E2E test for VoxCPM2 native AR offline inference.""" import os diff --git a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py index 4ac51a344..9594a0f99 100644 --- a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py @@ -156,7 +156,7 @@ def _generate_t2v_video(offline_client, prompt: str = "A cat sitting on a table" # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True) def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, offline_client): """Load the W4A16 quantized Wan2.2 I2V model and verify it produces a valid video.""" @@ -180,7 +180,7 @@ def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, offline_client): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True) def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, offline_client): """Load the W4A16 quantized Wan2.2 T2V model and verify it produces a valid video.""" @@ -202,7 +202,7 @@ def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, offline_client): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True) def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, offline_client): """Measure peak GPU memory of W4A16 quantized I2V model.""" @@ -218,7 +218,7 @@ def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, offline_client): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_runner", baseline_i2v_params, indirect=True) def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, offline_client): """Measure peak GPU memory of BF16 baseline I2V model.""" @@ -234,7 +234,7 @@ def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, offline_client): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) def test_wan22_i2v_autoround_w4a16_memory_savings(): """Assert quantized I2V model uses meaningfully less memory than BF16 baseline.""" quant_peak = _memory_results["quant_i2v"] @@ -259,7 +259,7 @@ def test_wan22_i2v_autoround_w4a16_memory_savings(): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True) def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, offline_client): """Measure peak GPU memory of W4A16 quantized T2V model.""" @@ -275,7 +275,7 @@ def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, offline_client): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_runner", baseline_t2v_params, indirect=True) def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, offline_client): """Measure peak GPU memory of BF16 baseline T2V model.""" @@ -291,7 +291,7 @@ def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, offline_client): # ------------------------------------------------------------------ -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) def test_wan22_t2v_autoround_w4a16_memory_savings(): """Assert quantized T2V model uses meaningfully less memory than BF16 baseline.""" quant_peak = _memory_results["quant_t2v"] diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py index 0bc3f449f..1d6df3317 100644 --- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py +++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py @@ -33,19 +33,13 @@ NEGATIVE_PROMPT = "low quality, blurry, distorted" MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" -SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) -PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) - -# Only the most basic single-card deployment row is cheap enough for PR (L2) / merge (L3) -# CI; the CacheDiT / parallel feature combinations run nightly (L4). -_CORE_MARKS = [pytest.mark.core_model, pytest.mark.advanced_model] -_FULL_MARKS = [pytest.mark.full_model] - def _get_diffusion_feature_cases(model: str): """Return diffusion feature cases for HunyuanVideo-1.5. Designed for 2x H100 environment per issue #1832. + Only the CPU-offload row is cheap enough for PR (L2) / merge (L3); + CacheDiT / parallel combinations run nightly (L4). """ return [ # (1 GPU) CPU offload @@ -57,7 +51,11 @@ def _get_diffusion_feature_cases(model: str): ], ), id="single_card_cpu_offload", - marks=SINGLE_CARD_MARKS + _CORE_MARKS, + marks=[ + *hardware_marks(res={"cuda": "H100"}), + pytest.mark.core_model, + pytest.mark.advanced_model, + ], ), # (1 GPU) CacheDiT + Layerwise CPU offloading pytest.param( @@ -70,7 +68,10 @@ def _get_diffusion_feature_cases(model: str): ], ), id="single_card_cachedit_layerwise", - marks=SINGLE_CARD_MARKS + _FULL_MARKS, + marks=[ + *hardware_marks(res={"cuda": ["H100", "B200"]}), + pytest.mark.full_model, + ], ), # (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2 pytest.param( @@ -87,7 +88,10 @@ def _get_diffusion_feature_cases(model: str): ], ), id="parallel_cachedit_tp2_vae2", - marks=PARALLEL_MARKS + _FULL_MARKS, + marks=[ + *hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2), + pytest.mark.full_model, + ], ), ] diff --git a/tests/e2e/online_serving/test_lingbot_video.py b/tests/e2e/online_serving/test_lingbot_video.py index 6ada319f5..e6f275bb9 100644 --- a/tests/e2e/online_serving/test_lingbot_video.py +++ b/tests/e2e/online_serving/test_lingbot_video.py @@ -20,7 +20,7 @@ PROMPT = "a robotic arm picks up a red block" NEGATIVE_PROMPT = "low quality, blurry, watermark, text" DEFAULT_SAMPLING_PARAMS = '{"0":{"num_frames":81,"num_inference_steps":40,"guidance_scale":6.0}}' -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}) def _get_diffusion_feature_cases(model: str): diff --git a/tests/e2e/online_serving/test_lingbot_video_expansion.py b/tests/e2e/online_serving/test_lingbot_video_expansion.py index fbf2bd5fa..4a411ddcf 100644 --- a/tests/e2e/online_serving/test_lingbot_video_expansion.py +++ b/tests/e2e/online_serving/test_lingbot_video_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``. @@ -26,7 +26,7 @@ MODEL = "robbyant/lingbot-video-dense-1.3b" PROMPT = "a robotic arm picks up a red block" NEGATIVE_PROMPT = "low quality, blurry, watermark, text" -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}) def _get_diffusion_feature_cases(model: str): diff --git a/tests/e2e/online_serving/test_lingbot_video_moe.py b/tests/e2e/online_serving/test_lingbot_video_moe.py index a75fa78e0..cc56161c9 100644 --- a/tests/e2e/online_serving/test_lingbot_video_moe.py +++ b/tests/e2e/online_serving/test_lingbot_video_moe.py @@ -20,7 +20,7 @@ PROMPT = "a robotic arm picks up a red block" NEGATIVE_PROMPT = "low quality, blurry, watermark, text" SMOKE_DEFAULT_SAMPLING_PARAMS = '{"0":{"num_frames":81,"num_inference_steps":2,"guidance_scale":3.0}}' -SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}) def _get_server_cases(model: str): diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py index c3c6fd90a..628040001 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -51,7 +51,7 @@ def _cached_server_vad_artifact() -> None: ) -@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) def test_duplex_admission_and_expiry_reaper(omni_server, tmp_path: Path) -> None: args = multi_session_args( @@ -79,7 +79,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, tmp_path: Path) -> None # still talking, which only holds when the duplex pipeline sustains real-time # throughput. The current NPU stack runs several times slower than real time, # so it never reaches a mid-stream decision point. -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) def test_duplex_soft_interrupt(omni_server, tmp_path: Path) -> None: input_wav = validated_soft_interrupt_wav() @@ -112,7 +112,7 @@ def test_duplex_soft_interrupt(omni_server, tmp_path: Path) -> None: assert result["followup_response_transcript_ok"] is True -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) def test_duplex_server_vad_hard_interrupt(_cached_server_vad_artifact, omni_server) -> None: result = asyncio.run( diff --git a/tests/e2e/online_serving/test_pi0_expansion.py b/tests/e2e/online_serving/test_pi0_expansion.py index 3a9c3b177..6a16edc71 100644 --- a/tests/e2e/online_serving/test_pi0_expansion.py +++ b/tests/e2e/online_serving/test_pi0_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving. @@ -58,7 +58,7 @@ test_params = [ @pytest.mark.full_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": ["H100", "B200"]}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_pi0_openpi_online(omni_server): try: diff --git a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py index 1a27a5449..a2d64de22 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py +++ b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py @@ -91,7 +91,7 @@ def _assert_audio_outputs(responses: list[OmniResponse], *, expect_text: bool) - @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_only_batch_uses_multi_replica_talker(omni_server, online_client) -> None: request_config = { @@ -108,7 +108,7 @@ def test_text_only_batch_uses_multi_replica_talker(omni_server, online_client) - @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, online_client) -> None: request_config = { @@ -125,7 +125,7 @@ def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, onli @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, online_client) -> None: request_config = { diff --git a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py index e126877db..abf80fb85 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py @@ -69,7 +69,7 @@ tts_no_async_chunk_server_params = [ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0 -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize( "omni_server", tts_async_chunk_server_params + tts_no_async_chunk_server_params, @@ -99,7 +99,7 @@ def test_voice_clone_streaming_001(omni_server, online_client) -> None: online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize( "omni_server", tts_async_chunk_server_params + tts_no_async_chunk_server_params, @@ -127,7 +127,7 @@ def test_response_format_001(omni_server, online_client) -> None: online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, online_client) -> None: """Regression for #5049: x-vector-only then ICL on the same ``ref_audio``. diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py index 15198870e..e292ea4e8 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py @@ -60,7 +60,7 @@ tts_server_params = [ ] -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_001(omni_server, online_client) -> None: """ @@ -92,7 +92,7 @@ def test_voice_001(omni_server, online_client) -> None: raise -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_002(omni_server, online_client) -> None: """ @@ -115,7 +115,7 @@ def test_voice_002(omni_server, online_client) -> None: online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_003(omni_server, online_client) -> None: """ @@ -138,7 +138,7 @@ def test_voice_003(omni_server, online_client) -> None: online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_sample_rate_001(omni_server, online_client) -> None: """The speech API resamples Qwen3-TTS output from native 24 kHz to 8 kHz.""" @@ -155,7 +155,7 @@ def test_sample_rate_001(omni_server, online_client) -> None: online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_language_001(omni_server, online_client) -> None: """ diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index 890b077d2..73e55789e 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -21,8 +21,8 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] T2I_PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." NEGATIVE_PROMPT = "blurry, low quality" -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) -PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}) +PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2) MODEL_IMAGE = "Qwen/Qwen-Image" MODEL_2512 = "Qwen/Qwen-Image-2512" diff --git a/tests/e2e/online_serving/test_sana_wm.py b/tests/e2e/online_serving/test_sana_wm.py index 883e4c367..38277f08e 100644 --- a/tests/e2e/online_serving/test_sana_wm.py +++ b/tests/e2e/online_serving/test_sana_wm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``). diff --git a/tests/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py index 322cde3cd..8734ddebb 100644 --- a/tests/e2e/online_serving/test_wan22_expansion.py +++ b/tests/e2e/online_serving/test_wan22_expansion.py @@ -30,8 +30,8 @@ PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight NEGATIVE_PROMPT = "low quality, blurry, distorted face, extra limbs, bad anatomy, watermark, logo, text, ugly, deformed, mutated, jpeg artifacts" # CUDA marks (original matrix, unchanged) -CUDA_SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) -CUDA_PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +CUDA_SINGLE_CARD_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}) +CUDA_PARALLEL_MARKS = hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2) # NPU marks NPU_TWO_CARD_MARKS = hardware_marks(res={"npu": "A2"}, num_cards=2) diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py index c37a947f4..680cc3910 100644 --- a/tests/e2e/online_serving/test_zimage_expansion.py +++ b/tests/e2e/online_serving/test_zimage_expansion.py @@ -19,7 +19,7 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard." diff --git a/tests/examples/offline_inference/test_image_to_video.py b/tests/examples/offline_inference/test_image_to_video.py index b53648cb3..182187298 100644 --- a/tests/examples/offline_inference/test_image_to_video.py +++ b/tests/examples/offline_inference/test_image_to_video.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Offline inference tests: image-to-video. See examples/offline_inference/image_to_video/README.md @@ -17,7 +20,7 @@ pytestmark = [ pytest.mark.usefixtures("clean_gpu_memory_between_tests"), pytest.mark.full_model, pytest.mark.example, - *hardware_marks(res={"cuda": "H100"}), + *hardware_marks(res={"cuda": ["H100", "B200"]}), ] I2V_SCRIPT = EXAMPLES / "offline_inference" / "image_to_video" / "image_to_video.py" diff --git a/tests/examples/offline_inference/test_internvla_a1.py b/tests/examples/offline_inference/test_internvla_a1.py index 791d380f0..f8ea50689 100644 --- a/tests/examples/offline_inference/test_internvla_a1.py +++ b/tests/examples/offline_inference/test_internvla_a1.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Offline E2E test for InternVLA-A1 open-loop inference.""" @@ -26,6 +26,7 @@ def _required_env(name: str) -> str: value = os.getenv(name) if not value: pytest.skip(f"{name} is not set") + assert value return value diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index 5ee2a2fae..6289dcd33 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -18,7 +18,7 @@ pytestmark = [ pytest.mark.usefixtures("clean_gpu_memory_between_tests"), pytest.mark.full_model, pytest.mark.example, - *hardware_marks(res={"cuda": "H100"}), + *hardware_marks(res={"cuda": ["H100", "B200"]}), ] T2I_SCRIPT = EXAMPLES / "offline_inference" / "text_to_image" / "text_to_image.py" diff --git a/tests/examples/offline_inference/test_vace_video_generation.py b/tests/examples/offline_inference/test_vace_video_generation.py index e4f6ca93d..fd23f0734 100644 --- a/tests/examples/offline_inference/test_vace_video_generation.py +++ b/tests/examples/offline_inference/test_vace_video_generation.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """VACE coverage through the shared text-to-video and image-to-video examples.""" from pathlib import Path @@ -14,7 +17,7 @@ pytestmark = [ pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.example, - *hardware_marks(res={"cuda": "H100"}), + *hardware_marks(res={"cuda": ["H100", "B200"]}), ] TEXT_TO_VIDEO_README = EXAMPLES / "offline_inference" / "text_to_video" / "text_to_video.md" diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index 0d0ec6f46..d56b6f958 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Online serving tests: Qwen2.5-Omni-7B. See examples/online_serving/qwen2_5_omni/README.md @@ -42,7 +45,7 @@ test_params = [ common_args = ["python", os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py")] -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: command = common_args + [ @@ -77,7 +80,7 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: command = common_args + [ @@ -112,7 +115,7 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: command = [ @@ -137,7 +140,7 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: command = common_args + [ @@ -165,7 +168,7 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: command = common_args + [ @@ -192,7 +195,7 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: command = common_args + [ @@ -228,7 +231,7 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": ["L4", "B200"], "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: command = common_args + [ diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index 711b7236a..23e2ba184 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Online serving tests: Qwen3-Omni. See examples/online_serving/qwen3_omni/README.md @@ -41,7 +44,7 @@ test_params = [ common_args = ["python", os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py")] -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: command = common_args + [ @@ -69,7 +72,7 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: command = common_args + [ @@ -100,7 +103,7 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: command = ["bash", os.path.join(example_dir, "qwen3_omni/run_curl_multimodal_generation.sh"), "use_image"] @@ -115,7 +118,7 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: command = common_args + [ @@ -137,7 +140,7 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: command = common_args + [ @@ -159,7 +162,7 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: command = common_args + [ @@ -189,7 +192,7 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: command = common_args + [ diff --git a/tests/examples/online_serving/test_text_to_image.py b/tests/examples/online_serving/test_text_to_image.py index 25970502b..f251cd16c 100644 --- a/tests/examples/online_serving/test_text_to_image.py +++ b/tests/examples/online_serving/test_text_to_image.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Online serving tests: text-to-image. See examples/online_serving/text_to_image/README.md @@ -18,7 +21,7 @@ from tests.helpers.assertions import assert_image_valid from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams -pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] +pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": ["H100", "B200"]})] T2I_ONLINE_CLIENT = EXAMPLES / "online_serving" / "text_to_image" / "openai_chat_client.py" EXAMPLE_OUTPUT_SUBFOLDER = "example_online_t2i" diff --git a/tests/model_tests/diffusion/case_filtering.py b/tests/model_tests/diffusion/case_filtering.py index da1aaa552..9f53edce0 100644 --- a/tests/model_tests/diffusion/case_filtering.py +++ b/tests/model_tests/diffusion/case_filtering.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py @@ -45,7 +48,6 @@ def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None f"but the max CI device count is {MAX_CI_DEVICES}. " ) - marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices)) # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case # directly here. This should probably be handled in a more common way later on. assert current_omni_platform is not None and current_omni_platform.device_count is not None @@ -53,8 +55,10 @@ def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None if current_omni_platform.is_cuda() and device_count < required_devices: marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}")) if required_devices > 1: + marks.extend(hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=required_devices)) marks.append(pytest.mark.full_model) else: + marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=1)) marks.append(pytest.mark.core_model) return marks From 2cbfd3e93b460462031a0339d9f1e8733e39e1e6 Mon Sep 17 00:00:00 2001 From: jiahengxie <13716265422@139.com> Date: Mon, 14 Sep 2026 15:47:39 +0800 Subject: [PATCH 055/209] [Performance][MammothModa2] Enable and validate FP8 KV cache for AR stage (#7436) Signed-off-by: JH <13716265422@139.com> --- recipes/MammothModa2/MammothModa2.md | 45 ++++++++++++++++++++++ tests/engine/test_stage_engine_args.py | 38 ++++++++++++++++++ vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml | 7 ++++ 3 files changed, 90 insertions(+) create mode 100644 vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 6dc5f12e9..02f8e30bf 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/recipes/MammothModa2/MammothModa2.md @@ -49,6 +49,51 @@ cache — see the note under *1x L40S 48GB*. ## GPU +### Optional FP8 AR KV cache + +For CUDA deployments, `mammoth_moda2_fp8_kv.yaml` is an opt-in preset that +stores the Stage 0 AR KV cache as FP8 E4M3. Stage 1 remains on +`kv_cache_dtype=auto`; its DiT execution is unaffected. This setting quantizes +only the autoregressive KV cache. It is neither FP8 weight/activation +quantization nor vLLM-Omni diffusion KV-cache quantization. + +Use the preset in place of the default deploy config: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ./MammothModa2-Preview \ + --deploy-config vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml \ + --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ + --height 1024 \ + --width 1024 \ + --seed 42 \ + --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ + --output mammoth_t2i.png +``` + +The preset was validated on one NVIDIA H800 80GB with CUDA and +FlashAttention 3. The native and FP8 runs used the same model, code revision, +and downstream configuration. + +| Metric | Native BF16 (`kv_cache_dtype=auto`) | FP8 E4M3 (`fp8_e4m3`) | Change | +| --- | ---: | ---: | ---: | +| KV cache memory | 15.59 GiB | 15.55 GiB | -0.04 GiB | +| GPU KV cache size | 145,904 tokens | 291,232 tokens | +99.6% | +| Maximum concurrency at 8,192 tokens | 17.81x | 35.55x | +99.6% | +| Steady-state AR median | 71.191 s | 79.816 s | +12.1% | +| Steady-state end-to-end median | 83.840 s | 92.473 s | +10.3% | +| Steady-state DiT median | 12.574 s | 12.561 s | effectively unchanged | + +The FP8 run used `kv_cache_dtype=fp8_e4m3` only for Stage 0; Stage 1 used +`kv_cache_dtype=auto`. The nearly unchanged reserved cache memory holds almost +twice as many tokens because FP8 reduces the bytes per cached token. This is a +capacity/concurrency tradeoff: the measured AR and end-to-end latencies were +higher than the native-BF16 baseline. + +1024x1024 fixed-seed smoke test completed successfully with no obvious visual +failure. FP8 is lossy, so numerical or image-quality equivalence with BF16 is +not implied. + ### 1x L40S 48GB > **48 GB config adjustment:** the committed diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index f20538284..7ac5b6b69 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -366,6 +366,44 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults(): assert inherited_defaults.isdisjoint(engine_args) +def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch): + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) + + pipeline = OMNI_PIPELINES["mammoth_moda2"] + deploy = load_deploy_config(_DEPLOY_DIR / "mammoth_moda2_fp8_kv.yaml") + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model="test-model", + ) + + assert deploy.stages[0].engine_extras["kv_cache_dtype"] == "fp8_e4m3" + assert "kv_cache_dtype" not in deploy.stages[1].engine_extras + + legacy_args = [build_legacy_engine_args_dict(stage, "test-model") for stage in legacy_stages] + typed_args = [ + build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(stage_id), + "test-model", + ) + for stage_id in (0, 1) + ] + + ar_stage = omni_config.stage_by_id(0) + dit_stage = omni_config.stage_by_id(1) + + assert ar_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR + assert legacy_args[0]["kv_cache_dtype"] == "fp8_e4m3" + assert ar_stage.cache_config.cache_dtype == "fp8_e4m3" + assert "cache_dtype" in ar_stage.cache_config._omni_explicit_fields + assert typed_args[0]["kv_cache_dtype"] == "fp8_e4m3" + + assert "kv_cache_dtype" not in legacy_args[1] + assert dit_stage.cache_config.cache_dtype == "auto" + assert "cache_dtype" not in dit_stage.cache_config._omni_explicit_fields + assert "kv_cache_dtype" not in typed_args[1] + + def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary(): stage_config = VllmOmniARStageConfig( stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), diff --git a/vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml b/vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml new file mode 100644 index 000000000..8d99226b5 --- /dev/null +++ b/vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml @@ -0,0 +1,7 @@ +# MammothModa2 AR→DiT text-to-image with opt-in FP8 KV cache for the AR stage. + +base_config: mammoth_moda2.yaml + +stages: + - stage_id: 0 + kv_cache_dtype: fp8_e4m3 From 446c2b5474dabffccf9a0034836eeab77c62b4b8 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:29:50 +0800 Subject: [PATCH 056/209] [Bugfix] Forward Qwen3-Omni streaming video sampling parameters (#7325) Signed-off-by: psv666 <2693925048@qq.com> --- docs/serving/video_stream_api.md | 16 ++- .../openai_api/test_serving_video_stream.py | 128 ++++++++++++++++++ .../entrypoints/openai/video_stream_base.py | 10 +- 3 files changed, 150 insertions(+), 4 deletions(-) diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index 78575e0fc..101289fc0 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -60,10 +60,24 @@ WebSocket /v1/video/chat/stream | `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. | | `system_prompt` | string or null | null | Optional custom system prompt. | | `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. | -| `sampling_params_list` | list or null | null | Optional per-stage sampling parameter overrides. | +| `sampling_params_list` | list or null | null | Optional per-stage parameter dictionaries. Each provided entry replaces that stage's deployment sampling settings. | | `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. | | `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. | +Sampling parameter bugfix: `sampling_params_list` is now forwarded to the engine; +earlier versions accepted this field but silently used deployment defaults instead. +Each provided entry constructs a fresh `SamplingParams` for that stage in pipeline +order; it is not merged with that stage's deployment defaults. Fields absent from a +provided entry use `SamplingParams` constructor defaults. Only omitted trailing +stages keep their deployment defaults. An omitted, `null`, or empty list keeps the +engine defaults. Invalid sampling parameters return an `error` when the query is +submitted. + +For example, `[{"temperature": 0.2, "max_tokens": 64}]` configures the thinker while +the talker and code2wav retain their deployment defaults. With only +`[{"temperature": 0.2}]`, the thinker's `max_tokens` and `top_p` use constructor +defaults, not its YAML values. + ### Legacy Aliases The server accepts these legacy field names and rewrites them before validation. New clients should send the canonical names above. diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index 58ea2b57f..f7267a061 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -12,11 +12,15 @@ import json import threading import weakref from typing import Any +from unittest.mock import AsyncMock, MagicMock import pytest import torch from PIL import Image +from vllm.sampling_params import RequestOutputKind, SamplingParams +from tests.helpers.serving_chat import build_serving_chat +from vllm_omni.config.stage_config import StageConfig from vllm_omni.entrypoints.openai import video_stream_base, video_stream_envs from vllm_omni.entrypoints.openai.serving_video_stream import ( QwenOmniStreamingVideoHandler, @@ -141,6 +145,130 @@ async def test_receive_config_accepts_client_legacy_aliases(): assert config.frame_filter_threshold == 0.87 +@pytest.fixture +def run_sampling_session(monkeypatch): + async def run(config_overrides): + engine = MagicMock() + engine.stage_configs = [ + StageConfig(stage_id=index, model_stage=name) + for index, name in enumerate(("thinker", "talker", "code2wav")) + ] + engine.default_sampling_params_list = [ + SamplingParams(temperature=0.4, max_tokens=64, top_p=0.85), + SamplingParams(temperature=0.7, max_tokens=96), + SamplingParams(temperature=0.8, max_tokens=128), + ] + + async def generate(**kwargs): + yield _text_result("answer") + + engine.generate.side_effect = generate + handler = QwenOmniStreamingVideoHandler( + chat_service=build_serving_chat(engine_client=engine), + engine_client=engine, + ) + monkeypatch.setattr(handler, "_preprocess_to_engine_prompt", AsyncMock(return_value={"prompt": "video"})) + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "enable_frame_filter": False, + **config_overrides, + } + ), + json.dumps({"type": "video.frame", "data": _b64(_make_jpeg())}), + json.dumps({"type": "video.query", "text": "describe"}), + json.dumps({"type": "video.done"}), + ] + ) + await handler.handle_session(ws) + return ws, engine + + return run + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_chunk_mode", ["on", "off"]) +@pytest.mark.parametrize("num_stage_overrides", [1, 3]) +async def test_video_sampling_params_reach_engine( + run_sampling_session, monkeypatch, async_chunk_mode, num_stage_overrides +): + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", async_chunk_mode) + overrides = [ + {"temperature": 0.2, "max_tokens": 1, "seed": 42}, + {"temperature": 0.6, "max_tokens": 12}, + {"temperature": 0.9, "max_tokens": 24}, + ][:num_stage_overrides] + + ws, engine = await run_sampling_session({"sampling_params_list": overrides}) + + assert not [msg for msg in ws.sent if msg["type"] == "error"] + engine.generate.assert_called_once() + params = engine.generate.call_args.kwargs["sampling_params_list"] + assert len(params) == 3 + for index, param in enumerate(params): + assert isinstance(param, SamplingParams) + default = engine.default_sampling_params_list[index] + expected = ( + overrides[index] + if index < len(overrides) + else { + "temperature": default.temperature, + "max_tokens": default.max_tokens, + } + ) + assert param.temperature == expected["temperature"] + assert param.max_tokens == expected["max_tokens"] + assert param.output_kind == RequestOutputKind.DELTA + assert param is not default + assert params[0].seed == 42 + assert all(param.output_kind == RequestOutputKind.CUMULATIVE for param in engine.default_sampling_params_list) + assert next(msg for msg in ws.sent if msg["type"] == "response.text.done")["text"] == "answer" + + +@pytest.mark.asyncio +async def test_video_sampling_params_provided_stage_uses_constructor_defaults(run_sampling_session): + ws, engine = await run_sampling_session({"sampling_params_list": [{"temperature": 0.2}]}) + + assert not [msg for msg in ws.sent if msg["type"] == "error"] + engine.generate.assert_called_once() + params = engine.generate.call_args.kwargs["sampling_params_list"] + constructor_defaults = SamplingParams() + assert params[0].temperature == 0.2 + for field in ("max_tokens", "top_p"): + assert getattr(params[0], field) == getattr(constructor_defaults, field) + assert getattr(params[0], field) != getattr(engine.default_sampling_params_list[0], field) + assert len(params) == 3 + for param, default in zip(params[1:], engine.default_sampling_params_list[1:]): + assert param.temperature == default.temperature + assert param.max_tokens == default.max_tokens + assert param.top_p == default.top_p + assert param is not default + + +@pytest.mark.asyncio +@pytest.mark.parametrize("config_overrides", [{}, {"sampling_params_list": None}, {"sampling_params_list": []}]) +async def test_video_sampling_params_omitted_preserve_engine_defaults(run_sampling_session, config_overrides): + ws, engine = await run_sampling_session(config_overrides) + + assert not [msg for msg in ws.sent if msg["type"] == "error"] + engine.generate.assert_called_once() + assert engine.generate.call_args.kwargs.get("sampling_params_list") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid_params", [{"temperature": -1}, {"max_tokens": 0}, {"unknown_sampling_option": 1}]) +async def test_video_sampling_params_invalid_do_not_reach_engine(run_sampling_session, invalid_params): + ws, engine = await run_sampling_session({"sampling_params_list": [invalid_params]}) + + engine.generate.assert_not_called() + assert any(msg["type"] == "error" for msg in ws.sent) + assert ws.sent[-1]["type"] == "session.done" + + @pytest.mark.asyncio async def test_video_frame_ack_reports_receiver_buffer_state(): ws = MockWebSocket( diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index 977a0b065..8885e9a33 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -51,6 +51,7 @@ from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilte from vllm_omni.entrypoints.openai.video_stream_context import ( text_only_message, ) +from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -660,11 +661,13 @@ class OmniStreamingVideoHandler: request_kwargs["mm_processor_kwargs"] = { "use_audio_in_video": True, } - if config.sampling_params_list: - request_kwargs["sampling_params_list"] = config.sampling_params_list - + sampling_kwargs: dict[str, Any] = {} try: chat_request = ChatCompletionRequest(**request_kwargs) + if config.sampling_params_list: + params = self._chat_service._to_sampling_params_list(config.sampling_params_list) + # Explicit params must retain the engine's streaming output mode. + sampling_kwargs["sampling_params_list"] = coerce_param_message_types(params, is_streaming=True) except Exception as e: await self._send_error(websocket, f"Failed to build request: {e}") return @@ -704,6 +707,7 @@ class OmniStreamingVideoHandler: prompt=engine_prompt, request_id=request_id, output_modalities=config.modalities, + **sampling_kwargs, ) async for output in result_gen: From 3d793481a04c9453b89508bbc67ecbc0b2d68357 Mon Sep 17 00:00:00 2001 From: BeatSeat <137616995+BeatSeat@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:11:15 +0800 Subject: [PATCH 057/209] [Bugfix][MiniCPMO45] Return bare tensor in Thinker forward to fix repetition loop (#7497) (#7517) Signed-off-by: BeatSeat --- .../test_minicpmo_4_5_omni_llm_forward.py | 123 ++++++++++++++++++ .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 6 +- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 15 +-- 3 files changed, 127 insertions(+), 17 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py new file mode 100644 index 000000000..b1cc48177 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_thinker_forward_returns_bare_hidden_states(): + """Verify that MiniCPMO45OmniLLMForConditionalGeneration.forward returns + a bare hidden_states Tensor (not a tuple), which stock vLLM GPUModelRunner + expects. + """ + model = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + torch.nn.Module.__init__(model) + model.llm = MagicMock() + mock_hidden_states = torch.randn(4, 4096) + model.llm.model = MagicMock(return_value=mock_hidden_states) + + out = model.forward( + input_ids=torch.tensor([1, 2, 3, 4]), + positions=torch.tensor([0, 1, 2, 3]), + inputs_embeds=torch.randn(4, 4096), + ) + + assert isinstance(out, torch.Tensor), f"Expected bare torch.Tensor, got {type(out)}" + assert out.shape == (4, 4096) + assert torch.equal(out, mock_hidden_states) + + +def test_thinker_forward_builds_embeddings_once_when_inputs_embeds_none(): + """Verify that when inputs_embeds is None, forward builds input embeddings + exactly once and forwards them to self.llm.model with input_ids=None.""" + model = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + torch.nn.Module.__init__(model) + model.llm = MagicMock() + mock_hidden_states = torch.randn(4, 4096) + model.llm.model = MagicMock(return_value=mock_hidden_states) + + mock_mm_embeds = [torch.randn(2, 4096)] + mock_input_embeds = torch.randn(4, 4096) + model.get_multimodal_embeddings = MagicMock(return_value=mock_mm_embeds) + model.get_input_embeddings = MagicMock(return_value=mock_input_embeds) + + input_ids = torch.tensor([1, 2, 3, 4]) + positions = torch.tensor([0, 1, 2, 3]) + + out = model.forward( + input_ids=input_ids, + positions=positions, + inputs_embeds=None, + ) + + # Assert input embeddings are built exactly once + model.get_multimodal_embeddings.assert_called_once() + model.get_input_embeddings.assert_called_once_with(input_ids, mock_mm_embeds) + + # Verify forwarded to self.llm.model with input_ids=None and inputs_embeds + model.llm.model.assert_called_once_with( + None, + positions, + None, + inputs_embeds=mock_input_embeds, + ) + + assert isinstance(out, torch.Tensor), f"Expected bare torch.Tensor, got {type(out)}" + assert torch.equal(out, mock_hidden_states) + + +@pytest.mark.parametrize( + ("thinker_shape", "expected_shape"), + [ + ((1, 4096), (1, 4096)), + ((5, 4096), (5, 4096)), + ((1, 1, 4096), (1, 4096)), + ((1, 5, 4096), (5, 4096)), + ], +) +def test_wrapper_shape_normalization(thinker_shape, expected_shape): + """Verify that MiniCPMO45OmniForConditionalGeneration normalizes 3D + hidden states of shape (1, seq_len, H) to (seq_len, H), while leaving + 2D hidden states intact (especially (1, H) for single-token decode).""" + wrapper = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + torch.nn.Module.__init__(wrapper) + wrapper.model_stage = "llm" + wrapper.thinker = torch.nn.Module() + mock_hidden = torch.randn(*thinker_shape) + wrapper.thinker.forward = MagicMock(return_value=mock_hidden) + + out = wrapper.forward( + input_ids=torch.tensor([1]), + positions=torch.tensor([0]), + ) + + assert out.text_hidden_states.shape == expected_shape, ( + f"Expected text_hidden_states shape {expected_shape}, got {out.text_hidden_states.shape}" + ) + assert out.multimodal_outputs["latent"].shape == expected_shape, ( + f"Expected latent shape {expected_shape}, got {out.multimodal_outputs['latent'].shape}" + ) + + +if __name__ == "__main__": + test_thinker_forward_returns_bare_hidden_states() + test_thinker_forward_builds_embeddings_once_when_inputs_embeds_none() + for shape, exp in [ + ((1, 4096), (1, 4096)), + ((5, 4096), (5, 4096)), + ((1, 1, 4096), (1, 4096)), + ((1, 5, 4096), (5, 4096)), + ]: + test_wrapper_shape_normalization(shape, exp) + print("ALL TESTS PASSED SUCCESSFULLY!") diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index e9ed05e75..60252e78c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -547,13 +547,13 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp ) if isinstance(thinker_output, tuple): - embeds, text_hidden_states = thinker_output + _, text_hidden_states = thinker_output else: text_hidden_states = thinker_output # Prepare hidden states for downstream stages - # Ensure correct shape: (batch_size, seq_len, hidden_dim) - if added_batch_dim: + # Ensure correct shape: (seq_len, hidden_dim) + if text_hidden_states.ndim == 3 and text_hidden_states.shape[0] == 1: text_hidden_states = text_hidden_states.squeeze(0) # Return hidden states with latent in multimodal_outputs for stage_input_processors diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 6d1c5de32..2781fcc30 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -4611,9 +4611,6 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S inputs_embeds: torch.Tensor | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: - """Forward pass through thinker model.""" - text_inputs_embeds = None - if intermediate_tensors is not None: inputs_embeds = None @@ -4621,22 +4618,12 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S elif inputs_embeds is None: multimodal_embeddings = self.get_multimodal_embeddings(**kwargs) inputs_embeds = self.get_input_embeddings(input_ids, multimodal_embeddings) - text_inputs_embeds = self.get_input_embeddings( - input_ids, - ( - [(torch.zeros_like(embeddings), "image") for embeddings in multimodal_embeddings] - if multimodal_embeddings is not None - else None - ), - ) input_ids = None - else: - text_inputs_embeds = inputs_embeds # Forward through language model hidden_states = self.llm.model(input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds) - return text_inputs_embeds, hidden_states.unsqueeze(0) if hidden_states.ndim == 2 else hidden_states + return hidden_states def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: """Compute logits from hidden states.""" From 1b6cd282f019fa8b49a8a34b4ae1b87f0ab4f0c2 Mon Sep 17 00:00:00 2001 From: Tai An Date: Mon, 14 Sep 2026 06:48:51 -0700 Subject: [PATCH 058/209] fix(qwen-image): drop the shadowed dead _get_qwen_prompt_embeds duplicate (#7461) Signed-off-by: Anai-Guo --- .../qwen_image/pipeline_qwen_image_edit.py | 45 ------------------- 1 file changed, 45 deletions(-) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 34e5733fd..a9b7ef95a 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -374,51 +374,6 @@ class QwenImageEditPipeline( return split_result - def _get_qwen_prompt_embeds( - self, - prompt: str | list[str] = None, - image: torch.Tensor | None = None, - dtype: torch.dtype | None = None, - ): - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - template = self.prompt_template_encode - drop_idx = self.prompt_template_encode_start_idx - txt = [template.format(e) for e in prompt] - - model_inputs = self.processor( - text=txt, - images=image, - padding=True, - return_tensors="pt", - ).to(self.device) - - outputs = self.text_encoder( - input_ids=model_inputs.input_ids, - attention_mask=model_inputs.attention_mask, - pixel_values=model_inputs.pixel_values, - image_grid_thw=model_inputs.image_grid_thw, - output_hidden_states=True, - ) - - hidden_states = outputs.hidden_states[-1] - split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask) - split_hidden_states = [e[drop_idx:] for e in split_hidden_states] - attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] - max_seq_len = max([e.size(0) for e in split_hidden_states]) - prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] - ) - encoder_attention_mask = torch.stack( - [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] - ) - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) - - return prompt_embeds, encoder_attention_mask - def _get_qwen_prompt_embeds( self, prompt: str | list[str] = None, From e0c4f97f326f895c7830bdbdcc8c6861b8fcb10a Mon Sep 17 00:00:00 2001 From: bcsdhjew Date: Tue, 15 Sep 2026 00:05:55 +0800 Subject: [PATCH 059/209] [Bugfix][Frontend] Preserve duplex chat audio format and duration (#7316) Signed-off-by: Nolen Liang --- docs/serving/realtime_duplex_api.md | 43 +++ .../openai/test_duplex_protocol.py | 48 ++++ .../openai/test_duplex_session_attachment.py | 84 ++++++ .../openai_api/test_audio_format.py | 96 ++++++- .../openai_api/test_duplex_handler.py | 260 +++++++++++++++++- vllm_omni/entrypoints/duplex/chat_fallback.py | 57 +++- vllm_omni/entrypoints/duplex/protocol.py | 40 ++- .../entrypoints/duplex/session_attachment.py | 10 + .../entrypoints/duplex/session_runner.py | 19 +- .../entrypoints/openai/audio_utils_mixin.py | 18 +- .../entrypoints/openai/protocol/audio.py | 14 + .../openai/protocol/chat_completion.py | 23 +- vllm_omni/entrypoints/openai/serving_chat.py | 12 +- 13 files changed, 684 insertions(+), 40 deletions(-) diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index c855c4491..288c934e8 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -387,6 +387,49 @@ Semantic divergences hidden behind shared names: the Tier 3 surface on top of the same Tier 1 events, so both kinds of client can share one deployment and one event log. +### Turn-based Chat audio metadata + +For turn-based models served through Chat fallback, non-empty audio choices +produced by the Omni Chat audio encoder include `audio_metadata`, in both +streaming and complete Chat responses. Text-only choices do not require it. + +| Field | Meaning | +| --- | --- | +| `format` | Encoding of the returned audio, such as `pcm` or `wav`; `pcm16` requests normalize to `pcm`. | +| `sample_rate_hz` | Sample rate used by the encoder, in samples per second. | +| `frame_count` | Samples per channel in this chunk, or in the complete waveform for a non-streaming choice. Not a byte count. | +| `channels` | Number of audio channels. | + +These values describe the waveform after audio transformations. Compressed +formats may add codec padding that is not included in `frame_count`. +The encoder uses a model-reported sample rate when available; Qwen3-Omni +currently omits it and uses the existing 24 kHz default. This metadata does +not independently determine a model's sample rate. + +Chat fallback requests the session's output format explicitly. Its internal +`response.output_audio.delta` carries `format`, `sample_rate_hz`, `channels`, +and `audio_duration_ms`. The Realtime endpoint exposes output `format` and +`sample_rate_hz` on `response.audio.delta`, with duration in +`metadata.audio_duration_ms`; it does not currently forward `channels`. +The duration is cumulative for the response: +integer frames are summed before converting to milliseconds and rounding down. +Non-empty audio without valid metadata, or a sample-rate change within one +response, fails that fallback response with `response_error`. + +In this Chat fallback path, `response.done.response.metadata.playback` reports: + +- `generated_ms`: cumulative source-waveform duration. +- `sent_ms`: cumulative audio accepted into the server output path, including + queued or journaled audio. A journaled event remains accepted if the + subsequent send fails; this is not proof of socket delivery. +- `played_ms`: cumulative playback reported by the client. +- `committed_ms`: playback position used for history reconciliation. + +The native model path still records audio before sending; it does not yet +share Chat fallback's output-acceptance timing. Only client playback reports +establish what was played. The existing duration-proportional text truncation +remains an approximation, not word-level audio alignment. + ### Event catalogue Every event type the Realtime route accepts or emits, with a worked example in the next section: 21 diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py index e5f3becbc..77c8f2cb0 100644 --- a/tests/entrypoints/openai/test_duplex_protocol.py +++ b/tests/entrypoints/openai/test_duplex_protocol.py @@ -462,6 +462,54 @@ def test_duplex_playback_ack_tracks_committed_cursor_separately(): assert session.playback.committed_ms == 2_000 +@pytest.mark.parametrize("sample_rate_hz,frame_count", [(22_050, 23), (24_000, 24)]) +def test_duplex_generated_audio_rounds_only_cumulative_frames(sample_rate_hz: int, frame_count: int): + session = DuplexSessionRegistry().create() + response_id = session.begin_response() + + durations = [ + session.record_generated_audio(response_id, frame_count=1, sample_rate_hz=sample_rate_hz) + for _ in range(frame_count) + ] + + assert durations[:-1] == [0] * (frame_count - 1) + assert durations[-1] == 1 + assert session.playback.generated_ms == 1 + assert session.playback.sent_ms == 0 + assert session.playback.played_ms == 0 + + +def test_duplex_generated_audio_rejects_rate_change_without_advancing(): + session = DuplexSessionRegistry().create() + response_id = session.begin_response() + session.record_generated_audio(response_id, frame_count=2205, sample_rate_hz=22_050) + + with pytest.raises(ValueError, match="sample rate changed"): + session.record_generated_audio(response_id, frame_count=2400, sample_rate_hz=24_000) + + assert session.playback.generated_ms == 100 + assert session.playback.sent_ms == 0 + + +def test_duplex_late_audio_accounting_cannot_change_new_response(): + session = DuplexSessionRegistry().create() + old_response = session.begin_response() + session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) + session.barge_in() + new_response = session.begin_response() + session.record_generated_audio(new_response, frame_count=2400, sample_rate_hz=48_000) + current_state = session.turn_state + + assert session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) is None + session.mark_audio_sent(100, response_id=old_response, text_chars=99) + + assert session.playback_for_response(old_response).sent_ms == 100 + assert session.playback.generated_ms == 50 + assert session.playback.sent_ms == 0 + assert session.assistant_audio_text_marks == () + assert session.turn_state == current_state + + def test_duplex_history_commit_uses_audio_text_alignment_marks(): registry = DuplexSessionRegistry() session = registry.create() diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index a73159064..13bd32761 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -122,6 +122,90 @@ def test_event_journal_overflow_is_explicit_and_does_not_evict_silently() -> Non journal.replay_after(0) +@pytest.mark.asyncio +@pytest.mark.parametrize("journal", [False, True]) +async def test_registry_failed_send_is_accepted_only_when_journaled(mocker, journal: bool) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + accepted = mocker.Mock() + expected_calls = [mocker.call()] if journal else [] + + async def failing_send(payload): + assert accepted.call_args_list == expected_calls + raise RuntimeError("transport lost") + + await registry.create("sid-send-failure", incarnation=0, send=failing_send, close=mocker.AsyncMock()) + + with pytest.raises(RuntimeError, match="transport lost"): + await registry.send_event( + "sid-send-failure", + {"type": "response.audio.delta", "delta": "AAAA"}, + journal=journal, + on_accepted=accepted, + ) + + assert accepted.call_args_list == expected_calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("journal", [False, True]) +async def test_registry_detached_event_is_accepted_only_when_journaled(mocker, journal: bool) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + send = mocker.AsyncMock() + accepted = mocker.Mock() + created = await registry.create("sid-detached-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + await registry.detach("sid-detached-acceptance", attachment_generation=created.attachment_generation) + + entry = await registry.send_event( + "sid-detached-acceptance", + {"type": "response.audio.delta", "delta": "AAAA"}, + journal=journal, + on_accepted=accepted, + ) + + send.assert_not_awaited() + assert (entry is not None) is journal + assert accepted.call_count == int(journal) + + +@pytest.mark.asyncio +async def test_registry_unjournaled_event_is_accepted_after_successful_send(mocker) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + accepted = mocker.Mock() + + async def send(payload): + accepted.assert_not_called() + + await registry.create("sid-live-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + + entry = await registry.send_event( + "sid-live-acceptance", + {"type": "response.audio.delta", "delta": "AAAA"}, + journal=False, + on_accepted=accepted, + ) + + assert entry is None + accepted.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_registry_overflow_does_not_accept_or_send_event(mocker) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=100) + accepted = mocker.Mock() + send = mocker.AsyncMock() + await registry.create("sid-overflow-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + + with pytest.raises(DuplexJournalOverflowError, match="byte limit"): + await registry.send_event( + "sid-overflow-acceptance", + {"type": "response.audio.delta", "delta": "A" * 200}, + on_accepted=accepted, + ) + + accepted.assert_not_called() + send.assert_not_awaited() + + @pytest.mark.asyncio async def test_registry_resume_rotates_token_replays_and_atomically_replaces_attachment() -> None: clock = _Clock() diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index 9c7eb9d89..c7fc4fbef 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -12,6 +12,8 @@ Covers: from __future__ import annotations +import base64 +import json from io import BytesIO import numpy as np @@ -27,6 +29,8 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( SUPPORTED_CHAT_AUDIO_FORMATS, CreateAudio, ) +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -65,6 +69,13 @@ class TestCreateAudio: ) response = mixin.create_audio(audio_obj) assert len(response.audio_data) > 0 + assert response.audio_metadata is not None + assert response.audio_metadata.model_dump() == { + "format": fmt, + "sample_rate_hz": 24000, + "frame_count": 24000, + "channels": 1, + } def test_wav_magic_bytes(self, mixin, audio_tensor): audio_obj = CreateAudio( @@ -113,10 +124,9 @@ class TestCreateAudio: response = mixin.create_audio(audio_obj) assert response.audio_data[:4] == b"RIFF" assert response.media_type == "audio/wav" + assert response.audio_metadata.format == "wav" def test_base64_encoding(self, mixin, audio_tensor): - import base64 - audio_obj = CreateAudio( audio_tensor=audio_tensor, sample_rate=24000, @@ -128,21 +138,33 @@ class TestCreateAudio: decoded = base64.b64decode(response.audio_data) assert decoded[:4] == b"RIFF" - def test_resamples_wav_to_requested_output_rate(self, mixin, audio_tensor): + @pytest.mark.parametrize( + "channels,speed,frame_count", + [(1, 1.0, 8000), (2, 2.0, 4000)], + ids=["mono-resampled", "stereo-resampled-faster"], + ) + def test_metadata_matches_transformed_waveform(self, mixin, audio_tensor, channels, speed, frame_count): response = mixin.create_audio( CreateAudio( - audio_tensor=audio_tensor, + audio_tensor=audio_tensor if channels == 1 else np.stack((audio_tensor, audio_tensor)), sample_rate=24000, output_sample_rate=8000, response_format="wav", - speed=1.0, + speed=speed, base64_encode=False, ) ) with soundfile.SoundFile(BytesIO(response.audio_data)) as audio_file: assert audio_file.samplerate == 8000 - assert audio_file.frames == 8000 + assert audio_file.frames == frame_count + assert audio_file.channels == channels + assert response.audio_metadata.model_dump() == { + "format": "wav", + "sample_rate_hz": audio_file.samplerate, + "frame_count": audio_file.frames, + "channels": audio_file.channels, + } class TestStreamingAudioResampler: @@ -262,3 +284,65 @@ class TestResolveAudioFormat: request = self._make_request({"format": fmt, "voice": "alloy"}) result = serving_chat._resolve_audio_format(request) assert not isinstance(result, ErrorResponse), f"Format {fmt} should be accepted" + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("fmt", ["pcm", "wav"]) +def test_chat_audio_metadata_survives_response_serialization(stream, fmt): + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo + + from vllm_omni.entrypoints.openai.protocol.chat_completion import ( + OmniChatCompletionResponse, + OmniChatCompletionStreamResponse, + ) + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + serving_chat = object.__new__(OmniOpenAIServingChat) + output = MultimodalCompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + multimodal_output=MultimodalPayload.from_dict( + {"audio": [torch.zeros(17), torch.zeros(23)], "sr": torch.tensor(22050)} + ), + finish_reason="stop", + stop_reason=None, + ) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + audio={"format": fmt, "voice": "alloy"}, + ) + choices = serving_chat._create_audio_choice( + OmniRequestOutput(outputs=[output], final_output_type="audio"), "assistant", request, stream=stream + ) + assert not isinstance(choices, ErrorResponse) + response_type = OmniChatCompletionStreamResponse if stream else OmniChatCompletionResponse + response = response_type( + id="chatcmpl-audio-metadata", + created=0, + model="test-model", + choices=choices, + usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0), + ) + + serialized = json.loads(response.model_dump_json(exclude_unset=True)) + choice = serialized["choices"][0] + frame_count = 23 if stream else 40 + assert choice["audio_metadata"] == { + "format": fmt, + "sample_rate_hz": 22050, + "frame_count": frame_count, + "channels": 1, + } + audio = choice["delta"]["content"] if stream else choice["message"]["audio"]["data"] + raw = base64.b64decode(audio) + if fmt == "wav": + with soundfile.SoundFile(BytesIO(raw)) as audio_file: + assert audio_file.frames == frame_count + assert audio_file.samplerate == 22050 + else: + assert len(raw) == frame_count * 2 diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 206a38917..b4a943873 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -9,6 +9,8 @@ import io import json import struct import wave +from collections.abc import Callable +from contextlib import nullcontext from types import SimpleNamespace from typing import Any @@ -2182,7 +2184,8 @@ async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip(): assert "session.updated" not in ws.sent_types() -def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): +@pytest.mark.parametrize("sample_rate_hz", [22_050, 24_000]) +def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(sample_rate_hz: int): ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] @@ -2192,7 +2195,8 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): "response_id": "resp-a", "audio": "AAAA", "format": "pcm", - "sample_rate_hz": 24000, + "sample_rate_hz": sample_rate_hz, + "audio_duration_ms": 135, } ) @@ -2201,7 +2205,8 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ] assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} assert {payload["format"] for payload in audio_events} == {"pcm16"} - assert {payload["sample_rate_hz"] for payload in audio_events} == {24000} + assert {payload["sample_rate_hz"] for payload in audio_events} == {sample_rate_hz} + assert audio_events[0]["metadata"]["audio_duration_ms"] == 135 def test_native_realtime_protocol_ignores_removed_legacy_event_switches(): @@ -5023,8 +5028,10 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): response_id = session.begin_response() sent: list[dict[str, Any]] = [] - async def send_json(data: dict[str, Any]) -> None: + async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: sent.append(data) + if on_accepted is not None: + on_accepted() await handler._emit_chat_payload( session, @@ -5034,7 +5041,13 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): { "delta": { "content": "AAAA", - } + }, + "audio_metadata": { + "format": "wav", + "sample_rate_hz": 22_050, + "frame_count": 2205, + "channels": 1, + }, } ], }, @@ -5051,8 +5064,222 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): "epoch": 0, "audio": "AAAA", "format": "wav", + "sample_rate_hz": 22_050, + "channels": 1, + "audio_duration_ms": 100, } ] + assert session.playback.generated_ms == 100 + assert session.playback.sent_ms == 100 + assert session.playback.played_ms == 0 + + +@pytest.mark.parametrize("response_format", ["pcm", "wav"]) +def test_duplex_chat_request_uses_response_audio_format(response_format: str): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + extra_audio = {"format": "flac", "voice": "alloy"} + session = DuplexSession( + session_id="sid-chat-audio-format", + config=DuplexSessionConfig(response_format=response_format, extra_body={"audio": extra_audio}), + ) + session.begin_response() + + request = handler._build_chat_request(session, "format-request") + + assert request.audio == {"format": response_format, "voice": "alloy"} + assert extra_audio == {"format": "flac", "voice": "alloy"} + + +def _chat_audio_completion_payload(*, frame_count: int = 2205, sample_rate_hz: int = 22_050) -> dict[str, object]: + return { + "modality": "audio", + "choices": [ + { + "delta": {"content": "AAAA"}, + "audio_metadata": { + "format": "pcm", + "sample_rate_hz": sample_rate_hz, + "frame_count": frame_count, + "channels": 1, + }, + } + ], + } + + +@pytest.mark.asyncio +async def test_duplex_chat_full_audio_message_preserves_actual_format_and_duration(): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-full-audio", config=DuplexSessionConfig(response_format="wav")) + response_id = session.begin_response() + sent: list[dict[str, object]] = [] + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + sent.append(data) + if on_accepted is not None: + on_accepted() + + await handler._emit_chat_payload( + session, + { + "choices": [ + { + "message": {"role": "assistant", "content": "hello", "audio": {"data": "AAAA"}}, + "audio_metadata": { + "format": "pcm", + "sample_rate_hz": 22_050, + "frame_count": 2205, + "channels": 1, + }, + } + ], + }, + session.epoch, + response_id, + send_json, + ) + + assert [event["type"] for event in sent] == ["response.text.delta", "response.output_audio.delta"] + assert sent[1]["audio"] == "AAAA" + assert sent[1]["format"] == "pcm" + assert sent[1]["sample_rate_hz"] == 22_050 + assert sent[1]["audio_duration_ms"] == 100 + assert session.assistant_text_buffer == ("hello",) + assert session.playback.sent_ms == 100 + + +@pytest.mark.asyncio +async def test_duplex_chat_audio_does_not_round_each_fragment(): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-sub-ms", config=DuplexSessionConfig()) + response_id = session.begin_response() + sent: list[dict[str, object]] = [] + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + sent.append(data) + if on_accepted is not None: + on_accepted() + + for _ in range(23): + await handler._emit_chat_payload( + session, + _chat_audio_completion_payload(frame_count=1), + session.epoch, + response_id, + send_json, + ) + + assert [event["audio_duration_ms"] for event in sent] == [0] * 22 + [1] + assert session.playback.generated_ms == 1 + assert session.playback.sent_ms == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("send_error", [None, RuntimeError("output unavailable")], ids=["rejected", "failed"]) +async def test_duplex_chat_unaccepted_audio_is_generated_but_not_sent(send_error: RuntimeError | None): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-unaccepted-audio", config=DuplexSessionConfig()) + response_id = session.begin_response() + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + if send_error is not None: + raise send_error + + with pytest.raises(RuntimeError, match="output unavailable") if send_error is not None else nullcontext(): + await handler._emit_chat_payload( + session, _chat_audio_completion_payload(), session.epoch, response_id, send_json + ) + + assert session.playback.generated_ms == 100 + assert session.playback.sent_ms == 0 + assert session.playback.played_ms == 0 + + +@pytest.mark.asyncio +async def test_duplex_chat_late_send_completion_only_updates_old_response(): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-late-audio", config=DuplexSessionConfig()) + old_response = session.begin_response() + send_started = asyncio.Event() + finish_send = asyncio.Event() + + async def delayed_send(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + send_started.set() + await finish_send.wait() + assert on_accepted is not None + on_accepted() + + emit_task = asyncio.create_task( + handler._emit_chat_payload(session, _chat_audio_completion_payload(), session.epoch, old_response, delayed_send) + ) + try: + await asyncio.wait_for(send_started.wait(), timeout=1) + session.barge_in() + new_response = session.begin_response() + current_state = session.turn_state + finish_send.set() + await asyncio.wait_for(emit_task, timeout=1) + + assert session.active_response_id == new_response + assert session.playback_for_response(old_response).sent_ms == 100 + assert session.playback.generated_ms == 0 + assert session.playback.sent_ms == 0 + assert session.turn_state == current_state + finally: + emit_task.cancel() + await asyncio.gather(emit_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_duplex_chat_missing_audio_metadata_reports_error(): + class MissingMetadataChatService(TurnBasedFakeChatService): + async def create_chat_completion(self, request, raw_request=None): + async def generate(): + payload = {"modality": "audio", "choices": [{"delta": {"content": "AAAA"}}]} + yield f"data: {json.dumps(payload)}\n\n" + + return generate() + + handler = OmniDuplexSessionHandler( + chat_service=MissingMetadataChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-missing-metadata", config=DuplexSessionConfig()) + sent: list[dict[str, object]] = [] + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + sent.append(data) + if on_accepted is not None: + on_accepted() + + await handler._run_response(session, send_json) + + assert [event["type"] for event in sent] == ["response.created", "error"] + assert sent[-1]["code"] == "response_error" + assert "AudioChunkMetadata" in sent[-1]["error"] + assert session.playback.generated_ms == 0 + assert session.playback.sent_ms == 0 @pytest.mark.asyncio @@ -5071,7 +5298,17 @@ async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): for index, snapshot in enumerate(snapshots): payload = { "modality": "audio", - "choices": [{"delta": {"content": f"audio-{index}"}}], + "choices": [ + { + "delta": {"content": f"audio-{index}"}, + "audio_metadata": { + "format": "pcm", + "sample_rate_hz": 24_000, + "frame_count": 2400, + "channels": 1, + }, + } + ], "metrics": {"stage_metrics": {"0": snapshot}}, } yield f"data: {json.dumps(payload)}\n\n" @@ -5089,13 +5326,22 @@ async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): ) sent: list[dict[str, Any]] = [] - async def send_json(data: dict[str, Any]) -> None: + async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: sent.append(data) + if on_accepted is not None: + on_accepted() await handler._run_response(session, send_json) audio_deltas = [event for event in sent if event.get("type") == "response.output_audio.delta"] done = next(event for event in sent if event.get("type") == "response.done") + assert [event["audio_duration_ms"] for event in audio_deltas] == [100, 200, 300] + assert done["playback"] == { + "generated_ms": 300, + "sent_ms": 300, + "played_ms": 0, + "committed_ms": 0, + } assert all("vllm_omni" not in event for event in audio_deltas) assert done["vllm_omni"]["stage_metrics"]["0"] == { "num_tokens_out": 3, diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py index f41d7205b..bd682c4f3 100644 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ b/vllm_omni/entrypoints/duplex/chat_fallback.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import json from collections.abc import AsyncGenerator +from functools import partial from typing import Any from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -115,6 +116,13 @@ class ChatFallbackProjectorMixin: ): model_extra.pop(protocol_key, None) kwargs.update(model_extra) + if "audio" in response_config.modalities: + audio_options = kwargs.get("audio") + audio_options = dict(audio_options) if isinstance(audio_options, dict) else {} + # Chat defaults to WAV. Request the session's internal output format + # explicitly rather than relabeling the returned bytes as PCM. + audio_options["format"] = response_config.response_format + kwargs["audio"] = audio_options if isinstance(tools, list): kwargs["tools"] = tools if isinstance(tool_choice, str | dict): @@ -190,6 +198,8 @@ class ChatFallbackProjectorMixin: response_id: str, send_json, ) -> None: + if session.epoch != epoch or session.active_response_id != response_id: + return metrics = payload.get("metrics") stage_metrics = metrics.get("stage_metrics") if isinstance(metrics, dict) else None if isinstance(stage_metrics, dict): @@ -233,19 +243,7 @@ class ChatFallbackProjectorMixin: content = message.get("content") if isinstance(content, str) and content: - if modality == "audio": - session.mark_audio_sent() - await send_json( - { - "type": "response.output_audio.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "audio": content, - "format": session.response_config.response_format, - } - ) - else: + if modality != "audio": session.append_assistant_text(content) await send_json( { @@ -257,6 +255,39 @@ class ChatFallbackProjectorMixin: } ) + audio_content = content if modality == "audio" else None + if isinstance(message, dict): + message_audio = message.get("audio") + if isinstance(message_audio, dict): + audio_content = message_audio.get("data") + if isinstance(audio_content, str) and audio_content: + # The OpenAI package imports duplex serving during initialization. + # Defer this import until chat fallback has finished loading. + from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata + + audio_metadata = AudioChunkMetadata.model_validate(choice.get("audio_metadata")) + duration_ms = session.record_generated_audio( + response_id, + frame_count=audio_metadata.frame_count, + sample_rate_hz=audio_metadata.sample_rate_hz, + ) + if duration_ms is None: + return + await send_json( + { + "type": "response.output_audio.delta", + "session_id": session.session_id, + "response_id": response_id, + "epoch": epoch, + "audio": audio_content, + "format": audio_metadata.format, + "sample_rate_hz": audio_metadata.sample_rate_hz, + "channels": audio_metadata.channels, + "audio_duration_ms": duration_ms, + }, + on_accepted=partial(session.mark_audio_sent, duration_ms, response_id=response_id), + ) + finish_reason = choice.get("finish_reason") if finish_reason is not None and modality != "audio": await send_json( diff --git a/vllm_omni/entrypoints/duplex/protocol.py b/vllm_omni/entrypoints/duplex/protocol.py index ff91e0db5..d28a7b009 100644 --- a/vllm_omni/entrypoints/duplex/protocol.py +++ b/vllm_omni/entrypoints/duplex/protocol.py @@ -178,6 +178,8 @@ class DuplexPlaybackCursor: sent_ms: int = 0 played_ms: int = 0 committed_ms: int = 0 + generated_frames: int = 0 + sample_rate_hz: int | None = None def acknowledge(self, played_ms: int, committed_ms: int | None = None) -> None: self.played_ms = max(self.played_ms, max(0, int(played_ms))) @@ -1103,25 +1105,54 @@ class DuplexSession: if text: self._response.assistant_text_buffer.append(text) + def record_generated_audio(self, response_id: str, *, frame_count: int, sample_rate_hz: int) -> int | None: + """Accumulate source frames, rounding only the cumulative duration. + + A late chunk cannot change a different response's playback cursor. + This records encoded audio, not transport delivery or client playback. + """ + if response_id != self.active_response_id: + return None + if type(frame_count) is not int or frame_count < 0: + raise ValueError("Audio frame_count must be a non-negative integer") + if type(sample_rate_hz) is not int or sample_rate_hz <= 0: + raise ValueError("Audio sample_rate_hz must be a positive integer") + playback = self._playback.current + if playback.sample_rate_hz is not None and playback.sample_rate_hz != sample_rate_hz: + raise ValueError("Audio sample rate changed within a response") + playback.sample_rate_hz = sample_rate_hz + playback.generated_frames += frame_count + playback.generated_ms = playback.generated_frames * 1000 // sample_rate_hz + return playback.generated_ms + def mark_audio_sent( self, duration_ms: int | None = None, *, + response_id: str | None = None, text_chars: int | None = None, audio_text_marks: list[dict[str, object]] | None = None, ) -> None: - playback = self._playback.current + """Record cumulative audio accepted by the output path. + + ``sent_ms`` includes queued/journaled output; only client acknowledgements + establish ``played_ms``. It is not confirmation of socket delivery. + """ + playback = self._playback.current if response_id is None else self._playback.by_response.get(response_id) + if playback is None: + return + is_current = response_id is None or response_id == self.active_response_id if duration_ms is not None: playback.generated_ms = max(playback.generated_ms, duration_ms) playback.sent_ms = max(playback.sent_ms, duration_ms) - if text_chars is not None and text_chars >= 0: + if is_current and text_chars is not None and text_chars >= 0: self._response.assistant_audio_text_marks.append( DuplexAssistantAudioTextMark( text_chars=int(text_chars), audio_end_ms=max(0, int(duration_ms)), ) ) - if audio_text_marks: + if is_current and audio_text_marks: for raw_mark in audio_text_marks: if not isinstance(raw_mark, dict): continue @@ -1135,7 +1166,8 @@ class DuplexSession: audio_end_ms=max(0, int(raw_audio_end_ms)), ) ) - self.turn_state = DuplexTurnState.ASSISTANT_PLAYING + if is_current: + self.turn_state = DuplexTurnState.ASSISTANT_PLAYING def _playback_cursor_for_response(self, response_id: str | None = None) -> DuplexPlaybackCursor: if response_id is None: diff --git a/vllm_omni/entrypoints/duplex/session_attachment.py b/vllm_omni/entrypoints/duplex/session_attachment.py index d68fb268b..242169a6f 100644 --- a/vllm_omni/entrypoints/duplex/session_attachment.py +++ b/vllm_omni/entrypoints/duplex/session_attachment.py @@ -275,12 +275,18 @@ class DuplexSessionAttachmentRegistry: payload: Mapping[str, object], *, journal: bool = True, + on_accepted: Callable[[], None] | None = None, ) -> JournalEntry | None: """Sequence and dispatch one event to the current attachment. The per-session lock keeps wire order equal to journal order without serializing unrelated sessions. A detached session still records replayable events, but has no transport side effect. + + ``on_accepted`` runs synchronously once the event is journaled, or + after a successful transport send when journaling is disabled. A later + send failure cannot undo acceptance into the journal. Detached, + non-journaled events do not invoke it. The callback must not raise. """ async with self._lock: state = self._require(session_id) @@ -291,8 +297,12 @@ class DuplexSessionAttachmentRegistry: entry = state.journal.record(payload) if journal else None attachment = state.attachment wire_payload = dict(entry.payload) if entry is not None else dict(payload) + if entry is not None and on_accepted is not None: + on_accepted() if attachment is not None: await attachment.send(wire_payload) + if entry is None and on_accepted is not None: + on_accepted() return entry async def acknowledge(self, session_id: str, sequence: int) -> int: diff --git a/vllm_omni/entrypoints/duplex/session_runner.py b/vllm_omni/entrypoints/duplex/session_runner.py index ed18d8591..8ee5060b0 100644 --- a/vllm_omni/entrypoints/duplex/session_runner.py +++ b/vllm_omni/entrypoints/duplex/session_runner.py @@ -7,6 +7,7 @@ import asyncio import binascii import json import uuid +from collections.abc import Callable from contextlib import suppress from copy import deepcopy @@ -99,7 +100,9 @@ class DuplexSessionRunnerMixin: if callable(close): await close(code=1000, reason=reason) - async def send_attachment_payload(payload: dict[str, object], *, journal: bool) -> None: + async def send_attachment_payload( + payload: dict[str, object], *, journal: bool, on_accepted: Callable[[], None] | None = None + ) -> None: assert session is not None session_id = session.session_id should_journal = journal and session_id not in self._resync_required_sessions @@ -108,6 +111,7 @@ class DuplexSessionRunnerMixin: session_id, payload, journal=should_journal, + on_accepted=on_accepted, ) except DuplexJournalOverflowError: first_overflow = session_id not in self._resync_required_sessions @@ -129,6 +133,7 @@ class DuplexSessionRunnerMixin: session_id, payload, journal=False, + on_accepted=on_accepted, ) if realtime_protocol is not None: @@ -155,23 +160,29 @@ class DuplexSessionRunnerMixin: event_emit_lock = asyncio.Lock() - async def send_outbound(payload: dict[str, object]) -> None: + async def send_outbound(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: if not attachment_ready or session is None: await actor.send_json(dict(payload)) + if on_accepted is not None: + on_accepted() return if realtime_protocol is None: await send_attachment_payload( payload, journal=payload.get("type") not in {"session.created", "session.resumed"}, + on_accepted=on_accepted, ) return for projected in realtime_protocol.encode_outbound_event(payload): await send_attachment_payload( projected, journal=projected.get("type") not in {"session.created", "session.resumed"}, + # Item/part announcements may precede the audio. Only the + # actual audio event establishes acceptance of this chunk. + on_accepted=on_accepted if projected.get("type") == "response.audio.delta" else None, ) - async def emit_event(payload: dict[str, object]) -> None: + async def emit_event(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: deferred_precreate_response = False async with event_emit_lock: accepted, deferred_overlap_payload = await self._apply_outbound_session_event( @@ -183,7 +194,7 @@ class DuplexSessionRunnerMixin: ) if not accepted: return - await send_outbound(payload) + await send_outbound(payload, on_accepted=on_accepted) if deferred_overlap_payload is not None: deferred_precreate_response = native.deferred_precreate_response native.deferred_precreate_response = False diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index 4382afc57..a13318ad4 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -29,7 +29,12 @@ import torch import torchaudio from vllm.logger import init_logger -from vllm_omni.entrypoints.openai.protocol.audio import DEFAULT_AUDIO_FORMAT, AudioResponse, CreateAudio +from vllm_omni.entrypoints.openai.protocol.audio import ( + DEFAULT_AUDIO_FORMAT, + AudioChunkMetadata, + AudioResponse, + CreateAudio, +) try: import soundfile @@ -244,7 +249,16 @@ class AudioMixin: audio_data = base64.b64encode(audio_data).decode("utf-8") - return AudioResponse(audio_data=audio_data, media_type=media_type) + return AudioResponse( + audio_data=audio_data, + media_type=media_type, + audio_metadata=AudioChunkMetadata( + format=response_format, + sample_rate_hz=int(sample_rate), + frame_count=int(audio_tensor.shape[0]), + channels=1 if audio_tensor.ndim == 1 else int(audio_tensor.shape[1]), + ), + ) @staticmethod def _resample_audio(audio_tensor: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 459f9d4ae..8a1b48304 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -421,9 +421,23 @@ class CreateAudio(BaseModel): arbitrary_types_allowed = True +class AudioChunkMetadata(BaseModel): + """Waveform dimensions after transforms, before encoding. + + Frames count samples per channel, not interleaved scalar samples or bytes. + For compressed formats this excludes any padding introduced by the codec. + """ + + format: str = Field(min_length=1, strict=True) + sample_rate_hz: int = Field(gt=0, strict=True) + frame_count: int = Field(ge=0, strict=True) + channels: int = Field(gt=0, strict=True) + + class AudioResponse(BaseModel): audio_data: bytes | str media_type: str + audio_metadata: AudioChunkMetadata | None = None # --- Batch Speech Models --- diff --git a/vllm_omni/entrypoints/openai/protocol/chat_completion.py b/vllm_omni/entrypoints/openai/protocol/chat_completion.py index 462e53af9..eaa15eea9 100644 --- a/vllm_omni/entrypoints/openai/protocol/chat_completion.py +++ b/vllm_omni/entrypoints/openai/protocol/chat_completion.py @@ -1,12 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from typing import Any -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse, ChatCompletionStreamResponse +from pydantic import SerializeAsAny +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, +) + +from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata + + +class OmniChatCompletionResponseStreamChoice(ChatCompletionResponseStreamChoice): + audio_metadata: AudioChunkMetadata | None = None + + +class OmniChatCompletionResponseChoice(ChatCompletionResponseChoice): + audio_metadata: AudioChunkMetadata | None = None class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse): + choices: list[SerializeAsAny[ChatCompletionResponseStreamChoice]] modality: str | None = "text" metrics: dict[str, Any] | None = None class OmniChatCompletionResponse(ChatCompletionResponse): + choices: list[SerializeAsAny[ChatCompletionResponseChoice]] metrics: dict[str, Any] | None = None diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 20bddaba6..484d5d170 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -34,7 +34,11 @@ from vllm_omni.entrypoints.openai.diffusion_request_utils import ( apply_normalized_diffusion_request_extra_args, normalize_diffusion_request_args, ) -from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse +from vllm_omni.entrypoints.openai.protocol.chat_completion import ( + OmniChatCompletionResponse, + OmniChatCompletionResponseChoice, + OmniChatCompletionResponseStreamChoice, +) from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.metrics import definitions as _metric_defs @@ -2860,18 +2864,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): for output in final_res.outputs: if stream: - choice_data = ChatCompletionResponseStreamChoice( + choice_data = OmniChatCompletionResponseStreamChoice( index=output.index, delta=DeltaMessage(role=role, content=audio_base64), + audio_metadata=audio_response.audio_metadata, logprobs=None, finish_reason=output.finish_reason, stop_reason=output.stop_reason, token_ids=(as_list(output.token_ids) if request.return_token_ids else None), ) else: - choice_data = ChatCompletionResponseChoice( + choice_data = OmniChatCompletionResponseChoice( index=output.index, message=ChatMessage(role=role, audio=audio_obj), + audio_metadata=audio_response.audio_metadata, logprobs=None, finish_reason="stop", stop_reason=output.stop_reason, From ea1a416a17849addf135e9aa74ef5c2d2574c130 Mon Sep 17 00:00:00 2001 From: noname <133508552+lzwnoname@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:55:53 +0800 Subject: [PATCH 060/209] [Model] Add Breeze-TTS-2 two-stage AR TTS support (talker + codec, streaming) (#7084) Signed-off-by: noname Signed-off-by: noname <133508552+lzwnoname@users.noreply.github.com> Signed-off-by: Yueqian Lin Co-authored-by: noname Co-authored-by: Yueqian Lin --- docs/models/supported_models.md | 1 + docs/serving/speech_api.md | 8 +- .../text_to_speech/README.md | 32 + .../text_to_speech/breeze_tts_2/end2end.py | 204 ++++++ .../online_serving/text_to_speech/README.md | 63 +- .../text_to_speech/breeze_tts_2/run_server.sh | 13 + recipes/BreezeBlue/Breeze-TTS-2.md | 193 ++++++ recipes/README.md | 1 + tests/config/test_environment_variables.py | 4 +- tests/e2e/online_serving/test_breeze_tts_2.py | 100 +++ .../openai_api/test_tts_adapter.py | 1 + .../openai_api/test_tts_detection.py | 3 + .../models/breeze_tts_2/__init__.py | 0 .../breeze_tts_2/test_audio_tokenizer.py | 128 ++++ .../models/breeze_tts_2/test_batching.py | 95 +++ .../models/breeze_tts_2/test_codec_loading.py | 68 ++ .../models/breeze_tts_2/test_configuration.py | 113 ++++ .../models/breeze_tts_2/test_instruction.py | 135 ++++ .../models/breeze_tts_2/test_local_modules.py | 152 +++++ .../breeze_tts_2/test_model_registration.py | 28 + .../test_pipeline_registration.py | 36 + .../breeze_tts_2/test_prompt_builder.py | 189 ++++++ .../breeze_tts_2/test_reference_voice.py | 187 ++++++ .../test_stage_input_processor.py | 133 ++++ .../models/breeze_tts_2/test_streaming.py | 151 +++++ .../breeze_tts_2/test_talker_state_machine.py | 243 +++++++ .../breeze_tts_2/test_weight_loading.py | 73 +++ tools/breeze_tts_2_compare_golden.py | 165 +++++ tools/breeze_tts_2_upstream_golden.py | 200 ++++++ .../config/environment_variable_inventory.py | 1 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/breeze_tts_2.yaml | 51 ++ vllm_omni/engine/arg_utils.py | 4 + .../openai/tts_adapters/__init__.py | 1 + .../openai/tts_adapters/breeze_tts_2.py | 242 +++++++ .../models/breeze_tts_2/__init__.py | 0 .../models/breeze_tts_2/audio_tokenizer.py | 149 +++++ .../configuration_breeze_tts_2.py | 366 +++++++++++ .../modeling_breeze_tts_2_codec.py | 404 ++++++++++++ .../modeling_breeze_tts_2_depth.py | 523 +++++++++++++++ .../modeling_breeze_tts_2_talker.py | 618 ++++++++++++++++++ .../modeling_breeze_tts_2_text.py | 350 ++++++++++ .../models/breeze_tts_2/pipeline.py | 51 ++ .../models/breeze_tts_2/prompt_builder.py | 336 ++++++++++ vllm_omni/model_executor/models/registry.py | 10 + .../stage_input_processors/breeze_tts_2.py | 244 +++++++ 46 files changed, 6067 insertions(+), 4 deletions(-) create mode 100644 examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py create mode 100755 examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh create mode 100644 recipes/BreezeBlue/Breeze-TTS-2.md create mode 100644 tests/e2e/online_serving/test_breeze_tts_2.py create mode 100644 tests/model_executor/models/breeze_tts_2/__init__.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_batching.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_codec_loading.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_configuration.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_instruction.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_local_modules.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_model_registration.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_prompt_builder.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_reference_voice.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_streaming.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_weight_loading.py create mode 100644 tools/breeze_tts_2_compare_golden.py create mode 100644 tools/breeze_tts_2_upstream_golden.py create mode 100644 vllm_omni/deploy/breeze_tts_2.yaml create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/__init__.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/pipeline.py create mode 100644 vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py create mode 100644 vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 1f4093828..a60eb6670 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -71,6 +71,7 @@ th { | `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | +| `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) | | `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/dots.tts-soar` | ✅︎ | | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index f0658f278..38cca54d6 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -7,7 +7,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz. - **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz. -See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano. +See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, MOSS-TTS-Nano, and Breeze-TTS-2. !!! tip "Deployment recipes" TTS deployment recipes are published at @@ -783,6 +783,12 @@ hidden-state dtype before the output projection. Consequently, a float32 stage c | ------- | ------------- | | `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. | +### Breeze-TTS-2 + +| Model | Description | +| ------- | ------------- | +| `BreezeBlue/Breeze-TTS-2` | Two-stage AR TTS (T5Gemma2 + Qwen3 talker with a depth decoder, bundled Qwen3-TTS codec) at 24 kHz. Four modes are selected from the request fields: plain (`input` + speaker tag `voice`, `S0`..`S9`), voice design (`instructions`), voice clone (`ref_audio` + `ref_text`, exactly one clip), and voice direction (reference + `instructions`). Greedy decoding only: `sample_rate` must be `24000`, `speed` must be `1.0`, and `guidance_scale`/`cfg_scale` other than `1.0`, `negative_prompt`, `temperature`/`top_p`/`top_k` overrides, `language`, and `speaker_embedding` are rejected. Streaming returns PCM `speech.audio.delta` events. See [`recipes/BreezeBlue/Breeze-TTS-2.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md). | + ## Error Responses ### 400 Bad Request diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 89ff2dbce..954b96bd3 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -16,6 +16,7 @@ list of supported architectures across all modalities, see |---|---|---|---|---|---|---| | AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | +| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | | Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz | @@ -662,6 +663,37 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). ### Notes + - `--num-prompts N` replicates the prompt for performance measurement. - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. + +--- + +## Breeze-TTS-2 + +Two-stage AR TTS (T5Gemma2 + Qwen3 talker with a depth decoder → bundled Qwen3-TTS codec) at 24 kHz mono. The prompt builder picks one of four templates automatically: plain, voice design (`--instruction`), clone (`--ref-audio` + `--ref-text`), or voice direction (reference + `--instruction`). + +### Quick start + +```bash +python examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py \ + --model BreezeBlue/Breeze-TTS-2 \ + --text "Hello, this is Breeze TTS 2 running on vLLM Omni." +``` + +### Voice direction + +```bash +python examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py \ + --text "We need to discuss what happened last night." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript of the reference audio." \ + --instruction "Speak slowly with a restrained, serious tone." +``` + +### Notes + +- Output: 24 kHz mono WAV; `--max-new-tokens` caps generated frames (80 ms each). +- Speaker tags `--voice S0`..`S9` (default `S0`) select the timbre in non-reference modes; with `--ref-audio` the timbre comes from the reference clip. +- Deploy config: `vllm_omni/deploy/breeze_tts_2.yaml` (auto-loaded by HF `model_type`); see [`recipes/BreezeBlue/Breeze-TTS-2.md`](../../../recipes/BreezeBlue/Breeze-TTS-2.md) for the online-serving recipe. diff --git a/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py b/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py new file mode 100644 index 000000000..942ead792 --- /dev/null +++ b/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Offline Breeze-TTS-2 inference example (two-stage talker + codec pipeline). + +Uses the async-chunk deploy config (vllm_omni/deploy/breeze_tts_2.yaml). +Reference checkpoint: BreezeBlue/Breeze-TTS-2. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import soundfile as sf +import torch +from transformers import AutoConfig + +from vllm_omni import Omni +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +REPO_ROOT = Path(__file__).resolve().parents[4] +SAMPLE_RATE = 24_000 + + +def parse_args(): + parser = TrackingArgumentParser(description="Offline Breeze-TTS-2 inference") + parser.add_argument( + "--model", + type=str, + default="BreezeBlue/Breeze-TTS-2", + help="Breeze-TTS-2 model path or HuggingFace repo ID.", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is Breeze TTS 2 running on vLLM Omni.", + help="Text to synthesize.", + ) + parser.add_argument( + "--voice", + type=str, + default="S0", + help="Speaker tag such as S0..S9.", + ) + parser.add_argument( + "--instruction", + type=str, + default=None, + help="Optional style instruction (voice design / voice direction).", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Optional reference wav path for cloning (requires --ref-text).", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Exact transcript of the reference audio.", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=None, + help="Cap on generated codec frames (one frame is 80 ms of audio).", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_audio", + help="Directory for output WAV files.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/breeze_tts_2.yaml based on the HF model_type.", + ) + return parser.parse_args() + + +def resolve_template(args) -> str: + if args.ref_audio is not None: + return "ref_edit_tata" if args.instruction else "ref_clone_tata" + return "tts_instruction" if args.instruction else "tts_plain" + + +def reference_payload(ref_audio_path: str | None) -> dict: + """Load the reference clip into the payload keys the prompt builder reads. + + ``BreezeTTS2PromptBuilder`` consumes ``ref_audio`` (a waveform) plus + ``ref_audio_sample_rate``; a bare path field is not read. + """ + if ref_audio_path is None: + return {} + waveform, sample_rate = sf.read(ref_audio_path, dtype="float32", always_2d=True) + # Down-mix to the single 1-D clip the reference encoder expects. + waveform = waveform.mean(axis=1) + return {"ref_audio": waveform, "ref_audio_sample_rate": int(sample_rate)} + + +def build_prompt(args) -> dict: + config = AutoConfig.from_pretrained(args.model, trust_remote_code=True) + reference_encoder = None + if args.ref_audio is not None: + from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + BreezeReferenceAudioTokenizer, + ) + + reference_encoder = BreezeReferenceAudioTokenizer.from_pretrained( + args.model, + num_codebooks=int(getattr(config, "num_codebooks", 16)), + device_map="cpu", + ) + builder = BreezeTTS2PromptBuilder.from_pretrained( + args.model, + config, + reference_audio_encoder=reference_encoder, + ) + payload = { + "text": args.text, + "speaker": args.voice, + **reference_payload(args.ref_audio), + } + if args.instruction: + payload["instruction"] = args.instruction + if args.ref_text: + payload["ref_text"] = args.ref_text + prompt = builder.build(payload, template=resolve_template(args)) + if args.max_new_tokens is not None: + prompt["additional_information"]["breeze_max_new_frames"] = int(args.max_new_tokens) + return prompt + + +def extract_audio(multimodal_output: dict) -> torch.Tensor: + """Extract the final complete audio tensor from multimodal output. + + The output processor concatenates per-step delta tensors under + ``model_outputs``. Falls back to ``audio`` for backwards compat. + """ + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") + + if isinstance(audio, list): + valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not valid: + raise ValueError("Audio list is empty or all elements are None.") + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def main(): + args = parse_args() + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + engine = Omni( + model=args.model, + deploy_config=args.deploy_config, + ) + + prompt = build_prompt(args) + + print(f"Model : {args.model}") + print(f"Text : {args.text}") + print(f"Template : {resolve_template(args)}") + print(f"Output dir : {output_dir}") + print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens") + + t_start = time.perf_counter() + outputs = engine.generate([prompt]) + elapsed = time.perf_counter() - t_start + + request_output = outputs[0] + mm = request_output.outputs[0].multimodal_output + audio = extract_audio(mm) + + duration = audio.numel() / SAMPLE_RATE + rtf = elapsed / duration if duration > 0 else float("inf") + + output_path = output_dir / "output.wav" + sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV") + + print(f"Saved : {output_path}") + print(f"Duration : {duration:.2f}s") + print(f"Inference : {elapsed:.2f}s") + print(f"RTF : {rtf:.3f}") + + +if __name__ == "__main__": + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + main() diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 5141a37c2..7ad939415 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -13,7 +13,8 @@ For the full list of supported architectures across all modalities, see ## Supported Models | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | -| --- | --- | --- | --- | --- | --- | +|---|---|---|---|---|---| +| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | @@ -1025,3 +1026,63 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice cloning is gated upstream and may require a recent `mistral_common`. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. + +## Breeze-TTS-2 + +Two-stage AR TTS (T5Gemma2 + Qwen3 talker with a depth decoder → bundled Qwen3-TTS codec) at 24 kHz mono, from [BreezeBlue](https://huggingface.co/BreezeBlue/Breeze-TTS-2). Four modes are selected automatically from the request fields: plain (`input`), voice design (`instructions`), clone (`ref_audio`+`ref_text`), and voice direction (reference trio + `instructions`). + +### Launch + +```bash +vllm serve BreezeBlue/Breeze-TTS-2 --omni --port 8091 +# or: +./breeze_tts_2/run_server.sh +``` + +The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming, 8-frame inter-stage chunks). + +### Sending requests + +```bash +# Plain synthesis (speaker tags S0..S9, default S0) +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Breeze TTS 2 on vLLM-Omni.", + "voice": "S0", + "response_format": "wav", + "sample_rate": 24000 + }' --output breeze_plain.wav + +# Voice direction (clone a reference, then steer the delivery) +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "We need to discuss what happened last night.", + "ref_audio": "file:///path/to/reference.wav", + "ref_text": "The exact transcript of the reference audio.", + "instructions": "Speak slowly with a restrained, serious tone.", + "response_format": "wav" + }' --output breeze_direction.wav +``` + +### Streaming PCM + +```bash +curl -N -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Streaming output from Breeze TTS 2.", + "stream": true, + "stream_format": "audio", + "response_format": "pcm", + "sample_rate": 24000 + }' --output breeze_stream.pcm +``` + +### Notes + +- Only `sample_rate=24000` is accepted; `speed` must be `1.0`. +- `guidance_scale`/`cfg_scale` must be `1.0`; `negative_prompt` is rejected (CFG companion support is a follow-up). +- `instructions` without `ref_audio` is voice design; with `ref_audio`+`ref_text` it is voice direction. +- See [`recipes/BreezeBlue/Breeze-TTS-2.md`](../../../recipes/BreezeBlue/Breeze-TTS-2.md) for the full recipe and the offline example under [`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../offline_inference/text_to_speech/breeze_tts_2/). diff --git a/examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh b/examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh new file mode 100755 index 000000000..83ca27084 --- /dev/null +++ b/examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Launch Breeze-TTS-2 online serving with the async-chunk deploy config. + +set -euo pipefail + +MODEL="${1:-BreezeBlue/Breeze-TTS-2}" +PORT="${PORT:-8091}" + +exec vllm-omni serve "${MODEL}" \ + --deploy-config vllm_omni/deploy/breeze_tts_2.yaml \ + --omni --port "${PORT}" diff --git a/recipes/BreezeBlue/Breeze-TTS-2.md b/recipes/BreezeBlue/Breeze-TTS-2.md new file mode 100644 index 000000000..50e8858af --- /dev/null +++ b/recipes/BreezeBlue/Breeze-TTS-2.md @@ -0,0 +1,193 @@ +# Breeze-TTS-2 + +> Text-to-speech serving (plain / voice design / voice clone / voice direction) + +## Summary + +- Vendor: BreezeBlue +- Model: `BreezeBlue/Breeze-TTS-2` (`BreezeForConditionalGeneration`) +- Task: Text-to-speech with speaker tags, natural-language style instructions, single-reference voice cloning, and reference+instruction "voice direction" +- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API (streaming and non-streaming); offline `Omni` example +- Hardware: 1x NVIDIA L20 48GB (CUDA) +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a known-good starting point for serving +Breeze-TTS-2 with vLLM-Omni. Breeze-TTS-2 is a two-stage AR TTS pipeline: + +| Stage | Components | Output | +| ----- | ---------- | ------ | +| 0 (talker, `LLM_AR`) | T5Gemma2 text encoder + Qwen3 backbone + 2052-way codebook-0 head + 12-layer depth decoder | 16-codebook codec frames | +| 1 (codec, `LLM_GENERATION`) | Bundled Qwen3-TTS audio tokenizer decoder (Mimi fallback) | 24 kHz mono waveform | + +Four prompt modes are selected automatically from the request fields: + +| Mode | Trigger | Fields | +| ---- | ------- | ------ | +| `tts_plain` | only `input` (default `voice=S0`) | `input`, `voice` | +| `tts_instruction` (voice design) | `input` + `instructions` | style/emotion/pace directives, not spoken | +| `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript | +| `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery | + +Current scope: greedy sampling with `cfg_scale=1.0`; CFG ≠ 1.0 and +`negative_prompt` are rejected with explicit client errors until companion +request support lands. + +## Supported model contract + +| Item | Value | +| ---- | ----- | +| Tasks | Text-to-speech through `/v1/audio/speech` (online) and the offline `Omni` example | +| Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) | +| Languages | English and Chinese (model card) | +| Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz | +| Output | 24 kHz mono; `wav` or `pcm`; streaming SSE `speech.audio.delta` | +| Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) | +| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0`; other guidance values and `negative_prompt` are rejected | +| Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` | + +## References + +- Issue: [#6656 [New Model]: Breeze TTS 2](https://github.com/vllm-project/vllm-omni/issues/6656) +- Related examples under `examples/`: + [`examples/online_serving/text_to_speech/breeze_tts_2/`](../../examples/online_serving/text_to_speech/breeze_tts_2/), + [`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../examples/offline_inference/text_to_speech/breeze_tts_2/) + +## Hardware + +- Accelerator model and per-device memory: 1x NVIDIA L20 48GB +- Number of devices: 1 (both stages are placed on device 0) +- Device interconnect: not applicable (single device) +- Host memory: no special requirement; per-request codec state is small +- Qualification scope: functional serving of all four prompt modes, streaming PCM, and 4 concurrent mixed-mode requests (greedy, `cfg_scale=1.0`). No other accelerator is qualified by this recipe. + +## Software environment + +- OS: Linux +- Python: 3.12.12 +- Driver / runtime: NVIDIA driver 580.82.07, CUDA 13.0 (PyTorch 2.13.0+cu130) +- vLLM version: 0.28.0 +- vLLM-Omni version or commit: the commit that introduced this recipe (branch based on `main` @ `67ef0b64`) +- Transformers: 5.14.1 + +## Command + +Start the server from the repository root. The deploy config +(`vllm_omni/deploy/breeze_tts_2.yaml`, async chunk streaming by default) +auto-loads from the checkpoint's `model_type`; pass it explicitly if you +customized it: + +```bash +vllm serve BreezeBlue/Breeze-TTS-2 \ + --deploy-config vllm_omni/deploy/breeze_tts_2.yaml \ + --omni --port 8091 +``` + +## Verification + +Plain synthesis (speaker tag `S0`..`S9`, default `S0`): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "BreezeBlue/Breeze-TTS-2", + "input": "Hello, this is Breeze TTS 2 running on vLLM-Omni.", + "voice": "S0", + "response_format": "wav", + "sample_rate": 24000 + }' --output breeze_plain.wav +``` + +Voice design (instruction only, no reference audio): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "[laughs] Welcome to tonight's story time.", + "instructions": "A warm young woman, clear voice, lively delivery.", + "response_format": "wav" + }' --output breeze_design.wav +``` + +Voice cloning (`ref_audio` and `ref_text` must be provided together): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "This is new text spoken in the cloned voice.", + "ref_audio": "file:///path/to/reference.wav", + "ref_text": "The exact transcript of the reference audio.", + "response_format": "wav" + }' --output breeze_clone.wav +``` + +Voice direction (reference + instruction): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "We need to discuss what happened last night.", + "ref_audio": "file:///path/to/reference.wav", + "ref_text": "The exact transcript of the reference audio.", + "instructions": "Speak slowly with a restrained, serious tone.", + "response_format": "wav" + }' --output breeze_direction.wav +``` + +Streaming PCM (SSE `speech.audio.delta` events): + +```bash +curl -N -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Streaming output from Breeze TTS 2.", + "stream": true, + "stream_format": "audio", + "response_format": "pcm", + "sample_rate": 24000 + }' --output breeze_stream.pcm +``` + +Expected results: finite, non-silent 24 kHz mono WAV files; streaming requests +return incremental PCM deltas. On 1x L20 (greedy, single request after +warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time. + +## Notes + +- **Sample rate**: only `24000` is accepted; other values fail validation + before inference. +- **Speed / seeds / multi-reference**: `speed` must be `1.0`; + `task_type=VoiceDesign`, `ref_audio_2`, and speaker embeddings are not + supported yet. +- **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and + `negative_prompt` return a client error. +- **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`; + synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the + 12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps. +- **Checkpoint layout**: the checkpoint must contain the `audio_tokenizer/` + subdirectory (bundled Qwen3-TTS codec) for the default streaming path; the + Mimi fallback is used only when that directory is absent and does not + support async-chunk streaming. +- **License**: the reference inference code is Apache 2.0; the checkpoints are + distributed under the separate BreezeBlue Research and Non-Commercial + License. Commercial use requires written authorization from BreezeBlue. + +## Supported features + +| Feature | Status | Notes | +| ------- | ------ | ----- | +| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks, SSE `speech.audio.delta` PCM | +| Voice cloning | ✓ | one reference clip plus transcript per request | +| Voice design / voice direction | ✓ | `instructions` field, with or without a reference | +| Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) | +| Non-greedy sampling | ✗ | greedy decoding only | +| Prefix caching | ✗ | disabled in the deploy config | +| Tensor / pipeline parallelism | untested | single-GPU profile only | +| Quantization | ✗ | not supported in this release | +| CUDA graphs | not tuned | depth decoder and stage-1 codec run eagerly | +| Platforms | CUDA | qualified on 1x NVIDIA L20 | diff --git a/recipes/README.md b/recipes/README.md index d9f4b6822..86eba6d87 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -45,6 +45,7 @@ recipes/ | [`black-forest-labs/FLUX.1-schnell.md`](./black-forest-labs/FLUX.1-schnell.md) | Offline text-to-image generation with CPU offload | 1x NVIDIA GeForce RTX 5090 32GB | | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ | | [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB | +| [`BreezeBlue/Breeze-TTS-2.md`](./BreezeBlue/Breeze-TTS-2.md) | Online + offline TTS (plain / voice design / voice clone / voice direction, streaming PCM) | 1x NVIDIA L20 48GB | | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB | | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB | | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index d958bbf18..ce608dca3 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 27, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 63, + EnvironmentVariableCategory.MODEL_SPECIFIC: 64, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ModelEnvironmentVariableDisposition.PROMOTE: 38, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 14, + ModelEnvironmentVariableDisposition.INTERNALIZE: 15, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, } diff --git a/tests/e2e/online_serving/test_breeze_tts_2.py b/tests/e2e/online_serving/test_breeze_tts_2.py new file mode 100644 index 000000000..29cd43b68 --- /dev/null +++ b/tests/e2e/online_serving/test_breeze_tts_2.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +""" +E2E Online tests for Breeze-TTS-2 with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +actual model inference, not mocks. Breeze-TTS-2 plain mode needs only +``input`` and a speaker tag (``voice`` defaults to ``S0``); reference +cloning requires ``ref_audio`` + ``ref_text`` and is not exercised here. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "BreezeBlue/Breeze-TTS-2" +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 +MAX_CONCURRENT = 4 + +# ~0.5 s of 24 kHz mono PCM_16 in WAV (~24k payload + header). +_MIN_AUDIO_BYTES = 40_000 + + +def get_prompt(prompt_type="text"): + """English prompt for zero-shot TTS.""" + prompts = { + "text": "The weather is nice today, perfect for a walk in the park.", + } + return prompts.get(prompt_type, prompts["text"]) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("breeze_tts_2.yaml"), + server_args=["--trust-remote-code", "--disable-log-stats"], + ), + id="breeze_tts_2", + ) +] + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_server, online_client) -> None: + """ + Test plain text-to-audio via OpenAI API. + Deploy Setting: breeze_tts_2.yaml + Input Modal: text + Output Modal: audio + Input Setting: stream=False + Datasets: few requests (max_num_seqs=8) + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "voice": "S0", + "sample_rate": 24000, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_002(omni_server, online_client) -> None: + """ + Test streaming text-to-audio via OpenAI API (async-chunk SSE deltas). + Deploy Setting: breeze_tts_2.yaml + Input Modal: text + Output Modal: audio + Input Setting: stream=True + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": True, + "stream_format": "audio", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "voice": "S0", + "sample_rate": 24000, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + online_client.send_audio_speech_request(request_config) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index a17a6ccbd..edd2b4cd4 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -65,6 +65,7 @@ EXPECTED_MODEL_TYPES = { "higgs_audio_v2", "higgs_audio_v3", "glm_tts", + "breeze_tts_2", "step_audio2", "indextts2", "indextts2_5", diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index d54ce36d5..6c7671659 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """TTS model detection: registry-driven stage -> model-type resolution. ``serving_speech.py`` used to carry a hand-written 20-branch ladder mapping @@ -138,6 +139,7 @@ _PIPELINE_STAGES = [ "audio_tokenizer", "audio_vae", "aura", + "breeze_tts_2", "code2wav", "cosyvoice3_code2wav", "cosyvoice3_talker", @@ -184,6 +186,7 @@ _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"] _ARCHS = [ None, "VoxCPM2TalkerForConditionalGeneration", + "BreezeForConditionalGeneration", "MingTTSForConditionalGeneration", "CovoAudioForConditionalGeneration", "MyCovoAudioThing", diff --git a/tests/model_executor/models/breeze_tts_2/__init__.py b/tests/model_executor/models/breeze_tts_2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py b/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py new file mode 100644 index 000000000..75dcbb1fb --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + BreezeReferenceAudioTokenizer, + resolve_audio_tokenizer_path, +) + + +def _patch_hf_api(monkeypatch, snapshot_download): + """Point the resolver's tagged HF helper at a fake API object.""" + from vllm.transformers_utils import repo_utils + + monkeypatch.setattr(repo_utils, "hf_api", lambda: SimpleNamespace(snapshot_download=snapshot_download)) + + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + def encode(self, audio, sr=None): + assert isinstance(audio, np.ndarray) + assert sr == 16000 + return {"audio_codes": [torch.arange(32).reshape(16, 2)]} + + +def test_reference_audio_tokenizer_normalizes_codebook_major_output(): + adapter = BreezeReferenceAudioTokenizer(_Tokenizer()) + + codes = adapter.encode(np.zeros(32, dtype=np.float32), 16000) + + assert tuple(codes.shape) == (2, 16) + assert codes.dtype == torch.int16 + assert codes.device.type == "cpu" + + +def test_reference_audio_tokenizer_rejects_waveform_without_sample_rate(): + adapter = BreezeReferenceAudioTokenizer(_Tokenizer()) + + try: + adapter.encode(np.zeros(32, dtype=np.float32)) + except ValueError as exc: + assert "sample_rate is required" in str(exc) + else: + raise AssertionError("missing sample rate should be rejected") + + +def test_resolve_prefers_local_directory_without_hub_access(tmp_path, monkeypatch): + bundled = tmp_path / "audio_tokenizer" + bundled.mkdir() + + def _no_hub(*_args, **_kwargs): + raise AssertionError("local directories must not touch the hub") + + _patch_hf_api(monkeypatch, _no_hub) + + assert resolve_audio_tokenizer_path(str(tmp_path)) == bundled + + +def test_resolve_local_directory_without_bundled_tokenizer_returns_none(tmp_path, monkeypatch): + def _no_hub(*_args, **_kwargs): + raise AssertionError("local directories must not touch the hub") + + _patch_hf_api(monkeypatch, _no_hub) + + assert resolve_audio_tokenizer_path(str(tmp_path)) is None + + +def test_resolve_repo_id_uses_hf_snapshot(tmp_path, monkeypatch): + snapshot_root = tmp_path / "snapshot" + (snapshot_root / "audio_tokenizer").mkdir(parents=True) + calls = [] + + def _fake_snapshot_download(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + if kwargs.get("local_files_only"): + raise FileNotFoundError("cold cache") + return str(snapshot_root) + + _patch_hf_api(monkeypatch, _fake_snapshot_download) + + resolved = resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") + + assert resolved == snapshot_root / "audio_tokenizer" + # The cold-cache probe runs first; the online fetch carries the pattern. + assert calls[0] == ( + "BreezeBlue/Breeze-TTS-2", + {"allow_patterns": ["audio_tokenizer/*"], "local_files_only": True}, + ) + assert calls[1] == ("BreezeBlue/Breeze-TTS-2", {"allow_patterns": ["audio_tokenizer/*"]}) + + +def test_resolve_repo_id_prefers_complete_cached_snapshot(tmp_path, monkeypatch): + snapshot_root = tmp_path / "snapshot" + (snapshot_root / "audio_tokenizer").mkdir(parents=True) + calls = [] + + def _fake_snapshot_download(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + assert kwargs.get("local_files_only"), "cached snapshot must not hit the network" + return str(snapshot_root) + + _patch_hf_api(monkeypatch, _fake_snapshot_download) + + resolved = resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") + + assert resolved == snapshot_root / "audio_tokenizer" + assert len(calls) == 1 + + +def test_resolve_repo_id_without_bundled_tokenizer_returns_none(tmp_path, monkeypatch): + snapshot_root = tmp_path / "snapshot" + snapshot_root.mkdir() + + def _fake_snapshot_download(repo_id, **kwargs): + if kwargs.get("local_files_only"): + raise FileNotFoundError("cold cache") + return str(snapshot_root) + + _patch_hf_api(monkeypatch, _fake_snapshot_download) + + assert resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") is None diff --git a/tests/model_executor/models/breeze_tts_2/test_batching.py b/tests/model_executor/models/breeze_tts_2/test_batching.py new file mode 100644 index 000000000..41fc0b9ba --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_batching.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( + BreezeTTS2TalkerForGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _RequestHead: + def __call__(self, hidden: torch.Tensor) -> torch.Tensor: + # Hidden row 0 selects codec id 2; row 1 selects codec id 3. + code0 = 2 + int(hidden.reshape(-1)[0].item()) + logits = torch.zeros((hidden.shape[0], 9), dtype=hidden.dtype) + logits[:, code0] = 1.0 + return logits + + +def _model() -> BreezeTTS2TalkerForGeneration: + model = object.__new__(BreezeTTS2TalkerForGeneration) + model.num_codebooks = 4 + model.codebook_vocab_size = 8 + model.codebook_size = 8 + model.codebook_eos_token_id = 0 + model.audio_token_id = 100 + model.audio_eos_token_id = 101 + model.lm_head = _RequestHead() + model.logits_processor = lambda head, hidden: head(hidden) + model._codec_disallowed_mask = None + model._batch_state = None + model._batch_state_spans = None + model._async_chunk = True + model._golden_dump_dir = None + model._generate_depth_codes = lambda _hidden, code0: _frame(code0).unsqueeze(0) + return model + + +def _frame(code0: int) -> torch.Tensor: + return torch.tensor([code0, code0 + 1, code0 + 2, code0 + 3], dtype=torch.long) + + +def _step(model: BreezeTTS2TalkerForGeneration, infos: list[dict[str, Any]], hidden: torch.Tensor): + return model.make_omni_output( + hidden, + model_intermediate_buffer=infos, + request_token_spans=[(0, 1), (1, 2)], + ) + + +def test_two_requests_accumulate_and_emit_independent_frames(): + model = _model() + infos = [ + { + "breeze_generated_frames": 1, + "breeze_max_new_frames": 2, + "breeze_audio_codes": _frame(2).reshape(1, 4), + }, + {"breeze_generated_frames": 0, "breeze_max_new_frames": 3}, + ] + hidden = torch.tensor([[0.0], [1.0]], dtype=torch.float32) + + first = _step(model, infos, hidden) + assert [tuple(item.shape) for item in first.multimodal_outputs["codes"]["audio"]] == [(1, 4), (1, 4)] + assert torch.equal(first.multimodal_outputs["codes"]["audio"][0][0], _frame(2)) + assert torch.equal(first.multimodal_outputs["codes"]["audio"][1][0], _frame(3)) + + second = _step(model, infos, hidden) + # Request 0 reaches its two-frame budget and emits only a finish marker; + # request 1 still has budget and emits its next tail frame. + assert second.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert torch.equal(second.multimodal_outputs["codes"]["audio"][1][0], _frame(3)) + # Async streaming only maintains the per-frame tail; the cumulative + # buffer is deliberately not grown (request 0 keeps its pre-seeded + # entry untouched, request 1 never gains one). + assert infos[0]["breeze_audio_codes"].shape == (1, 4) + assert "breeze_audio_codes" not in infos[1] + assert infos[0]["breeze_force_eos"] is True + assert infos[1]["breeze_force_eos"] is False + + +def test_finished_requests_only_clear_model_owned_batch_mapping(): + model = _model() + model._batch_state = [{"id": "a"}, {"id": "b"}] + model._batch_state_spans = [(0, 1), (1, 2)] + + model.on_requests_finished(["a", "b"]) + + assert model._batch_state is None + assert model._batch_state_spans is None diff --git a/tests/model_executor/models/breeze_tts_2/test_codec_loading.py b/tests/model_executor/models/breeze_tts_2/test_codec_loading.py new file mode 100644 index 000000000..f558d067a --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_codec_loading.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_codec import ( + BreezeTTS2MimiCodec, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _BundledTokenizer: + def __init__(self) -> None: + self.model = nn.Linear(2, 2) + + def named_parameters(self): + for name, parameter in self.model.named_parameters(): + yield f"model.{name}", parameter + + def decode(self, _encoded: Any) -> tuple[list[torch.Tensor], int]: + return [torch.zeros(4)], 24_000 + + +def test_bundled_qwen_tokenizer_consumes_unused_root_fallback_weights(tmp_path): + bundled = tmp_path / "audio_tokenizer" + bundled.mkdir() + codec = object.__new__(BreezeTTS2MimiCodec) + codec.model_path = str(tmp_path) + codec.vllm_config = None + # load_weights reuses the path resolved once in __init__. + codec._tokenizer_path = bundled + codec._codec = None + codec._audio_tokenizer = _BundledTokenizer() + codec._loaded_local_weights = False + + weights = [ + ("backbone_model.layers.0.norm.weight", torch.ones(2)), + ("codec_model.decoder.layers.0.conv.weight", torch.ones(2)), + ("lm_head.weight", torch.ones(2, 2)), + ] + + loaded = codec.load_weights(iter(weights)) + + assert "_audio_tokenizer.model.weight" in loaded + assert "_audio_tokenizer.model.bias" in loaded + # The bundled tokenizer was preloaded, so load_weights must not try to copy + # the root Mimi fallback into the nonexistent ``self._codec`` module. + assert codec._codec is None + + +def test_codec_warmup_with_single_token_returns_empty_audio(): + codec = object.__new__(BreezeTTS2MimiCodec) + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = None + codec._codec = None + codec._async_chunk = False + + output = codec.forward(torch.tensor([1], dtype=torch.long), seq_token_counts=[1]) + + assert output.multimodal_outputs["model_outputs"][0].numel() == 0 + assert output.multimodal_outputs["sr"][0].item() == 24_000 diff --git a/tests/model_executor/models/breeze_tts_2/test_configuration.py b/tests/model_executor/models/breeze_tts_2/test_configuration.py new file mode 100644 index 000000000..950f420cf --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_configuration.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +from transformers.models.qwen3 import Qwen3Config + +from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2Config, + BreezeTTS2DepthDecoderConfig, + BreezeTTS2TextEncoderConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _config() -> BreezeTTS2Config: + return BreezeTTS2Config( + backbone_config={ + "model_type": "qwen3", + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "vocab_size": 128, + "max_position_embeddings": 512, + }, + text_encoder_config={ + "model_type": "t5gemma2_text", + "hidden_size": 24, + "intermediate_size": 48, + "num_hidden_layers": 6, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "vocab_size": 256, + }, + depth_decoder_config={ + "model_type": "breeze_depth_decoder_model", + "hidden_size": 32, + "backbone_hidden_size": 32, + "audio_embed_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "num_codebooks": 16, + "vocab_size": 2051, + "rope_scaling": { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 0.001953125, + "high_freq_factor": 0.0078125, + "original_max_position_embeddings": 16, + }, + }, + codec_config={"model_type": "mimi", "codebook_size": 2048, "num_quantizers": 32}, + num_codebooks=16, + vocab_size=2051, + text_vocab_size=262158, + ) + + +def test_nested_configs_and_text_config_use_qwen_backbone(): + config = _config() + + assert isinstance(config.backbone_config, Qwen3Config) + assert config.get_text_config() is config.backbone_config + assert config.get_text_config().vocab_size == 128 + assert isinstance(config.text_encoder_config, BreezeTTS2TextEncoderConfig) + assert isinstance(config.depth_decoder_config, BreezeTTS2DepthDecoderConfig) + assert config.num_codebooks == config.depth_decoder_config.num_codebooks == 16 + assert config.codec_config["codebook_size"] == 2048 + + +def test_transformers_five_roundtrip_preserves_nested_configs_and_rope(): + config = _config() + + roundtrip = BreezeTTS2Config(**dict(config.to_dict())) + + assert isinstance(roundtrip.backbone_config, Qwen3Config) + assert isinstance(roundtrip.text_encoder_config, BreezeTTS2TextEncoderConfig) + assert isinstance(roundtrip.depth_decoder_config, BreezeTTS2DepthDecoderConfig) + assert roundtrip.depth_decoder_config.rope_parameters["rope_theta"] == 500_000.0 + + +def test_nested_autoconfig_registration_is_idempotent(): + import importlib + + module = importlib.import_module("vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2") + importlib.reload(module) + + assert module.BreezeTTS2Config.model_type == "breeze" + + +def test_text_encoder_defaults_to_upstream_layer_pattern(): + config = BreezeTTS2TextEncoderConfig(num_hidden_layers=6) + + assert config.layer_types == [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] + + +def test_invalid_codebook_layout_is_rejected(): + with pytest.raises(ValueError, match="positive"): + BreezeTTS2DepthDecoderConfig(num_codebooks=0) diff --git a/tests/model_executor/models/breeze_tts_2/test_instruction.py b/tests/model_executor/models/breeze_tts_2/test_instruction.py new file mode 100644 index 000000000..3817ccb30 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_instruction.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, token): + return {"<|AUDIO|>": 262144, "<|audio_eos|>": 262145}.get(token, self.unk_token_id) + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [[1000 + ord(char) % 1000 for char in text] for text in texts] + if add_special_tokens: + ids = [[2, *row, 1] for row in ids] + return {"input_ids": ids} + + def decode(self, ids, *, skip_special_tokens=False): + del skip_special_tokens + return "".join(chr(int(token_id) - 1000) for token_id in ids if token_id not in (1, 2)) + + +def test_instruction_prompt_remains_one_independent_text_segment(): + builder = BreezeTTS2PromptBuilder( + _Tokenizer(), + SimpleNamespace( + audio_token_id=262144, + audio_eos_token_id=262145, + pad_token_id=0, + backbone_config=SimpleNamespace(vocab_size=100000), + num_codebooks=16, + codec_config={"codebook_size": 2048}, + ), + ) + + prompt = builder.build( + { + "template": "tts_instruction", + "speaker": "S2", + "instruction": "Use a bright and friendly delivery.", + "text": "Hello from Breeze.", + } + ) + + info = prompt["additional_information"] + assert info["template"] == "tts_instruction" + assert "input_values" not in info + assert bool(info["text_ids_mask"].all()) + assert info["text_ids_len"].tolist() == [len(info["prompt_ids"])] + + +def _request(extra_params=None, **overrides): + base = dict( + input="Hello.", + ref_text=None, + ref_audio=None, + task_type=None, + speed=1.0, + max_new_tokens=None, + language=None, + speaker_embedding=None, + x_vector_only_mode=None, + extra_params=extra_params or {}, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def _adapter() -> BreezeTTS2Adapter: + adapter = object.__new__(BreezeTTS2Adapter) + adapter.ctx = SimpleNamespace(server=SimpleNamespace()) + return adapter + + +def test_cfg_scale_one_is_explicitly_supported_without_companion_request(): + adapter = _adapter() + + assert adapter.validate(_request({"guidance_scale": 1.0})) is None + assert adapter.validate(_request({"cfg_scale": 1.0})) is None + + +def test_non_unit_cfg_is_rejected_until_negative_branch_exists(): + adapter = _adapter() + + assert "guidance_scale=1.0" in adapter.validate(_request({"guidance_scale": 2.0})) + assert "guidance_scale=1.0" in adapter.validate(_request({"cfg_scale": 0.5})) + assert "negative_prompt" in adapter.validate(_request({"negative_prompt": "noisy"})) + + +def test_unsupported_conditioning_fields_are_rejected_not_ignored(): + adapter = _adapter() + + assert "language" in adapter.validate(_request(language="en")) + assert "speaker_embedding" in adapter.validate(_request(speaker_embedding=[0.1, 0.2])) + assert "x_vector_only_mode" in adapter.validate(_request(x_vector_only_mode=True)) + # Absent fields must keep the request valid. + assert adapter.validate(_request()) is None + assert adapter.validate(_request(x_vector_only_mode=False)) is None + + +def test_multiple_reference_clips_are_rejected_not_silently_truncated(): + adapter = _adapter() + adapter.ctx.server._validate_ref_audio_format = lambda _uri: None + + error = adapter.validate(_request(ref_audio=["a.wav", "b.wav"], ref_text="hello")) + assert "exactly one reference clip" in error + assert "2" in error + # A single-element list stays valid and reaches format validation. + assert adapter.validate(_request(ref_audio=["a.wav"], ref_text="hello")) is None + + +def test_non_greedy_sampling_overrides_are_rejected(): + adapter = _adapter() + + # Greedy values, including the deploy-config defaults, stay valid. + assert adapter.validate(_request({"temperature": 0.0, "top_p": 1.0, "top_k": -1})) is None + assert adapter.validate(_request({"temperature": 0})) is None + # The talker ignores the sampled code0, so any non-greedy override must be + # refused instead of silently producing greedy audio. + assert "greedy" in adapter.validate(_request({"temperature": 0.7})) + assert "greedy" in adapter.validate(_request({"top_p": 0.9})) + assert "greedy" in adapter.validate(_request({"top_k": 50})) + assert "temperature must be a number" in adapter.validate(_request({"temperature": "warm"})) diff --git a/tests/model_executor/models/breeze_tts_2/test_local_modules.py b/tests/model_executor/models/breeze_tts_2/test_local_modules.py new file mode 100644 index 000000000..2ce54b09a --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_local_modules.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2DepthDecoderConfig, + BreezeTTS2TextEncoderConfig, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_depth import ( + BreezeDepthDecoderForCausalLM, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text import ( + BreezeTTS2TextEncoder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _depth_config() -> BreezeTTS2DepthDecoderConfig: + return BreezeTTS2DepthDecoderConfig( + hidden_size=16, + backbone_hidden_size=16, + audio_embed_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + max_position_embeddings=8, + num_codebooks=4, + vocab_size=8, + rope_scaling={ + "rope_type": "llama3", + "rope_theta": 100.0, + "factor": 1.0, + "low_freq_factor": 0.5, + "high_freq_factor": 1.0, + "original_max_position_embeddings": 4, + }, + ) + + +def test_depth_decoder_position_specific_logits_and_offsets(): + config = _depth_config() + decoder = BreezeDepthDecoderForCausalLM(config) + + output = decoder( + input_ids=torch.tensor([[0, 2]]), + backbone_last_hidden_state=torch.randn(1, 16), + use_cache=False, + ) + + assert output.logits.shape == (1, 1, 8) + with torch.no_grad(): + embedding = decoder.model.embed_tokens.weight + expected_first = embedding[2] @ decoder.model.inputs_embeds_projector.weight.T + # The first hidden row is replaced by the backbone state, so compare the + # second input row (codebook 0, id 2 + 0 * vocab_size). + actual_input = decoder.model.embed_tokens(torch.tensor([[2]]))[0, 0] + actual_projected = actual_input @ decoder.model.inputs_embeds_projector.weight.T + assert torch.allclose(expected_first, actual_projected) + + +def _text_config() -> BreezeTTS2TextEncoderConfig: + return BreezeTTS2TextEncoderConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + vocab_size=32, + layer_types=["sliding_attention", "full_attention"], + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 100.0}, + "full_attention": {"rope_type": "default", "rope_theta": 100.0}, + }, + ) + + +def test_text_encoder_masks_right_padding_independently(monkeypatch): + import vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text as text_module + + class _ParallelEmbedding(torch.nn.Embedding): + def __init__(self, num_embeddings: int, embedding_dim: int): + super().__init__(num_embeddings, embedding_dim) + + class _QKVLinear(torch.nn.Module): + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int, + **_: object, + ): + super().__init__() + output_size = (total_num_heads + 2 * total_num_kv_heads) * head_size + self.linear = torch.nn.Linear(hidden_size, output_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.linear(hidden_states), None + + class _RowLinear(torch.nn.Module): + def __init__(self, input_size: int, output_size: int, **_: object): + super().__init__() + self.linear = torch.nn.Linear(input_size, output_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.linear(hidden_states), None + + class _MergedLinear(torch.nn.Module): + def __init__(self, input_size: int, output_sizes: list[int], **_: object): + super().__init__() + self.linear = torch.nn.Linear(input_size, sum(output_sizes), bias=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.linear(hidden_states), None + + monkeypatch.setattr(text_module, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(text_module, "VocabParallelEmbedding", _ParallelEmbedding) + monkeypatch.setattr(text_module, "QKVParallelLinear", _QKVLinear) + monkeypatch.setattr(text_module, "RowParallelLinear", _RowLinear) + monkeypatch.setattr(text_module, "MergedColumnParallelLinear", _MergedLinear) + + config = _text_config() + encoder = BreezeTTS2TextEncoder(config).eval() + ids = torch.tensor([[3, 4]]) + padded_ids = torch.tensor([[3, 4, 0, 0]]) + mask = torch.tensor([[True, True, False, False]]) + + with torch.inference_mode(): + short = encoder(ids) + padded = encoder(padded_ids, attention_mask=mask) + + assert short.shape == (1, 2, 16) + assert padded.shape == (1, 4, 16) + assert torch.allclose(short[0, :2], padded[0, :2], atol=1e-5) + + +def test_text_encoder_weight_loader_reports_fused_target_names(): + encoder = object.__new__(BreezeTTS2TextEncoder) + parameter = torch.nn.Parameter(torch.zeros(2, 2)) + parameter.weight_loader = lambda param, tensor, shard_id: param.data.copy_(tensor) + encoder.__dict__["named_parameters"] = lambda: [("layers.0.qkv_proj.weight", parameter)] + + loaded = encoder.load_weights([("layers.0.q_proj.weight", torch.ones(2, 2))]) + + assert loaded == {"layers.0.qkv_proj.weight"} + assert torch.equal(parameter.detach(), torch.ones(2, 2)) diff --git a/tests/model_executor/models/breeze_tts_2/test_model_registration.py b/tests/model_executor/models/breeze_tts_2/test_model_registration.py new file mode 100644 index 000000000..12a2b31f3 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_model_registration.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +from vllm.model_executor.models.registry import ModelRegistry + +from vllm_omni.model_executor.models.registry import OmniModelRegistry + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_breeze_architectures_resolve_to_omni_classes(): + assert "BreezeForConditionalGeneration" in OmniModelRegistry.get_supported_archs() + assert "BreezeTTS2MimiCodec" in OmniModelRegistry.get_supported_archs() + + talker = OmniModelRegistry._try_load_model_cls("BreezeForConditionalGeneration") + codec = OmniModelRegistry._try_load_model_cls("BreezeTTS2MimiCodec") + assert talker.__name__ == "BreezeTTS2TalkerForGeneration" + assert codec.__name__ == "BreezeTTS2MimiCodec" + + +def test_plugin_registration_publishes_architectures_to_upstream_registry(): + from vllm_omni.engine.arg_utils import register_omni_models_to_vllm + + register_omni_models_to_vllm() + + assert "BreezeForConditionalGeneration" in ModelRegistry.get_supported_archs() + assert "BreezeTTS2MimiCodec" in ModelRegistry.get_supported_archs() diff --git a/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py b/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py new file mode 100644 index 000000000..069296942 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest +import yaml + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_breeze_pipeline_is_registered_with_talker_and_codec_stages(): + assert OMNI_PIPELINES["breeze"] is BREEZE_TTS_2_PIPELINE + assert [stage.model_stage for stage in BREEZE_TTS_2_PIPELINE.stages] == [ + "breeze_tts_2", + "breeze_tts_2_codec", + ] + assert BREEZE_TTS_2_PIPELINE.stages[1].requires_full_payload_input is True + assert BREEZE_TTS_2_PIPELINE.stages[0].async_chunk_process_next_stage_input_func.endswith( + "talker2codec_async_chunk" + ) + assert BREEZE_TTS_2_PIPELINE.stages[1].retains_state_across_chunks is True + + +def test_breeze_default_deploy_is_streaming(): + deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml" + config = yaml.safe_load(deploy_path.read_text(encoding="utf-8")) + assert config["async_chunk"] is True + assert len(config["stages"]) == 2 + # Stage 0 must prefill each prompt in one step: the talker builds a codec + # frame from the last hidden row of every scheduled request. + assert config["stages"][0]["enable_chunked_prefill"] is False + assert config["connectors"]["connector_of_shared_memory"]["extra"]["breeze_codec_chunk_frames"] == 8 diff --git a/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py b/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py new file mode 100644 index 000000000..3bd26fce6 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + pad_token_id = 0 + eos_token_id = 1 + unk_token_id = -1 + _special = { + "<|AUDIO|>": 262144, + "<|audio_eos|>": 262145, + } + + def convert_tokens_to_ids(self, token): + return self._special.get(token, self.unk_token_id) + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [] + for text in texts: + values = [1000 + ord(char) % 1000 for char in text] + if add_special_tokens: + values = [2, *values, 1] + ids.append(values) + return {"input_ids": ids} + + def decode(self, ids, *, skip_special_tokens=False): + del skip_special_tokens + rendered = [] + for token_id in ids: + if token_id == 262144: + rendered.append("<|AUDIO|>") + elif token_id == 262145: + rendered.append("<|audio_eos|>") + elif token_id in (2, 1): + continue + else: + rendered.append(chr(int(token_id) - 1000)) + return "".join(rendered) + + +class _RenderingTokenizer(_Tokenizer): + """Tiny tokenizer that preserves Breeze placeholders after decode.""" + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [] + for text in texts: + row = [] + cursor = 0 + while cursor < len(text): + if text.startswith("<|AUDIO|>", cursor): + row.append(262144) + cursor += len("<|AUDIO|>") + elif text.startswith("<|audio_eos|>", cursor): + row.append(262145) + cursor += len("<|audio_eos|>") + else: + row.append(1000 + ord(text[cursor]) % 1000) + cursor += 1 + if add_special_tokens: + row = [2, *row, 1] + ids.append(row) + return {"input_ids": ids} + + +class _ReferenceEncoder: + def __init__(self): + self.calls = [] + + def encode(self, audio, sample_rate=None): + self.calls.append((audio, sample_rate)) + return torch.arange(32, dtype=torch.int16).reshape(2, 16) + + +def _config(): + return SimpleNamespace( + audio_token_id=262144, + audio_eos_token_id=262145, + pad_token_id=0, + backbone_config=SimpleNamespace(vocab_size=100000), + num_codebooks=16, + codec_config={"codebook_size": 2048}, + ) + + +def test_builder_accepts_missing_codec_config(): + config = _config() + config.codec_config = None + + builder = BreezeTTS2PromptBuilder(_Tokenizer(), config) + + assert builder.codebook_size == 2048 + + +def test_build_instruction_prompt(): + builder = BreezeTTS2PromptBuilder(_Tokenizer(), _config()) + + prompt = builder.build({"text": "hello", "instruction": "happy", "speaker": "S0"}) + + assert prompt["prompt_token_ids"] + info = prompt["additional_information"] + assert "input_values" not in info + assert set(prompt["prompt_token_ids"]) == {0} + assert info["prompt_ids"].numel() == len(prompt["prompt_token_ids"]) + assert info["text_ids_mask"].dtype == torch.bool + assert info["text_ids_len"].tolist() == [len(prompt["prompt_token_ids"])] + assert bool(info["text_ids_mask"].all()) + + +def test_build_reference_prompt_and_reuse_encoder(): + encoder = _ReferenceEncoder() + builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config(), encoder) + + prompt = builder.build( + { + "template": "ref_edit_tata", + "text": "target", + "instruction": "calm", + "ref_text": "reference", + "ref_audio": "ref.wav", + "ref_audio_sample_rate": 16000, + } + ) + + info = prompt["additional_information"] + codes = info["input_values"] + assert tuple(codes.shape) == (2, 16) + assert codes.dtype == torch.int16 + assert set(prompt["prompt_token_ids"]) == {0} + assert info["prompt_ids"].tolist().count(262144) == 2 + assert info["prompt_ids"].tolist().count(262145) == 1 + assert int((~info["text_ids_mask"]).sum()) == 3 + assert info["text_ids_len"].numel() == 2 + assert encoder.calls == [("ref.wav", 16000)] + + +def test_build_reference_prompt_preserves_special_tokens_after_rendering(): + encoder = _ReferenceEncoder() + builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config(), encoder) + + prompt = builder.build( + { + "template": "ref_clone_tata", + "text": "target", + "ref_text": "reference", + "ref_audio": "ref.wav", + "ref_audio_sample_rate": 16000, + } + ) + + info = prompt["additional_information"] + prompt_ids = info["prompt_ids"].tolist() + assert prompt_ids.count(262144) == 2 + assert prompt_ids.count(262145) == 1 + assert int((~info["text_ids_mask"]).sum()) == 3 + assert len(info["text_ids_len"]) == 2 + assert sum(info["text_ids_len"].tolist()) == int(info["text_ids_mask"].sum()) + + +def test_build_reference_prompt_accepts_preencoded_codes_mapping(): + builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config()) + codes = torch.arange(32, dtype=torch.int16).reshape(2, 16) + + prompt = builder.build( + { + "template": "ref_clone_tata", + "text": "target", + "ref_text": "reference", + "codes": {"ref": codes}, + } + ) + + info = prompt["additional_information"] + assert torch.equal(info["input_values"], codes) + assert info["ref_code_len"] == 2 diff --git a/tests/model_executor/models/breeze_tts_2/test_reference_voice.py b/tests/model_executor/models/breeze_tts_2/test_reference_voice.py new file mode 100644 index 000000000..d29b863c4 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_reference_voice.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import soundfile as sf +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_EXAMPLE_PATH = ( + Path(__file__).resolve().parents[4] + / "examples" + / "offline_inference" + / "text_to_speech" + / "breeze_tts_2" + / "end2end.py" +) + + +class _Tokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, token): + return { + "<|AUDIO|>": 262144, + "<|audio_eos|>": 262145, + }.get(token, self.unk_token_id) + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [[1000 + ord(char) % 1000 for char in text] for text in texts] + for row_index, text in enumerate(texts): + if "<|AUDIO|>" in text or "<|audio_eos|>" in text: + cursor = 0 + row = [] + while cursor < len(text): + if text.startswith("<|AUDIO|>", cursor): + row.append(262144) + cursor += len("<|AUDIO|>") + elif text.startswith("<|audio_eos|>", cursor): + row.append(262145) + cursor += len("<|audio_eos|>") + else: + row.append(1000 + ord(text[cursor]) % 1000) + cursor += 1 + ids[row_index] = row + if add_special_tokens: + ids = [[2, *row, 1] for row in ids] + return {"input_ids": ids} + + def decode(self, ids, *, skip_special_tokens=False): + del skip_special_tokens + rendered = [] + for token_id in ids: + if token_id == 262144: + rendered.append("<|AUDIO|>") + elif token_id == 262145: + rendered.append("<|audio_eos|>") + elif token_id in (1, 2): + continue + else: + rendered.append(chr(int(token_id) - 1000)) + return "".join(rendered) + + +class _ReferenceEncoder: + def __init__(self) -> None: + self.calls: list[tuple[object, int | None]] = [] + + def encode(self, audio, sample_rate=None): + self.calls.append((audio, sample_rate)) + return torch.arange(32, dtype=torch.int16).reshape(2, 16) + + +def _builder() -> tuple[BreezeTTS2PromptBuilder, _ReferenceEncoder]: + encoder = _ReferenceEncoder() + builder = BreezeTTS2PromptBuilder( + _Tokenizer(), + SimpleNamespace( + audio_token_id=262144, + audio_eos_token_id=262145, + pad_token_id=0, + backbone_config=SimpleNamespace(vocab_size=100000), + num_codebooks=16, + codec_config={"codebook_size": 2048}, + ), + encoder, + ) + return builder, encoder + + +def test_reference_voice_uses_exact_transcript_and_audio_codes(): + builder, encoder = _builder() + + prompt = builder.build( + { + "template": "ref_clone_tata", + "speaker": "S1", + "ref_text": "exact reference transcript", + "ref_audio": "reference.wav", + "ref_audio_sample_rate": 24_000, + "text": "new target text", + } + ) + + info = prompt["additional_information"] + assert tuple(info["input_values"].shape) == (2, 16) + assert info["ref_code_len"] == 2 + assert info["prompt_ids"].tolist().count(262144) == 2 + assert info["text_ids_len"].numel() == 2 + assert int(info["text_ids_mask"].sum()) == int(info["text_ids_len"].sum()) + assert encoder.calls == [("reference.wav", 24_000)] + + +def test_reference_plus_instruction_composes_voice_direction_prompt(): + builder, _ = _builder() + + prompt = builder.build( + { + "template": "ref_edit_tata", + "speaker": "S0", + "ref_text": "exact reference transcript", + "ref_audio": "reference.wav", + "instruction": "speak calmly and warmly", + "text": "new target text", + } + ) + + info = prompt["additional_information"] + assert info["template"] == "ref_edit_tata" + assert tuple(info["input_values"].shape) == (2, 16) + assert info["prompt_ids"].tolist().count(262144) == 2 + assert info["text_ids_len"].numel() == 2 + + +def test_offline_example_reference_payload_feeds_prompt_builder(tmp_path): + """Smoke: the offline example's payload keys flow into the real builder. + + Loads ``end2end.py`` (voxtral-style importlib), reads a real stereo wav + through its ``reference_payload`` helper, and feeds the result into the + real ``BreezeTTS2PromptBuilder`` — guarding the example↔builder field + contract that previously broke via an unread ``ref_audio_path`` key. + """ + spec = importlib.util.spec_from_file_location("breeze_tts_2_end2end", _EXAMPLE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + sample_rate = 16_000 + t = np.arange(sample_rate, dtype=np.float32) / sample_rate + tone = np.sin(2 * np.pi * 220.0 * t) + ref_path = tmp_path / "reference.wav" + sf.write(str(ref_path), np.stack([tone, tone], axis=1), sample_rate) + + payload = module.reference_payload(str(ref_path)) + + # The helper emits exactly the keys the prompt builder consumes. + assert set(payload) == {"ref_audio", "ref_audio_sample_rate"} + assert payload["ref_audio_sample_rate"] == sample_rate + assert payload["ref_audio"].ndim == 1 # stereo down-mixed to one clip + + builder, encoder = _builder() + prompt = builder.build( + { + "template": "ref_clone_tata", + "speaker": "S0", + "ref_text": "exact reference transcript", + "text": "new target text", + **payload, + } + ) + + info = prompt["additional_information"] + assert tuple(info["input_values"].shape) == (2, 16) + assert info["ref_code_len"] == 2 + assert info["prompt_ids"].tolist().count(262144) == 2 + # The builder received the example's waveform and sample rate verbatim. + assert encoder.calls == [(payload["ref_audio"], sample_rate)] diff --git a/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py b/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py new file mode 100644 index 000000000..b9795f798 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import ( + talker2codec, + talker2codec_full_payload, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_talker2codec_flattens_codebook_major(): + codes = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int16) + output = SimpleNamespace(multimodal_output={"codes": {"audio": codes}}) + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [1, 3, 5, 2, 4, 6] + assert torch.equal(prompts[0]["multi_modal_data"]["codes"]["audio"], codes.transpose(0, 1)) + + +def test_talker2codec_keeps_empty_request_terminal(): + output = SimpleNamespace(multimodal_output={"codes": {}}) + + prompts = talker2codec([output]) + + assert prompts == [{"prompt_token_ids": []}] + + +def test_talker2codec_accepts_engine_core_output_wrapper(): + codes = torch.tensor([[7, 8]], dtype=torch.int16) + output = SimpleNamespace(outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": codes}})]) + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [7, 8] + + +def test_talker2codec_unwraps_per_request_audio_container(): + codes = torch.tensor([[7, 8], [9, 10]], dtype=torch.int16) + output = SimpleNamespace(multimodal_output={"codes": {"audio": [codes]}}) + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [7, 9, 8, 10] + + +def test_talker2codec_accepts_flat_full_payload(): + output = {"codes.audio": torch.tensor([1, 3, 2, 4], dtype=torch.int16)} + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [1, 3, 2, 4] + assert torch.equal( + prompts[0]["multi_modal_data"]["codes"]["audio"], + torch.tensor([1, 3, 2, 4], dtype=torch.int16), + ) + + +def test_talker2codec_full_payload_transposes_accumulated_frames(): + payload = talker2codec_full_payload( + transfer_manager=None, + pooling_output={"codes.audio": torch.tensor([[1, 2], [3, 4], [5, 6]])}, + request=None, + ) + + assert payload is not None + assert torch.equal( + payload.codes.audio, + torch.tensor([1, 3, 5, 2, 4, 6]), + ) + + +def test_talker2codec_full_payload_accepts_runner_pooling_output(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.tensor([[9, 8]], dtype=torch.int16)}, + request=None, + ) + + assert payload is not None + assert torch.equal(payload.codes.audio, torch.tensor([9, 8])) + + +def test_full_payload_zero_frame_finish_sends_explicit_terminal_marker(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.empty(0, 16, dtype=torch.long)}, + request=None, + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert payload.meta.code_flat_numel == 0 + assert bool(payload.meta.finished.item()) is True + + +def test_full_payload_missing_codes_finish_sends_explicit_terminal_marker(): + payload = talker2codec_full_payload( + pooling_output={}, + request=None, + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert bool(payload.meta.finished.item()) is True + + +def test_full_payload_empty_without_finish_stays_silent(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.empty(0, 16, dtype=torch.long)}, + request=None, + is_finished=False, + ) + + assert payload is None + + +def test_full_payload_carries_finished_flag_on_nonempty_frames(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}, + request=None, + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.numel() == 4 + assert bool(payload.meta.finished.item()) is True diff --git a/tests/model_executor/models/breeze_tts_2/test_streaming.py b/tests/model_executor/models/breeze_tts_2/test_streaming.py new file mode 100644 index 000000000..045444ec3 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_streaming.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_codec import ( + BreezeTTS2MimiCodec, +) +from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import ( + talker2codec_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _TransferManager: + def __init__(self, chunk_frames: int = 4): + self.connector = SimpleNamespace(config={"breeze_codec_chunk_frames": chunk_frames}) + self.code_prompt_token_ids = defaultdict(list) + + +class _Request: + def __init__(self, request_id: str, finished: bool = False): + self.external_req_id = request_id + self._finished = finished + + def is_finished(self): + return self._finished + + +def test_async_processor_sends_only_unemitted_tail_and_flushes_finish(): + manager = _TransferManager(chunk_frames=4) + request = _Request("req-1") + + for code0 in range(3): + frame = torch.tensor([code0, code0 + 4, code0 + 8, code0 + 12], dtype=torch.long) + assert talker2codec_async_chunk(manager, {"codes": {"audio": frame}}, request) is None + + payload = talker2codec_async_chunk( + manager, + {"codes": {"audio": torch.tensor([3, 7, 11, 15], dtype=torch.long)}}, + request, + ) + assert payload is not None + assert payload.codes.audio.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + assert bool(payload.meta.finished.item()) is False + assert manager.code_prompt_token_ids["req-1"] == [] + + assert talker2codec_async_chunk(manager, {"codes": {"audio": torch.arange(4)}}, request) is None + payload = talker2codec_async_chunk(manager, None, _Request("req-1", finished=True)) + assert payload.codes.audio.tolist() == [0, 1, 2, 3] + assert bool(payload.meta.finished.item()) is True + assert bool(payload.meta.stream_finished.item()) is True + assert payload.meta.codec_streaming is True + + +class _Decoder: + def __init__(self): + self.calls = [] + + def batched_chunked_decode(self, codes, lengths, caches=None, **kwargs): + self.calls.append((codes.shape, tuple(lengths), tuple(caches), kwargs)) + return [torch.ones(length * 1920, dtype=torch.float32) for length in lengths] + + +def test_stateful_codec_batches_request_chunks_and_releases_finished_state(): + decoder = _Decoder() + codec = object.__new__(BreezeTTS2MimiCodec) + codec._async_chunk = True + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + codec._decoder_state_cache = {} + + input_ids = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.long) + output = codec.forward( + input_ids, + runtime_additional_information=[ + # ``stream_finished`` is the model-visible flag: the connector + # receiver strips ``meta.finished`` while merging runtime info. + {"meta": {"request_id": "live", "stream_finished": False}}, + {"meta": {"request_id": "done", "stream_finished": True}}, + ], + seq_token_counts=[4, 4], + request_ids=["scheduler-live", "scheduler-done"], + ) + + assert decoder.calls[0][0] == (2, 4, 1) + assert [len(item) for item in output.multimodal_outputs["model_outputs"]] == [1920, 1920] + assert set(codec._decoder_state_cache) == {"scheduler-live"} + + codec.on_requests_finished(["scheduler-live"]) + assert codec._decoder_state_cache == {} + + +def test_stateful_codec_rejects_partial_frame_chunk(): + decoder = _Decoder() + codec = object.__new__(BreezeTTS2MimiCodec) + codec._async_chunk = True + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + codec._decoder_state_cache = {} + + # 5 codes cannot form a whole frame; a silent state reset would corrupt + # the rest of the stream, so the codec must fail loudly. + with pytest.raises(ValueError, match="not divisible"): + codec.forward( + torch.tensor([0, 1, 2, 3, 4], dtype=torch.long), + runtime_additional_information=[{"meta": {"request_id": "bad", "finished": False}}], + seq_token_counts=[5], + request_ids=["scheduler-bad"], + ) + assert decoder.calls == [] + + +def test_stateful_codec_terminal_marker_pops_state_without_decoding(): + decoder = _Decoder() + codec = object.__new__(BreezeTTS2MimiCodec) + codec._async_chunk = True + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + codec._decoder_state_cache = {"scheduler-live": {}, "scheduler-done": {}} + + output = codec.forward( + torch.tensor([0, 1, 2, 3], dtype=torch.long), + runtime_additional_information=[ + # Plain ``finished`` covers the fallback for direct in-process + # callers that do not go through the connector receiver. + {"meta": {"request_id": "live", "finished": False}}, + {"meta": {"request_id": "done", "finished": True}}, + ], + seq_token_counts=[4, 0], + request_ids=["scheduler-live", "scheduler-done"], + ) + + # The empty terminal marker decodes nothing and releases its state; the + # live request still decodes its own chunk. + assert len(decoder.calls) == 1 + assert decoder.calls[0][0][0] == 1 + assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [1920, 0] + assert list(codec._decoder_state_cache) == ["scheduler-live"] + assert codec._decoder_state_cache["scheduler-live"]["prefix_frames"] == 0 diff --git a/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py b/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py new file mode 100644 index 000000000..2f4a13df1 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( + BreezeTTS2TalkerForGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Head: + def __init__(self) -> None: + self.code0 = 2 + + def __call__(self, hidden: torch.Tensor) -> torch.Tensor: + logits = torch.zeros((hidden.shape[0], 9), dtype=hidden.dtype) + logits[:, self.code0] = 1.0 + return logits + + +class _Embedding: + def __init__(self) -> None: + self.weight = torch.randn(4 * 8, 16) + + def __call__(self, ids: torch.Tensor) -> torch.Tensor: + return self.weight[ids.reshape(-1)].reshape(*ids.shape, -1) + + +def _talker() -> BreezeTTS2TalkerForGeneration: + model = object.__new__(BreezeTTS2TalkerForGeneration) + model.num_codebooks = 4 + model.codebook_vocab_size = 8 + model.codebook_size = 8 + model.hidden_size = 16 + model.codebook_eos_token_id = 0 + model.audio_token_id = 100 + model.audio_eos_token_id = 101 + # Bypass nn.Module registration; this fake is created with object.__new__. + model.__dict__["embed_audio_tokens"] = _Embedding() + model.audio_token_offsets = torch.arange(4) * 8 + model.lm_head = _Head() + model.logits_processor = lambda head, hidden: head(hidden) + model._codec_disallowed_mask = None + model._batch_state = None + model._batch_state_spans = None + model._async_chunk = False + model._golden_dump_dir = None + return model + + +def _depth_codes(_model: Any, _hidden: torch.Tensor, code0: int) -> torch.Tensor: + if code0 >= 8: + return torch.empty(0, dtype=torch.long) + return torch.tensor([code0, code0 + 1, code0 + 2, code0 + 3], dtype=torch.long).unsqueeze(0) + + +def test_prefill_hidden_generates_and_stores_first_complete_frame(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"_omni_is_prefill": True, "breeze_generated_frames": 0, "breeze_max_new_frames": 3} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + codes = output.multimodal_outputs["codes"]["audio"][0] + assert torch.equal(codes, torch.tensor([[2, 3, 4, 5]])) + assert torch.equal(info["breeze_current_frame"], codes[0]) + assert info["breeze_generated_frames"] == 1 + assert info["breeze_force_eos"] is False + + +def test_first_decode_consumes_full_frame_instead_of_scalar_eos_fallback(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"_omni_is_prefill": True, "breeze_generated_frames": 0, "breeze_max_new_frames": 3} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + frame = output.multimodal_outputs["codes"]["audio"][0][0] + decode_info = dict(info) + decode_info.pop("_omni_is_prefill", None) + + safe_id, decode_embeds, update = model.preprocess( + torch.tensor([2], dtype=torch.long), + None, + **decode_info, + ) + + assert safe_id.item() == 2 + assert update == {} + expected = model.embed_input_ids(frame.reshape(1, 1, -1)).reshape(1, -1) + assert torch.equal(decode_embeds, expected) + + +def test_real_decode_rejects_missing_current_frame(): + model = _talker() + info = { + "breeze_generated_frames": 1, + "breeze_audio_codes": torch.tensor([[2, 3, 4, 5]]), + } + + try: + model.preprocess(torch.tensor([2], dtype=torch.long), None, **info) + except RuntimeError as exc: + assert "breeze_current_frame" in str(exc) + assert "codebooks 1..15" in str(exc) + else: + raise AssertionError("missing breeze_current_frame must fail explicitly") + + +def test_embed_input_ids_rejects_scalar_codec_id(): + model = _talker() + + try: + model.embed_input_ids(torch.tensor([[2]])) + except ValueError as exc: + assert "complete codec frames" in str(exc) + else: + raise AssertionError("scalar codec ids must not be embedded") + + +def test_decode_accumulates_frames_then_natural_eos_stops_without_fake_frame(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 1, "breeze_max_new_frames": 3} + info["breeze_audio_codes"] = torch.tensor([[2, 3, 4, 5]]) + model.lm_head.code0 = 3 + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) + + # Main-head class 8 is the Breeze backbone EOS (one past vocab_size=8). + model.lm_head.code0 = 8 + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) + assert info["breeze_force_eos"] is True + assert info["breeze_force_eos"] is True + + logits = model.compute_logits(hidden) + assert torch.equal(logits[0, :8], torch.full((8,), float("-inf"))) + assert logits[0, 8].item() == 0.0 + + +def test_frame_budget_forces_eos_on_followup_step(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 2, "breeze_max_new_frames": 2} + info["breeze_audio_codes"] = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 0]]) + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) + assert info["breeze_generated_frames"] == 2 + assert info["breeze_force_eos"] is True + + +def test_async_chunk_mode_skips_full_accumulation(): + model = _talker() + model._async_chunk = True + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 0, "breeze_max_new_frames": 8} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + for _ in range(3): + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 4) + + # Streaming keeps only the per-frame tail; the cumulative buffer (and its + # per-step torch.cat) must not exist without the opt-in golden dump. + assert "breeze_audio_codes" not in info + assert info["breeze_generated_frames"] == 3 + assert torch.equal(info["breeze_current_frame"], torch.tensor([2, 3, 4, 5])) + + +def test_async_chunk_mode_keeps_accumulation_for_golden_dump(tmp_path): + model = _talker() + model._async_chunk = True + model._golden_dump_dir = str(tmp_path) + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 0, "breeze_max_new_frames": 8} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + for _ in range(3): + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 4) + + assert info["breeze_audio_codes"].shape == (3, 4) + assert torch.equal(info["breeze_audio_codes"][-1], torch.tensor([2, 3, 4, 5])) + + +def test_opt_in_golden_dump_writes_terminal_frames(tmp_path): + model = _talker() + model._golden_dump_dir = str(tmp_path) + info = { + "template": "tts_plain", + "prompt_ids": torch.tensor([1, 2, 3]), + "breeze_audio_codes": torch.tensor([[1, 2, 3, 4]]), + "breeze_generated_frames": 1, + "breeze_force_eos": True, + } + + model._dump_golden_frames(info) + + dump_path = next(tmp_path.glob("breeze_*.pt")) + payload = torch.load(dump_path, weights_only=False) + assert payload["template"] == "tts_plain" + assert torch.equal(payload["prompt_ids"], torch.tensor([1, 2, 3])) + assert torch.equal(payload["codes"], torch.tensor([[1, 2, 3, 4]])) + assert payload["finished"] is True diff --git a/tests/model_executor/models/breeze_tts_2/test_weight_loading.py b/tests/model_executor/models/breeze_tts_2/test_weight_loading.py new file mode 100644 index 000000000..f4c0da429 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_weight_loading.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( + BreezeTTS2TalkerForGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _PaddedHead(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(2_112, 32), requires_grad=False) + + +class _DepthDecoder: + def __init__(self) -> None: + self.loaded: list[tuple[str, torch.Tensor]] = [] + + def load_weights(self, weights): + self.loaded.extend(weights) + return {"model.embed_tokens.weight"} + + +class _TalkerHost(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(tie_codebooks_embeddings=True) + self.codebook_vocab_size = 2_051 + self.hidden_size = 32 + self.lm_head = _PaddedHead() + self.embed_audio_tokens = nn.Embedding(4 * 2_051, 32) + self.depth_decoder = _DepthDecoder() + + +def _padded_loader(parameter: nn.Parameter, weight: torch.Tensor) -> None: + assert parameter.shape != weight.shape + parameter[: weight.shape[0]].copy_(weight) + + +def test_lm_head_loader_accepts_vocab_padding(): + host = _TalkerHost() + host.lm_head.weight.weight_loader = _padded_loader + checkpoint_weight = torch.full((2_052, 32), 1.5) + + loaded = BreezeTTS2TalkerForGeneration.load_weights( + host, + iter([("lm_head.weight", checkpoint_weight)]), + ) + + assert "lm_head.weight" in loaded + assert torch.equal(host.lm_head.weight[:2_052], checkpoint_weight) + assert torch.count_nonzero(host.lm_head.weight[2_052:]).item() == 0 + + +def test_tied_audio_embedding_reports_both_strict_loading_aliases(): + host = _TalkerHost() + checkpoint_weight = torch.full((4 * 2_051, 32), 2.0) + + loaded = BreezeTTS2TalkerForGeneration.load_weights( + host, + iter([("depth_decoder.model.embed_tokens.weight", checkpoint_weight)]), + ) + + assert "embed_audio_tokens.weight" in loaded + assert "depth_decoder.model.embed_tokens.weight" in loaded diff --git a/tools/breeze_tts_2_compare_golden.py b/tools/breeze_tts_2_compare_golden.py new file mode 100644 index 000000000..23306896e --- /dev/null +++ b/tools/breeze_tts_2_compare_golden.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Compare upstream and vLLM-Omni Breeze-TTS-2 golden artifacts.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import soundfile as sf +import torch + +TEMPLATES = { + "plain": "tts_plain", + "instruction": "tts_instruction", + "reference_edit": "ref_edit_tata", +} + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _audio_metrics(path: Path) -> dict[str, object]: + data, sample_rate = sf.read(path, dtype="float32") + return { + "path": str(path), + "bytes": path.stat().st_size, + "sha256": _sha256(path), + "samples": int(data.shape[0]), + "sample_rate": int(sample_rate), + "channels": int(1 if data.ndim == 1 else data.shape[1]), + "duration_s": round(float(data.shape[0]) / sample_rate, 6), + "finite": bool(torch.isfinite(torch.from_numpy(data)).all().item()), + "peak": round(float(abs(data).max()), 6), + "rms": round(float((data**2).mean() ** 0.5), 6), + } + + +def _vector_metrics(upstream: torch.Tensor, actual: torch.Tensor) -> dict[str, float]: + difference = upstream.double() - actual.double() + return { + "mae": round(float(difference.abs().mean()), 8), + "rmse": round(float(difference.pow(2).mean().sqrt()), 8), + "max_abs_error": round(float(difference.abs().max()), 8), + "relative_l2": round(float(difference.norm() / upstream.norm()), 8), + "cosine": round( + float(torch.dot(upstream.double(), actual.double()) / (upstream.norm() * actual.norm())), + 8, + ), + } + + +def _load_vllm_payload(vllm_dir: Path, template: str) -> dict[str, object]: + matches = [] + for path in vllm_dir.glob("breeze_*.pt"): + payload = torch.load(path, weights_only=False) + if payload.get("template") == template: + matches.append(payload) + if len(matches) != 1: + raise ValueError(f"expected one vLLM dump for {template!r}, found {len(matches)}") + return matches[0] + + +def _longest_exact_frame_prefix(upstream: torch.Tensor, actual: torch.Tensor) -> int: + count = min(int(upstream.shape[0]), int(actual.shape[0])) + equal = (upstream[:count] == actual[:count]).all(dim=-1) + prefix = 0 + for is_equal in equal.tolist(): + if not is_equal: + break + prefix += 1 + return prefix + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--upstream-dir", type=Path, required=True) + parser.add_argument("--vllm-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--plain-audio", type=Path) + parser.add_argument("--instruction-audio", type=Path) + parser.add_argument("--reference-audio", type=Path) + args = parser.parse_args() + + audio_paths = { + "plain": args.plain_audio, + "instruction": args.instruction_audio, + "reference_edit": args.reference_audio, + } + report: dict[str, object] = { + "upstream_dir": str(args.upstream_dir.resolve()), + "vllm_dir": str(args.vllm_dir.resolve()), + "cases": [], + } + for name, template in TEMPLATES.items(): + upstream_codes = torch.load(args.upstream_dir / name / "codes.pt", weights_only=False) + upstream_prompt = torch.load(args.upstream_dir / name / "prompt_ids.pt", weights_only=False) + actual_payload = _load_vllm_payload(args.vllm_dir, template) + actual_codes = actual_payload["codes"] + actual_prompt = actual_payload["prompt_ids"] + common_frames = min(int(upstream_codes.shape[0]), int(actual_codes.shape[0])) + code_equal = upstream_codes[:common_frames] == actual_codes[:common_frames] + frame_equal = code_equal.all(dim=-1) + prefill_metrics = None + upstream_hidden_path = args.upstream_dir / name / "prefill_hidden.pt" + if upstream_hidden_path.exists(): + prefill_files = [ + path + for path in args.vllm_dir.glob("prefill_*.pt") + if torch.load(path, weights_only=False).get("template") == template + ] + if len(prefill_files) == 1: + actual_hidden = torch.load(prefill_files[0], weights_only=False)["hidden"] + prefill_metrics = _vector_metrics( + torch.load(upstream_hidden_path, weights_only=False), + actual_hidden, + ) + entry: dict[str, object] = { + "name": name, + "template": template, + "prompt": { + "upstream_tokens": int(upstream_prompt.numel()), + "vllm_tokens": int(actual_prompt.numel()), + "exact": bool(torch.equal(upstream_prompt, actual_prompt)), + }, + "frames": { + "upstream": int(upstream_codes.shape[0]), + "vllm": int(actual_codes.shape[0]), + "codebooks": int(upstream_codes.shape[-1]), + "exact_sequence": bool(torch.equal(upstream_codes, actual_codes)), + "common_frames": common_frames, + "exact_common_frames": int(frame_equal.sum()), + "longest_exact_prefix": _longest_exact_frame_prefix(upstream_codes, actual_codes), + "exact_codes": int(code_equal.sum()), + "total_common_codes": int(code_equal.numel()), + "code_accuracy": round(float(code_equal.float().mean()), 8), + "codebook0_accuracy": round( + float((upstream_codes[:common_frames, 0] == actual_codes[:common_frames, 0]).float().mean()), + 8, + ), + "range_valid": bool( + upstream_codes.min() >= 0 + and upstream_codes.max() < 2048 + and actual_codes.min() >= 0 + and actual_codes.max() < 2048 + ), + }, + "prefill_hidden": prefill_metrics, + } + if audio_paths[name] is not None and Path(audio_paths[name]).exists(): # type: ignore[arg-type] + entry["audio"] = _audio_metrics(Path(audio_paths[name])) # type: ignore[arg-type] + report["cases"].append(entry) + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/breeze_tts_2_upstream_golden.py b/tools/breeze_tts_2_upstream_golden.py new file mode 100644 index 000000000..337a0444f --- /dev/null +++ b/tools/breeze_tts_2_upstream_golden.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Generate deterministic Breeze-TTS-2 golden frames with the upstream runtime. + +This tool is intentionally separate from vLLM-Omni: it imports the unmodified +``breezeblue-ai/breeze-tts`` reference implementation, fixes greedy decoding, +and writes both generated codec frames and decoded audio for offline parity +checks. It is a diagnostic tool, not a serving path. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import soundfile as sf +import torch + + +def _plain_template(upstream_templates): + """The upstream CLI does not register tts_plain; reuse its segment builder.""" + return SimpleNamespace( + name="tts_plain", + required_fields=("text",), + build_segments=upstream_templates._tts_plain_segments, + build_negative_segments=None, + build_dual_branches=None, + ) + + +def _greedy(model) -> None: + model.generation_config.do_sample = False + model.generation_config.temperature = None + model.generation_config.top_p = None + model.generation_config.top_k = None + depth = model.depth_decoder.generation_config + depth.depth_decoder_do_sample = False + depth.do_sample = False + depth.temperature = None + depth.top_p = None + depth.top_k = None + + +def _valid_frames(sequences: torch.Tensor, pad_token_id: int) -> torch.Tensor: + frames = sequences[0].detach().cpu().to(torch.long) + if frames.ndim == 1: + frames = frames.unsqueeze(0) + is_pad = (frames == pad_token_id).all(dim=-1) + if bool(is_pad.any()): + frames = frames[: int(torch.argmax(is_pad.int()).item())] + return frames.contiguous() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--upstream", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--reference-audio", type=Path) + parser.add_argument("--max-new-tokens", type=int, default=32) + args = parser.parse_args() + + args.output_dir.mkdir(parents=True, exist_ok=True) + upstream_posix = str(args.upstream.resolve()) + if upstream_posix not in sys.path: + sys.path.insert(0, upstream_posix) + + from breeze_infer.runtime import load_runtime, resolve_device + from breeze_infer.templates import get_template, prepare_inputs + + device = resolve_device() + # ``torch.manual_seed`` seeds CUDA RNGs as well, so no device-specific + # seeding call is needed here. + torch.manual_seed(0) + tokenizer, model, audio_tokenizer = load_runtime( + args.model, + device=device, + attn_implementation="eager", + ) + # Do not enable Transformers 4.x's ``fix_mistral_regex`` compatibility + # path: vLLM-Omni runs with Transformers 5.x, and direct tokenization checks + # showed both runtimes produce the same token IDs without that patch. + _greedy(model) + model.eval() + first_backbone_hidden: list[torch.Tensor] = [] + + def _capture_first_backbone_hidden(_module, _inputs, output): + if first_backbone_hidden: + return + hidden = getattr(output, "last_hidden_state", None) + if isinstance(hidden, torch.Tensor): + first_backbone_hidden.append(hidden[0, -1].detach().cpu().to(torch.float32).contiguous()) + + model.backbone_model.register_forward_hook(_capture_first_backbone_hidden) + + cases = [ + { + "name": "plain", + "template": _plain_template(sys.modules["breeze_infer.templates"]), + "request": { + "id": "golden-plain", + "text": "Golden plain alignment check.", + "speaker": "S0", + }, + }, + { + "name": "instruction", + "template": get_template("tts_instruction"), + "request": { + "id": "golden-instruction", + "text": "Golden instruction alignment check.", + "instruction": "Speak calmly, warmly, and clearly.", + "speaker": "S0", + }, + }, + ] + if args.reference_audio is not None: + cases.append( + { + "name": "reference_edit", + "template": get_template("ref_edit_tata"), + "request": { + "id": "golden-reference-edit", + "text": "Golden reference alignment check.", + "instruction": "Speak calmly and clearly.", + "ref_audio_path": str(args.reference_audio), + "ref_text": "Hello from Breeze.", + "speaker": "S0", + }, + } + ) + + manifest = { + "runtime": "upstream breezeblue-ai/breeze-tts", + "model": str(args.model.resolve()), + "upstream": upstream_posix, + "device": device, + "seed": 0, + "max_new_tokens": args.max_new_tokens, + "decoding": "greedy", + "cfg_scale": 1.0, + "cases": [], + } + for case in cases: + first_backbone_hidden.clear() + inputs = prepare_inputs( + tokenizer, + audio_tokenizer, + model, + [case["request"]], + case["template"], + guidance_scale=1.0, + guidance_scale_ref=None, + guidance_scale_ins=None, + ) + output = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + return_dict_in_generate=True, + output_audio=True, + audio_tokenizer=audio_tokenizer, + ) + frames = _valid_frames(output.sequences, model.config.codebook_pad_token_id) + audio = output.audio[0].detach().cpu().to(torch.float32).reshape(-1) + case_dir = args.output_dir / case["name"] + case_dir.mkdir(parents=True, exist_ok=True) + torch.save(frames, case_dir / "codes.pt") + prompt_ids = inputs["input_ids"].detach().cpu().to(torch.long).reshape(-1) + torch.save(prompt_ids, case_dir / "prompt_ids.pt") + if first_backbone_hidden: + torch.save(first_backbone_hidden[0], case_dir / "prefill_hidden.pt") + sf.write(case_dir / "upstream.wav", audio.numpy(), 24_000, subtype="PCM_16") + entry = { + "name": case["name"], + "request": case["request"], + "prompt_tokens": int(prompt_ids.numel()), + "frames": int(frames.shape[0]), + "codebooks": int(frames.shape[-1]) if frames.ndim == 2 else 0, + "min_code": int(frames.min().item()) if frames.numel() else None, + "max_code": int(frames.max().item()) if frames.numel() else None, + "audio_samples": int(audio.numel()), + "audio_sha256": __import__("hashlib").sha256((case_dir / "upstream.wav").read_bytes()).hexdigest(), + } + manifest["cases"].append(entry) + print(json.dumps(entry), flush=True) + + (args.output_dir / "manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index fd7617144..b5214720a 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -201,6 +201,7 @@ _MODEL_INTERNALIZE = ( "MING_TTS_STAGE1_FINAL_LOG", "MOSS_TTS_DEBUG_STOP", "NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", + "VLLM_OMNI_BREEZE_GOLDEN_DIR", "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS", "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY", "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 2378ce5b9..92aad9f4d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -47,6 +47,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_SINGLE_STAGE_PIPELINE, BAGEL_THINK_PIPELINE, ) +from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE @@ -128,6 +129,7 @@ PipelineResolverFunc: TypeAlias = Callable[[PretrainedConfig | None], PipelineCo # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) --- OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "aura_omni": AURA_OMNI_PIPELINE, + "breeze": BREEZE_TTS_2_PIPELINE, "joyai_vl_interaction": JOYAI_VL_INTERACTION_PIPELINE, "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE, "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, diff --git a/vllm_omni/deploy/breeze_tts_2.yaml b/vllm_omni/deploy/breeze_tts_2.yaml new file mode 100644 index 000000000..1d1b85391 --- /dev/null +++ b/vllm_omni/deploy/breeze_tts_2.yaml @@ -0,0 +1,51 @@ +# Breeze-TTS-2: talker -> stateful Qwen3-TTS codec decoder. +# Stage 0 batches codec frames and sends incremental codebook-major chunks. +async_chunk: true +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + breeze_codec_chunk_frames: 8 + +stages: + - stage_id: 0 + max_num_seqs: 8 + gpu_memory_utilization: 0.80 + enable_prefix_caching: false + # The talker derives one codec frame from the last hidden row of every + # scheduled request. A partially scheduled prompt would yield a frame from + # a mid-prompt hidden state, so keep each prefill in a single step. + enable_chunked_prefill: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + min_tokens: 1 + + - stage_id: 1 + max_num_seqs: 8 + gpu_memory_utilization: 0.15 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 20319de86..d1e53775a 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -54,6 +54,9 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig + from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2Config, + ) from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import ( IndexTTS2Config, IndexTTS25Config, @@ -105,6 +108,7 @@ def _register_omni_hf_configs() -> None: ("glm_tts", GLMTTSConfig), ("omnivoice", OmniVoiceConfig), ("voxcpm2", VoxCPM2Config), + ("breeze", BreezeTTS2Config), ]: try: AutoConfig.register(model_type, config_cls) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index 39882274d..8ed54afd8 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -105,6 +105,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 audex, audex_tta, audio8_tts, + breeze_tts_2, cosyvoice3, covo_audio, dots_tts, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py new file mode 100644 index 000000000..d2851e32c --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 serving adapter for synchronous full-payload generation.""" + +import asyncio +from typing import TYPE_CHECKING, Any + +import numpy as np + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) +from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + BreezeReferenceAudioTokenizer, +) +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + + +@register_tts_adapter +class BreezeTTS2Adapter(ARTTSAdapter): + """Build Breeze prompts before scheduler submission. + + The adapter owns only request-facing concerns. Model workers still own the + stage-0 embeddings and stage-1 waveform decode. + """ + + stage_keys = frozenset({"breeze_tts_2"}) + model_archs = frozenset({"BreezeForConditionalGeneration"}) + name = "breeze_tts_2" + # Breeze's architecture and stage key are unique, so this detector cannot + # collide with another adapter. Lower priority runs first; 9 keeps it + # ahead of the default-priority (100) detectors. + detect_priority = 9 + supported_output_sample_rates = frozenset({24000}) + + def _load_supported_speakers(self) -> set[str]: + # Breeze accepts arbitrary speaker tags such as S0/S1 in its prompt; + # there is no finite checkpoint speaker inventory to advertise. + return set() + + def _load_codec_frame_rate(self) -> float | None: + config = self.ctx.engine_client.model_config.hf_config + codec = getattr(config, "codec_config", None) + if isinstance(codec, dict): + frame_rate = codec.get("_frame_rate") + if frame_rate is not None: + return float(frame_rate) + return 12.5 + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + if not request.input or not request.input.strip(): + return "Breeze-TTS-2 input text cannot be empty" + if request.ref_text is not None and not request.ref_text.strip(): + return "Breeze-TTS-2 ref_text cannot be empty" + if request.ref_text is not None and request.ref_audio is None: + return "Breeze-TTS-2 ref_text requires ref_audio" + if request.ref_audio is not None and (request.ref_text is None or not request.ref_text.strip()): + return "Breeze-TTS-2 ref_audio requires ref_text" + if isinstance(request.ref_audio, list) and len(request.ref_audio) != 1: + # The prompt path conditions on exactly one reference waveform; + # extra clips would be silently discarded otherwise. + return f"Breeze-TTS-2 supports exactly one reference clip, got {len(request.ref_audio)}" + if request.task_type == "Base" and request.ref_audio is None: + return "Breeze-TTS-2 Base task requires ref_audio and ref_text" + if request.task_type == "VoiceDesign": + return "Breeze-TTS-2 does not support task_type=VoiceDesign" + if request.speed is not None and request.speed != 1.0: + return "Breeze-TTS-2 does not support speed adjustment" + if request.language is not None: + return "Breeze-TTS-2 does not support 'language'; the prompt language follows the input text" + if request.speaker_embedding is not None: + return "Breeze-TTS-2 does not support 'speaker_embedding' (voice-cloning helper field)" + if request.x_vector_only_mode: + return "Breeze-TTS-2 does not support 'x_vector_only_mode' (voice-cloning helper field)" + extra_params = request.extra_params or {} + guidance_scale = extra_params.get("guidance_scale", extra_params.get("cfg_scale", 1.0)) + try: + unsupported_guidance = guidance_scale is not None and float(guidance_scale) != 1.0 + except (TypeError, ValueError): + return "Breeze-TTS-2 guidance_scale must be a number" + if unsupported_guidance: + return "Breeze-TTS-2 currently supports only guidance_scale=1.0" + if extra_params.get("negative_prompt"): + return "Breeze-TTS-2 does not support negative_prompt/CFG yet" + # The talker selects codebook 0 by argmax inside make_omni_output and + # ignores the scheduler's sampled token, so non-greedy overrides would + # not change the audio while still letting the sampler draw an early + # EOS. Reject them until the talker consumes the sampled code. + greedy_error = _validate_greedy_sampling(extra_params) + if greedy_error is not None: + return greedy_error + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + if request.ref_audio is not None: + ref_audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio + return self.ctx.server._validate_ref_audio_format(ref_audio) + return None + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del has_inline_ref_audio + server = self.ctx.server + builder = await self._get_builder() + template = self._resolve_template(request) + payload: dict[str, Any] = { + "text": request.input, + "template": template, + "speaker": request.voice or "S0", + } + if request.instructions: + payload["instruction"] = request.instructions + if request.ref_text: + payload["ref_text"] = request.ref_text + + if request.ref_audio is not None: + ref_audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio + wav_list, sr, cache_key = await server._resolve_ref_audio(ref_audio) + payload["ref_audio"] = np.asarray(wav_list, dtype=np.float32) + payload["ref_audio_sample_rate"] = int(sr) + else: + cache_key = None + + prompt = await asyncio.to_thread(builder.build, payload, template=template) + # The AR scheduler's completion budget is measured in codebook-0 + # frames for Breeze. Preserve it in the mutable request metadata so + # the talker can emit its own EOS before a scheduler length cutoff. + max_new_frames = request.max_new_tokens + if max_new_frames is None and sampling_params_list: + max_new_frames = getattr(sampling_params_list[0], "max_tokens", None) + if max_new_frames is not None: + prompt["additional_information"]["breeze_max_new_frames"] = int(max_new_frames) + tts_params = { + "template": [template], + "text": [request.input], + } + if cache_key: + tts_params["ref_audio_cache_key"] = cache_key + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) + return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + + async def _get_builder(self) -> BreezeTTS2PromptBuilder: + server = self.ctx.server + cached = getattr(server, "_breeze_tts_2_prompt_builder", None) + if cached is not None: + return cached + # Multiple speech requests can arrive while the first tokenizer load + # is still in progress. Serialize the one-time CPU initialization so + # model files are not opened and decoded repeatedly. + lock = getattr(server, "_breeze_tts_2_prompt_builder_lock", None) + if lock is None: + lock = asyncio.Lock() + server._breeze_tts_2_prompt_builder_lock = lock + async with lock: + cached = getattr(server, "_breeze_tts_2_prompt_builder", None) + if cached is not None: + return cached + model_path = server.engine_client.model_config.model + config = server.engine_client.model_config.hf_config + audio_tokenizer = await asyncio.to_thread( + BreezeReferenceAudioTokenizer.from_pretrained, + model_path, + num_codebooks=int(getattr(config, "num_codebooks", 16)), + codebook_size=self._codebook_size(config), + device_map="cpu", + ) + builder = await asyncio.to_thread( + BreezeTTS2PromptBuilder.from_pretrained, + model_path, + config, + reference_audio_encoder=audio_tokenizer, + ) + server._breeze_tts_2_prompt_builder = builder + return builder + + @staticmethod + def _codebook_size(config: Any) -> int: + codec = getattr(config, "codec_config", None) + if isinstance(codec, dict): + return int(codec.get("codebook_size", 2048)) + return int(getattr(codec, "codebook_size", 2048)) + + @staticmethod + def _resolve_template(request: "OpenAICreateSpeechRequest") -> str: + if request.ref_audio is not None: + return "ref_edit_tata" if request.instructions else "ref_clone_tata" + return "tts_instruction" if request.instructions else "tts_plain" + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + del prompt, request_id + return apply_max_new_tokens(sampling_params_list, request) + + +def _validate_greedy_sampling(extra_params: dict[str, Any]) -> str | None: + """Return an error unless the sampling overrides keep decoding greedy.""" + temperature = extra_params.get("temperature") + if temperature is not None: + try: + if float(temperature) != 0.0: + return "Breeze-TTS-2 currently supports only greedy decoding (temperature=0)" + except (TypeError, ValueError): + return "Breeze-TTS-2 temperature must be a number" + top_p = extra_params.get("top_p") + if top_p is not None: + try: + if float(top_p) != 1.0: + return "Breeze-TTS-2 currently supports only greedy decoding (top_p=1.0)" + except (TypeError, ValueError): + return "Breeze-TTS-2 top_p must be a number" + top_k = extra_params.get("top_k") + if top_k is not None: + try: + if int(top_k) not in (-1, 0, 1): + return "Breeze-TTS-2 currently supports only greedy decoding (top_k=-1)" + except (TypeError, ValueError): + return "Breeze-TTS-2 top_k must be an integer" + return None + + +__all__ = ["BreezeTTS2Adapter"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/__init__.py b/vllm_omni/model_executor/models/breeze_tts_2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py b/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py new file mode 100644 index 000000000..ea2796a53 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Reference-audio encoding helpers for Breeze-TTS-2. + +The Breeze prompt uses a Qwen3-TTS tokenizer only for reference audio. This +module deliberately does not contain text tokenization or waveform decoding; +the latter belongs to the pipeline's codec stage. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch + + +def resolve_audio_tokenizer_path(model_path: str) -> Path | None: + """Resolve the bundled ``audio_tokenizer`` dir for a local path or HF repo id. + + Local directories (dev checkouts) resolve directly: present returns the + subdirectory, absent returns ``None`` so callers can fall back. Repo ids + resolve through the vLLM-tagged HF helpers: an already-complete cached + snapshot is preferred (offline deployments keep working), and a partial + cache falls through to a targeted online fetch restricted to the bundled + tokenizer subtree. Returns ``None`` when the checkpoint carries no + ``audio_tokenizer`` subdirectory. + """ + root = Path(model_path) + if root.is_dir(): + bundled = root / "audio_tokenizer" + return bundled if bundled.is_dir() else None + + from vllm.transformers_utils.repo_utils import hf_api + + kwargs: dict[str, object] = {"allow_patterns": ["audio_tokenizer/*"]} + try: + cached = Path(hf_api().snapshot_download(model_path, local_files_only=True, **kwargs)) + except Exception: + cached = None + if cached is not None and (cached / "audio_tokenizer").is_dir(): + return cached / "audio_tokenizer" + online = Path(hf_api().snapshot_download(model_path, **kwargs)) + bundled = online / "audio_tokenizer" + return bundled if bundled.is_dir() else None + + +class BreezeReferenceAudioTokenizer: + """Small adapter that normalizes Qwen3-TTS encoder output to ``(T, Q)``. + + The underlying tokenizer is created once per serving worker and is reused + for all requests. Keeping this wrapper independent from the + prompt builder makes it straightforward to inject a fake encoder in tests + and to replace the upstream Qwen3 implementation later. + """ + + def __init__(self, tokenizer: Any, *, num_codebooks: int = 16, codebook_size: int = 2048) -> None: + self.tokenizer = tokenizer + self.num_codebooks = int(num_codebooks) + self.codebook_size = int(codebook_size) + if self.num_codebooks <= 0 or self.codebook_size <= 0: + raise ValueError("num_codebooks and codebook_size must be positive") + + @classmethod + def from_pretrained( + cls, + model_path: str, + *, + num_codebooks: int = 16, + codebook_size: int = 2048, + **kwargs: Any, + ) -> BreezeReferenceAudioTokenizer: + """Load the bundled ``audio_tokenizer`` with the existing Qwen wrapper.""" + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import ( + Qwen3TTSTokenizer, + ) + + resolved = resolve_audio_tokenizer_path(model_path) + if resolved is None: + raise FileNotFoundError( + "Breeze reference audio tokenizer was not found for " + f"{model_path} (expected an audio_tokenizer/ subdirectory)" + ) + tokenizer = Qwen3TTSTokenizer.from_pretrained(str(resolved), **kwargs) + return cls(tokenizer, num_codebooks=num_codebooks, codebook_size=codebook_size) + + @torch.inference_mode() + def encode(self, audio: Any, sample_rate: int | None = None) -> torch.Tensor: + """Encode one waveform or audio source into ``(T, Q)`` codes.""" + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu().numpy() + elif isinstance(audio, list) and audio and all(isinstance(x, (int, float)) for x in audio): + audio = np.asarray(audio, dtype=np.float32) + if isinstance(audio, np.ndarray) and audio.ndim > 1: + # Accept both common layouts: (samples, channels) and + # (channels, samples). Qwen3TTSTokenizer expects one 1-D clip. + channel_first = audio.shape[0] <= 8 and audio.shape[-1] > audio.shape[0] + audio = audio.mean(axis=0 if channel_first else -1) + + if isinstance(audio, np.ndarray) and sample_rate is None: + raise ValueError("sample_rate is required when ref_audio is a numpy waveform") + + encoded = self.tokenizer.encode(audio, sr=sample_rate) + return self._normalize_codes(encoded) + + def _normalize_codes(self, encoded: Any) -> torch.Tensor: + codes = getattr(encoded, "audio_codes", None) + if codes is None and isinstance(encoded, Mapping): + codes = encoded.get("audio_codes") + if codes is None and isinstance(encoded, (tuple, list)): + codes = encoded[0] if encoded else None + if codes is None: + raise ValueError("reference audio tokenizer returned no audio_codes") + + if isinstance(codes, (tuple, list)) and codes: + # Qwen3TTSTokenizer returns one tensor per input waveform. + if len(codes) != 1: + raise ValueError(f"only one reference waveform is supported, got {len(codes)}") + codes = codes[0] + codes = torch.as_tensor(codes) + if codes.ndim == 3: + if codes.shape[0] != 1: + raise ValueError(f"only one reference waveform is supported, got {tuple(codes.shape)}") + codes = codes[0] + if codes.ndim != 2: + raise ValueError(f"reference audio codes must be 2D, got {tuple(codes.shape)}") + if codes.shape[0] == 0: + raise ValueError("reference audio tokenizer returned no frames") + if codes.shape[-1] != self.num_codebooks: + if codes.shape[0] == self.num_codebooks: + codes = codes.transpose(0, 1) + else: + raise ValueError( + f"reference audio codes must have {self.num_codebooks} codebooks, got {tuple(codes.shape)}" + ) + codes = codes.to(device="cpu", dtype=torch.long).contiguous() + if codes.numel() and (int(codes.min()) < 0 or int(codes.max()) >= self.codebook_size): + raise ValueError( + f"reference audio code outside [0, {self.codebook_size}): " + f"min={int(codes.min())}, max={int(codes.max())}" + ) + # int16 is enough for Mimi/Qwen3-TTS codes and halves IPC/cache memory. + return codes.to(dtype=torch.int16) + + +__all__ = ["BreezeReferenceAudioTokenizer", "resolve_audio_tokenizer_path"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py new file mode 100644 index 000000000..2c04768a9 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Configuration classes for Breeze-TTS-2.""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig +from transformers.models.qwen3 import Qwen3Config + + +class BreezeTTS2Config(PretrainedConfig): + """Configuration for ``BreezeForConditionalGeneration``.""" + + model_type = "breeze" + + def __init__( + self, + backbone_config: dict | Qwen3Config | None = None, + text_encoder_config: dict | BreezeTTS2TextEncoderConfig | None = None, + depth_decoder_config: dict | BreezeTTS2DepthDecoderConfig | None = None, + codec_config: dict | None = None, + architectures: list[str] | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + audio_embed_size: int = 2048, + audio_eos_token_id: int = 262145, + audio_num_codebooks: int = 16, + audio_token_id: int = 262144, + audio_vocab_size: int = 2051, + backbone_flavor: str = "llama-1B", + backbone_model_name_or_path: str | None = None, + backbone_model_type: str = "qwen3", + bos_token_id: int = 2, + codebook_eos_token_id: int = 0, + codebook_pad_token_id: int = 2050, + decoder_flavor: str = "llama-100M", + depth_header_loss_weight: float = 1.0, + dtype: str = "bfloat16", + eos_token_id: int = 1, + head_dim: int = 128, + hidden_act: str = "silu", + hidden_size: int = 2048, + initializer_range: float = 0.02, + intermediate_size: int = 6144, + max_position_embeddings: int = 2048, + mlp_bias: bool = False, + num_attention_heads: int = 16, + num_codebooks: int = 16, + num_hidden_layers: int = 28, + num_key_value_heads: int = 8, + pad_token_id: int = 0, + rms_norm_eps: float = 1e-5, + rope_scaling: dict | None = None, + rope_theta: float = 500_000.0, + text_encoder_bucket_max_length_ratio: float = 4.0, + text_encoder_lora_config: dict | None = None, + text_encoder_proj_type: str = "linear", + text_encoder_special_tokens_config: dict | None = None, + text_vocab_size: int = 262158, + tie_codebooks_embeddings: bool = True, + tie_word_embeddings: bool = False, + use_cache: bool = True, + vocab_size: int = 2051, + sampling_rate: int = 24000, + codec_model_name_or_path: str = "kyutai/mimi", + **kwargs: Any, + ) -> None: + serialized_rope_parameters = kwargs.pop("rope_parameters", None) + if rope_scaling is None and serialized_rope_parameters is not None: + rope_scaling = serialized_rope_parameters + # Breeze stores the Qwen3 backbone in backbone_config. This follows the + # MOSS-TTS pattern and deliberately does not declare sub_configs. + if backbone_config is None: + backbone_config = {} + if isinstance(backbone_config, dict): + backbone_config = dict(backbone_config) + backbone_config.pop("model_type", None) + self.backbone_config = Qwen3Config(**backbone_config) + else: + self.backbone_config = backbone_config + + if text_encoder_config is None: + text_encoder_config = {} + if isinstance(text_encoder_config, dict): + text_encoder_config = dict(text_encoder_config) + text_encoder_config.pop("model_type", None) + self.text_encoder_config = BreezeTTS2TextEncoderConfig(**text_encoder_config) + else: + self.text_encoder_config = text_encoder_config + + if depth_decoder_config is None: + depth_decoder_config = {} + if isinstance(depth_decoder_config, dict): + depth_decoder_config = dict(depth_decoder_config) + depth_decoder_config.pop("model_type", None) + self.depth_decoder_config = BreezeTTS2DepthDecoderConfig(**depth_decoder_config) + else: + self.depth_decoder_config = depth_decoder_config + + # Nested configs must exist before the parent constructor: Transformers + # 5.x may call get_text_config() during config validation. + super().__init__( + architectures=architectures, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.audio_embed_size = audio_embed_size + self.audio_eos_token_id = audio_eos_token_id + self.audio_num_codebooks = audio_num_codebooks + self.audio_token_id = audio_token_id + self.audio_vocab_size = audio_vocab_size + # Keep Mimi's nested config as a plain checkpoint dictionary. Stage 1 + # owns codec construction; stage 0 only needs codebook_size here. + self.codec_config = codec_config + self.backbone_flavor = backbone_flavor + self.backbone_model_name_or_path = backbone_model_name_or_path + self.backbone_model_type = backbone_model_type + self.codebook_eos_token_id = codebook_eos_token_id + self.codebook_pad_token_id = codebook_pad_token_id + self.decoder_flavor = decoder_flavor + self.depth_header_loss_weight = depth_header_loss_weight + self.dtype = dtype + self.head_dim = head_dim + self.hidden_act = hidden_act + self.hidden_size = hidden_size + self.initializer_range = initializer_range + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_position_embeddings + self.mlp_bias = mlp_bias + self.num_attention_heads = num_attention_heads + self.num_codebooks = num_codebooks + self.num_hidden_layers = num_hidden_layers + self.num_key_value_heads = num_key_value_heads + self.rms_norm_eps = rms_norm_eps + if rope_scaling is not None: + # Preserve the upstream legacy field while also satisfying the + # normalized ``rope_parameters`` validation added in Transformers 5. + rope_scaling = dict(rope_scaling) + rope_scaling.setdefault("rope_type", rope_scaling.get("type", "default")) + rope_scaling.setdefault("rope_theta", rope_theta) + self.rope_scaling = rope_scaling + self.rope_theta = rope_theta + self.text_encoder_bucket_max_length_ratio = text_encoder_bucket_max_length_ratio + self.text_encoder_lora_config = text_encoder_lora_config or {} + self.text_encoder_proj_type = text_encoder_proj_type + self.text_encoder_special_tokens_config = text_encoder_special_tokens_config or {} + self.text_vocab_size = text_vocab_size + self.tie_codebooks_embeddings = tie_codebooks_embeddings + self.use_cache = use_cache + self.vocab_size = vocab_size + self.sampling_rate = sampling_rate + self.codec_model_name_or_path = codec_model_name_or_path + + if not getattr(self, "architectures", None): + self.architectures = ["BreezeForConditionalGeneration"] + + def get_text_config(self, **_: Any) -> Qwen3Config: + return self.backbone_config + + +class BreezeTTS2TextEncoderConfig(PretrainedConfig): + """Configuration for Breeze's T5Gemma2 text encoder.""" + + model_type = "t5gemma2_text" + + def __init__( + self, + hidden_size: int = 1152, + intermediate_size: int = 6912, + num_hidden_layers: int = 26, + num_attention_heads: int = 4, + num_key_value_heads: int = 1, + head_dim: int = 256, + max_position_embeddings: int = 32768, + hidden_activation: str = "gelu_pytorch_tanh", + rms_norm_eps: float = 1e-6, + rope_theta: float = 10_000.0, + sliding_window: int = 512, + query_pre_attn_scalar: int = 256, + eoi_token_index: int = 256000, + vocab_size: int = 262158, + layer_types: list[str] | None = None, + rope_parameters: dict | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + dropout_rate: float = 0.0, + initializer_range: float = 0.02, + tie_word_embeddings: bool = True, + use_bidirectional_attention: bool = False, + use_cache: bool = True, + requires_grad: bool = False, + preferred_attn_implementation: str = "flash_attention_2", + dtype: str = "bfloat16", + is_encoder_decoder: bool = False, + bos_token_id: int = 2, + eos_token_id: int = 1, + pad_token_id: int = 0, + **kwargs: Any, + ) -> None: + kwargs.pop("model_type", None) + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + is_encoder_decoder=is_encoder_decoder, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.hidden_activation = hidden_activation + self.hidden_act = hidden_activation + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.sliding_window = sliding_window + self.query_pre_attn_scalar = query_pre_attn_scalar + self.eoi_token_index = eoi_token_index + self.vocab_size = vocab_size + if layer_types is None: + # T5Gemma2 uses sliding attention except every sixth layer. + layer_types = [ + "full_attention" if (index + 1) % 6 == 0 else "sliding_attention" for index in range(num_hidden_layers) + ] + self.layer_types = list(layer_types) + if rope_parameters is None: + rope_parameters = { + "full_attention": { + "factor": 8.0, + "rope_theta": 1_000_000.0, + "rope_type": "linear", + }, + "sliding_attention": { + "rope_theta": 10_000.0, + "rope_type": "default", + }, + } + self.rope_parameters = dict(rope_parameters) + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.dropout_rate = dropout_rate + self.initializer_range = initializer_range + self.use_bidirectional_attention = use_bidirectional_attention + self.use_cache = use_cache + self.requires_grad = requires_grad + self.preferred_attn_implementation = preferred_attn_implementation + self.dtype = dtype + + +class BreezeTTS2DepthDecoderConfig(PretrainedConfig): + """Configuration for Breeze's per-frame depth decoder.""" + + model_type = "breeze_depth_decoder_model" + + def __init__( + self, + hidden_size: int = 1024, + backbone_hidden_size: int = 2048, + audio_embed_size: int = 2048, + intermediate_size: int = 8192, + num_hidden_layers: int = 12, + num_attention_heads: int = 8, + num_key_value_heads: int = 2, + head_dim: int = 128, + max_position_embeddings: int = 33, + hidden_act: str = "silu", + rms_norm_eps: float = 1e-5, + rope_theta: float = 500_000.0, + rope_scaling: dict | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + mlp_bias: bool = False, + initializer_range: float = 0.02, + use_cache: bool = True, + num_codebooks: int = 16, + vocab_size: int = 2051, + codebook_loss_weights: list[int] | None = None, + dtype: str = "bfloat16", + tie_word_embeddings: bool = False, + pad_token_id: int | None = 0, + **kwargs: Any, + ) -> None: + kwargs.pop("model_type", None) + # Transformers 5 serializes the normalized field as + # ``rope_parameters`` while the upstream Breeze checkpoint calls it + # ``rope_scaling``. Normalize before ``PretrainedConfig`` validates + # so ``to_dict() -> BreezeTTS2DepthDecoderConfig(**dict)`` round-trips. + serialized_rope_parameters = kwargs.pop("rope_parameters", None) + if rope_scaling is None: + rope_scaling = { + "factor": 32.0, + "high_freq_factor": 0.0078125, + "low_freq_factor": 0.001953125, + "original_max_position_embeddings": 16, + "rope_type": "llama3", + } + if codebook_loss_weights is None: + codebook_loss_weights = [3] + if num_codebooks <= 0 or vocab_size <= 0: + raise ValueError("num_codebooks and vocab_size must be positive") + rope_parameters = dict(serialized_rope_parameters or rope_scaling) + rope_parameters.setdefault("rope_type", rope_parameters.get("type", "default")) + rope_parameters.setdefault("rope_theta", rope_theta) + + super().__init__( + vocab_size=vocab_size, + tie_word_embeddings=tie_word_embeddings, + pad_token_id=pad_token_id, + **kwargs, + ) + # The vLLM depth decoder is a plain ``nn.Module``, so it does not get + # Transformers' PreTrainedModel default attention initialization. + self._attn_implementation = getattr(self, "_attn_implementation", None) or "eager" + self.hidden_size = hidden_size + self.backbone_hidden_size = backbone_hidden_size + self.audio_embed_size = audio_embed_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.hidden_act = hidden_act + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.rope_parameters = dict(getattr(self, "rope_parameters", None) or rope_parameters) + self.rope_parameters.setdefault("rope_theta", rope_theta) + # Transformers 4 exposes this as a plain attribute; Transformers 5 + # implements ``rope_scaling`` as a property backed by rope_parameters. + if not hasattr(type(self), "rope_scaling"): + self.rope_scaling = self.rope_parameters + self.rope_type = rope_scaling.get("rope_type", "default") + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + self.initializer_range = initializer_range + self.use_cache = use_cache + self.num_codebooks = num_codebooks + self.vocab_size = vocab_size + self.codebook_loss_weights = list(codebook_loss_weights) + self.dtype = dtype + + +AutoConfig.register("breeze", BreezeTTS2Config, exist_ok=True) + + +__all__ = [ + "BreezeTTS2Config", + "BreezeTTS2DepthDecoderConfig", + "BreezeTTS2TextEncoderConfig", +] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py new file mode 100644 index 000000000..7addb3934 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 stage-1 codec (codec codes -> waveform). + +The talker emits a frame-major ``(T, 16)`` matrix. The stage input +processor flattens it codebook-major; this module restores the ``(T, 16)`` +layout and delegates decoding to Breeze's bundled Qwen3-TTS tokenizer. A +Mimi fallback remains available for checkpoints that omit ``audio_tokenizer``. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +import torch +from torch import nn +from vllm.config import VllmConfig + +from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + resolve_audio_tokenizer_path, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + + +def _split_flat_ids(ids: torch.Tensor, counts: list[int] | None) -> list[torch.Tensor]: + if counts is None or len(counts) <= 1: + return [ids] + boundaries = [0] + for count in counts: + boundaries.append(boundaries[-1] + int(count)) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(counts))] + + +class BreezeTTS2MimiCodec(nn.Module): + """Synchronous Breeze codec stage used by ``GenerationModelRunner``.""" + + input_modalities = "audio" + requires_raw_input_tokens = True + # Stateful chunk decoding needs the scheduler-side request id, which is + # also used by on_requests_finished to release decoder state. + requires_request_ids = True + have_multimodal_outputs = True + has_preprocess = False + has_postprocess = False + enable_update_additional_information = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + del prefix + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + self.model_path = vllm_config.model_config.model + self._async_chunk = bool(getattr(vllm_config.model_config, "async_chunk", False)) + # Resolve once: a local directory works directly, an HF repo id + # resolves to its cached snapshot. ``None`` means the checkpoint has + # no bundled tokenizer and the Mimi fallback applies. + self._tokenizer_path = resolve_audio_tokenizer_path(self.model_path) + if self._async_chunk and self._tokenizer_path is None: + raise RuntimeError( + "Breeze-TTS-2 async_chunk requires the bundled Qwen3-TTS audio_tokenizer; " + "the stateless Mimi fallback cannot preserve chunk decoder state" + ) + + self._audio_tokenizer = None + codec_config = getattr(self.config, "codec_config", None) + if isinstance(codec_config, Mapping): + codec_kwargs = dict(codec_config) + elif codec_config is not None: + codec_kwargs = codec_config.to_dict() + else: + codec_kwargs = {} + codec_kwargs.pop("model_type", None) + codec_kwargs.pop("architectures", None) + + # Breeze checkpoints normally carry a Qwen3-TTS decoder in this + # subdirectory. Avoid constructing a second, unused Mimi network in + # that case; this saves both startup time and worker memory. + self._codec = None + if self._tokenizer_path is None: + # Import lazily: stage 0 does not need the Mimi implementation. + from transformers import MimiConfig, MimiModel + + self._codec = MimiModel(MimiConfig(**codec_kwargs)).eval() + for parameter in self._codec.parameters(): + parameter.requires_grad_(False) + + self._num_codebooks = int(getattr(self.config, "num_codebooks", 16)) + self._codebook_size = int(codec_kwargs.get("codebook_size", 2048)) + self._sample_rate = int(codec_kwargs.get("sampling_rate", getattr(self.config, "sampling_rate", 24000))) + self._decoder_state_cache: dict[str, dict[str, Any]] = {} + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Return stable dummy embeddings; this stage never samples.""" + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> None: + del hidden_states, sampling_metadata + return None + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + empty = torch.empty(0, dtype=torch.float32) + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr]}, + ) + # The generation runner's profiling/dummy invocation has no request ids + # or runtime payload. Real requests always provide both. + if runtime_additional_information is None and kwargs.get("request_ids") is None: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + counts = kwargs.get("seq_token_counts") + if counts is not None: + counts = [int(item) for item in counts] + requests = _split_flat_ids(ids, counts) + waveforms: list[torch.Tensor] = [] + sample_rates: list[torch.Tensor] = [] + if self._async_chunk: + return self._forward_stateful_chunks( + ids=ids, + requests=requests, + runtime_additional_information=runtime_additional_information, + empty=empty, + sr=sr, + **kwargs, + ) + for index, request_ids in enumerate(requests): + runtime_info = ( + runtime_additional_information[index] + if runtime_additional_information is not None and index < len(runtime_additional_information) + else None + ) + flat = self._payload_codes(request_ids, runtime_info) + if flat.numel() == 0: + waveforms.append(empty.to(device=ids.device)) + sample_rates.append(sr) + continue + if flat.numel() % self._num_codebooks != 0: + raise ValueError( + f"Breeze codec input length {int(flat.numel())} is not divisible by " + f"num_codebooks {self._num_codebooks}" + ) + if int(flat.min()) < 0 or int(flat.max()) >= self._codebook_size: + raise ValueError( + f"Breeze codec id outside [0, {self._codebook_size}): min={int(flat.min())}, max={int(flat.max())}" + ) + frames = flat.numel() // self._num_codebooks + codes = flat.reshape(self._num_codebooks, frames).unsqueeze(0) + waveform = self._decode_codes(codes) + if waveform.ndim == 3: + waveform = waveform[0, 0] + elif waveform.ndim == 2: + waveform = waveform[0] + if waveform.ndim != 1: + raise ValueError(f"Breeze Mimi waveform must be 1D, got {tuple(waveform.shape)}") + # Serving/output serialization consumes CPU tensors. Move only + # the final waveform off the worker device; codec inference itself + # remains on the accelerator. + waveforms.append(waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()) + sample_rates.append(sr) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, + ) + + def _forward_stateful_chunks( + self, + *, + ids: torch.Tensor, + requests: list[torch.Tensor], + runtime_additional_information: list[dict[str, Any]] | None, + empty: torch.Tensor, + sr: torch.Tensor, + **kwargs: Any, + ) -> OmniOutput: + """Decode incremental codebook-major chunks with per-request state.""" + decoder = self._audio_tokenizer.model.decoder + scheduler_request_ids = kwargs.get("request_ids") + if scheduler_request_ids is None: + raise ValueError("Breeze async_chunk codec inputs require scheduler request_ids") + request_ids = [str(request_id) for request_id in scheduler_request_ids] + if len(request_ids) != len(requests): + raise ValueError( + f"Breeze codec request id count {len(request_ids)} does not match input request count {len(requests)}" + ) + valid: list[tuple[str, torch.Tensor, bool]] = [] + for index, request_ids_tensor in enumerate(requests): + runtime_info = ( + runtime_additional_information[index] + if runtime_additional_information is not None and index < len(runtime_additional_information) + else {} + ) + state_id = request_ids[index] + # The connector receiver strips ``meta.finished`` while merging + # runtime metadata, so the stream-terminated signal arrives on + # the model-visible ``stream_finished`` field (the same contract + # CosyVoice3 / Audex use). Fall back to ``finished`` for direct + # in-process callers that still set only the transport flag. + finished = self._meta_bool(runtime_info, "stream_finished") or self._meta_bool(runtime_info, "finished") + flat = self._payload_codes(request_ids_tensor, runtime_info) + if flat.numel() == 0: + # Terminal marker: the stage-0 processor sends an empty + # payload with ``finished=True`` once generation ends. + self._decoder_state_cache.pop(state_id, None) + continue + if flat.numel() % self._num_codebooks != 0: + # A partial frame means the transport payload is corrupt. + # Silently resetting decoder state would corrupt the rest of + # the stream; fail loudly instead. The scheduler aborts the + # request and ``on_requests_finished`` releases the state. + raise ValueError( + f"Breeze codec chunk length {int(flat.numel())} is not divisible by " + f"num_codebooks {self._num_codebooks}" + ) + if int(flat.min()) < 0 or int(flat.max()) >= self._codebook_size: + raise ValueError( + f"Breeze codec id outside [0, {self._codebook_size}): min={int(flat.min())}, max={int(flat.max())}" + ) + frames = int(flat.numel()) // self._num_codebooks + valid.append((state_id, flat.reshape(self._num_codebooks, frames), finished)) + + waveforms = [empty.to(device=ids.device)] * len(requests) + sample_rates = [sr] * len(requests) + if not valid: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, + ) + states: list[dict[str, Any]] = [] + valid_indices: list[int] = [] + max_frames = max(codes.shape[-1] for _, codes, _ in valid) + batched_codes = valid[0][1].new_zeros((len(valid), self._num_codebooks, max_frames)) + for row, (state_id, row_codes, _) in enumerate(valid): + batched_codes[row, :, : row_codes.shape[-1]].copy_(row_codes) + state = self._decoder_state_cache.get(state_id) + if state is None: + state = {} + self._decoder_state_cache[state_id] = state + state.setdefault("prefix_frames", 0) + states.append(state) + valid_indices.append(row) + + decoded = decoder.batched_chunked_decode( + batched_codes, + [int(codes.shape[-1]) for _, codes, _ in valid], + caches=states, + chunk_size=300, + left_context_size=25, + ) + if len(decoded) != len(valid): + raise RuntimeError(f"Breeze codec returned {len(decoded)} chunks for {len(valid)} requests") + for request_index, waveform in zip(valid_indices, decoded, strict=True): + if waveform.ndim == 2 and waveform.shape[0] == 1: + waveform = waveform[0] + elif waveform.ndim != 1: + waveform = waveform.reshape(-1) + waveforms[request_index] = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous() + + for state_id, _, finished in valid: + if finished: + self._decoder_state_cache.pop(state_id, None) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, + ) + + @staticmethod + def _meta_bool(info: Mapping[str, Any] | None, key: str) -> bool: + if not isinstance(info, Mapping): + return False + meta = info.get("meta", {}) + if not isinstance(meta, Mapping): + return False + value = meta.get(key, False) + if isinstance(value, torch.Tensor): + return bool(value.reshape(-1)[0].item()) if value.numel() else False + return bool(value) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + for request_id in finished_req_ids: + self._decoder_state_cache.pop(str(request_id), None) + + def _decode_codes(self, codes: torch.Tensor) -> torch.Tensor: + """Decode one ``(1, Q, T)`` sequence using Breeze's bundled codec.""" + if self._audio_tokenizer is not None: + wavs, sample_rate = self._audio_tokenizer.decode({"audio_codes": [codes[0].transpose(0, 1).contiguous()]}) + self._sample_rate = int(sample_rate) + if not wavs: + raise RuntimeError("Breeze Qwen3-TTS tokenizer returned no waveform") + return torch.as_tensor(wavs[0], device=codes.device, dtype=torch.float32).reshape(-1) + + if self._codec is None: + raise RuntimeError("Breeze codec is not initialized; load_weights() was not completed") + decoded = self._codec.decode(codes) + waveform = getattr(decoded, "audio_values", decoded) + if isinstance(waveform, Mapping): + waveform = waveform.get("audio_values") + if waveform is None: + raise RuntimeError("MimiModel.decode returned no audio_values") + return torch.as_tensor(waveform, device=codes.device, dtype=torch.float32) + + @staticmethod + def _payload_codes(input_ids: torch.Tensor, runtime_info: Mapping[str, Any] | None) -> torch.Tensor: + if isinstance(runtime_info, Mapping): + codes = runtime_info.get("codes") + if isinstance(codes, Mapping): + audio = codes.get("audio") + if audio is not None: + return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) + # Full-payload transport may flatten nested OmniPayload keys at + # the wire boundary and restore them only after this hook runs. + audio = runtime_info.get("codes.audio") + if audio is not None: + return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) + return input_ids + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load Breeze's bundled Qwen3-TTS codec, or fallback Mimi weights.""" + if self._tokenizer_path is not None: + # The root checkpoint contains both the talker weights and a Mimi + # ``codec_model.*`` fallback. When the bundled Qwen3-TTS tokenizer + # is present, it owns decoding; consume the root iterator without + # copying the unused fallback into worker memory. + if self._audio_tokenizer is None: + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import ( + Qwen3TTSTokenizer, + ) + + device = self.vllm_config.device_config.device + self._audio_tokenizer = Qwen3TTSTokenizer.from_pretrained( + str(self._tokenizer_path), + device_map=str(device), + ) + self._sample_rate = int(self._audio_tokenizer.model.get_output_sample_rate()) + for _ in weights: + pass + tokenizer_module = getattr(self._audio_tokenizer, "model", None) + named_parameters = getattr(tokenizer_module, "named_parameters", None) + if callable(named_parameters): + return {f"_audio_tokenizer.model.{name}" for name, _ in named_parameters()} + return set() + + params = dict(self._codec.named_parameters()) if self._codec is not None else {} + loaded: set[str] = set() + for name, tensor in weights: + if name.startswith("codec_model."): + target_name = name[len("codec_model.") :] + else: + continue + parameter = params.get(target_name) + if parameter is None: + raise ValueError(f"Unknown Breeze Mimi weight: {target_name}") + loader = getattr(parameter, "weight_loader", None) + if loader is None: + parameter.data.copy_(tensor.to(device=parameter.device, dtype=parameter.dtype)) + else: + loader(parameter, tensor) + loaded.add(target_name) + + if loaded: + return {f"_codec.{name}" for name in loaded} + + # Breeze inference uses the bundled Qwen3-TTS tokenizer for both + # reference encoding and generated-code decoding. Load it once per + # Stage-1 worker; unlike the request-side wrapper this instance is + # placed on the worker device and reused for the entire batch. + # A few Breeze distributions omit codec tensors and point to the + # standalone kyutai/mimi repository instead. Load it once, never per + # request, and fail loudly if it cannot be resolved. + codec_name = getattr(self.config, "codec_model_name_or_path", None) or "kyutai/mimi" + from transformers import AutoModel + + self._codec = AutoModel.from_pretrained(codec_name).eval() + for parameter in self._codec.parameters(): + parameter.requires_grad_(False) + self._sample_rate = int(getattr(self._codec.config, "sampling_rate", self._sample_rate)) + self._codebook_size = int( + getattr(getattr(self._codec.config, "quantizer", None), "cardinality", self._codebook_size) + ) + return {f"_codec.{name}" for name, _ in self._codec.named_parameters()} + + +__all__ = ["BreezeTTS2MimiCodec"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py new file mode 100644 index 000000000..559c88b07 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py @@ -0,0 +1,523 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections.abc import Callable +from inspect import signature +from typing import Any + +import torch +from torch import nn +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.integrations import use_kernel_forward_from_hub +from transformers.masking_utils import create_causal_mask +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, logging +from transformers.utils.deprecation import deprecate_kwarg + +logger = logging.get_logger(__name__) + + +def _create_depth_causal_mask( + *, + config: Any, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.LongTensor | None, + cache_position: torch.LongTensor, +) -> torch.Tensor | None: + """Bridge Transformers 4.57 and 5.x ``create_causal_mask`` contracts.""" + parameters = signature(create_causal_mask).parameters + kwargs: dict[str, Any] = { + "config": config, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + if "inputs_embeds" in parameters: + kwargs["inputs_embeds"] = inputs_embeds + else: + kwargs["input_embeds"] = inputs_embeds + kwargs["cache_position"] = cache_position + return create_causal_mask(**kwargs) + + +class BreezeDepthDecoderModel(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + if hasattr(config, "audio_embed_size") and config.audio_embed_size: + self.audio_embed_size = config.audio_embed_size + else: + self.audio_embed_size = config.backbone_hidden_size + self.embed_tokens = nn.Embedding((config.num_codebooks * config.vocab_size), self.audio_embed_size) + self.layers = nn.ModuleList( + [BreezeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = BreezeRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.inputs_embeds_projector = nn.Linear(self.audio_embed_size, config.hidden_size, bias=False) + + if config.backbone_hidden_size != self.audio_embed_size: + self.backbone_hidden_state_projector = nn.Linear( + config.backbone_hidden_size, self.audio_embed_size, bias=False + ) + else: + self.backbone_hidden_state_projector = None + + # Weights are loaded from the Breeze checkpoint by the parent talker. + # Avoid PreTrainedModel.post_init here because this module is a plain + # nn.Module used for eager per-frame inference. + + def forward( + self, + input_ids: torch.LongTensor = None, + backbone_last_hidden_state: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + r""" + backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*): + The last hidden state of the backbone model. This is required when + the first codebook token (generated by the backbone model) is + provided in the `input_ids` argument. + """ + if position_ids is not None and not torch.compiler.is_compiling(): + logger.warning_once( + "Custom `position_ids` were provided but will be ignored. " + "Breeze depth decoder derives them from `cache_position`; " + "they must be identical across the batch." + ) + position_ids = None + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds.") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + inputs_seq_length = inputs_embeds.shape[1] if inputs_embeds is not None else input_ids.shape[1] + device = inputs_embeds.device if inputs_embeds is not None else input_ids.device + cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_seq_length, device=device) + + if inputs_embeds is None: + codebook_idxs = torch.clamp(cache_position - 1, min=0) + offset = codebook_idxs * self.vocab_size + inputs_embeds = self.embed_tokens(input_ids + offset) + + input_ids_are_first_codebook = cache_position[0] == 0 + if backbone_last_hidden_state is not None: + if self.backbone_hidden_state_projector is not None: + inputs_embeds[:, 0] = self.backbone_hidden_state_projector(backbone_last_hidden_state) + else: + inputs_embeds[:, 0] = backbone_last_hidden_state + else: + if not torch.compiler.is_compiling() and input_ids_are_first_codebook: + logger.warning( + "When the first codebook token is provided, " + "`backbone_last_hidden_state` is required for correct inference." + ) + + inputs_embeds = self.inputs_embeds_projector(inputs_embeds) + + causal_mask = _create_depth_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + cache_position=cache_position, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_ids = cache_position.unsqueeze(0) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +@use_kernel_forward_from_hub("RMSNorm") +class BreezeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + BreezeRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class BreezeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class BreezeMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class BreezeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class BreezeDecoderLayer(nn.Module): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = BreezeAttention(config=config, layer_idx=layer_idx) + + self.mlp = BreezeMLP(config) + self.input_layernorm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + cache_position: torch.LongTensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, # necessary, but kept here for BC + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class BreezeCodebooksHead(nn.Module): + """Position-specific heads for codebooks 1..15.""" + + def __init__(self, hidden_size: int, num_codebooks: int, vocab_size: int) -> None: + super().__init__() + self.num_codebooks = int(num_codebooks) + self.weight = nn.Parameter(torch.empty(self.num_codebooks - 1, hidden_size, vocab_size)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # ``hidden_states`` contains the placeholder at position 0 followed + # by codebook-0..N tokens. Return logits for codebook-1..N. + depth_states = hidden_states[:, 1 : self.num_codebooks, :] + if depth_states.shape[1] > self.num_codebooks - 1: + raise ValueError(f"depth decoder received too many input positions: {hidden_states.shape[1]}") + return torch.einsum("bkh,khv->bkv", depth_states, self.weight[: depth_states.shape[1]]) + + +class BreezeDepthDecoderForCausalLM(nn.Module): + """Eager depth decoder used by the vLLM stage-0 talker.""" + + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.model = BreezeDepthDecoderModel(config) + self.codebooks_head = BreezeCodebooksHead( + config.hidden_size, + config.num_codebooks, + config.vocab_size, + ) + + def forward( + self, + input_ids: torch.Tensor, + backbone_last_hidden_state: torch.Tensor, + *, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + outputs = self.model( + input_ids=input_ids, + backbone_last_hidden_state=backbone_last_hidden_state, + use_cache=use_cache, + **kwargs, + ) + logits = self.codebooks_head(outputs.last_hidden_state) + return CausalLMOutputWithPast( + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=getattr(outputs, "hidden_states", None), + attentions=getattr(outputs, "attentions", None), + ) + + def load_weights(self, weights): + params_dict = dict(self.named_parameters()) + loaded: set[str] = set() + for name, tensor in weights: + param = params_dict.get(name) + if param is None: + continue + if name == "codebooks_head.weight": + expected = ( + int(self.config.num_codebooks) - 1, + int(self.config.hidden_size), + int(self.config.vocab_size), + ) + actual = tuple(int(dim) for dim in tensor.shape) + if actual != expected: + raise ValueError(f"Breeze depth weight {name!r} has shape {actual}, expected {expected}") + loader = getattr(param, "weight_loader", None) + if loader is None: + param.data.copy_(tensor.to(device=param.device, dtype=param.dtype)) + else: + loader(param, tensor) + loaded.add(name) + return loaded + + +__all__ = [ + "BreezeCodebooksHead", + "BreezeDepthDecoderForCausalLM", + "BreezeDepthDecoderModel", +] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py new file mode 100644 index 000000000..beaa390c8 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py @@ -0,0 +1,618 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 stage-0 talker for vLLM-Omni. + +The scheduler samples Breeze codebook-0 tokens one at a time. The other 15 +codebooks are kept in the request-local Omni buffer: after each backbone step, +the eager depth decoder completes the frame and the next backbone step +consumes the summed 16-codebook embedding. This preserves the upstream Breeze +generation order without importing HuggingFace GenerationMixin. +""" + +from __future__ import annotations + +import hashlib +import os +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import torch +from torch import nn +from torch.nn.utils.rnn import pad_sequence +from vllm.config import VllmConfig +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2Config, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_depth import ( + BreezeDepthDecoderForCausalLM, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text import ( + BreezeTTS2TextEncoder, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + + +def _prefixed(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +def _normalize_request_infos(infos: Any) -> list[dict[str, Any]] | None: + """Normalize runner metadata to one mutable dict per request.""" + if not isinstance(infos, Sequence) or isinstance(infos, (str, bytes, bytearray)): + return None + if not all(isinstance(info, dict) for info in infos): + return None + return infos + + +def _iter_request_rows( + states: Sequence[dict[str, Any]], + spans: Sequence[tuple[int, int]] | None, + num_rows: int, +) -> Iterable[tuple[dict[str, Any], int, int]]: + """Map request state to logits rows for full or sampled hidden layouts.""" + if num_rows <= 0 or not states: + return + if spans is not None and len(spans) == len(states): + if all(0 <= int(start) <= int(end) <= num_rows for start, end in spans): + for state, (start, end) in zip(states, spans, strict=True): + if int(start) < int(end): + yield state, int(start), int(end) + return + # The runner passes full-sequence spans to make_omni_output but only + # one sampled row per request to compute_logits. + if num_rows == len(states): + for index, state in enumerate(states): + yield state, index, index + 1 + return + if num_rows == len(states): + for index, state in enumerate(states): + yield state, index, index + 1 + return + raise RuntimeError( + "Breeze talker cannot align request state with logits rows: " + f"states={len(states)}, rows={num_rows}, spans={spans}" + ) + + +def _check_weight_shape(name: str, tensor: torch.Tensor, expected: tuple[int, ...]) -> None: + """Fail early for Breeze-owned, non-sharded checkpoint tensors.""" + actual = tuple(int(dim) for dim in tensor.shape) + if actual != expected: + raise ValueError(f"Breeze weight {name!r} has shape {actual}, expected {expected}") + + +class BreezeTTS2TalkerForGeneration(nn.Module): + """Stage-0 Breeze talker: text/audio conditioning -> codec frames.""" + + have_multimodal_outputs = True + has_preprocess = True + has_postprocess = True + # Stage 0 emits a tiny per-frame codec payload. Let the AR runner snapshot + # it through the async-output path instead of the final-only client path. + use_async_omni_output = True + eager_omni_postprocess_before_async_output = True + omni_pooler_payload_include_hidden = False + input_modalities = "text" + gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("hidden_states", "last"), + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config: BreezeTTS2Config = vllm_config.model_config.hf_config + self._async_chunk = bool(vllm_config.model_config.async_chunk) + self._golden_dump_dir = os.environ.get("VLLM_OMNI_BREEZE_GOLDEN_DIR") + + backbone_config = self.config.backbone_config + backbone_vllm_config = vllm_config.with_hf_config( + backbone_config, + architectures=["Qwen3ForCausalLM"], + ) + backbone_vllm_config.model_config.hf_text_config = backbone_config + self.model = Qwen3Model( + vllm_config=backbone_vllm_config, + prefix=_prefixed(prefix, "model"), + ) + # Upstream's Qwen3 adapter keeps only decoder layers/norm/RoPE and + # replaces the standard text embedding with Breeze's summed audio/text + # tables. Breeze always supplies ``inputs_embeds`` through preprocess, + # so remove the unused Qwen embedding instead of allocating and then + # reporting an uninitialized 151936x2048 parameter at load time. + self.model.embed_tokens = nn.Identity() + + self.hidden_size = int(backbone_config.hidden_size) + self.num_codebooks = int(self.config.num_codebooks) + self.codebook_vocab_size = int(self.config.vocab_size) + codec_config = self.config.codec_config + self.codebook_size = ( + int(codec_config.get("codebook_size", self.codebook_vocab_size)) + if isinstance(codec_config, Mapping) + else self.codebook_vocab_size + ) + self.codebook_eos_token_id = int(self.config.codebook_eos_token_id) + self.audio_token_id = int(self.config.audio_token_id) + self.audio_eos_token_id = int(self.config.audio_eos_token_id) + self.text_pad_token_id = int(self.config.pad_token_id) + + self.lm_head = ParallelLMHead( + self.codebook_vocab_size + 1, + self.hidden_size, + bias=False, + prefix=_prefixed(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.codebook_vocab_size + 1) + + # Upstream Breeze uses embedding(code_i + i * vocab_size) and sums + # over codebooks before entering the Qwen3 backbone. + self.embed_audio_tokens = nn.Embedding( + self.num_codebooks * self.codebook_vocab_size, + self.hidden_size, + ) + # Keep Breeze's standalone text table even though the normal path + # uses T5Gemma2 outputs. It is part of the checkpoint contract and is + # also the correct fallback when a checkpoint omits text_encoder. + self.embed_text_tokens = nn.Embedding(int(self.config.text_vocab_size), self.hidden_size) + self.register_buffer( + "audio_token_offsets", + torch.arange(self.num_codebooks, dtype=torch.long) * self.codebook_vocab_size, + persistent=False, + ) + + self.text_encoder = BreezeTTS2TextEncoder(self.config.text_encoder_config) + self.text_encoder_proj = nn.Linear( + int(self.config.text_encoder_config.hidden_size), + self.hidden_size, + bias=False, + ) + self.depth_decoder = BreezeDepthDecoderForCausalLM(self.config.depth_decoder_config) + if bool(self.config.tie_codebooks_embeddings): + # The upstream HF model ties these tables. Sharing the Parameter + # also prevents a tied checkpoint from leaving the depth input + # embedding randomly initialized when only one key is serialized. + self.depth_decoder.model.embed_tokens.weight = self.embed_audio_tokens.weight + + self._codec_disallowed_mask: torch.Tensor | None = None + self._batch_state: list[dict[str, Any]] | None = None + self._batch_state_spans: Sequence[tuple[int, int]] | None = None + + # ------------------------------------------------------------------ + # Embedding and backbone hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Embed complete generated codec frames; scalar ids are invalid.""" + if input_ids.ndim < 1 or input_ids.shape[-1] != self.num_codebooks: + raise ValueError( + "Breeze embed_input_ids requires complete codec frames with shape " + f"(..., {self.num_codebooks}), got {tuple(input_ids.shape)}" + ) + codes = input_ids.to(dtype=torch.long) + offsets = self.audio_token_offsets.to(device=codes.device) + ids = codes.clamp(0, self.codebook_vocab_size - 1) + offsets + return self.embed_audio_tokens(ids).sum(dim=-2) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + # ------------------------------------------------------------------ + # Prompt/text conditioning + # ------------------------------------------------------------------ + + def _encode_text_segments( + self, + input_ids: torch.Tensor, + text_mask: torch.Tensor, + text_lengths: torch.Tensor, + ) -> torch.Tensor: + text_ids = input_ids[text_mask.to(device=input_ids.device, dtype=torch.bool)].to(torch.long) + lengths = [int(x) for x in text_lengths.reshape(-1).tolist() if int(x) > 0] + if int(text_ids.numel()) != sum(lengths): + raise ValueError(f"Breeze text mask/length mismatch: {text_ids.numel()} != {sum(lengths)}") + if not lengths: + return input_ids.new_empty((0, self.hidden_size), dtype=self.text_encoder.dtype) + segments = list(torch.split(text_ids, lengths)) + padded = pad_sequence(segments, batch_first=True, padding_value=0) + attn = torch.zeros_like(padded, dtype=torch.bool) + for row, length in enumerate(lengths): + attn[row, :length] = True + hidden = self.text_encoder(padded, attention_mask=attn) + projected = self.text_encoder_proj(hidden) + return torch.cat([projected[row, :length] for row, length in enumerate(lengths)], dim=0) + + def _build_prompt_embeddings(self, prompt_ids: torch.Tensor, info: Mapping[str, Any]) -> torch.Tensor: + device = next(self.parameters()).device + prompt_ids = prompt_ids.reshape(-1).to(device=device, dtype=torch.long) + text_mask = torch.as_tensor(info.get("text_ids_mask"), device=device, dtype=torch.bool).reshape(-1) + text_lengths = torch.as_tensor(info.get("text_ids_len"), device=device, dtype=torch.long) + if text_mask.numel() != prompt_ids.numel(): + raise ValueError("Breeze text_ids_mask length does not match prompt length") + output = torch.zeros( + (prompt_ids.numel(), self.hidden_size), + device=device, + dtype=self.embed_text_tokens.weight.dtype, + ) + text_embeds = self._encode_text_segments(prompt_ids, text_mask, text_lengths) + output[text_mask] = text_embeds.to(dtype=output.dtype) + + audio_mask = prompt_ids == self.audio_token_id + input_values = info.get("input_values") + if bool(audio_mask.any()): + if input_values is None: + raise ValueError("Breeze audio placeholders require input_values") + codes = torch.as_tensor(input_values, device=device, dtype=torch.long) + if codes.ndim != 2 or codes.shape[-1] != self.num_codebooks: + raise ValueError(f"Breeze input_values must be (T, {self.num_codebooks})") + if int(codes.shape[0]) != int(audio_mask.sum()): + raise ValueError("Breeze reference code count does not match audio placeholders") + output[audio_mask] = self.embed_input_ids(codes).to(dtype=output.dtype) + + eos_mask = prompt_ids == self.audio_eos_token_id + if bool(eos_mask.any()): + eos_codes = torch.full( + (1, self.num_codebooks), + self.codebook_eos_token_id, + dtype=torch.long, + device=device, + ) + output[eos_mask] = self.embed_input_ids(eos_codes).expand(int(eos_mask.sum()), -1) + return output + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build prefill embeddings and consume one completed frame per decode.""" + del input_embeds + info = dict(info_dict) + span_len = int(input_ids.shape[0]) + is_prefill = bool(info.get("_omni_is_prefill", span_len > 1)) + if is_prefill: + full_ids = info.get("prompt_ids") + if not isinstance(full_ids, torch.Tensor): + raise RuntimeError("Breeze prefill is missing prompt_ids") + full_embeds = info.get("breeze_prompt_embeds") + if not isinstance(full_embeds, torch.Tensor): + full_embeds = self._build_prompt_embeddings(full_ids, info).detach().cpu().contiguous() + offset = int(info.get("breeze_prefill_offset", info.get("_omni_num_computed_tokens", 0)) or 0) + take = full_embeds[offset : offset + span_len].to(device=input_ids.device) + if take.shape[0] != span_len: + raise RuntimeError("Breeze prompt embedding slice is shorter than scheduled span") + return ( + torch.full_like(input_ids, self.text_pad_token_id), + take, + { + "breeze_prompt_embeds": full_embeds, + "breeze_prefill_offset": offset + span_len, + "breeze_generated_frames": int(info.get("breeze_generated_frames", 0) or 0), + "breeze_max_new_frames": int(info.get("breeze_max_new_frames", -1)), + }, + ) + + frame = info.get("breeze_current_frame") + if isinstance(frame, torch.Tensor) and frame.numel() == self.num_codebooks: + frame = frame.to(device=input_ids.device, dtype=torch.long).reshape(1, 1, -1) + embeds = self.embed_input_ids(frame).reshape(1, -1) + safe_id = frame[..., 0].reshape_as(input_ids).to(dtype=input_ids.dtype) + return safe_id, embeds, {} + has_breeze_state = any( + key in info + for key in ( + "breeze_current_frame", + "breeze_audio_codes", + "breeze_generated_frames", + "breeze_prompt_embeds", + ) + ) + if has_breeze_state: + # A real decode must consume the full frame produced from the + # preceding backbone hidden state. A scalar fallback would silently + # fill codebooks 1..15 with EOS and corrupt AR state. + raise RuntimeError( + "Breeze decode is missing breeze_current_frame; refusing scalar fallback " + "because it would fill codebooks 1..15 with EOS" + ) + # Runner-only dummy execution has no Breeze codec state. It must not + # invent a semantically invalid audio frame; a stable zero embedding is + # sufficient for shape/profiling warmup. + dummy_embeds = torch.zeros( + (*input_ids.shape, self.hidden_size), + device=input_ids.device, + dtype=self.embed_audio_tokens.weight.dtype, + ) + return input_ids, dummy_embeds, {} + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + return {} + return {"hidden_states": {"last": hidden_states[-1].detach()}} + + # ------------------------------------------------------------------ + # Main head, depth decoder and Omni output + # ------------------------------------------------------------------ + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: SamplingMetadata | None = None, + ) -> torch.Tensor | None: + del sampling_metadata + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self._main_head_logits(hidden_states) + + # ``make_omni_output`` runs immediately before this hook. It records + # requests whose previous hidden state selected EOS or reached the + # frame budget, so the scheduler receives a real stop token instead + # of depending on an arbitrary deploy-level max_tokens timeout. + states = self._batch_state or [] + if states: + for state, row_start, row_end in _iter_request_rows( + states, + getattr(self, "_batch_state_spans", None), + int(logits.shape[0]), + ): + if state.get("breeze_force_eos", False): + row = logits[row_start:row_end] + row.fill_(float("-inf")) + row[:, self.codebook_vocab_size] = 0.0 + return logits + + def _main_head_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + raise RuntimeError("Breeze main head returned no logits") + if self._codec_disallowed_mask is None or self._codec_disallowed_mask.device != logits.device: + mask = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) + if self.codebook_size < self.codebook_vocab_size: + mask[self.codebook_size : self.codebook_vocab_size] = True + self._codec_disallowed_mask = mask + return logits.masked_fill(self._codec_disallowed_mask, float("-inf")) + + @torch.no_grad() + def _generate_depth_codes(self, backbone_hidden: torch.Tensor, code0: int) -> torch.Tensor: + sequence = torch.tensor([[0, code0]], dtype=torch.long, device=backbone_hidden.device) + generated = [int(code0)] + for _ in range(self.num_codebooks - 1): + depth_out = self.depth_decoder( + input_ids=sequence, + backbone_last_hidden_state=backbone_hidden.reshape(1, -1), + use_cache=False, + ) + logits = depth_out.logits[:, -1, :].clone() + if self.codebook_size < logits.shape[-1]: + logits[:, self.codebook_size :] = float("-inf") + next_code = int(torch.argmax(logits, dim=-1).item()) + generated.append(next_code) + sequence = torch.cat([sequence, torch.tensor([[next_code]], device=sequence.device)], dim=1) + return torch.tensor(generated, dtype=torch.long, device=backbone_hidden.device).unsqueeze(0) + + @staticmethod + def _accumulate_codes(info: dict[str, Any], frame: torch.Tensor) -> torch.Tensor: + """Append one ``(1, Q)`` frame to the request's cumulative codes.""" + accumulated = info.get("breeze_audio_codes") + accumulated = ( + torch.cat([accumulated.to(frame.device), frame], dim=0) + if isinstance(accumulated, torch.Tensor) and accumulated.numel() + else frame + ) + info["breeze_audio_codes"] = accumulated + return accumulated + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + self._batch_state = None + self._batch_state_spans = None + return model_outputs + hidden = model_outputs + infos = kwargs.get("model_intermediate_buffer", []) + spans = kwargs.get("request_token_spans") + info_items = _normalize_request_infos(infos) + if info_items is None: + self._batch_state = None + self._batch_state_spans = None + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + # The runner invokes compute_logits after this method, so retain the + # same request ordering and spans for per-request EOS masking. + self._batch_state = info_items + self._batch_state_spans = spans + per_request_codes: list[torch.Tensor] = [] + + for index, info in enumerate(info_items): + if spans is not None and index < len(spans): + row_start = max(0, int(spans[index][0])) + row_end = min(int(spans[index][1]), int(hidden.shape[0])) + elif len(info_items) == 1: + row_start, row_end = 0, int(hidden.shape[0]) + else: + raise RuntimeError("Breeze make_omni_output requires request_token_spans for batched decode") + if row_end <= row_start: + per_request_codes.append(hidden.new_empty((0, self.num_codebooks), dtype=torch.long)) + continue + row_hidden = hidden[row_end - 1] + if self._golden_dump_dir and bool(info.get("_omni_is_prefill", False)): + self._dump_golden_prefill(info, row_hidden) + main_logits = self._main_head_logits(row_hidden.reshape(1, -1)) + code0 = int(torch.argmax(main_logits, dim=-1).item()) + generated_frames = int(info.get("breeze_generated_frames", 0) or 0) + max_new_frames = int(info.get("breeze_max_new_frames", -1)) + terminal = code0 == self.codebook_vocab_size or (max_new_frames > 0 and generated_frames >= max_new_frames) + if terminal: + info["breeze_force_eos"] = True + existing = info.get("breeze_audio_codes") + per_request_codes.append( + existing + if isinstance(existing, torch.Tensor) and existing.numel() and not self._async_chunk + else hidden.new_empty((0, self.num_codebooks), dtype=torch.long) + ) + continue + if code0 < 0 or code0 >= self.codebook_size: + raise RuntimeError(f"Breeze main head produced invalid codec id {code0}") + + frame = self._generate_depth_codes(row_hidden, code0) + info["breeze_current_frame"] = frame[0].detach() + info["breeze_generated_frames"] = generated_frames + 1 + info["breeze_force_eos"] = False + if self._async_chunk: + # Streaming consumers only ever read the per-frame tail. The + # cumulative buffer exists for full-payload transport and the + # opt-in golden dump; maintaining it here would also cost an + # O(T^2) tensor-cat per request. + if self._golden_dump_dir: + self._accumulate_codes(info, frame) + per_request_codes.append(frame) + else: + per_request_codes.append(self._accumulate_codes(info, frame)) + if self._golden_dump_dir: + for info in info_items: + self._dump_golden_frames(info) + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"codes": {"audio": per_request_codes}}) + + def _dump_golden_frames(self, info: dict[str, Any]) -> None: + """Write terminal codec frames for opt-in golden-parity diagnostics.""" + codes = info.get("breeze_audio_codes") + if not info.get("breeze_force_eos", False) or not isinstance(codes, torch.Tensor) or codes.numel() == 0: + return + prompt_ids = info.get("prompt_ids") + if not isinstance(prompt_ids, torch.Tensor): + raise RuntimeError("Breeze golden dump is missing prompt_ids") + digest = hashlib.sha256(prompt_ids.detach().cpu().numpy().tobytes()).hexdigest()[:16] + output_dir = Path(self._golden_dump_dir) # type: ignore[arg-type] + output_dir.mkdir(parents=True, exist_ok=True) + torch.save( + { + "template": info.get("template"), + "prompt_ids": prompt_ids.cpu(), + "codes": codes.detach().cpu().to(torch.long).contiguous(), + "generated_frames": int(info.get("breeze_generated_frames", 0) or 0), + "finished": True, + }, + output_dir / f"breeze_{digest}.pt", + ) + + def _dump_golden_prefill(self, info: dict[str, Any], hidden: torch.Tensor) -> None: + """Write the terminal-prompt hidden row for opt-in parity diagnostics.""" + prompt_ids = info.get("prompt_ids") + if not isinstance(prompt_ids, torch.Tensor): + raise RuntimeError("Breeze golden dump is missing prompt_ids") + digest = hashlib.sha256(prompt_ids.detach().cpu().numpy().tobytes()).hexdigest()[:16] + output_dir = Path(self._golden_dump_dir) # type: ignore[arg-type] + output_dir.mkdir(parents=True, exist_ok=True) + torch.save( + { + "template": info.get("template"), + "prompt_ids": prompt_ids.cpu(), + "hidden": hidden.detach().reshape(-1).cpu().to(torch.float32).contiguous(), + }, + output_dir / f"prefill_{digest}.pt", + ) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Drop cross-step references retained for the finished AR requests.""" + del finished_req_ids + # Codec/frame state lives in request-local additional_information and is + # freed with the Request. Only the short logits-to-state mapping is + # model-owned and must not leak into the next micro-batch. + self._batch_state = None + self._batch_state_spans = None + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + params = dict(self.named_parameters()) + backbone_weights: list[tuple[str, torch.Tensor]] = [] + text_weights: list[tuple[str, torch.Tensor]] = [] + depth_weights: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + # This table belongs to Breeze's custom audio embedding module, + # although its checkpoint key is nested under backbone_model. + # Handle it before forwarding the remaining backbone weights to + # Qwen3Model. + if name == "backbone_model.embed_tokens.embed_audio_tokens.weight": + target = "embed_audio_tokens.weight" + param = params.get(target) + if param is None: + raise ValueError(f"Missing Breeze parameter: {target}") + _check_weight_shape(name, tensor, tuple(param.shape)) + default_weight_loader(param, tensor) + loaded.add(target) + continue + if name.startswith("backbone_model."): + # Qwen3Model.load_weights() is called on the nested module and + # therefore expects names relative to that module (the same + # convention used by MOSS-TTS), not the outer ``model.`` + # namespace used by ``named_parameters()``. + backbone_weights.append((name[len("backbone_model.") :], tensor)) + continue + if name.startswith("text_encoder."): + text_name = name[len("text_encoder.") :] + text_weights.append((text_name, tensor)) + continue + if name.startswith("depth_decoder."): + depth_weights.append((name[len("depth_decoder.") :], tensor)) + continue + if name in {"lm_head.weight", "embed_text_tokens.weight"} or name.startswith("text_encoder_proj."): + target = name + else: + continue + param = params.get(target) + if param is None: + continue + # ParallelLMHead stores a TP-local shard while checkpoint tensors + # are global; validate against the global Breeze head shape. + expected_shape = ( + (self.codebook_vocab_size + 1, self.hidden_size) if name == "lm_head.weight" else tuple(param.shape) + ) + _check_weight_shape(name, tensor, expected_shape) + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, tensor) + loaded.add(target) + + if backbone_weights: + loaded.update("model." + key for key in self.model.load_weights(backbone_weights)) + if text_weights: + loaded.update("text_encoder." + key for key in self.text_encoder.load_weights(text_weights)) + if depth_weights: + loaded.update("depth_decoder." + key for key in self.depth_decoder.load_weights(depth_weights)) + if bool(getattr(self.config, "tie_codebooks_embeddings", False)): + # The checkpoint serializes only one side of the tied table. + # ``named_parameters(remove_duplicate=True)`` exposes it as the + # talker-owned ``embed_audio_tokens.weight``, so report both + # aliases to the strict weight tracker. + if "depth_decoder.model.embed_tokens.weight" in loaded: + loaded.add("embed_audio_tokens.weight") + if "embed_audio_tokens.weight" in loaded: + loaded.add("depth_decoder.model.embed_tokens.weight") + return loaded + + +__all__ = ["BreezeTTS2TalkerForGeneration"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py new file mode 100644 index 000000000..0276e8278 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze T5Gemma2 text encoder. + +The Breeze checkpoint contains a T5Gemma2 encoder, not the older T5Gemma +implementation used by some diffusion models. This module keeps the encoder +local to Breeze so checkpoint names and attention semantics stay explicit, +while using vLLM tensor-parallel layers for the expensive projections. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + +class _RMSNorm(nn.Module): + """T5Gemma2 RMSNorm, whose checkpoint stores a zero-centered weight.""" + + def __init__(self, dim: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) + self.eps = float(eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + normalized = hidden_states.float() * torch.rsqrt( + hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + self.eps + ) + return (normalized * (1.0 + self.weight.float())).to(hidden_states.dtype) + + +class _RotaryEmbedding(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + for layer_type in dict.fromkeys(config.layer_types): + params = dict(config.rope_parameters.get(layer_type, {})) + base = float(params.get("rope_theta", 10_000.0)) + factor = float(params.get("factor", 1.0)) + dim = int(config.head_dim) + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if params.get("rope_type") == "linear": + inv_freq = inv_freq / factor + elif params.get("rope_type", "default") != "default": + raise ValueError(f"Unsupported T5Gemma2 rope type: {params.get('rope_type')!r}") + self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + layer_type: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = getattr(self, f"{layer_type}_inv_freq") + expanded = inv_freq[None, :, None].to(device=hidden_states.device) + positions = position_ids[:, None, :].float() + freqs = (expanded.float() @ positions).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().to(hidden_states.dtype), emb.sin().to(hidden_states.dtype) + + +def _rotate_half(value: torch.Tensor) -> torch.Tensor: + first, second = value[..., : value.shape[-1] // 2], value[..., value.shape[-1] // 2 :] + return torch.cat((-second, first), dim=-1) + + +def _apply_rotary( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (query * cos) + (_rotate_half(query) * sin), (key * cos) + (_rotate_half(key) * sin) + + +def _repeat_kv(hidden_states: torch.Tensor, repeats: int) -> torch.Tensor: + if repeats == 1: + return hidden_states + batch, heads, seq_len, head_dim = hidden_states.shape + return ( + hidden_states[:, :, None, :, :] + .expand(batch, heads, repeats, seq_len, head_dim) + .reshape(batch, heads * repeats, seq_len, head_dim) + ) + + +def _attention_mask( + padding_mask: torch.Tensor | None, + *, + batch_size: int, + seq_len: int, + layer_type: str, + sliding_window: int | None, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + if layer_type == "full_attention" and padding_mask is None: + return None + valid = ( + torch.ones((batch_size, seq_len), device=device, dtype=torch.bool) + if padding_mask is None + else padding_mask.to(device=device, dtype=torch.bool) + ) + allowed = valid[:, None, None, :].expand(batch_size, 1, seq_len, seq_len) + if layer_type == "sliding_attention": + if sliding_window is None: + raise ValueError("T5Gemma2 sliding_attention requires sliding_window") + q_idx = torch.arange(seq_len, device=device)[:, None] + kv_idx = torch.arange(seq_len, device=device)[None, :] + left = (int(sliding_window) + 1) // 2 + right = int(sliding_window) // 2 + 1 + distance = q_idx - kv_idx + local = ((distance >= 0) & (distance < left)) | ((distance < 0) & (-distance < right)) + allowed = allowed & local[None, None, :, :] + mask = torch.zeros((batch_size, 1, seq_len, seq_len), device=device, dtype=dtype) + mask.masked_fill_(~allowed, torch.finfo(dtype).min) + return mask + + +class _SelfAttention(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + if config.num_attention_heads % tp_size != 0: + raise ValueError("T5Gemma2 attention heads must be divisible by tensor parallel size") + self.layer_type = config.layer_type + self.head_dim = int(config.head_dim) + self.num_heads = int(config.num_attention_heads) // tp_size + self.num_kv_heads = max(1, int(config.num_key_value_heads) // tp_size) + self.num_key_value_groups = self.num_heads // self.num_kv_heads + self.scaling = float(config.query_pre_attn_scalar) ** -0.5 + self.qkv_proj = QKVParallelLinear( + hidden_size=config.hidden_size, + head_size=self.head_dim, + total_num_heads=config.num_attention_heads, + total_num_kv_heads=config.num_key_value_heads, + bias=config.attention_bias, + ) + self.o_proj = RowParallelLinear( + input_size=config.num_attention_heads * self.head_dim, + output_size=config.hidden_size, + bias=config.attention_bias, + input_is_parallel=True, + ) + self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split((q_size, kv_size, kv_size), dim=-1) + shape = (*hidden_states.shape[:-1], -1, self.head_dim) + query = self.q_norm(query.view(shape).transpose(1, 2)) + key = self.k_norm(key.view(shape).transpose(1, 2)) + value = value.view(shape).transpose(1, 2) + query, key = _apply_rotary(query, key, *position_embeddings) + key = _repeat_kv(key, self.num_key_value_groups) + value = _repeat_kv(value, self.num_key_value_groups) + # SDPA scales by head_dim**-0.5; adjust Q to T5Gemma2's scalar. + query = query * (self.scaling / (self.head_dim**-0.5)) + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + ) + output = output.transpose(1, 2).contiguous().reshape(*hidden_states.shape[:-1], -1) + output, _ = self.o_proj(output) + return output + + +class _MLP(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size, config.intermediate_size], + bias=False, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + input_is_parallel=True, + ) + self.act_fn = get_act_fn(config.hidden_activation) + self.dropout = nn.Dropout(float(config.dropout_rate)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(hidden_states) + gate, up = gate_up.chunk(2, dim=-1) + output, _ = self.down_proj(self.dropout(self.act_fn(gate) * up)) + return output + + +class _EncoderLayer(nn.Module): + def __init__(self, config: Any, layer_idx: int) -> None: + super().__init__() + layer_config = type("LayerConfig", (), dict(vars(config)))() + layer_config.layer_type = config.layer_types[layer_idx] + self.attention_type = layer_config.layer_type + self.pre_self_attn_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = _SelfAttention(layer_config) + self.post_self_attn_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.pre_feedforward_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = _MLP(config) + self.post_feedforward_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.dropout = nn.Dropout(float(config.dropout_rate)) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.pre_self_attn_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask) + hidden_states = residual + self.dropout(self.post_self_attn_layernorm(hidden_states)) + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + self.dropout(self.post_feedforward_layernorm(hidden_states)) + return hidden_states + + +class BreezeTTS2TextEncoder(nn.Module): + """Frozen T5Gemma2 encoder returning ``(B, S, hidden_size)`` states.""" + + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.embed_scale = float(config.hidden_size**0.5) + self.eoi_token_index = int(getattr(config, "eoi_token_index", 256000)) + # T5Gemma2 stores this learned replacement vector under the input + # embedding module (``embed_tokens.eoi_embedding``), so keep the same + # parameter path for checkpoint loading and tied-weight inspection. + self.embed_tokens.eoi_embedding = nn.Parameter(torch.zeros(config.hidden_size)) + self.layers = nn.ModuleList([_EncoderLayer(config, i) for i in range(config.num_hidden_layers)]) + self.norm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.rotary_emb = _RotaryEmbedding(config) + self.dropout = nn.Dropout(float(config.dropout_rate)) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if input_ids.ndim != 2: + raise ValueError(f"Breeze text encoder expects 2D input_ids, got {tuple(input_ids.shape)}") + if attention_mask is not None and attention_mask.shape != input_ids.shape: + raise ValueError( + "Breeze text encoder attention_mask must match input_ids: " + f"{tuple(attention_mask.shape)} != {tuple(input_ids.shape)}" + ) + hidden_states = self.embed_tokens(input_ids) * self.embed_scale + eoi_mask = input_ids == self.eoi_token_index + if bool(eoi_mask.any()): + hidden_states = torch.where( + eoi_mask.unsqueeze(-1), + self.embed_tokens.eoi_embedding.to(dtype=hidden_states.dtype), + hidden_states, + ) + hidden_states = self.dropout(hidden_states) + positions = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long).unsqueeze(0) + position_embeddings = { + layer_type: self.rotary_emb(hidden_states, positions, layer_type) + for layer_type in set(self.config.layer_types) + } + masks = { + layer_type: _attention_mask( + attention_mask, + batch_size=input_ids.shape[0], + seq_len=input_ids.shape[1], + layer_type=layer_type, + sliding_window=self.config.sliding_window, + dtype=hidden_states.dtype, + device=input_ids.device, + ) + for layer_type in set(self.config.layer_types) + } + for layer in self.layers: + hidden_states = layer( + hidden_states, + position_embeddings[layer.attention_type], + masks[layer.attention_type], + ) + return self.dropout(self.norm(hidden_states)) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = ( + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ) + params = dict(self.named_parameters()) + loaded: set[str] = set() + for name, tensor in weights: + target = name + shard_id: str | int | None = None + for packed_name, source_name, source_id in stacked_params_mapping: + marker = f".{source_name}." + if marker in name: + target = name.replace(marker, f".{packed_name}.") + shard_id = source_id + break + parameter = params.get(target) + if parameter is None: + continue + loader = getattr(parameter, "weight_loader", default_weight_loader) + if shard_id is None: + loader(parameter, tensor) + else: + loader(parameter, tensor, shard_id) + loaded.add(target) + return loaded + + +__all__ = ["BreezeTTS2TextEncoder"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py new file mode 100644 index 000000000..d1aaeb13e --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Synchronous Breeze-TTS-2 pipeline: talker -> Mimi codec.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.breeze_tts_2" + +BREEZE_TTS_2_PIPELINE = PipelineConfig( + model_type="breeze", + default_deploy_config_name="breeze_tts_2.yaml", + model_arch="BreezeForConditionalGeneration", + hf_architectures=("BreezeForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="breeze_tts_2", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=f"{_PROC}.talker2codec_async_chunk", + custom_process_next_stage_input_func=f"{_PROC}.talker2codec_full_payload", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2051], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="breeze_tts_2_codec", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="BreezeTTS2MimiCodec", + sync_process_input_func=f"{_PROC}.talker2codec", + sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, + retains_state_across_chunks=True, + ), + ), +) + +__all__ = ["BREEZE_TTS_2_PIPELINE"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py b/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py new file mode 100644 index 000000000..32c8145bb --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 request-to-prompt conversion. + +This module is intentionally limited to request preparation. It creates the +real prompt token ids needed by the scheduler and stores the text/audio masks +and reference codes in ``additional_information``. The stage-0 model later +turns those values into embeddings; this keeps tokenization independent from +the Qwen3/text-encoder implementation and avoids changing sequence length in +the model ``preprocess`` hook. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from transformers import AutoTokenizer + +from vllm_omni.inputs.data import OmniTokensPrompt + +from .audio_tokenizer import BreezeReferenceAudioTokenizer + +AUDIO_TAG = "<|AUDIO|>" +AUDIO_EOS = "<|audio_eos|>" +INSTRUCTION_BOS = "" +INSTRUCTION_EOS = "" + + +@dataclass(frozen=True) +class _TextSegment: + text: str + + +@dataclass(frozen=True) +class _AudioSegment: + source: Any + sample_rate: int | None + + +class BreezeTTS2PromptBuilder: + """Build one Breeze prompt without depending on model modules. + + ``tokenizer`` is a normal HuggingFace tokenizer. ``reference_audio_encoder`` + is injected so the builder can use the serving worker's cached audio + encoder and never reload it per request. The returned ``OmniTokensPrompt`` + contains scheduler token ids and compact CPU metadata for stage-0 + ``preprocess``. + """ + + def __init__( + self, + tokenizer: Any, + config: Any, + reference_audio_encoder: BreezeReferenceAudioTokenizer | None = None, + ) -> None: + self.tokenizer = tokenizer + self.audio_encoder = reference_audio_encoder + self.audio_token_id = int(config.audio_token_id) + self.audio_eos_token_id = int(config.audio_eos_token_id) + # The Breeze tokenizer vocabulary is larger than the Qwen3 backbone + # vocabulary. ``prompt_token_ids`` are scheduler bookkeeping only, so + # keep them in the backbone vocabulary and retain the real Breeze ids + # in ``additional_information['prompt_ids']`` for the talker. + self.scheduler_token_id = int(config.pad_token_id) + scheduler_vocab_size = int(config.backbone_config.vocab_size) + if self.scheduler_token_id >= scheduler_vocab_size: + raise ValueError( + "prompt_pad_token_id must be smaller than the backbone vocabulary: " + f"{self.scheduler_token_id} >= {scheduler_vocab_size}" + ) + self.num_codebooks = int(config.num_codebooks) + codec_config = getattr(config, "codec_config", None) + if isinstance(codec_config, Mapping): + self.codebook_size = int(codec_config.get("codebook_size", 2048)) + elif codec_config is not None: + self.codebook_size = int(getattr(codec_config, "codebook_size", 2048)) + else: + # ``codec_config`` is optional in BreezeTTS2Config; a minimal + # checkpoint may omit it. Breeze's codec uses 2048-entry + # codebooks, the same fallback the talker and codec stages apply. + self.codebook_size = 2048 + if self.num_codebooks <= 0 or self.codebook_size <= 0: + raise ValueError("num_codebooks and codec codebook_size must be positive") + self._validate_audio_special_tokens() + + @classmethod + def from_pretrained( + cls, + model_path: str, + config: Any, + *, + reference_audio_encoder: BreezeReferenceAudioTokenizer | None = None, + tokenizer_kwargs: Mapping[str, Any] | None = None, + ) -> BreezeTTS2PromptBuilder: + """Load the text tokenizer once and construct a reusable builder.""" + kwargs = dict(tokenizer_kwargs or {}) + kwargs.setdefault("use_fast", True) + kwargs.setdefault("trust_remote_code", True) + tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs) + return cls(tokenizer, config, reference_audio_encoder) + + def build( + self, + request: Mapping[str, Any], + *, + template: str | None = None, + ) -> OmniTokensPrompt: + """Convert one request to a vLLM-Omni token prompt. + + Supported templates are ``tts_plain``, ``tts_instruction``, + ``ref_clone_tata`` and ``ref_edit_tata``. When omitted, the template + is inferred from the presence of instruction/reference fields. + """ + template_name = template or request.get("template") + if not isinstance(template_name, str) or not template_name: + template_name = _infer_template(request) + segments = self._build_segments(template_name, request) + + # Breeze's reference implementation first adds tokenizer special + # tokens to every text segment, decodes those segments back to text, + # inserts the audio placeholders, and tokenizes the complete string + # once. This matters at segment boundaries: tokenizing each raw text + # fragment independently can produce different boundary tokens. + rendered_segments: list[tuple[str, bool, int]] = [] + reference_codes: list[torch.Tensor] = [] + + for segment in segments: + if isinstance(segment, _TextSegment): + ids = self._encode_one_text(segment.text, add_special_tokens=True) + if not ids: + raise ValueError("Breeze text tokenizer returned an empty segment") + rendered = str(self.tokenizer.decode(ids, skip_special_tokens=False)) + rendered_segments.append((rendered, True, 0)) + continue + + codes = self._resolve_reference_codes(segment) + reference_codes.append(codes) + audio_text = AUDIO_TAG * int(codes.shape[0]) + AUDIO_EOS + rendered_segments.append((audio_text, False, int(codes.shape[0]) + 1)) + + prompt_ids: list[int] = [] + text_mask: list[bool] = [] + text_lengths: list[int] = [] + rendered_prompt = "".join(item[0] for item in rendered_segments) + prompt_ids = self._encode_one_text(rendered_prompt, add_special_tokens=False) + offset = 0 + for rendered, is_text, expected_len in rendered_segments: + segment_ids = self._encode_one_text(rendered, add_special_tokens=False) + actual_len = len(segment_ids) + if not is_text and actual_len != expected_len: + raise ValueError( + "Breeze audio placeholder token length changed after rendering: " + f"expected {expected_len}, got {actual_len}" + ) + if not is_text and segment_ids != [ + *([self.audio_token_id] * (expected_len - 1)), + self.audio_eos_token_id, + ]: + raise ValueError("Breeze tokenizer did not preserve audio placeholder ids") + text_mask.extend([is_text] * actual_len) + if is_text: + text_lengths.append(actual_len) + offset += actual_len + if offset != len(prompt_ids) or len(text_mask) != len(prompt_ids): + raise RuntimeError("Breeze rendered prompt length accounting error") + if not prompt_ids: + raise ValueError("Breeze prompt is empty") + + input_values = torch.cat(reference_codes, dim=0) if reference_codes else None + if input_values is not None and int(input_values.shape[0]) != prompt_ids.count(self.audio_token_id): + raise RuntimeError("reference code count does not match audio placeholder count") + + additional_information: dict[str, Any] = { + # The model runner may invoke preprocess on a chunk of a long + # prompt. Keep one compact CPU copy so the talker can reconstruct + # the full embedding buffer before slicing that chunk. + "prompt_ids": torch.tensor(prompt_ids, dtype=torch.long), + "attention_mask": torch.ones(len(prompt_ids), dtype=torch.bool), + "text_ids_mask": torch.tensor(text_mask, dtype=torch.bool), + "text_ids_len": torch.tensor(text_lengths, dtype=torch.long), + "template": template_name, + } + if input_values is not None: + # ``input_values`` is the name used by the upstream Breeze model. + # Keep the tensor unbatched here: Omni transports one request at a + # time and the talker adds the batch dimension after collation. + additional_information["input_values"] = input_values + additional_information["ref_code_len"] = int(input_values.shape[0]) + + if self.scheduler_token_id < 0: + raise ValueError("prompt_pad_token_id must be non-negative") + + return OmniTokensPrompt( + # Never expose Breeze's high text/audio ids to the Qwen3 scheduler + # or its token-range validation. The real ids remain available in + # additional_information and are consumed by ``preprocess``. + prompt_token_ids=[self.scheduler_token_id] * len(prompt_ids), + additional_information=additional_information, + ) + + def _build_segments(self, template: str, request: Mapping[str, Any]) -> list[_TextSegment | _AudioSegment]: + speaker = _speaker_prefix(request.get("speaker", "S0")) + text = _required_text(request, "text") + instruction = request.get("instruction") + ref_text = request.get("ref_text") + + if template == "tts_plain": + return [_TextSegment(f"{speaker}{text}")] + if template == "tts_instruction": + if not isinstance(instruction, str) or not instruction: + raise ValueError("tts_instruction requires a non-empty instruction") + return [_TextSegment(f"{speaker}{INSTRUCTION_BOS}{instruction}{INSTRUCTION_EOS}{text}")] + if template in ("ref_clone_tata", "ref_edit_tata"): + if not isinstance(ref_text, str) or not ref_text: + raise ValueError(f"{template} requires a non-empty ref_text") + audio = self._audio_segment(request) + if template == "ref_clone_tata": + return [_TextSegment(f"{speaker}{ref_text}"), audio, _TextSegment(f"{speaker}{text}")] + if not isinstance(instruction, str) or not instruction: + raise ValueError("ref_edit_tata requires a non-empty instruction") + return [ + _TextSegment(f"{speaker}{ref_text}"), + audio, + _TextSegment(f"{speaker}{INSTRUCTION_BOS}{instruction}{INSTRUCTION_EOS}{text}"), + ] + raise ValueError( + f"unknown Breeze template {template!r}; expected tts_plain, tts_instruction, " + "ref_clone_tata or ref_edit_tata" + ) + + def _audio_segment(self, request: Mapping[str, Any]) -> _AudioSegment: + source = request.get("ref_audio") + if source is None: + source = request.get("ref_audio_codes") + if source is None: + raw_codes = request.get("codes") + source = raw_codes.get("ref") if isinstance(raw_codes, Mapping) else None + if source is None: + raise ValueError("reference template requires ref_audio or ref_audio_codes") + sample_rate = request.get("ref_audio_sample_rate") + return _AudioSegment(source, int(sample_rate) if sample_rate is not None else None) + + def _resolve_reference_codes(self, segment: _AudioSegment) -> torch.Tensor: + source = segment.source + if _looks_like_codes(source, self.num_codebooks): + codes = torch.as_tensor(source) + return _normalize_codes(codes, self.num_codebooks, self.codebook_size) + if self.audio_encoder is None: + raise RuntimeError("reference audio was provided but no reference_audio_encoder is configured") + codes = self.audio_encoder.encode(source, segment.sample_rate) + return _normalize_codes(codes, self.num_codebooks, self.codebook_size) + + def _encode_one_text(self, text: str, *, add_special_tokens: bool) -> list[int]: + encoded = self.tokenizer(text, add_special_tokens=add_special_tokens, padding=False) + ids = encoded["input_ids"] if isinstance(encoded, Mapping) else encoded.input_ids + if isinstance(ids, torch.Tensor): + ids = ids.tolist() + return _as_int_list(ids) + + def _validate_audio_special_tokens(self) -> None: + convert = self.tokenizer.convert_tokens_to_ids + for token, expected in ((AUDIO_TAG, self.audio_token_id), (AUDIO_EOS, self.audio_eos_token_id)): + actual = convert(token) + unk = getattr(self.tokenizer, "unk_token_id", None) + if actual is None or (unk is not None and actual == unk) or int(actual) != expected: + raise ValueError(f"tokenizer id for {token!r} is {actual}, expected config value {expected}") + + +def _infer_template(request: Mapping[str, Any]) -> str: + has_ref = ( + request.get("ref_audio") is not None + or request.get("ref_audio_codes") is not None + or bool(request.get("ref_text")) + ) + has_instruction = bool(request.get("instruction")) + if has_ref and has_instruction: + return "ref_edit_tata" + if has_ref: + return "ref_clone_tata" + return "tts_instruction" if has_instruction else "tts_plain" + + +def _required_text(request: Mapping[str, Any], key: str) -> str: + value = request.get(key) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Breeze request requires non-empty {key}") + return value + + +def _speaker_prefix(value: Any) -> str: + if value in (None, ""): + return "" + value = str(value) + return value if value.startswith("[") and value.endswith("]") else f"[{value}]" + + +def _as_int_list(value: Any) -> list[int]: + if isinstance(value, torch.Tensor): + value = value.tolist() + elif isinstance(value, np.ndarray): + value = value.tolist() + if not isinstance(value, list): + raise ValueError("tokenizer input_ids must be a list") + if value and isinstance(value[0], (list, tuple, np.ndarray, torch.Tensor)): + if len(value) != 1: + raise ValueError("tokenizer input_ids contains multiple sequences") + return _as_int_list(value[0]) + return [int(item) for item in value] + + +def _looks_like_codes(value: Any, num_codebooks: int) -> bool: + if isinstance(value, (torch.Tensor, np.ndarray)): + return value.ndim == 2 and value.shape[-1] == num_codebooks + return False + + +def _normalize_codes(codes: Any, num_codebooks: int, codebook_size: int) -> torch.Tensor: + codes = torch.as_tensor(codes) + if codes.ndim != 2: + raise ValueError(f"reference audio codes must be 2D, got {tuple(codes.shape)}") + if codes.shape[0] == 0: + raise ValueError("reference audio codes must contain at least one frame") + if codes.shape[-1] != num_codebooks: + raise ValueError(f"expected {num_codebooks} codebooks, got {tuple(codes.shape)}") + codes = codes.to(device="cpu", dtype=torch.long).contiguous() + if int(codes.min()) < 0 or int(codes.max()) >= codebook_size: + raise ValueError(f"reference code outside [0, {codebook_size})") + return codes.to(dtype=torch.int16) + + +__all__ = ["BreezeTTS2PromptBuilder"] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index febda68af..e5d7b4115 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -492,6 +492,16 @@ _OMNI_MODELS = { "acoustic", "MiniMaxMusic3AcousticForConditionalGeneration", ), + "BreezeForConditionalGeneration": ( + "breeze_tts_2", + "modeling_breeze_tts_2_talker", + "BreezeTTS2TalkerForGeneration", + ), + "BreezeTTS2MimiCodec": ( + "breeze_tts_2", + "modeling_breeze_tts_2_codec", + "BreezeTTS2MimiCodec", + ), } diff --git a/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py b/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py new file mode 100644 index 000000000..ae393948a --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Stage-0 to stage-1 input conversion for Breeze-TTS-2. + +The prompt builder and the reference-audio tokenizer live under the Breeze +model directory. This module intentionally contains only the inter-stage +contract: flatten the generated ``(T, Q)`` codec matrix into the codebook-major +token sequence consumed by the stage-1 codec model. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + +# ``BreezeTTS2TalkerForGeneration`` emits the complete sequence accumulated +# so far on every decode step. The generic full-payload accumulator normally +# concatenates rank-2 tensors, which would duplicate all earlier frames. +# Replace the latest complete sequence instead. +_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes.audio"}) +_DEFAULT_CODEC_CHUNK_FRAMES = 8 + + +def _extract_audio_codes(output: Any) -> torch.Tensor | None: + """Extract ``codes.audio`` from the supported Omni output shapes.""" + # Depending on the engine path, source_outputs contains either the + # per-request output itself or an EngineCoreOutput with ``outputs[0]``. + if isinstance(output, Mapping): + multimodal = output.get("multimodal_output", output.get("multimodal_outputs")) + if multimodal is None and "codes" in output: + multimodal = output + if multimodal is None and "codes.audio" in output: + audio = output["codes.audio"] + if isinstance(audio, (list, tuple)): + if len(audio) != 1: + raise ValueError(f"expected one Breeze audio-code tensor, got {len(audio)}") + audio = audio[0] + audio = torch.as_tensor(audio) + if audio.ndim not in (1, 2): + raise ValueError(f"Breeze audio codes must be rank 1 or 2, got {tuple(audio.shape)}") + return audio.to(device="cpu", dtype=torch.long).contiguous() + else: + multimodal = getattr(output, "multimodal_output", None) + if multimodal is None: + multimodal = getattr(output, "multimodal_outputs", None) + if multimodal is None: + candidates = getattr(output, "outputs", None) + if isinstance(candidates, Sequence) and candidates: + return _extract_audio_codes(candidates[0]) + if not isinstance(multimodal, Mapping): + return None + codes = multimodal.get("codes") + audio = codes.get("audio") if isinstance(codes, Mapping) else multimodal.get("codes.audio") + if audio is None: + return None + # ``OmniOutput`` stores one tensor per request under ``codes.audio``. + # ``process_engine_inputs`` invokes this function once per source output, + # so unwrap that single-item container before tensor conversion. + if isinstance(audio, (list, tuple)): + if len(audio) != 1: + raise ValueError(f"expected one Breeze audio-code tensor, got {len(audio)}") + audio = audio[0] + audio = torch.as_tensor(audio) + if audio.ndim == 3 and audio.shape[0] == 1: + audio = audio[0] + if audio.ndim not in (1, 2): + raise ValueError( + f"Breeze audio codes must have shape (T, Q) or codebook-major (Q*T,), got {tuple(audio.shape)}" + ) + return audio.to(device="cpu", dtype=torch.long).contiguous() + + +def talker2codec( + source_outputs: Sequence[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """Build synchronous stage-1 prompts from completed stage-0 outputs. + + Breeze's Mimi decoder receives one flat codebook-major sequence. The + generated matrix is therefore transposed from ``(T, Q)`` to ``(Q, T)`` + before flattening. No audio decoding, text tokenization, or streaming + state is handled here. + """ + del prompt, _requires_multimodal_data + results: list[OmniTokensPrompt] = [] + for source_index, output in enumerate(source_outputs): + codes = _extract_audio_codes(output) + if codes is None or codes.numel() == 0: + logger.warning("Breeze stage output %d contains no audio codes", source_index) + results.append(OmniTokensPrompt(prompt_token_ids=[])) + continue + # Connector full-payload builders send a codebook-major flat tensor; + # direct in-process stage outputs carry frame-major ``(T, Q)``. + if codes.ndim == 1: + codebook_major_data = codes.contiguous() + codebook_major = codebook_major_data + else: + codebook_major_data = codes.transpose(0, 1).contiguous() + codebook_major = codebook_major_data.reshape(-1) + results.append( + OmniTokensPrompt( + prompt_token_ids=codebook_major.tolist(), + multi_modal_data={"codes": {"audio": codebook_major_data}}, + ) + ) + return results + + +def talker2codec_full_payload( + transfer_manager: Any = None, + pooling_output: Mapping[str, Any] | None = None, + request: Any = None, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Build the final Stage-1 payload from accumulated talker codes. + + ``pooling_output`` is the AR runner's flattened per-request payload. Mimi + consumes codebook-major IDs, so the frame-major talker matrix is transposed + before flattening. A request that finishes without any codec frame still + gets an explicit empty payload with ``finished=True``: returning ``None`` + would leave Stage 1 without a terminal marker, and the connector receiver + requeues marker-less empty chunks forever. + """ + del transfer_manager, request + audio = None + if isinstance(pooling_output, Mapping): + audio = pooling_output.get("codes.audio") + if audio is None: + nested = pooling_output.get("codes") + if isinstance(nested, Mapping): + audio = nested.get("audio") + + codes = torch.empty(0, dtype=torch.long) if audio is None else torch.as_tensor(audio, dtype=torch.long) + if codes.numel() == 0: + if not is_finished: + return None + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + code_flat_numel=0, + finished=torch.tensor(True, dtype=torch.bool), + ), + ) + if codes.ndim == 2: + flat = codes.transpose(0, 1).contiguous().reshape(-1) + elif codes.ndim == 1: + flat = codes.contiguous() + else: + raise ValueError(f"Breeze audio codes must be rank 1 or 2, got {tuple(codes.shape)}") + return OmniPayloadStruct( + codes=CodesStruct(audio=flat), + meta=MetaStruct( + code_flat_numel=int(flat.numel()), + finished=torch.tensor(bool(is_finished), dtype=torch.bool), + ), + ) + + +def _extract_current_audio_frame( + multimodal_output: Mapping[str, Any] | None, +) -> torch.Tensor | None: + """Extract the single ``(Q,)`` frame emitted by async Stage 0.""" + if not isinstance(multimodal_output, Mapping): + return None + codes = multimodal_output.get("codes") + audio = codes.get("audio") if isinstance(codes, Mapping) else None + if audio is None: + return None + frame = torch.as_tensor(audio, dtype=torch.long).reshape(-1) + if frame.numel() == 0: + return None + return frame.cpu().contiguous() + + +def _codec_chunk_frames(transfer_manager: Any) -> int: + connector = getattr(transfer_manager, "connector", None) + # SharedMemoryConnector stores its ConnectorSpec.extra directly in + # ``connector.config``; there is no second ``extra`` wrapper at runtime. + config = getattr(connector, "config", {}) or {} + value = config.get("breeze_codec_chunk_frames") + try: + chunk_frames = int(value) if value is not None else _DEFAULT_CODEC_CHUNK_FRAMES + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid Breeze codec_chunk_frames={value!r}") from exc + if chunk_frames <= 0: + raise ValueError(f"Breeze codec_chunk_frames must be positive, got {chunk_frames}") + return chunk_frames + + +def talker2codec_async_chunk( + transfer_manager: Any, + multimodal_output: Mapping[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Batch Stage-0 frames and send only the un-emitted codebook-major tail.""" + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + frame = _extract_current_audio_frame(multimodal_output) + if frame is not None: + transfer_manager.code_prompt_token_ids[request_id].append(frame) + + pending = transfer_manager.code_prompt_token_ids[request_id] + chunk_frames = _codec_chunk_frames(transfer_manager) + if not pending and not finished: + return None + if len(pending) < chunk_frames and not finished: + return None + + if pending: + # Stage 0 supplies frame-major ``(T, Q)`` rows. The Qwen3-TTS decoder + # consumes ``(Q, T)``, so transpose before flattening the emitted tail. + tail = torch.stack(pending, dim=0) + flat = tail.transpose(0, 1).contiguous().reshape(-1) + else: + flat = torch.empty(0, dtype=torch.long) + pending.clear() + + return OmniPayloadStruct( + codes=CodesStruct(audio=flat), + meta=MetaStruct( + request_id=request_id, + code_flat_numel=int(flat.numel()), + codec_streaming=True, + finished=torch.tensor(finished, dtype=torch.bool), + # The chunk adapter strips ``meta.finished`` before it reaches the + # stage-1 model; ``stream_finished`` survives and is what the + # codec uses to release per-request decoder state. + stream_finished=torch.tensor(finished, dtype=torch.bool), + ), + ) + + +__all__ = ["talker2codec", "talker2codec_async_chunk", "talker2codec_full_payload"] From 909b121e0a5485ed86a604ce091a6b153d6c7257 Mon Sep 17 00:00:00 2001 From: nodeeeeee Date: Tue, 15 Sep 2026 01:02:11 +0800 Subject: [PATCH 061/209] [Qwen2.5-Omni] Avoid redundant embedding computation (#7477) Signed-off-by: nodeeeeee --- benchmarks/README.md | 5 + benchmarks/qwen2_5_omni/README.md | 66 ++++ .../qwen2_5_omni/benchmark_embeddings.py | 306 ++++++++++++++++++ .../qwen2_5_omni/test_qwen2_5_omni_embed.py | 34 +- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 14 +- 5 files changed, 412 insertions(+), 13 deletions(-) create mode 100644 benchmarks/qwen2_5_omni/README.md create mode 100644 benchmarks/qwen2_5_omni/benchmark_embeddings.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 0a8fb8058..182bba02d 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -4,6 +4,11 @@ This directory contains benchmark suites for evaluating different model families ## Benchmark families +### [Qwen2.5-Omni](qwen2_5_omni/README.md) — Embedding Wrapper Latency + +CPU-observed A/B timings for redundant text embedding removal, with real checkpoint +embeddings, warmup, paired blocks, and embedding call-count validation. + ### [TTS](tts/README.md) — Text-to-Speech Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2. diff --git a/benchmarks/qwen2_5_omni/README.md b/benchmarks/qwen2_5_omni/README.md new file mode 100644 index 000000000..c8aea8980 --- /dev/null +++ b/benchmarks/qwen2_5_omni/README.md @@ -0,0 +1,66 @@ +# Qwen2.5-Omni embedding microbenchmark + +Measure the host-observed latency of `embed_input_ids` before and after the +redundant text embedding removal in PR #7477. Both variants use the same real +checkpoint embedding table, dependencies, inputs and batch shapes. The baseline +wrapper is extracted from the specified Git commit; the fixed wrapper comes from +the current checkout. No full model is initialized. + +## Run + +Use Linux, a CUDA GPU, the project dependencies, and a local +`Qwen/Qwen2.5-Omni-7B` checkpoint with its safetensors index and weights. +From the repository root: + +```bash +git fetch origin 01a2f93256975c7ff9565c5414b0fe0225bc4765 +PYTHONPATH="$PWD" OMP_NUM_THREADS=1 python benchmarks/qwen2_5_omni/benchmark_embeddings.py \ + --model /path/to/Qwen2.5-Omni-7B \ + --baseline-ref 01a2f93256975c7ff9565c5414b0fe0225bc4765 \ + --cpu-core 2 --lengths 1172 --warmup 50 --blocks 20 --iterations 20 \ + --output-dir /tmp/qwen25-embedding-results +``` + +Choose a CPU core in your allowed affinity set. For the original CUDA length +sweep, use `--lengths 512 1172 4096 16384`. + +## Method + +- BF16 checkpoint embeddings, TP=1, CUDA int32 token IDs, CPU multimodal masks. +- One pinned CPU core, one PyTorch thread, garbage collection disabled. +- 50 warmups per variant, then 20 paired blocks of 20 calls with balanced, + randomized variant order. Summary: median of block medians. +- Host latency includes synchronization inside the wrapper. Explicit GPU + synchronization before and after the timed call is excluded from host latency. + Calling-thread CPU time and synchronized completion time are recorded separately. +- CPU frequency is sampled at block boundaries; clocks are not locked. +- Validation checks one text embedding call for the fixed wrapper, two for the + multimodal baseline, and `torch.equal` for identical inputs and batch shapes. + This does not establish batch invariance or full-model output equality. + +Interleaved audio/video uses a valid synthetic multimodal mask. This benchmark +measures wrapper latency, not E2E throughput. + +## Reference results + +Original environment: H100 80GB; Python 3.12.3; vLLM 0.29.0+cu129; +PyTorch 2.13.0+cu129; Transformers 5.14.0; Qwen2.5-Omni-7B BF16. +At 1172 tokens, the September 13 measurements were: + +| Input | Before (μs) | After (μs) | Host latency reduction | +| --- | ---: | ---: | ---: | +| Text | 18.81 | 18.66 | 0.80% | +| Audio | 104.32 | 85.30 | 18.23% | +| Image | 104.72 | 86.27 | 17.62% | +| Video | 104.73 | 85.76 | 18.11% | +| Audio + image + video | 567.92 | 549.54 | 3.24% | +| Interleaved audio/video | 622.76 | 597.84 | 4.00% | + +Results vary between runs and environments. These local savings did not produce +a consistent E2E throughput improvement in the separate repeated E2E experiment. + +## Output + +The output directory contains `metadata.json`, `micro-raw.json`, +`validation.json`, and `summary.json`. Raw records include per-call timings, +block timestamps and CPU frequency samples. Timer overhead is measured separately. diff --git a/benchmarks/qwen2_5_omni/benchmark_embeddings.py b/benchmarks/qwen2_5_omni/benchmark_embeddings.py new file mode 100644 index 000000000..1a5b32435 --- /dev/null +++ b/benchmarks/qwen2_5_omni/benchmark_embeddings.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reproduce PR #7477 host-observed embedding-wrapper latency on one CUDA GPU. + +Run from the PR checkout with PYTHONPATH=$PWD. Baseline and fixed methods share +weights and inputs. Only the historical wrapper is loaded from --baseline-ref. +""" + +import argparse +import ast +import gc +import json +import os +import random +import statistics +import subprocess +import textwrap +import time +from pathlib import Path + +import torch +from safetensors import safe_open +from torch import nn +from transformers import Qwen2_5OmniThinkerConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.multimodal.utils import set_mm_embedding_modality + +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerForConditionalGeneration as Thinker, +) +from vllm_omni.platforms import current_omni_platform + + +def baseline_method(cls, ref): + path = "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py" + source = subprocess.check_output( + ["git", "-C", str(Path(__file__).resolve().parents[2]), "show", f"{ref}:{path}"], text=True + ) + tree = ast.parse(source) + owner = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == cls.__name__) + node = next(n for n in owner.body if isinstance(n, ast.FunctionDef) and n.name == "embed_input_ids") + method = textwrap.dedent(ast.get_source_segment(source, node)) + # Preserve zero-argument super() against the real model class and MRO. + factory = "from __future__ import annotations\ndef make_method(cls):\n __class__ = cls\n" + factory += textwrap.indent(method, " ") + "\n return embed_input_ids\n" + namespace = dict(cls.embed_input_ids.__globals__) + exec(compile(factory, f"{ref}:{path}", "exec"), namespace) + return namespace["make_method"](cls) + + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--model", type=Path, required=True) +parser.add_argument("--baseline-ref", default="01a2f93256975c7ff9565c5414b0fe0225bc4765") +parser.add_argument("--cpu-core", type=int, default=sorted(os.sched_getaffinity(0))[0]) +parser.add_argument("--lengths", type=int, nargs="+", default=[1172]) +parser.add_argument("--warmup", type=int, default=50) +parser.add_argument("--blocks", type=int, default=20) +parser.add_argument("--iterations", type=int, default=20) +parser.add_argument("--output-dir", type=Path, default=Path("micro-results")) +args = parser.parse_args() +if min(args.lengths) < 300 or min(args.blocks, args.iterations) < 1 or args.warmup < 0: + parser.error("Require lengths >= 300, positive blocks/iterations, and nonnegative warmup") +if args.cpu_core not in os.sched_getaffinity(0): + parser.error("CPU core is outside the allowed affinity set") +ROOT, MODEL, CORE = args.output_dir, args.model, args.cpu_core +ROOT.mkdir(parents=True, exist_ok=True) +CONFIG = json.loads((MODEL / "config.json").read_text())["thinker_config"] +HIDDEN = CONFIG["text_config"]["hidden_size"] +VOCAB = CONFIG["text_config"]["vocab_size"] +IDS = {m: CONFIG[m + "_token_index"] for m in ("audio", "image", "video")} +METHODS = {"baseline": baseline_method(Thinker, args.baseline_ref), "fixed": Thinker.embed_input_ids} +RESULTS, VALIDATIONS = [], [] +os.sched_setaffinity(0, {CORE}) +torch.set_num_threads(1) +torch.set_num_interop_threads(1) +gc.disable() + + +def cpu_frequency(): + path = Path(f"/sys/devices/system/cpu/cpu{CORE}/cpufreq/scaling_cur_freq") + if path.exists(): + return float(path.read_text()) / 1000 + for section in Path("/proc/cpuinfo").read_text().split("\n\n"): + fields = dict(line.split(":", 1) for line in section.splitlines() if ":" in line) + fields = {k.strip(): v.strip() for k, v in fields.items()} + if fields.get("processor") == str(CORE): + return float(fields.get("cpu MHz", "nan")) + + +def payload(case, length, device): + generator = torch.Generator().manual_seed(7451 + length) + ids = torch.randint(2, 10000, (length,), generator=generator) + specs = { + "text": [], + "audio": [("audio", 100)], + "image": [("image", 16)], + "video": [("video", 128)], + "mixed": [("audio", 100), ("image", 16), ("video", 128)], + "interleaved": [("video", 64), ("audio", 50), ("video", 64), ("audio", 50)], + }[case] + offset = 44 + counts = {} + for modality, count in specs: + ids[offset : offset + count] = IDS[modality] + counts[modality] = counts.get(modality, 0) + count + offset += count + (0 if case == "interleaved" else 2) + mask = torch.zeros(length, dtype=torch.bool) + for token in IDS.values(): + mask |= ids == token + mm = [ + set_mm_embedding_modality(torch.full((count, HIDDEN), i + 1, device=device, dtype=torch.bfloat16), modality) + for i, (modality, count) in enumerate(counts.items()) + ] + return ids.to(device=device, dtype=torch.int32), None if case == "text" else mm, None if case == "text" else mask + + +def once(fn, cuda): + # Drain preceding work before the host timer; the production function's + # own D2H synchronizations remain inside the timed call. + if cuda: + current_omni_platform.synchronize() + w0 = time.perf_counter_ns() + c0 = time.thread_time_ns() + value = fn() + c1 = time.thread_time_ns() + w1 = time.perf_counter_ns() + if cuda: + current_omni_platform.synchronize() + w2 = time.perf_counter_ns() if cuda else w1 + del value + return [(w1 - w0) / 1000, (c1 - c0) / 1000, (w2 - w0) / 1000] + + +def measure(functions, device, case, length, scope, rounds=20, iterations=20, warmups=50): + cuda = device == "cuda" + for fn in functions.values(): + for _ in range(warmups): + once(fn, cuda) + order_rng = random.Random(7451 + length + len(case) + len(scope)) + directions = [i % 2 for i in range(rounds)] + order_rng.shuffle(directions) + names = list(functions) + for block, direction in enumerate(directions): + order = names if direction == 0 else names[::-1] + for name in order: + frequency_start = cpu_frequency() + wall_start = time.time() + samples = [once(functions[name], cuda) for _ in range(iterations)] + wall_end = time.time() + RESULTS.append( + dict( + device=device, + case=case, + tokens=length, + scope=scope, + variant=name, + block=block, + iterations=iterations, + samples_us=samples, + wall_start=wall_start, + wall_end=wall_end, + cpu_mhz_start=frequency_start, + cpu_mhz_end=cpu_frequency(), + ) + ) + (ROOT / "micro-raw.json").write_text(json.dumps(RESULTS)) + current = [r for r in RESULTS if (r["device"], r["case"], r["tokens"], r["scope"]) == (device, case, length, scope)] + print( + "RESULT", + json.dumps( + { + "device": device, + "case": case, + "tokens": length, + "scope": scope, + "host_median_us": { + v: statistics.median(t[0] for r in current if r["variant"] == v for t in r["samples_us"]) + for v in names + }, + } + ), + flush=True, + ) + + +def build_model(weights, device): + model = Thinker.__new__(Thinker) + nn.Module.__init__(model) + model.config = Qwen2_5OmniThinkerConfig(**CONFIG) + model._has_oov_mm_tokens = False + table = VocabParallelEmbedding(VOCAB, HIDDEN, params_dtype=torch.bfloat16, disable_tp=True) + table.weight = nn.Parameter(weights.to(device), requires_grad=False) + + class LM(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = table + + def embed_input_ids(self, ids): + return self.embed_tokens(ids) + + model.language_model = LM() + return model + + +def validate(model, args, case, length, device): + ids, mm, mask = args + lm = model.get_language_model() + original = lm.embed_input_ids + outputs, counts = {}, {} + for variant, method in METHODS.items(): + calls = [] + + def counted(tokens): + calls.append(tokens.numel()) + return original(tokens) + + lm.embed_input_ids = counted + try: + outputs[variant] = method(model, ids, mm, is_multimodal=mask) + finally: + lm.embed_input_ids = original + counts[variant] = len(calls) + assert torch.equal(outputs["baseline"], outputs["fixed"]) + assert counts == {"baseline": 1 if case == "text" else 2, "fixed": 1} + VALIDATIONS.append(dict(case=case, tokens=length, device=device, calls=counts, same_input_wrapper_torch_equal=True)) + (ROOT / "validation.json").write_text(json.dumps(VALIDATIONS, indent=2)) + + +@torch.inference_mode() +def main(): + import transformers + import vllm + + metadata = dict( + baseline_ref=args.baseline_ref, + checkout=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + torch=torch.__version__, + vllm=vllm.__version__, + transformers=transformers.__version__, + gpu=torch.cuda.get_device_name(), + cpu_core=CORE, + torch_cpu_threads=1, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + lengths=args.lengths, + embedding_shape=[VOCAB, HIDDEN], + dtype="bfloat16", + input_ids_dtype="int32", + mask_device="cpu", + start_time=time.time(), + metric_order=["host_observed_us", "calling_thread_cpu_us", "synchronized_completion_us"], + ) + index = json.loads((MODEL / "model.safetensors.index.json").read_text())["weight_map"] + key = "thinker.model.embed_tokens.weight" + with safe_open(MODEL / index[key], framework="pt", device="cpu") as f: + weights = f.get_tensor(key) + assert weights.shape == (VOCAB, HIDDEN) and weights.dtype == torch.bfloat16 + model = build_model(weights, "cuda") + measure({"timer_floor": lambda: None}, "cuda", "timer", 0, "overhead", rounds=10, iterations=100) + for length in args.lengths: + for case in ("text", "audio", "image", "video", "mixed", "interleaved"): + ids, mm, mask = payload(case, length, "cuda") + validate(model, (ids, mm, mask), case, length, "cuda") + functions = { + variant: (lambda method=method: method(model, ids, mm, is_multimodal=mask)) + for variant, method in METHODS.items() + } + measure( + functions, + "cuda", + case, + length, + "whole_wrapper", + rounds=args.blocks, + iterations=args.iterations, + warmups=args.warmup, + ) + summary = [] + for length in args.lengths: + for case in ("text", "audio", "image", "video", "mixed", "interleaved"): + medians = { + variant: statistics.median( + statistics.median(sample[0] for sample in record["samples_us"]) + for record in RESULTS + if record["tokens"] == length and record["case"] == case and record["variant"] == variant + ) + for variant in METHODS + } + summary.append( + dict( + tokens=length, + case=case, + host_us=medians, + reduction_pct=100 * (1 - medians["fixed"] / medians["baseline"]), + ) + ) + metadata["end_time"] = time.time() + (ROOT / "metadata.json").write_text(json.dumps(metadata, indent=2)) + (ROOT / "summary.json").write_text(json.dumps(summary, indent=2)) + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py index 13899d47d..7b624ba11 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py +++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py @@ -11,6 +11,7 @@ Regression test for: https://github.com/vllm-project/vllm/issues/34506 """ import functools +from types import SimpleNamespace import pytest import torch @@ -134,16 +135,15 @@ def make_mock_model(mocker: MockerFixture, hidden: int = 8, thinker_cls=None): that has enough structure to run embed_input_ids. ``thinker_cls`` selects whose ``embed_input_ids`` the mock dispatches to. - It defaults to upstream's. Only the interleaved case passes omni's override: - omni's non-interleaved path falls through to a zero-arg ``super()`` call, - which raises TypeError when ``self`` is a Mock rather than an instance. + It defaults to upstream's. Use the selected class as the mock spec so + zero-argument ``super()`` in omni's override resolves correctly. """ m = _qwen2_5_omni_thinker_mod() Qwen2_5OmniThinkerForConditionalGeneration = m.Qwen2_5OmniThinkerForConditionalGeneration if thinker_cls is None: thinker_cls = Qwen2_5OmniThinkerForConditionalGeneration - model = mocker.Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration) + model = mocker.Mock(spec=thinker_cls) # Config with token IDs cfg = mocker.Mock() @@ -157,8 +157,7 @@ def make_mock_model(mocker: MockerFixture, hidden: int = 8, thinker_cls=None): # view with shared memory, which masked_scatter_ cannot handle). return ids.float().unsqueeze(-1).expand(-1, hidden).clone() - lang_model = mocker.Mock() - lang_model.embed_input_ids = fake_lm_embed + lang_model = SimpleNamespace(embed_input_ids=fake_lm_embed) model.get_language_model = mocker.Mock(return_value=lang_model) from vllm.model_executor.models.interfaces import SupportsMultiModal @@ -271,6 +270,27 @@ class TestEmbedInputIds: "Text positions should keep text embeddings" ) + @pytest.mark.parametrize( + "audio_n,image_n,video_n", + [(5, 0, 0), (0, 4, 0), (0, 0, 6), (5, 4, 6)], + ids=["audio", "image", "video", "mixed"], + ) + def test_non_interleaved_multimodal_single_embedding(self, mocker: MockerFixture, audio_n, image_n, video_n): + """Omni's regular path embeds once through super() and preserves placement.""" + hidden = 8 + input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n) + mm_embeds = build_mm_embeds(audio_n, image_n, video_n, hidden) + model, _ = make_mock_model(mocker, hidden, thinker_cls=_omni_qwen2_5_thinker_cls()) + embed = mocker.spy(model.get_language_model(), "embed_input_ids") + + result = model.embed_input_ids(input_ids, mm_embeds, is_multimodal=is_multimodal) + + embed.assert_called_once() + expected = input_ids.float().unsqueeze(-1).expand(-1, hidden).clone() + for token_id, value in ((AUDIO_TOKEN_ID, 10.0), (IMAGE_TOKEN_ID, 20.0), (VIDEO_TOKEN_ID, 30.0)): + expected[input_ids == token_id] = value + assert torch.equal(result, expected) + def test_interleaved_use_audio_in_video(self, mocker: MockerFixture): """ Interleaved (use_audio_in_video): video chunks interleaved with audio. @@ -296,7 +316,9 @@ class TestEmbedInputIds: ] model, _ = make_mock_model(mocker, hidden, thinker_cls=_omni_qwen2_5_thinker_cls()) + embed = mocker.spy(model.get_language_model(), "embed_input_ids") result = model.embed_input_ids(input_ids, mm_embeds, is_multimodal=is_multimodal) + embed.assert_called_once() video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0] audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0] diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 6bbefceab..263ffa2ab 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -1326,14 +1326,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration( if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) - inputs_embeds = self._embed_text_input_ids( - input_ids, - self.get_language_model().embed_input_ids, - is_multimodal=is_multimodal, - ) - if len(multimodal_embeddings) == 0: - return inputs_embeds + return self._embed_text_input_ids( + input_ids, + self.get_language_model().embed_input_ids, + is_multimodal=is_multimodal, + ) # Check for audio-in-video: interleaved video and audio tokens # in the multimodal region. Only use the interleaved path when @@ -1341,6 +1339,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration( video_token_id = self.config.video_token_index audio_token_id = self.config.audio_token_index + # The interleaving check reads scalar positions in a Python loop. + # Keep its masks on CPU to avoid repeated CUDA synchronization. input_ids_cpu = input_ids.cpu() is_video = is_multimodal & (input_ids_cpu == video_token_id) is_audio = is_multimodal & (input_ids_cpu == audio_token_id) From 2f1a18e69fad43e36b3dbc58e698fa142dba207b Mon Sep 17 00:00:00 2001 From: Deep Shah Date: Mon, 14 Sep 2026 10:02:21 -0700 Subject: [PATCH 062/209] [Bugfix][Qwen3-Omni] Fix synchronous Thinker pipeline parallelism (#7345) Signed-off-by: Deep Shah --- .../examples/online_serving/qwen3_omni.md | 4 + .../qwen3_omni/fixtures/thinker_pp2.yaml | 66 ++++++++++ .../qwen3_omni/test_qwen3_omni_pp_capture.py | 115 ++++++++++++++++++ tests/worker/test_pp_intermediate_output.py | 48 ++++++++ .../models/qwen3_omni/qwen3_omni.py | 20 ++- .../qwen3_omni/qwen3_omni_moe_thinker.py | 23 +++- vllm_omni/worker/gpu_model_runner.py | 12 +- 7 files changed, 279 insertions(+), 9 deletions(-) create mode 100644 tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py create mode 100644 tests/worker/test_pp_intermediate_output.py diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index d4e13fc89..4e01ef4c2 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -22,6 +22,10 @@ both stages on vLLM's Triton unquantized MoE path and avoids the unstable FlashI some GPUs. Asynchronous chunk streaming is **enabled by default** within the bundled configuration. +For Thinker pipeline parallelism, configure `pipeline_parallel_size` and `devices` +on the Thinker stage and set `async_scheduling: false`. Keep Talker and Code2Wav +at `pipeline_parallel_size: 1`. Audio chunk streaming can remain enabled. + To explicitly utilize a custom deployment YAML, specify the configuration path: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ diff --git a/tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml b/tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml new file mode 100644 index 000000000..56720b49c --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml @@ -0,0 +1,66 @@ +async_chunk: true +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + initial_codec_chunk_frames: 4 + codec_chunk_frames: 25 + codec_left_context_frames: 25 +stages: +- stage_id: 0 + max_num_batched_tokens: 4096 + max_num_seqs: 2 + gpu_memory_utilization: 0.65 + moe_backend: triton + trust_remote_code: true + enable_prefix_caching: false + devices: 0,1 + default_sampling_params: + temperature: 0.0 + max_tokens: 128 + enforce_eager: true + max_model_len: 4096 + async_scheduling: false + tensor_parallel_size: 1 + pipeline_parallel_size: 2 +- stage_id: 1 + max_num_batched_tokens: 4096 + max_num_seqs: 2 + gpu_memory_utilization: 0.2 + moe_backend: triton + trust_remote_code: true + enable_prefix_caching: false + devices: '1' + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_k: 50 + max_tokens: 1024 + repetition_penalty: 1.05 + enforce_eager: true + max_model_len: 4096 + async_scheduling: false + tensor_parallel_size: 1 + pipeline_parallel_size: 1 +- stage_id: 2 + max_num_batched_tokens: 4096 + max_num_seqs: 2 + gpu_memory_utilization: 0.05 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + enable_chunked_prefill: false + async_scheduling: false + devices: '1' + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + repetition_penalty: 1.1 + max_model_len: 4096 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py new file mode 100644 index 000000000..de8d02e3e --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Required Thinker snapshots must survive consecutive pipeline ranks. + +Actual forward methods with deterministic decoder doubles test the contract +without model weights. Distributed serving remains a separate integration gate. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.qwen3_omni import qwen3_omni as combined_module +from vllm_omni.model_executor.models.qwen3_omni import qwen3_omni_moe_thinker as thinker_module + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel] + + +class _Layer(nn.Module): + def forward(self, positions, hidden_states, residual): + total = hidden_states if residual is None else hidden_states + residual + return total * 0.01 + 0.125, total + + +class _Norm(nn.Module): + def forward(self, hidden_states, residual): + return hidden_states + residual, None + + +def _forward(monkeypatch, start, end, embeddings, intermediate=None, *, capture=True): + rank = SimpleNamespace(is_first_rank=start == 0, is_last_rank=end == 48) + monkeypatch.setattr(thinker_module, "get_pp_group", lambda: rank) + model = object.__new__(thinker_module.Qwen3MoeLLMModel) + nn.Module.__init__(model) + model.start_layer = start + model.end_layer = end + model.layers = nn.ModuleList([_Layer() for _ in range(48)]) + model.norm = _Norm() + return model.forward( + input_ids=None, + positions=torch.arange(embeddings.shape[0]), + inputs_embeds=embeddings, + intermediate_tensors=intermediate, + capture_layer_indices=[0, 24] if capture else None, + return_hidden_states=capture, + ) + + +@pytest.mark.parametrize("tokens", [1, 4]) +@pytest.mark.parametrize("boundaries", [(0, 12, 48), (0, 24, 48), (0, 36, 48), (0, 16, 32, 48)]) +def test_required_captures_match_unsplit_forward(monkeypatch, tokens, boundaries): + embeddings = torch.arange(tokens * 8, dtype=torch.float32).reshape(tokens, 8) / 8 + reference, captures = _forward(monkeypatch, 0, 48, embeddings) + intermediate = None + for start, end in zip(boundaries, boundaries[1:]): + result = _forward(monkeypatch, start, end, embeddings, intermediate) + if end < 48: + assert isinstance(result, IntermediateTensors) + assert all(isinstance(value, torch.Tensor) for value in result.tensors.values()) + intermediate = result + actual, actual_captures = result + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + actual_layers = actual_captures["hidden_states"]["layers"] + assert set(actual_layers) == {0, 24} + for key in (0, 24): + torch.testing.assert_close(actual_layers[key], captures["hidden_states"]["layers"][key], rtol=0, atol=0) + # A later receive into the rank's persistent buffers cannot rewrite snapshots. + assert intermediate is not None + for value in intermediate.tensors.values(): + value.zero_() + for key in (0, 24): + torch.testing.assert_close(actual_layers[key], captures["hidden_states"]["layers"][key], rtol=0, atol=0) + + +def test_capture_disabled_keeps_original_intermediate_schema(monkeypatch): + embeddings = torch.ones(4, 8) + first = _forward(monkeypatch, 0, 24, embeddings, capture=False) + assert set(first.tensors) == {"hidden_states", "residual"} + actual, captures = _forward(monkeypatch, 24, 48, embeddings, first, capture=False) + reference, _ = _forward(monkeypatch, 0, 48, embeddings, capture=False) + assert captures is None + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "start,expected", [(0, set()), (16, {"capture_0"}), (24, {"capture_0"}), (32, {"capture_0", "capture_24"})] +) +@pytest.mark.parametrize("staged", [False, True]) +def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, expected, staged): + thinker_config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=8)) + config = SimpleNamespace( + thinker_config=thinker_config, + talker_config=SimpleNamespace(accept_hidden_layer=24), + code2wav_config=SimpleNamespace(), + tts_bos_token_id=1, + tts_eos_token_id=2, + tts_pad_token_id=3, + ) + model_config = SimpleNamespace(hf_config=config, multimodal_config=None) + if staged: + model_config.model_stage = "thinker" + vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None) + thinker = nn.Module() + thinker.weight = nn.Parameter(torch.ones(1)) + thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start)) + thinker.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states", "residual"], 8) + monkeypatch.setattr(combined_module, "init_vllm_registered_model", lambda **kwargs: thinker) + model = combined_module.Qwen3OmniMoeForConditionalGeneration(vllm_config=vllm_config) + buffers = model.make_empty_intermediate_tensors(4, dtype=torch.float32, device="cpu") + assert set(buffers.tensors) == {"hidden_states", "residual"} | (expected if staged else set()) + assert all(value.shape == (4, 8) for value in buffers.tensors.values()) diff --git a/tests/worker/test_pp_intermediate_output.py b/tests/worker/test_pp_intermediate_output.py new file mode 100644 index 000000000..a9cf02bba --- /dev/null +++ b/tests/worker/test_pp_intermediate_output.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Non-final pipeline outputs must reach vLLM's existing PP sender unchanged.""" + +from types import SimpleNamespace + +import pytest +import torch +from vllm.sequence import IntermediateTensors +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel] + + +def test_nonfinal_forward_bypasses_final_output_conversion(monkeypatch): + intermediate = IntermediateTensors({"hidden_states": torch.ones(4, 8), "residual": torch.ones(4, 8)}) + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: intermediate) + + def final_only_converter(*args, **kwargs): + pytest.fail("A non-final rank must not convert intermediate tensors to OmniOutput") + + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(make_omni_output=final_only_converter, have_multimodal_outputs=True) + runner._build_model_kwargs_extra = lambda: {} + result = runner._model_forward() + assert result is intermediate + hidden, multimodal = runner.extract_multimodal_outputs(result) + assert hidden is intermediate + assert multimodal == {} + + +def test_final_tensor_still_uses_model_converter(monkeypatch): + tensor = torch.ones(4, 8) + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: tensor) + converted = object() + calls = [] + + def convert(value, **kwargs): + calls.append(value) + return converted + + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(make_omni_output=convert) + runner._build_model_kwargs_extra = lambda: {} + assert runner._model_forward() is converted + assert calls == [tensor] diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index a5856718d..28accbc51 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -34,6 +34,7 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import ( from vllm.model_executor.models.utils import ( WeightsMapper, init_vllm_registered_model, + make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY @@ -53,6 +54,7 @@ from vllm_omni.model_executor.models.qwen3_omni.quantization import ( apply_outer_quant_config_mapping, ) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + PP_CAPTURE_PREFIX, Qwen3OmniMoeThinkerDummyInputsBuilder, Qwen3OmniMoeThinkerForConditionalGeneration, Qwen3OmniMoeThinkerMultiModalProcessor, @@ -274,6 +276,18 @@ class Qwen3OmniMoeForConditionalGeneration( self.make_empty_intermediate_tensors = ( self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None ) + if self.model_stage == "thinker": + capture_indices = self._thinker_capture_layer_indices() + start_layer = self.thinker.language_model.model.start_layer + incoming_captures = [f"{PP_CAPTURE_PREFIX}{index}" for index in capture_indices if index < start_layer] + if incoming_captures: + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size + ) + + def _thinker_capture_layer_indices(self) -> list[int]: + accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) + return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else [] @classmethod async def buffer_realtime_audio( @@ -441,11 +455,11 @@ class Qwen3OmniMoeForConditionalGeneration( # Only staged runs have a talker stage to consume the capture; in a # plain vLLM run capturing would waste a clone per layer per step # and make the thinker return a tuple stock vLLM cannot handle. - accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) + capture_indices = self._thinker_capture_layer_indices() capture_kwargs = {} - if accept_layer is not None and self.is_staged_run: + if capture_indices: capture_kwargs = { - "capture_layer_indices": [0, int(accept_layer)], + "capture_layer_indices": capture_indices, "return_hidden_states": True, } diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index b2a7d7abe..8d0389d63 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -118,6 +118,7 @@ from vllm.multimodal.utils import set_mm_embedding_modality from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config +from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerMultiModalDataParser, @@ -140,6 +141,8 @@ except (ImportError, ModuleNotFoundError): logger = init_logger(__name__) +PP_CAPTURE_PREFIX = "capture_" + _THINKER_ARCHITECTURE = "Qwen3OmniMoeThinkerForConditionalGeneration" @@ -568,7 +571,15 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] capture_set = set(capture_layer_indices) if capture_layer_indices else None - captured_hidden_states: dict[str, torch.Tensor] | None = {} if return_hidden_states else None + captured_hidden_states: OmniPayload | None = {} if return_hidden_states else None + + if captured_hidden_states is not None and capture_set and intermediate_tensors is not None: + for layer_idx in capture_set: + if layer_idx < self.start_layer: + hs = captured_hidden_states.setdefault("hidden_states", {}) + layers = hs.setdefault("layers", {}) + # Receive buffers are reused on the next step; retain an independent snapshot. + layers[layer_idx] = intermediate_tensors[f"{PP_CAPTURE_PREFIX}{layer_idx}"].clone() for layer_idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): layer_idx = layer_idx + self.start_layer @@ -589,7 +600,15 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"] if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + tensors = {"hidden_states": hidden_states, "residual": residual} + if captured_hidden_states: + tensors.update( + { + f"{PP_CAPTURE_PREFIX}{index}": value + for index, value in captured_hidden_states["hidden_states"]["layers"].items() + } + ) + return IntermediateTensors(tensors) hidden_states, _ = self.norm(hidden_states, residual) if captured_hidden_states is not None: return hidden_states, captured_hidden_states diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 8dd9f5f58..d3ee8de02 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -34,6 +34,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs @@ -892,8 +893,8 @@ class OmniGPUModelRunner(GPUModelRunner): @torch.inference_mode() def extract_multimodal_outputs( - self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput - ) -> tuple[Any, Any]: + self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput | IntermediateTensors + ) -> tuple[torch.Tensor | IntermediateTensors, OmniPayload | None]: if ( hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs @@ -902,7 +903,7 @@ class OmniGPUModelRunner(GPUModelRunner): text_hidden_states = hidden_states.text_hidden_states multimodal_outputs = hidden_states.multimodal_outputs - elif isinstance(hidden_states, torch.Tensor): + elif isinstance(hidden_states, (torch.Tensor, IntermediateTensors)): text_hidden_states = hidden_states multimodal_outputs = {} elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple): @@ -1225,6 +1226,9 @@ class OmniGPUModelRunner(GPUModelRunner): else: hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + if isinstance(hidden_states, IntermediateTensors): + # Profiling/graph capture needs a tensor to select dummy sampler rows. + hidden_states = hidden_states["hidden_states"] if self.speculative_config and ( self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() @@ -2072,7 +2076,7 @@ class OmniGPUModelRunner(GPUModelRunner): **model_kwargs, **model_kwargs_extra, ) - if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + if not isinstance(model_output, (OmniOutput, IntermediateTensors)) and hasattr(self.model, "make_omni_output"): model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra) # Cache model output so later sample_tokens can consume multimodal results. self._omni_last_model_output = model_output From 40f8e94400fca12ddcefbb46b9043f6a6d6dba96 Mon Sep 17 00:00:00 2001 From: heyuanliu-intel <90166842+heyuanliu-intel@users.noreply.github.com> Date: Tue, 15 Sep 2026 01:02:30 +0800 Subject: [PATCH 063/209] [Model][MiniMax-H3] Multi-GPU serving fixes and layer-wise offload component selection (#7047) Signed-off-by: heyuanliu-intel --- docs/design/feature/sequence_parallel.md | 6 +- .../diffusion/regional_compilation.md | 29 ++ recipes/MiniMaxAI/MiniMax-H3.md | 5 +- .../diffusion/attention/test_ring_padding.py | 221 ++++++++ .../minimax_h3/test_minimax_h3_contract.py | 306 ++++++++++- .../minimax_h3/test_minimax_h3_offload.py | 244 ++++++++- .../minimax_h3/test_minimax_h3_packing.py | 89 +++- .../test_minimax_h3_vae_gather_reuse.py | 473 ++++++++++++++++++ .../minimax_h3/test_minimax_h3_vae_ops.py | 2 +- .../openai_api/test_video_server.py | 26 + tests/model_extras/test_model_extras.py | 10 + .../diffusion/attention/backends/abstract.py | 5 +- .../attention/backends/ring/ring_utils.py | 29 +- .../attention/backends/ring_flash_attn.py | 80 ++- .../attention/backends/ring_pytorch_attn.py | 55 +- .../diffusion/attention/parallel/ring.py | 4 + .../minimax_h3/minimax_h3_transformer.py | 25 +- .../models/minimax_h3/packed_sequence.py | 48 +- .../models/minimax_h3/pipeline_minimax_h3.py | 81 ++- vllm_omni/diffusion/models/minimax_h3/vae.py | 160 ++++++ vllm_omni/model_extras/registry.py | 19 + 21 files changed, 1857 insertions(+), 60 deletions(-) create mode 100644 tests/diffusion/attention/test_ring_padding.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py diff --git a/docs/design/feature/sequence_parallel.md b/docs/design/feature/sequence_parallel.md index d0328bcf6..7ae263b7d 100644 --- a/docs/design/feature/sequence_parallel.md +++ b/docs/design/feature/sequence_parallel.md @@ -77,7 +77,9 @@ Use it when plain Ulysses-SP would otherwise fail because: - `auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`. - `advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes. -- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`. +- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume arbitrary + `attention_mask` values. A model with contiguous suffix padding may instead publish `valid_kv_length`; Ring trims + each circulated K/V block against that global valid prefix. - `advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement. --- @@ -416,7 +418,7 @@ _sp_plan = { | Constraint | Description | |------------|-------------| | **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. | -| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. | +| **Ring Attention Limitation** | Ring attention does not support arbitrary `attention_mask` values, so `auto_pad=True` remains incompatible. Contiguous suffix padding is supported when the model publishes the global `valid_kv_length` prefix. | 1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`: ```python diff --git a/docs/user_guide/diffusion/regional_compilation.md b/docs/user_guide/diffusion/regional_compilation.md index 448f19518..364900ff5 100644 --- a/docs/user_guide/diffusion/regional_compilation.md +++ b/docs/user_guide/diffusion/regional_compilation.md @@ -34,5 +34,34 @@ that provide their own `setup_compile()` implementation manage their compilation policy independently. Compilation is lazy, so backend or graph errors can first surface on the initial request. +## Pinning one packed shape (MiniMax-H3) + +Compiled regions are keyed by input shape. MiniMax-H3 packs a request into a +row count that follows the request -- a longer prompt, more frames, a larger +frame -- and rounds it up to the next 64-row boundary, so two requests can land +on two different packed lengths. + +Under the default `diffusion_compile_dynamic: true` a new packed length is +absorbed by the dynamic shape, so it is not by itself a recompilation. It still +hands the backend a shape it has not autotuned, which is the risk this knob +exists to remove: a fixed-shape deployment (`--no-diffusion-compile-dynamic`) +recompiles the transformer blocks per shape, and a first-seen large shape can +push autotuning into its worst-case memory use. + +A request pins the length with `extra_args["pad_seq_len"]`: + +```python +sampling_params = SamplingParams(extra_args={"pad_seq_len": 54080}) +``` + +The value must be a positive multiple of 64 and must cover the rows the request +actually uses; the packed sequence is then padded to it instead of to the next +64-row boundary. Pick a bucket that covers every request dimension the +deployment accepts -- prompt length, frame count, frame size and reference +blocks all change the used row count. The server logs the effective length as +`MiniMax H3 packed sequence: ... pad_seq_len=... used=... seq_len=...` whenever +a request pins it. The padding rows are masked out, so their cost is the +attention and feed-forward work on those rows. + Use `--enforce-eager` to disable the model runner's generic compile setup. Pipelines that compile internally define their own eager-mode behavior. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 812ed6f30..c966b6b2e 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -1240,8 +1240,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ - `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch). - VAE patch parallelism requires size 1 or the full DiT group size and supports the H3 native `tile` mode only. -- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use - pure Ulysses. +- Ulysses x Ring hybrid attention supports H3's single-request contiguous suffix + padding. Arbitrary attention masks and multi-request packed batches remain + unsupported on the Ring path. - Online FP8 with DLO AllGather temporarily materializes the complete FP8 model in host memory on every rank during startup before retaining only each rank's shard. Size startup host memory for that transient peak. diff --git a/tests/diffusion/attention/test_ring_padding.py b/tests/diffusion/attention/test_ring_padding.py new file mode 100644 index 000000000..8994e3fc1 --- /dev/null +++ b/tests/diffusion/attention/test_ring_padding.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from dataclasses import dataclass +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends import ring_flash_attn, ring_pytorch_attn +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType +from vllm_omni.diffusion.attention.backends.ring.ring_utils import ring_kv_block_valid_length +from vllm_omni.diffusion.attention.parallel import ring as ring_parallel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _SingleRankComm: + rank = 0 + world_size = 1 + + +@dataclass(frozen=True) +class _FakeSequenceParallelGroup: + ring_group: object + + +def _fake_attention_result(q, _k, _v, **_kwargs): + return torch.zeros_like(q), torch.zeros(q.shape[0], q.shape[2], q.shape[1]) + + +@pytest.mark.parametrize(("block_rank", "expected"), [(0, 4), (1, 1), (2, 0)]) +def test_global_prefix_is_partitioned_across_ring_blocks(block_rank, expected): + assert ring_kv_block_valid_length(5, 4, block_rank, 3) == expected + + +def test_absent_prefix_length_keeps_the_full_ring_block(): + assert ring_kv_block_valid_length(None, 4, 2, 3) == 4 + + +def test_pytorch_ring_trims_padding_from_each_kv_block(monkeypatch): + captured: dict[str, Any] = {} + + def fake_attention(q, k, v, **kwargs): + captured.update(k=k, v=v, **kwargs) + return _fake_attention_result(q, k, v, **kwargs) + + monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention) + qkv = torch.randn(1, 4, 2, 8) + + output = ring_pytorch_attn.ring_pytorch_attn_func( + qkv, + qkv, + qkv, + group=object(), + valid_kv_length=3, + ) + + assert output.shape == qkv.shape + assert captured["k"].shape[1] == 3 + assert captured["v"].shape[1] == 3 + + +def test_pytorch_ring_trims_circulated_blocks_by_source_rank(monkeypatch): + observed_lengths = [] + # Each ring step sends K and V separately. + blocks = [torch.randn(1, 4, 2, 8) for _ in range(4)] + + class FakeRingComm: + rank = 2 + world_size = 3 + + def __init__(self, _group): + self._next_block = 0 + + def send_recv(self, _tensor): + block = blocks[self._next_block] + self._next_block += 1 + return block + + def commit(self): + pass + + def wait(self): + pass + + def fake_attention(q, k, v, **kwargs): + observed_lengths.append(k.shape[1]) + return _fake_attention_result(q, k, v, **kwargs) + + monkeypatch.setattr(ring_pytorch_attn, "RingComm", FakeRingComm) + monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention) + qkv = torch.randn(1, 4, 2, 8) + + ring_pytorch_attn.ring_pytorch_attn_func( + qkv, + qkv, + qkv, + group=object(), + valid_kv_length=5, + ) + + # Rank 2 starts with global block 2, then receives blocks 1 and 0. + # The first block is empty and therefore does not launch attention. + assert observed_lengths == [1, 4] + + +def test_flash_ring_trims_padding_from_each_kv_block(monkeypatch): + captured: dict[str, Any] = {} + + def fake_attention(q, k, v, **kwargs): + captured.update(k=k, v=v, **kwargs) + return _fake_attention_result(q, k, v, **kwargs) + + monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr(ring_flash_attn, "select_flash_attn_impl", lambda *_args, **_kwargs: fake_attention) + qkv = torch.randn(1, 4, 2, 8) + + output, _ = ring_flash_attn.ring_flash_attn_forward( + object(), + qkv, + qkv, + qkv, + softmax_scale=0.25, + causal=False, + attn_type=AttnType.FA, + valid_kv_length=3, + ) + + assert output.shape == qkv.shape + assert captured["k"].shape[1] == 3 + assert captured["v"].shape[1] == 3 + + +def test_ring_dispatch_forwards_global_valid_kv_length(monkeypatch): + captured: dict[str, Any] = {} + + def fake_ring_attention(*_args, **kwargs): + captured.update(kwargs) + return "output" + + monkeypatch.setattr(ring_pytorch_attn, "ring_pytorch_attn_func", fake_ring_attention) + strategy = ring_parallel.RingParallelAttention( + _FakeSequenceParallelGroup(ring_group=object()), + attn_backend_pref="sdpa", + ) + metadata = AttentionMetadata(extra={"valid_kv_length": 3}) + query = torch.randn(1, 4, 2, 8) + + assert strategy.run_attention(query, query, query, metadata) == "output" + assert captured["valid_kv_length"] == 3 + + +def test_pytorch_ring_rejects_causal_with_a_valid_prefix(monkeypatch): + """Trimming a K/V block under causal=True would silently shift the diagonal. + + The query block keeps the full padded length while the circulated K/V block + is trimmed, and a causal mask over a rectangular tile is bottom-right + aligned, so the combination must be refused rather than computed. + """ + monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", _fake_attention_result) + qkv = torch.randn(1, 4, 2, 8) + + with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"): + ring_pytorch_attn.ring_pytorch_attn_func( + qkv, + qkv, + qkv, + causal=True, + group=object(), + valid_kv_length=3, + ) + + +def test_flash_ring_rejects_causal_with_a_valid_prefix(monkeypatch): + """Same refusal on the Flash-Attention ring path.""" + monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr( + ring_flash_attn, + "select_flash_attn_impl", + lambda *_args, **_kwargs: _fake_attention_result, + ) + qkv = torch.randn(1, 4, 2, 8) + + with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"): + ring_flash_attn.ring_flash_attn_forward( + object(), + qkv, + qkv, + qkv, + softmax_scale=0.25, + causal=True, + attn_type=AttnType.FA, + valid_kv_length=3, + ) + + +def test_ring_paths_still_accept_causal_without_a_valid_prefix(monkeypatch): + """The guard is scoped to the padded case: plain causal Ring is unchanged.""" + monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr( + ring_flash_attn, + "select_flash_attn_impl", + lambda *_args, **_kwargs: _fake_attention_result, + ) + qkv = torch.randn(1, 4, 2, 8) + + output, _ = ring_flash_attn.ring_flash_attn_forward( + object(), + qkv, + qkv, + qkv, + softmax_scale=0.25, + causal=True, + attn_type=AttnType.FA, + ) + + assert output.shape == qkv.shape diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 7868c255b..67498abc5 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -459,6 +459,7 @@ def test_combined_task_inference_and_transformer_routing(): ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, False), ], ) +@pytest.mark.parametrize("trust_remote_code", [False, True]) def test_pipeline_loads_task_selected_components_with_encoder_ownership( monkeypatch, tmp_path, @@ -469,6 +470,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( source_partitions, expected_tasks, load_text_encoder, + trust_remote_code, ): from vllm_omni.diffusion.data import ( DiffusionParallelConfig, @@ -500,6 +502,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( (tmp_path / partition_name / component).mkdir() created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} + created_kwargs: dict[str, list[Any]] = {"text_encoder": [], "video_vae": [], "audio_vae": []} class FakeModule(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -513,6 +516,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( def component_factory(name): def create(path, *args, **kwargs): created[name].append(str(path)) + created_kwargs[name].append(kwargs) return FakeModule() return create @@ -589,6 +593,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( cfg_parallel_size=1, text_encoder_tp_size=1, ), + trust_remote_code=trust_remote_code, ) pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config) @@ -598,6 +603,13 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( component_path = tmp_path / component_partition assert created["video_vae"] == [str(component_path / "video_vae")] assert created["audio_vae"] == [str(component_path / "audio_vae")] + # The checkpoint VAEs execute the modeling code shipped with the weights, so + # the pipeline must hand them the engine's trust_remote_code decision. Both + # VAEs are created for every parametrization, so this stays unconditional. + for component in ("video_vae", "audio_vae"): + assert [kwargs["trust_remote_code"] for kwargs in created_kwargs[component]] == [trust_remote_code] + # The tokenizer/processor "loaded once" assertions are asserted per + # load_text_encoder branch below, matching upstream's own expectations. expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"] assert pipeline._dit_modules == expected_dit_modules assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2) @@ -1122,6 +1134,64 @@ def test_packed_attention_keeps_padding_mask_for_other_backends(): ) +def test_ring_packed_attention_publishes_prefix_length_without_global_mask(): + attention = _fake_packed_attention("CUDNN_ATTN") + attention.attention.use_ring = True + q = torch.randn(4, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + metadata = attention.attention.metadata + assert metadata.attn_mask is None + assert metadata.extra["valid_kv_length"] == 5 + + +def test_ring_packed_attention_keeps_mask_when_runtime_sp_is_inactive(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer + + class InactiveContext: + sp_active = False + + attention = _fake_packed_attention("CUDNN_ATTN") + attention.attention.use_ring = True + monkeypatch.setattr( + minimax_h3_transformer, + "is_forward_context_available", + lambda: True, + raising=False, + ) + monkeypatch.setattr( + minimax_h3_transformer, + "get_forward_context", + InactiveContext, + raising=False, + ) + q = torch.randn(8, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + metadata = attention.attention.metadata + assert torch.equal( + metadata.attn_mask, + torch.tensor([[True, True, True, True, True, False, False, False]]), + ) + assert metadata.extra["npu_attn_varlen"] is True + + def _fake_packed_attention( backend_name: str, *, @@ -1143,6 +1213,10 @@ def _fake_packed_attention( def supports_multi_doc_packed_varlen(cls) -> bool: return supports_multi_doc + @classmethod + def supports_packed_mask_free(cls) -> bool: + return False + class FakeAttention(torch.nn.Module): attn_backend = FakeBackend @@ -1712,7 +1786,7 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, *, trust_remote_code: FakeRemote(), ) video_vae = vae_module.MiniMaxH3VideoVAE( @@ -1834,7 +1908,7 @@ def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, *, trust_remote_code: FakeRemote(), ) video_vae = vae_module.MiniMaxH3VideoVAE( @@ -1869,7 +1943,7 @@ def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, *, trust_remote_code: FakeRemote(), ) monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls) @@ -1890,6 +1964,123 @@ def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch): stager.offload.assert_called_once_with() +def test_remote_component_refuses_untrusted_checkpoint_code(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + loaded = [] + + def spy_get_class(class_reference, component_path): + loaded.append((class_reference, component_path)) + raise AssertionError("checkpoint code must not be resolved when untrusted") + + monkeypatch.setattr(vae_module, "get_class_from_dynamic_module", spy_get_class) + + with pytest.raises(ValueError, match="--trust-remote-code"): + vae_module._load_remote_component( + "/checkpoints/video_vae", + {"auto_map": {"AutoModel": "modeling_vae.VideoVAE"}}, + trust_remote_code=False, + ) + + assert loaded == [] + + +def test_remote_component_resolves_checkpoint_code_when_trusted(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeComponent(torch.nn.Module): + @classmethod + def from_pretrained(cls, component_path): + module = cls() + module.loaded_from = component_path + return module + + loaded = [] + + def spy_get_class(class_reference, component_path): + loaded.append((class_reference, component_path)) + return FakeComponent + + monkeypatch.setattr(vae_module, "get_class_from_dynamic_module", spy_get_class) + + module = vae_module._load_remote_component( + "/checkpoints/video_vae", + {"auto_map": {"AutoModel": "modeling_vae.VideoVAE"}}, + trust_remote_code=True, + ) + + # Negative control: the trusted path stays byte-for-byte the pre-gate call. + assert loaded == [("modeling_vae.VideoVAE", "/checkpoints/video_vae")] + assert module.loaded_from == "/checkpoints/video_vae" + + +@pytest.mark.parametrize("vae_cls_name", ["MiniMaxH3VideoVAE", "MiniMaxH3AudioVAE"]) +@pytest.mark.parametrize("trusted", [True, False]) +def test_vae_forwards_trust_remote_code_to_loader(monkeypatch, vae_cls_name, trusted): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1) + + forwarded = [] + + def spy_load(_component_path, _config, *, trust_remote_code): + forwarded.append(trust_remote_code) + return FakeRemote() + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + "sample_rate": 32000, + }, + ) + monkeypatch.setattr(vae_module, "_load_remote_component", spy_load) + + vae_cls = getattr(vae_module, vae_cls_name) + vae_cls( + "unused", + device=torch.device("cpu"), + trust_remote_code=trusted, + ) + + assert forwarded == [trusted] + + +@pytest.mark.parametrize("vae_cls_name", ["MiniMaxH3VideoVAE", "MiniMaxH3AudioVAE"]) +def test_vae_defaults_to_refusing_checkpoint_code(monkeypatch, vae_cls_name): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + forwarded = [] + + def spy_load(_component_path, _config, *, trust_remote_code): + forwarded.append(trust_remote_code) + raise AssertionError("unreachable") + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + "sample_rate": 32000, + }, + ) + monkeypatch.setattr(vae_module, "_load_remote_component", spy_load) + + vae_cls = getattr(vae_module, vae_cls_name) + with pytest.raises(AssertionError, match="unreachable"): + vae_cls("unused", device=torch.device("cpu")) + + assert forwarded == [False] + + def test_video_vae_encode_uses_configured_parallel_tiling(): from vllm_omni.diffusion.models.minimax_h3.vae import ( MiniMaxH3VideoVAE, @@ -2889,3 +3080,112 @@ def test_encoder_load_places_every_source_in_its_own_rows(tmp_path): } for source, rows in slices.items(): assert torch.equal(rows, torch.full_like(rows, _SOURCE_FILL[source])), source + + +def _pad_seq_len_batch(pad_seq_len): + """A t2va request that pins the packed length through ``extra_args``.""" + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + extra_args = {"task": "t2va", "aspect_ratio": "16:9"} + if pad_seq_len is not None: + extra_args["pad_seq_len"] = pad_seq_len + sampling = OmniDiffusionSamplingParams( + quality="lossless", + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=None, + extra_args=extra_args, + ) + return DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="pinned packed length", + sampling_params=sampling, + request_id="pad-seq-len", + ) + ] + ) + + +@pytest.mark.parametrize("pinned, expected", [(None, None), (54080, 54080)]) +def test_pad_seq_len_reaches_diffuse_from_extra_args(pinned, expected): + """The request field must arrive at ``diffuse``, not stop at validation.""" + diffuse_calls: list[dict[str, Any]] = [] + pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": None, "ref2va": None}) + + pipeline.forward(_pad_seq_len_batch(pinned)) + + assert [call["pad_seq_len"] for call in diffuse_calls] == [expected] + + +@pytest.mark.parametrize("task", ["t2va", "ref2va"]) +def test_pad_seq_len_reaches_the_packer_for_every_task_family(task, monkeypatch): + """``diffuse`` must hand the pin to whichever packer the task uses.""" + import vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 as module + + class _CapturedError(Exception): + pass + + recorded = {} + + def _capture(**kwargs): + recorded.update(kwargs) + raise _CapturedError + + monkeypatch.setattr(module, "minimax_h3_packed_sequence", _capture) + monkeypatch.setattr(module, "minimax_h3_packed_sequence_ref2va_blocks", _capture) + + pipeline = _distilled_pipeline([], {"fl2va": None, "ref2va": None}) + pipeline._initial_noise = lambda **kwargs: (torch.zeros(1), torch.zeros(1)) + + with pytest.raises(_CapturedError): + pipeline._build_denoise_inputs( + task=task, + text_embeddings=torch.ones(7, 2), + text_tags=torch.ones(7, dtype=torch.long), + seed=0, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + num_frames=5, + num_steps=2, + video_shift=12.0, + audio_shift=3.0, + base_schedule=None, + visual_condition=None, + visual_condition_shape=(1, 4, 6), + audio_condition=None, + ref_audio_t=1, + pad_seq_len=192, + ) + + assert recorded["seq_len"] == 192 + + +@pytest.mark.parametrize("value", [0, -64, 100, "128", True, 1.5, (1 << 20) + 64]) +def test_pad_seq_len_request_validation_rejects_bad_values(value): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _resolve_pad_seq_len, + ) + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError): + _resolve_pad_seq_len(value) + + +def test_pad_seq_len_request_validation_accepts_aligned_values(): + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + MINIMAX_H3_MAX_PAD_SEQ_LEN, + ) + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _resolve_pad_seq_len, + ) + + assert _resolve_pad_seq_len(None) is None + assert _resolve_pad_seq_len(54080) == 54080 + assert _resolve_pad_seq_len(MINIMAX_H3_MAX_PAD_SEQ_LEN) == MINIMAX_H3_MAX_PAD_SEQ_LEN diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index 865b2248a..31a8916a6 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from contextlib import contextmanager +import weakref +from contextlib import contextmanager, nullcontext from types import SimpleNamespace from typing import Any from unittest.mock import Mock @@ -76,6 +77,247 @@ def test_requested_model_offload_without_backend_uses_resident_components(): pipeline.text_encoder.encode_ids.assert_called_once() +def test_decoded_output_is_unchanged_without_active_model_offload(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [] + output = Mock() + + assert pipeline._offload_model_cpu_stage_output(output) is output + output.cpu.assert_not_called() + + +def test_reply_rank_copies_decoded_output_to_host(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [Mock()] + monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: True) + output = Mock() + host_output = Mock() + output.cpu.return_value = host_output + + assert pipeline._offload_model_cpu_stage_output(output) is host_output + output.cpu.assert_called_once_with() + + +def test_non_reply_rank_drops_decoded_output_storage(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [Mock()] + monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: False) + output = torch.ones((2, 3), dtype=torch.float16) + + placeholder = pipeline._offload_model_cpu_stage_output(output) + + assert placeholder.device.type == "meta" + assert placeholder.shape == output.shape + assert placeholder.dtype == output.dtype + + +def test_decode_releases_only_audio_before_return(monkeypatch): + """Video is released by the callers of ``decode`` -- after quantization.""" + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + video = torch.ones((1, 1, 4, 4)) + audio = torch.ones((1, 8)) + pipeline = SimpleNamespace( + device=torch.device("cpu"), + video_vae=Mock(decode_latent=Mock(return_value=video)), + audio_vae=Mock(decode_latent=Mock(return_value=audio)), + _component_on_device=lambda component: nullcontext(), + _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value), + ) + monkeypatch.setattr( + module.current_omni_platform, + "create_autocast_context", + lambda **kwargs: nullcontext(), + ) + + actual_video, actual_audio = MiniMaxH3Pipeline.decode( + pipeline, + torch.zeros(1), + torch.zeros(1), + height=2, + width=3, + ) + + assert actual_video.shape[-2:] == (2, 3) + assert actual_audio is audio + calls = pipeline._offload_model_cpu_stage_output.call_args_list + assert len(calls) == 1 + assert calls[0].args[0] is audio + + +def _recording_release(events, releases): + """Stand-in for ``_offload_model_cpu_stage_output`` that records what it got.""" + + def release(value): + events.append("release") + releases.append(value) + return value + + return release + + +def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + # MiniMaxH3VideoVAE.decode_latent ends in frames.float(), so this is what + # the release hook would see if it ran before the quantizer. + decoded_dtype = torch.float32 + events: list[str] = [] + releases: list[torch.Tensor] = [] + + def diffuse(**kwargs): + events.append(f"diffuse:{kwargs['seed']}") + return torch.zeros(1), torch.zeros(1) + + def decode(video_latent, audio_latent, *, height, width): + events.append("decode") + return ( + torch.zeros((1, 3, 2, height, width), dtype=decoded_dtype), + torch.zeros((1, 8)), + ) + + # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the + # tensor-decode branch, which is the one the release hook wraps. + context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False} + pipeline = SimpleNamespace( + od_config=SimpleNamespace(), + _extract_prompt=lambda raw: ("a prompt", {}), + _extract_text_conditioning=lambda raw: None, + _extract_prepared_reference_videos=lambda raw: None, + _prepare_request_inputs=lambda **kwargs: context, + _denoise_kwargs=lambda ctx: {}, + diffuse=diffuse, + decode=decode, + _offload_model_cpu_stage_output=_recording_release(events, releases), + ) + request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace()) + + output = MiniMaxH3Pipeline.forward(pipeline, request) + + # The release must see the quantized frames, not the decoded floats. + assert len(releases) == 2 + assert [tensor.dtype for tensor in releases] == [torch.uint8, torch.uint8] + assert decoded_dtype is not torch.uint8 + # ... and it must still happen before the next seed reloads the DiT. + assert events == ["diffuse:41", "decode", "release", "diffuse:42", "decode", "release"] + assert output.output[0].dtype is torch.uint8 + + +def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypatch): + """Releasing the copy is not enough if a local still pins the original. + + ``_offload_model_cpu_stage_output`` hands back a *new* tensor, so the + decoded frames stay alive for as long as anything still references them. + The whole point of this branch is that they are gone before the next seed + reloads the DiT, and ``diffuse`` is where that reload happens. + """ + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + decoded_refs: list[weakref.ref] = [] + alive_at_diffuse: list[int] = [] + + def diffuse(**kwargs): + alive_at_diffuse.append(sum(1 for ref in decoded_refs if ref() is not None)) + return torch.zeros(1), torch.zeros(1) + + def decode(video_latent, audio_latent, *, height, width): + video = torch.zeros((1, 3, 2, height, width), dtype=torch.float32) + decoded_refs.append(weakref.ref(video)) + return video, torch.zeros((1, 8)) + + # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the + # tensor-decode branch, which is the one the release hook wraps. + context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False} + pipeline = SimpleNamespace( + od_config=SimpleNamespace(), + _extract_prompt=lambda raw: ("a prompt", {}), + _extract_text_conditioning=lambda raw: None, + _extract_prepared_reference_videos=lambda raw: None, + _prepare_request_inputs=lambda **kwargs: context, + _denoise_kwargs=lambda ctx: {}, + diffuse=diffuse, + decode=decode, + _offload_model_cpu_stage_output=lambda value: value.clone(), + ) + request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace()) + + MiniMaxH3Pipeline.forward(pipeline, request) + + assert len(decoded_refs) == 2 + # Seed 1's diffuse sees nothing decoded yet; seed 2's must not see seed 1's + # decoded frames still resident. + assert alive_at_diffuse == [0, 0] + + +def test_post_decode_releases_video_after_quantization(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + events: list[str] = [] + releases: list[torch.Tensor] = [] + + def decode(video_latent, audio_latent, *, height, width): + events.append("decode") + return ( + torch.zeros((1, 3, 2, height, width), dtype=torch.float16), + torch.zeros((1, 8)), + ) + + pipeline = SimpleNamespace( + od_config=SimpleNamespace(), + _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)), + decode=decode, + _offload_model_cpu_stage_output=_recording_release(events, releases), + ) + state = SimpleNamespace( + latents=torch.zeros(1), + extra={ + module._STEP_BRANCH: object(), + module._STEP_AUDIO_ROWS: torch.zeros(1), + module._STEP_SHAPE: { + "latent_t": 1, + "latent_h": 1, + "latent_w": 1, + "audio_t": 1, + "height": 4, + "width": 4, + }, + }, + ) + + output = MiniMaxH3Pipeline.post_decode(pipeline, state) + + assert events == ["decode", "release"] + assert len(releases) == 1 + assert releases[0].dtype is torch.uint8 + assert output.output[0] is releases[0] + + +@pytest.mark.parametrize( + ("available", "initialized", "rank", "expected"), + [(False, False, 7, True), (True, False, 7, True), (True, True, 0, True), (True, True, 1, False)], +) +def test_output_owner_rank_matches_executor_contract(monkeypatch, available, initialized, rank, expected): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + monkeypatch.setattr(module.dist, "is_available", lambda: available) + monkeypatch.setattr(module.dist, "is_initialized", lambda: initialized) + monkeypatch.setattr(module.dist, "get_rank", lambda: rank) + + assert MiniMaxH3Pipeline._is_output_owner_rank() is expected + + def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py index 53348f579..4340e093f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch @@ -190,3 +190,90 @@ def test_condition_noise_accepts_a_reference_video_longer_than_the_target(): noise_aug=0.5, ) assert result.shape == rows.shape + + +def test_explicit_seq_len_pins_one_shape_across_prompt_lengths(): + """Prompts of different token counts must land on the same packed length. + + Without the pin they fall into different 64-row buckets, and each bucket is + a packed shape the compiled transformer has not seen before. + """ + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + ) + + common = dict( + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + include_keyframe_cond=False, + ) + # used = text_len + 18 rows: 22 -> bucket 64; 118 -> bucket 128. + short = minimax_h3_packed_sequence(text_len=4, **common) + long = minimax_h3_packed_sequence(text_len=100, **common) + assert int(short["seq_len"]) == 64 + assert int(long["seq_len"]) == 128 + + short_pinned = minimax_h3_packed_sequence(text_len=4, seq_len=192, **common) + long_pinned = minimax_h3_packed_sequence(text_len=100, seq_len=192, **common) + assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192 + # The pin only adds padding rows; the used prefix is untouched. + assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1]) == 22 + assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1]) == 118 + assert short_pinned["img_pos"].tolist() == short["img_pos"].tolist() + + +def test_the_ref2va_packer_pins_one_shape_across_prompt_lengths_too(): + """Both packers share one resolver, so both honour the pin.""" + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence_ref2va_blocks, + ) + + common = dict( + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ref_blocks=[{"kind": "audio", "ref_audio_t": 1}], + ) + short = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, **common) + long = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, **common) + assert int(short["seq_len"]) != int(long["seq_len"]) + + short_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, seq_len=192, **common) + long_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, seq_len=192, **common) + assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192 + assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1]) + assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1]) + + +def test_explicit_seq_len_below_used_rows_is_a_client_error(): + """Both packers reject an under-sized pin as a 4xx, not a server error.""" + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + minimax_h3_packed_sequence_ref2va_blocks, + ) + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match="used rows"): + minimax_h3_packed_sequence( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + include_keyframe_cond=False, + seq_len=8, + ) + + with pytest.raises(OmniClientError, match="used rows"): + minimax_h3_packed_sequence_ref2va_blocks( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ref_blocks=[{"kind": "audio", "ref_audio_t": 1}], + seq_len=8, + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py new file mode 100644 index 000000000..116f44ef9 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py @@ -0,0 +1,473 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for the pinned tiled-VAE gather buffer in the MiniMax-H3 video VAE. + +The checkpoint's ``_all_gather_tiled_results`` allocates its gather output on +every call, so with a per-request ``empty_cache()`` the collective lands on a +new device address every request and XCCL's per-address registrations pile up. +On XPU the adapter replaces that method with an equal-shape gather into a +buffer it keeps alive. These tests pin the properties that make it a fix +rather than a rewrite: the address is stable, the tiles come back +bit-identical and in the same order, the caller owns what it gets back, and no +other device is touched. +""" + +import logging +import sys +import types +from typing import Any + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module +from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +TILE = (1, 3, 2, 4, 4) + + +class _CheckpointGatherSentinel: + """Stands in for the checkpoint's own bound method. + + Identity is the point: a test asserts this exact object is still installed + on the devices the adapter must leave alone. + """ + + def __call__(self, tasks, num_tiles): # pragma: no cover - never invoked + raise AssertionError("the checkpoint gather must not run in these tests") + + +class _FakeCheckpointModel: + def __init__(self): + self._all_gather_tiled_results = _CheckpointGatherSentinel() + + +def _local_tiles(all_tiles, num_tiles, rank, sp_size): + """One rank's round-robin share, as ``klvae.py:275`` computes it.""" + return [all_tiles[index] for index in range(rank, num_tiles, sp_size)] + + +def _reference_gather(all_tiles, num_tiles, sp_size): + """The checkpoint's variable-shape list gather, reimplemented verbatim. + + Written from the control flow of ``_all_gather_tiled_results`` in the + checkpoint's ``klvae.py`` (lines 252-273) -- stack each rank's share, gather + the list, walk it back into global tile order -- so that it is an + independent oracle rather than a restatement of the code under test. + """ + gathered = [torch.stack(_local_tiles(all_tiles, num_tiles, rank, sp_size), dim=0) for rank in range(sp_size)] + results = [None] * num_tiles + for rank, rank_tensors in enumerate(gathered): + for k in range(rank_tensors.shape[0]): + global_index = k * sp_size + rank + if global_index >= num_tiles: + break + results[global_index] = rank_tensors[k] + return results + + +def _peer_stacks(all_tiles, num_tiles, sp_size): + """What every rank hands to the collective, padded the way the adapter pads.""" + max_tasks = -(-num_tiles // sp_size) + stacks = [] + for rank in range(sp_size): + share = _local_tiles(all_tiles, num_tiles, rank, sp_size) + stacked = share[0].new_zeros((max_tasks, *share[0].shape)) + stacked[: len(share)] = torch.stack(share, dim=0) + stacks.append(stacked) + return stacks + + +def _tiles(num_tiles, seed=0, dtype=torch.float32): + generator = torch.Generator().manual_seed(seed) + return [torch.rand(TILE, generator=generator).to(dtype) for _ in range(num_tiles)] + + +def _workspace_bytes(num_tiles, sp_size, dtype=torch.float32): + """Bytes one gather needs: ``sp_size`` copies of the padded local stack.""" + max_tasks = -(-num_tiles // sp_size) + element_size = torch.empty(0, dtype=dtype).element_size() + numel = 1 + for dim in TILE: + numel *= dim + return sp_size * max_tasks * numel * element_size + + +@pytest.fixture +def fake_collective(monkeypatch): + """Replace the collective with a single-process stand-in. + + ``all_gather_into_tensor`` is the only distributed call on this path, so + filling the output buffer from a caller-supplied peer list exercises the + real padding, keying, addressing and unpacking code without a process + group. + """ + calls: list[dict[str, Any]] = [] + + def fake_all_gather_into_tensor(output, input_tensor, group=None, async_op=False): + peers = calls[-1]["peers"] + assert output.shape == (len(peers), *input_tensor.shape) + calls[-1]["sent"] = input_tensor.clone() + calls[-1]["buf_ptr"] = output.data_ptr() + for rank, peer in enumerate(peers): + output[rank].copy_(peer) + + monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor) + return calls + + +class _Records(logging.Handler): + """vllm's logger sets ``propagate=False``, so caplog never sees these.""" + + def __init__(self): + super().__init__(level=logging.DEBUG) + self.messages = [] + self.records = [] + + def emit(self, record): + self.messages.append(record.getMessage()) + self.records.append((record.levelno, record.getMessage())) + + +@pytest.fixture +def gather_log(): + handler = _Records() + logger = vae_module.logger + previous = logger.level + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + try: + yield handler + finally: + logger.removeHandler(handler) + logger.setLevel(previous) + + +def _vae(sp_size, sp_rank=0): + """An adapter instance with the override installed, no checkpoint needed.""" + state = { + "sp_size": sp_size, + "sp_rank": sp_rank, + "sp_enabled": sp_size > 1, + "sp_process_group": object(), + } + module = types.ModuleType("fake_gather_ckpt.parallel") + module.get_parallel_state = lambda: state # type: ignore[attr-defined] + sys.modules.setdefault("fake_gather_ckpt", types.ModuleType("fake_gather_ckpt")) + sys.modules["fake_gather_ckpt.parallel"] = module + + vae = object.__new__(MiniMaxH3VideoVAE) + vae.remote = type("Remote", (), {"__module__": "fake_gather_ckpt.klvae"})() + vae.model = _FakeCheckpointModel() + vae._device_target = torch.device("xpu") + vae._tile_gather_workspace = None + vae._tile_gather_stats = {"hits": 0, "allocs": 0, "workspace_bytes": 0} + vae._checkpoint_tile_gather = None + vae._install_persistent_tile_gather() + return vae + + +def _run(vae, calls, all_tiles, num_tiles, sp_size, sp_rank): + calls.append({"peers": _peer_stacks(all_tiles, num_tiles, sp_size)}) + return vae.model._all_gather_tiled_results(_local_tiles(all_tiles, num_tiles, sp_rank, sp_size), num_tiles) + + +# -------------------------------------------------------------------------- +# The address is stable across calls. This is what the accumulation is about: +# drop the workspace and every call registers a new receive address. And the +# workspace is one grow-only block, so varying request geometries cannot keep +# adding resident buffers the way a per-geometry cache would. +# -------------------------------------------------------------------------- + + +def test_the_landing_buffer_address_is_stable_across_calls(fake_collective): + vae = _vae(4) + pointers = [] + for request in range(3): + _run(vae, fake_collective, _tiles(8, seed=request), 8, 4, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + + assert len(set(pointers)) == 1, f"gather buffer moved across calls: {pointers}" + assert vae._tile_gather_stats == { + "hits": 2, + "allocs": 1, + "workspace_bytes": _workspace_bytes(8, 4), + } + assert vae._tile_gather_workspace.numel() == _workspace_bytes(8, 4) + assert vae._tile_gather_workspace.dtype is torch.uint8 + + +def test_distinct_shapes_share_the_one_workspace(fake_collective): + """One decode gathers several temporal clips, and the head/tail clips differ.""" + vae = _vae(4) + _run(vae, fake_collective, _tiles(8), 8, 4, 0) + first = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(12), 12, 4, 0) + second = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(8, seed=5), 8, 4, 0) + third = fake_collective[-1]["buf_ptr"] + + # The larger geometry did not fit, so the block was replaced -- not added to. + assert second == third + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["hits"] == 1 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(12, 4) + # ``first`` is deliberately not compared: the replaced block is freed, so + # the allocator is free to hand its address back for the larger one. + assert isinstance(first, int) + + +def test_workspace_is_bounded_across_geometries(fake_collective): + """Alternating geometries must not each retain a buffer of their own.""" + sp_size = 4 + large, small, largest = 12, 8, 20 + vae = _vae(sp_size) + + pointers = [] + for round_index in range(3): + for num_tiles in (large, small): + _run(vae, fake_collective, _tiles(num_tiles, seed=round_index), num_tiles, sp_size, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + + assert len(pointers) == 6 + assert len(set(pointers)) == 1, f"workspace moved across geometries: {pointers}" + assert vae._tile_gather_stats["allocs"] == 1 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(large, sp_size) + assert vae._tile_gather_workspace.numel() == _workspace_bytes(large, sp_size) + + # A geometry that does not fit grows the single block; it does not add one. + _run(vae, fake_collective, _tiles(largest), largest, sp_size, 0) + grown = fake_collective[-1]["buf_ptr"] + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(largest, sp_size) + assert vae._tile_gather_workspace.numel() == _workspace_bytes(largest, sp_size) + + # ...and the geometries seen before it now fit, so nothing reallocates again. + for num_tiles in (large, small): + _run(vae, fake_collective, _tiles(num_tiles, seed=7), num_tiles, sp_size, 0) + assert fake_collective[-1]["buf_ptr"] == grown + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(largest, sp_size) + + +def test_workspace_grows_only_upward(fake_collective): + """Resident bytes track the largest geometry seen, not the sum of them.""" + sp_size = 4 + small, large = 8, 12 + vae = _vae(sp_size) + + _run(vae, fake_collective, _tiles(small), small, sp_size, 0) + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(small, sp_size) + + _run(vae, fake_collective, _tiles(large), large, sp_size, 0) + + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["hits"] == 0 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(large, sp_size) + assert vae._tile_gather_stats["workspace_bytes"] < _workspace_bytes(small, sp_size) + _workspace_bytes( + large, sp_size + ) + assert vae._tile_gather_workspace.numel() == _workspace_bytes(large, sp_size) + + +def test_mixed_dtype_shares_one_workspace(fake_collective): + """The workspace is bytes, so a narrower dtype reuses it instead of adding one.""" + sp_size, num_tiles = 4, 8 + vae = _vae(sp_size) + fp32_bytes = _workspace_bytes(num_tiles, sp_size, dtype=torch.float32) + + pointers = [] + for round_index in range(2): + for dtype in (torch.float32, torch.bfloat16): + tiles = _tiles(num_tiles, seed=round_index, dtype=dtype) + produced = _run(vae, fake_collective, tiles, num_tiles, sp_size, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + assert all(tile.dtype is dtype for tile in produced) + for index, tile in enumerate(produced): + assert torch.equal(tile, tiles[index]) + + assert len(set(pointers)) == 1, f"workspace moved across dtypes: {pointers}" + assert vae._tile_gather_stats["allocs"] == 1 + assert vae._tile_gather_stats["hits"] == 3 + assert vae._tile_gather_stats["workspace_bytes"] == fp32_bytes + assert vae._tile_gather_workspace.numel() == fp32_bytes + + +# -------------------------------------------------------------------------- +# Equivalence with the checkpoint's own gather, padded branch included. A +# reordered tile grid is silent at runtime, so order is asserted as well. +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_tiles", [5, 8, 9, 28]) +@pytest.mark.parametrize("sp_rank", [0, 3]) +def test_gathered_tiles_match_the_checkpoint_gather(fake_collective, num_tiles, sp_rank): + sp_size = 4 + all_tiles = _tiles(num_tiles, seed=num_tiles) + vae = _vae(sp_size, sp_rank=sp_rank) + + produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, sp_rank) + expected = _reference_gather(all_tiles, num_tiles, sp_size) + + assert len(produced) == num_tiles + assert all(tile is not None for tile in produced) + for index, (got, want) in enumerate(zip(produced, expected)): + assert torch.equal(got, want), f"tile {index} differs from the checkpoint gather" + # ...and it is the tile that belongs at that index, not a permutation that + # happens to agree with an equally permuted oracle. + for index, tile in enumerate(produced): + assert torch.equal(tile, all_tiles[index]) + + +def test_the_padded_branch_is_actually_exercised(fake_collective): + """9 tiles over 4 ranks: ranks 1-3 are one short, so the padding runs.""" + sp_size, num_tiles = 4, 9 + all_tiles = _tiles(num_tiles, seed=9) + assert len(_local_tiles(all_tiles, num_tiles, 1, sp_size)) < -(-num_tiles // sp_size) + + vae = _vae(sp_size, sp_rank=1) + produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, 1) + + sent = fake_collective[-1]["sent"] + assert sent.shape[0] == 3 + assert torch.equal(sent[2], torch.zeros_like(sent[2])) + for index, tile in enumerate(produced): + assert torch.equal(tile, all_tiles[index]) + + +def test_an_empty_local_share_still_raises(fake_collective): + """A rank with no tiles fails loudly, exactly as the checkpoint does.""" + vae = _vae(4) + with pytest.raises(ValueError, match="empty tasks on sp rank"): + vae.model._all_gather_tiled_results([], 8) + + +def test_more_tiles_than_round_robin_allows_is_refused(fake_collective): + """The unpacking arithmetic is only valid for the checkpoint's ownership.""" + vae = _vae(4) + with pytest.raises(ValueError, match="round-robin"): + vae.model._all_gather_tiled_results(_tiles(5), 8) + + +# -------------------------------------------------------------------------- +# Ownership: the buffer is overwritten by the next gather, so returned tiles +# must not alias it. +# -------------------------------------------------------------------------- + + +def test_returned_tiles_do_not_alias_the_pinned_buffer(fake_collective): + sp_size, num_tiles = 4, 8 + first_tiles = _tiles(num_tiles, seed=1) + vae = _vae(sp_size) + produced = _run(vae, fake_collective, first_tiles, num_tiles, sp_size, 0) + + buffer_storage = vae._tile_gather_workspace.untyped_storage().data_ptr() + for tile in produced: + assert tile.untyped_storage().data_ptr() != buffer_storage + + # The next gather reuses that memory; the first call's tiles must survive it. + _run(vae, fake_collective, _tiles(num_tiles, seed=2), num_tiles, sp_size, 0) + for index, tile in enumerate(produced): + assert torch.equal(tile, first_tiles[index]) + + +# -------------------------------------------------------------------------- +# Non-XPU behaviour is unchanged: the checkpoint's own method stays installed. +# -------------------------------------------------------------------------- + + +def _construct(monkeypatch, device, model): + """Run the real ``__init__`` against a stub checkpoint component.""" + + class _Remote(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1)) + self.model = model + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda path: {"latent_channels": 1, "latents_mean": [0.0], "latents_std": [1.0]}, + ) + # **kwargs: the loader gained a trust_remote_code argument at some point, + # and this stub does not care either way. + monkeypatch.setattr(vae_module, "_load_remote_component", lambda path, config, **kwargs: _Remote()) + # The adapter only stores this handle; resolving it for real needs an + # accelerator, which a CPU test runner does not have. + monkeypatch.setattr(torch, "get_device_module", lambda *args, **kwargs: types.SimpleNamespace()) + return MiniMaxH3VideoVAE("unused", device=torch.device(device)) + + +@pytest.mark.parametrize("device", ["cpu", "meta"]) +def test_non_xpu_devices_keep_the_checkpoint_gather(monkeypatch, device): + model = _FakeCheckpointModel() + original = model._all_gather_tiled_results + + vae = _construct(monkeypatch, device, model) + + assert vae._tile_gather_reuse_enabled() is False + assert vae.model._all_gather_tiled_results is original + assert vae._checkpoint_tile_gather is None + + +@pytest.mark.parametrize( + ("device", "enabled"), + [("xpu", True), ("cpu", False), ("cuda", False), ("meta", False)], +) +def test_only_xpu_asks_for_the_pinned_buffer(device, enabled): + vae = object.__new__(MiniMaxH3VideoVAE) + vae._device_target = torch.device(device) + assert vae._tile_gather_reuse_enabled() is enabled + + +def test_construction_installs_the_override_when_the_device_asks_for_it(monkeypatch): + """Close the loop: ``__init__`` must consult the predicate, not ignore it.""" + monkeypatch.setattr(MiniMaxH3VideoVAE, "_tile_gather_reuse_enabled", lambda self: True) + model = _FakeCheckpointModel() + original = model._all_gather_tiled_results + + vae = _construct(monkeypatch, "cpu", model) + + assert vae.model._all_gather_tiled_results is not original + assert vae._checkpoint_tile_gather is original + + +# -------------------------------------------------------------------------- +# The decision-point receipt carries quantities, not just presence: a line +# that only says "installed" cannot tell reuse from reallocation. +# -------------------------------------------------------------------------- + + +def test_the_receipt_line_reports_hits_allocs_and_the_address(fake_collective, gather_log): + vae = _vae(4) + install = [ + (level, line) for level, line in gather_log.records if "persistent tile gather installed device=xpu" in line + ] + # The one-shot install line is what an operator needs at info level. + assert install and all(level == logging.INFO for level, _ in install) + gather_log.messages.clear() + gather_log.records.clear() + + _run(vae, fake_collective, _tiles(8), 8, 4, 0) + first_ptr = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(8, seed=3), 8, 4, 0) + + records = [(level, line) for level, line in gather_log.records if "[H3_VAE_GATHER] reuse=" in line] + # The per-gather line fires once per decoder tile batch, so it is debug. + assert all(level == logging.DEBUG for level, _ in records) + lines = [line for _, line in records] + assert len(lines) == 2 + assert "reuse=alloc" in lines[0] and "hits=0 allocs=1" in lines[0] + assert "reuse=hit" in lines[1] and "hits=1 allocs=1" in lines[1] + # Resident bytes are reported, not just the bytes this one gather needed. + needed = _workspace_bytes(8, 4) / (1024.0 * 1024.0) + for line in lines: + assert f"need_mib={needed:.2f}" in line + assert f"workspace_mib={needed:.2f}" in line + assert f"buf_ptr=0x{first_ptr:x}" in lines[0] + assert f"buf_ptr=0x{first_ptr:x}" in lines[1] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py index 7a1c54f7b..86ccaa27b 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py @@ -198,7 +198,7 @@ def test_video_vae_installs_exact_optimizations(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, **_kwargs: FakeRemote(), ) install = Mock(return_value=True) monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 01332beec..57b2daaf8 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -635,6 +635,32 @@ def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): } +def test_i2v_minimax_h3_preserves_reference_geometry(): + engine = FakeAsyncOmni() + engine.model_class_name = "MiniMaxH3Pipeline" + handler = OmniOpenAIServingVideo.for_diffusion( + diffusion_engine=engine, + model_name="MiniMaxAI/MiniMax-H3", + ) + image = Image.new("RGB", (48, 32)) + + asyncio.run( + handler._run_and_extract( + VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64), + "minimax-h3-reference-geometry", + reference_image=ReferenceImage(image), + ) + ) + + assert engine.captured_prompt is not None + assert engine.captured_sampling_params_list is not None + input_image = engine.captured_prompt["multi_modal_data"]["image"] + assert isinstance(input_image, Image.Image) + assert input_image.size == (48, 32) + sampling_params = engine.captured_sampling_params_list[0] + assert (sampling_params.width, sampling_params.height) == (96, 64) + + def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture): image_bytes = _make_test_image_bytes((48, 48)) mocker.patch( diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index dcdf13678..a187c570c 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -290,6 +290,16 @@ def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.Monkey assert captured == {"model": "org/model", "revision": "pinned-revision"} +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("model_class_name", ["MiniMaxH3Pipeline", "MiniMaxH3ModularPipeline"]) +def test_minimax_h3_preserves_reference_image_size(model_class_name: str) -> None: + assert should_preserve_reference_image_size( + model_class_name, + model="MiniMaxAI/MiniMax-H3", + ) + + @pytest.mark.core_model @pytest.mark.cpu def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None: diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index e903960d6..592f5ba9b 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -215,8 +215,9 @@ class AttentionMetadata: # variable-length query/key sequences for FlashAttention. # "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with # the packed cu_seqlens tensors. - # "valid_kv_length": int — contiguous valid K/V prefix length for a - # backend that advertises supports_prefix_kv_slicing. + # "valid_kv_length": int — contiguous valid K/V prefix length. Local + # backends may slice it directly; ring backends trim each circulated + # K/V block from the same global prefix. # "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path # (TND npu_fusion_attention driven by cu_seqlens, mask never read). # Requires the [real, pad] two-document packing contract; see diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py index b8743fe13..a898c7729 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention @@ -9,7 +10,33 @@ import torch.nn.functional as F LseLayout = Literal["bhs", "bsh"] -__all__ = ["update_out_and_lse", "flatten_varlen_lse", "unflatten_varlen_lse"] +__all__ = [ + "ring_kv_block_valid_length", + "update_out_and_lse", + "flatten_varlen_lse", + "unflatten_varlen_lse", +] + + +def ring_kv_block_valid_length( + valid_kv_length: int | None, + block_size: int, + block_rank: int, + world_size: int, +) -> int: + """Return the valid prefix length in one globally ordered ring K/V block.""" + if valid_kv_length is None: + return block_size + if isinstance(valid_kv_length, bool) or not isinstance(valid_kv_length, int): + raise ValueError("valid_kv_length must be a Python integer") + total_length = block_size * world_size + if not 0 < valid_kv_length <= total_length: + raise ValueError( + "valid_kv_length must be within the global ring K/V sequence, " + f"got {valid_kv_length} for length {total_length}" + ) + block_start = block_rank * block_size + return max(0, min(block_size, valid_kv_length - block_start)) def _normalize_lse(block_lse: torch.Tensor, block_out: torch.Tensor, lse_layout: LseLayout) -> torch.Tensor: diff --git a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py b/vllm_omni/diffusion/attention/backends/ring_flash_attn.py index 1fce4f1b4..87c3b9e3d 100644 --- a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/ring_flash_attn.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention @@ -6,7 +7,10 @@ import torch from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType, select_flash_attn_impl -from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse +from vllm_omni.diffusion.attention.backends.ring.ring_utils import ( + ring_kv_block_valid_length, + update_out_and_lse, +) from vllm_omni.diffusion.distributed.comm import RingComm @@ -27,6 +31,7 @@ def ring_flash_attn_forward( joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): # Validate causal + joint_strategy combination # When causal=True and joint_strategy="rear", the causal mask would incorrectly @@ -42,6 +47,21 @@ def ring_flash_attn_forward( "from seeing the joint tokens." ) + # Trimming the circulated K/V blocks changes their length while the query + # block keeps the full padded length. Flash-Attention aligns a causal mask + # to the bottom-right of a rectangular (seqlen_q, seqlen_k) tile, so a + # trimmed block would shift the diagonal and let queries attend to the + # wrong key positions -- silently, with no shape error. The two are + # therefore an unsupported combination rather than a slow path: the caller + # must either drop the padding before the ring or run non-causal. + if causal and valid_kv_length is not None: + raise ValueError( + "valid_kv_length is not supported with causal=True in Ring Attention. " + "Trimming a circulated K/V block shortens seqlen_k while seqlen_q keeps " + "the padded length, and the causal diagonal is bottom-right aligned, so " + "the mask would silently shift. Unpad before the ring, or use causal=False." + ) + comm = RingComm(process_group) out = None @@ -66,8 +86,15 @@ def ring_flash_attn_forward( comm.commit() if not causal or step <= comm.rank: - step_k = k - step_v = v + block_rank = (comm.rank - step) % comm.world_size + block_valid_length = ring_kv_block_valid_length( + valid_kv_length, + k.shape[1], + block_rank, + comm.world_size, + ) + step_k = k[:, :block_valid_length] + step_v = v[:, :block_valid_length] if step == 0 and joint_tensor_key is not None: if joint_strategy == "front": step_k = torch.cat([joint_tensor_key, step_k], dim=1) @@ -76,27 +103,28 @@ def ring_flash_attn_forward( step_k = torch.cat([step_k, joint_tensor_key], dim=1) step_v = torch.cat([step_v, joint_tensor_value], dim=1) - fn = select_flash_attn_impl(attn_type, stage="fwd-only", attn_processor=attn_processor) - block_out, block_lse = fn( - q, - step_k, - step_v, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal and step == 0, - window_size=window_size, - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=True and dropout_p > 0, - ) + if step_k.shape[1] > 0: + fn = select_flash_attn_impl(attn_type, stage="fwd-only", attn_processor=attn_processor) + block_out, block_lse = fn( + q, + step_k, + step_v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal and step == 0, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=True and dropout_p > 0, + ) - # Ensure block_out is contiguous if needed, though usually it is from FA + # Ensure block_out is contiguous if needed, though usually it is from FA - if attn_type == AttnType.SPARSE_SAGE: - out, lse = block_out, block_lse - else: - # Ring kernel wrappers canonicalize LSE to (B, H, S). - out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") + if attn_type == AttnType.SPARSE_SAGE: + out, lse = block_out, block_lse + else: + # Ring kernel wrappers canonicalize LSE to (B, H, S). + out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") if step + 1 != comm.world_size: comm.wait() @@ -132,6 +160,7 @@ class RingFlashAttnFunc(torch.autograd.Function): joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -158,6 +187,7 @@ class RingFlashAttnFunc(torch.autograd.Function): joint_tensor_key=joint_tensor_key, joint_tensor_value=joint_tensor_value, joint_strategy=joint_strategy, + valid_kv_length=valid_kv_length, ) return out if not return_softmax else (out, softmax_lse, None) @@ -193,6 +223,7 @@ def ring_flash_attn_qkvpacked_func( None, # joint_tensor_key None, # joint_tensor_value "front", # joint_strategy + None, # valid_kv_length ) @@ -228,6 +259,7 @@ def ring_flash_attn_kvpacked_func( None, # joint_tensor_key None, # joint_tensor_value "front", # joint_strategy + None, # valid_kv_length ) @@ -249,6 +281,7 @@ def ring_flash_attn_func( joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, None]: """Ring Attention forward pass using Flash Attention backend. @@ -290,6 +323,8 @@ def ring_flash_attn_func( Defaults to None. joint_strategy (str): Concatenation strategy ("front" or "back"). Defaults to "front". + valid_kv_length (int | None): Length of the contiguous valid prefix in + the global ring K/V sequence. Defaults to the full sequence. Returns: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, None]]: @@ -314,4 +349,5 @@ def ring_flash_attn_func( joint_tensor_key, joint_tensor_value, joint_strategy, + valid_kv_length, ) diff --git a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py index 7ab60b784..65f5d8ad7 100644 --- a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py +++ b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention @@ -10,7 +11,10 @@ import torch from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.ring.ring_kernels import pytorch_attn_forward -from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse +from vllm_omni.diffusion.attention.backends.ring.ring_utils import ( + ring_kv_block_valid_length, + update_out_and_lse, +) from vllm_omni.diffusion.distributed.comm import RingComm logger = init_logger(__name__) @@ -33,6 +37,7 @@ def ring_pytorch_attn_func( joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): return RingAttentionFunc.apply( group, @@ -45,6 +50,7 @@ def ring_pytorch_attn_func( joint_tensor_key, joint_tensor_value, joint_strategy, + valid_kv_length, ) @@ -64,6 +70,7 @@ class RingAttentionFunc(torch.autograd.Function): joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): # Validate causal + joint_strategy combination # When causal=True and joint_strategy="rear", the causal mask would incorrectly @@ -79,6 +86,22 @@ class RingAttentionFunc(torch.autograd.Function): "from seeing the joint tokens." ) + # Trimming the circulated K/V blocks changes their length while the + # query block keeps the full padded length. A causal SDPA mask over a + # rectangular (seqlen_q, seqlen_k) tile is aligned to the bottom right, + # so a trimmed block would shift the diagonal and let queries attend to + # the wrong key positions -- silently, with no shape error. The two are + # therefore an unsupported combination rather than a slow path: the + # caller must either drop the padding before the ring or run + # non-causal. + if is_causal and valid_kv_length is not None: + raise ValueError( + "valid_kv_length is not supported with causal=True in Ring Attention. " + "Trimming a circulated K/V block shortens seqlen_k while seqlen_q keeps " + "the padded length, and the causal diagonal is bottom-right aligned, so " + "the mask would silently shift. Unpad before the ring, or use causal=False." + ) + comm = RingComm(group) # Ensure tensors are contiguous for P2P communication q = q.contiguous() @@ -98,8 +121,15 @@ class RingAttentionFunc(torch.autograd.Function): comm.commit() if not is_causal or step <= comm.rank: - step_k = k - step_v = v + block_rank = (comm.rank - step) % comm.world_size + block_valid_length = ring_kv_block_valid_length( + valid_kv_length, + k.shape[1], + block_rank, + comm.world_size, + ) + step_k = k[:, :block_valid_length] + step_v = v[:, :block_valid_length] if step == 0 and joint_tensor_key is not None: if joint_strategy == "front": step_k = torch.cat([joint_tensor_key, step_k], dim=1) @@ -108,15 +138,16 @@ class RingAttentionFunc(torch.autograd.Function): step_k = torch.cat([step_k, joint_tensor_key], dim=1) step_v = torch.cat([step_v, joint_tensor_value], dim=1) - block_out, block_lse = pytorch_attn_forward( - q, - step_k, - step_v, - softmax_scale=sm_scale, - causal=is_causal and step == 0, - op_type=op_type, - ) - out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") + if step_k.shape[1] > 0: + block_out, block_lse = pytorch_attn_forward( + q, + step_k, + step_v, + softmax_scale=sm_scale, + causal=is_causal and step == 0, + op_type=op_type, + ) + out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") if step + 1 != comm.world_size: comm.wait() diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index d71359ae7..1eb811636 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/vllm_omni/diffusion/attention/parallel/ring.py @@ -156,9 +156,11 @@ class RingParallelAttention: # Extract joint tensors joint_key, joint_value = None, None joint_strategy = "front" + valid_kv_length = None if attn_metadata is not None: joint_key = attn_metadata.joint_key joint_value = attn_metadata.joint_value + valid_kv_length = attn_metadata.extra.get("valid_kv_length") if attn_metadata.joint_strategy is not None: joint_strategy = attn_metadata.joint_strategy @@ -176,6 +178,7 @@ class RingParallelAttention: joint_tensor_key=joint_key, joint_tensor_value=joint_value, joint_strategy=joint_strategy, + valid_kv_length=valid_kv_length, ) # Ring only implements local Flash/AITER and SDPA kernels. HuggingFace @@ -261,4 +264,5 @@ class RingParallelAttention: joint_tensor_key=joint_key, joint_tensor_value=joint_value, joint_strategy=joint_strategy, + valid_kv_length=valid_kv_length, ) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 3abc90476..f0839429b 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -44,6 +44,10 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, +) from vllm_omni.diffusion.layers.activation import SiluAndMul from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm @@ -67,6 +71,15 @@ logger = init_logger(__name__) # packing, and ``_run_packed_attention`` re-checks it per forward; a name-only # gate would let FLASH_ATTN's NPU/XPU code paths through even though those # variants would silently attend across request boundaries. +def _ring_sequence_parallel_is_active(attention_layer: Attention) -> bool: + """Match :meth:`Attention._get_active_parallel_strategy` for Ring.""" + if not getattr(attention_layer, "use_ring", False) or getattr(attention_layer, "skip_sequence_parallel", False): + return False + if is_forward_context_available() and not get_forward_context().sp_active: + return False + return True + + def _attention_isolates_packed_requests(attention_layer: Any) -> bool: """True if this attention layer keeps N-document packed boundaries. @@ -78,7 +91,7 @@ def _attention_isolates_packed_requests(attention_layer: Any) -> bool: backend = getattr(attention_layer, "attn_backend", None) if backend is None or not backend.supports_multi_doc_packed_varlen(): return False - return not getattr(attention_layer, "use_ring", False) + return not _ring_sequence_parallel_is_active(attention_layer) @dataclass @@ -508,6 +521,7 @@ class MiniMaxH3Attention(nn.Module): ) attn_mask = None mask_free_packed_padding = False + use_ring = _ring_sequence_parallel_is_active(self.attention) if num_requests > 1: # A step-mode batch packs one document per request, so its valid # rows are block-diagonal rather than a prefix: neither a KV prefix @@ -534,12 +548,15 @@ class MiniMaxH3Attention(nn.Module): # supports_packed_mask_free: backend consumes the packed metadata # without ever reading attn_mask (CUDA packed varlen, NPU # npu_attn_varlen opt-in with its own fallback rebuild). - use_ring = getattr(self.attention, "use_ring", False) mask_free_packed_padding = not use_ring and self.attention.attn_backend.supports_packed_mask_free() no_mask = not use_ring and ( self.attention.attn_backend.supports_prefix_kv_slicing or mask_free_packed_padding ) - if used < packed_total and not no_mask: + # Hybrid Ulysses reshards Q to one ring partition before the ring + # kernel runs, so a global [packed_total] mask cannot pass its + # query-length check. Ring consumes valid_kv_length directly and + # trims the circulated K/V blocks instead. + if used < packed_total and not no_mask and not use_ring: attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, @@ -563,7 +580,7 @@ class MiniMaxH3Attention(nn.Module): # quadratic full_qk mask is never materialized. Ring attention # is excluded: it keeps the aligned padding rows for its # fixed-size P2P buffers and still needs the mask. - "npu_attn_varlen": not getattr(self.attention, "use_ring", False), + "npu_attn_varlen": not use_ring, # fp16-range protection for the ascend_laser_attention kernel # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other # backend/path. diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py index 0dcbd19b8..ecf797a9c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py +++ b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 packed-sequence materialization for FL2VA and Ref2VA tasks. Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P]. @@ -18,6 +19,8 @@ from collections.abc import Mapping, Sequence import numpy as np import torch +from vllm_omni.errors import OmniClientError + MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: tuple[tuple[int, ...], ...] = ( (0,), (-1,), @@ -38,13 +41,42 @@ _INTERP = 32 _T_GROUP = 5 _FRAME_PER_TOKEN = (1, 4, 4, 4, 4) _FRAME_RESCALE = 5.0 / 3.0 -_SEQ_ALIGN = 64 +MINIMAX_H3_SEQ_ALIGN = 64 +"""Row alignment of the packed sequence. + +An explicit ``seq_len`` must be a multiple of this, so that pinning a length +lands on a bucket the default rounding could also produce. +""" + +MINIMAX_H3_MAX_PAD_SEQ_LEN = 1 << 20 +"""Ceiling for a request-pinned packed length. + +Every packed row costs a fixed number of bytes across the structural tensors +built below, so an unbounded request value would size those allocations. The +largest shape this pipeline is documented against -- 672x384 at 209 frames -- +packs about 16k rows, so this ceiling leaves roughly two orders of magnitude of +headroom while keeping the structural tensors in the tens of megabytes. +""" # MiniMax H3 packs video latents with a fixed [1, 2, 2] (t, h, w) patch. _PATCH_T = 1 _PATCH_H = 2 _PATCH_W = 2 +def _resolve_padded_seq_len(used: int, seq_len: int | None) -> int: + """Resolve the padded row count for a packed sequence. + + Without an explicit ``seq_len`` the used rows are rounded up to the next + :data:`MINIMAX_H3_SEQ_ALIGN` boundary. An explicit value is honored as-is + and must cover the used rows. + """ + if seq_len is None: + return ((used + MINIMAX_H3_SEQ_ALIGN - 1) // MINIMAX_H3_SEQ_ALIGN) * MINIMAX_H3_SEQ_ALIGN + if seq_len < used: + raise OmniClientError(f"seq_len {seq_len} < used rows {used}") + return seq_len + + def _keyframe_cond_frame_indices( *, include_keyframe_cond: bool, @@ -124,10 +156,15 @@ def minimax_h3_packed_sequence( include_keyframe_cond: bool, keyframe_frame_indices: list[int] | tuple[int, ...] | None = None, frame_count: int | None = None, + seq_len: int | None = None, ) -> dict[str, object]: """Build the packed-sequence structural fields for one CFG branch. - The used length is padded up to a multiple of 64. + The used length is padded up to a multiple of + :data:`MINIMAX_H3_SEQ_ALIGN`, or to an explicit ``seq_len`` that covers the + used rows -- the same contract + :func:`minimax_h3_packed_sequence_ref2va_blocks` already offers. Pinning + the length keeps requests of different prompt lengths on one packed shape. """ ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W frame_rows = ph * pw @@ -143,7 +180,7 @@ def minimax_h3_packed_sequence( video_rows = latent_t * frame_rows audio_rows = audio_t * audio_channel used = text_len + cond_rows + audio_rows + video_rows - seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN + seq_len = _resolve_padded_seq_len(used, seq_len) text_sl = slice(0, text_len) cond_sl = slice(text_len, text_len + cond_rows) @@ -376,10 +413,7 @@ def minimax_h3_packed_sequence_ref2va_blocks( audio_rows = audio_t * audio_channel ref_rows = ref_visual_rows + ref_audio_rows used = text_len + ref_rows + audio_rows + video_rows - if seq_len is None: - seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN - if seq_len < used: - raise ValueError(f"seq_len {seq_len} < used rows {used}") + seq_len = _resolve_padded_seq_len(used, seq_len) text_sl = slice(0, text_len) cursor = text_len diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index f80513d4e..8f814a27a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -137,6 +137,8 @@ from .npu.lora import ( load_minimax_h3_native_lora, ) from .packed_sequence import ( + MINIMAX_H3_MAX_PAD_SEQ_LEN, + MINIMAX_H3_SEQ_ALIGN, minimax_h3_packed_sequence, minimax_h3_packed_sequence_ref2va_blocks, ) @@ -288,6 +290,7 @@ _MINIMAX_H3_DENOISE_INPUT_KEYS = ( "visual_condition_shapes", "audio_condition_lengths", "keyframe_frame_indices", + "pad_seq_len", ) # ``StepRequestState.extra`` keys owned by the step-execution path. @@ -543,6 +546,28 @@ def _resolve_minimax_h3_num_outputs(value: Any) -> int: return value +def _resolve_pad_seq_len(value: object) -> int | None: + """Validate the optional packed-length pin from ``extra_args``. + + The packer itself only needs a value that covers the used rows. A request + field needs two more guards: an unaligned value would silently open yet + another compiled shape, and an unbounded one would size every structural + tensor the packer allocates. + """ + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, np.integer)): + raise OmniClientError("MiniMax H3 pad_seq_len must be an integer") + pinned = int(value) + if pinned <= 0: + raise OmniClientError(f"MiniMax H3 pad_seq_len must be positive, got {pinned}") + if pinned % MINIMAX_H3_SEQ_ALIGN: + raise OmniClientError(f"MiniMax H3 pad_seq_len must be a multiple of {MINIMAX_H3_SEQ_ALIGN}, got {pinned}") + if pinned > MINIMAX_H3_MAX_PAD_SEQ_LEN: + raise OmniClientError(f"MiniMax H3 pad_seq_len must be at most {MINIMAX_H3_MAX_PAD_SEQ_LEN}, got {pinned}") + return pinned + + def _minimax_h3_output_seeds(seed: int, num_outputs: int) -> list[int]: return [int(seed) + output_index for output_index in range(int(num_outputs))] @@ -1077,11 +1102,13 @@ class MiniMaxH3Pipeline( os.path.join(model_path, "video_vae"), device=self.device, load_device=component_load_device, + trust_remote_code=od_config.trust_remote_code, ) self.audio_vae = MiniMaxH3AudioVAE( os.path.join(model_path, "audio_vae"), device=self.device, load_device=component_load_device, + trust_remote_code=od_config.trust_remote_code, ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae @@ -1674,6 +1701,31 @@ class MiniMaxH3Pipeline( logger.exception("Failed to release retained allocator cache after offload failure") raise + @staticmethod + def _is_output_owner_rank() -> bool: + """Whether this rank's output is returned by the diffusion executor.""" + return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + + def _offload_model_cpu_stage_output(self, tensor: torch.Tensor) -> torch.Tensor: + """Release a decoded output's storage before a later seed reloads the DiT. + + Model-level CPU offload makes the VAE hook evict the DiT before decode. + For a multi-output request, however, the next seed reloads the DiT while + the preceding decoded output would otherwise still occupy the device. + The reply rank performs the D2H copy early; ranks whose outputs are not + consumed retain only a metadata placeholder for the final concatenation. + + Call this on the tensor that is actually returned to the engine. Audio is + released inside ``decode``; video is released by the callers of ``decode`` + only after ``_prepare_minimax_h3_video_output`` has quantized it, so the + early D2H copy carries ``uint8`` frames rather than the decoded floats. + """ + if not getattr(self, "_model_cpu_offload_modules", None): + return tensor + if self._is_output_owner_rank(): + return tensor.cpu() + return torch.empty(tensor.shape, dtype=tensor.dtype, device="meta") + def _encode_visual_conditions( self, images: list[Image.Image], @@ -1932,6 +1984,7 @@ class MiniMaxH3Pipeline( visual_condition_shapes: list[tuple[int, int, int]] | None = None, audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, + pad_seq_len: int | None = None, ) -> dict[str, Any]: """Build the packed layout, initial rows, anchors, and sigma schedules. @@ -1961,6 +2014,7 @@ class MiniMaxH3Pipeline( latent_w=latent_w, audio_t=audio_t, ref_blocks=ref_blocks, + seq_len=pad_seq_len, ) else: packed = minimax_h3_packed_sequence( @@ -1972,7 +2026,19 @@ class MiniMaxH3Pipeline( include_keyframe_cond=task == "fl2va", keyframe_frame_indices=keyframe_frame_indices if task == "fl2va" else None, frame_count=num_frames if task == "fl2va" else None, + seq_len=pad_seq_len, ) + # Report the effective shape at info level exactly when a request pins + # it, so a deployment can confirm the pin landed without raising the + # log level for every request. + log = logger.info if pad_seq_len is not None else logger.debug + log( + "MiniMax H3 packed sequence: task=%s pad_seq_len=%s used=%d seq_len=%d", + task, + pad_seq_len, + int(packed["cu_seqlens"][1]), + int(packed["seq_len"]), + ) tags = packed["token_tags"].clone() tags[packed["text_pos"]] = text_tags.cpu() @@ -2108,6 +2174,7 @@ class MiniMaxH3Pipeline( visual_condition_shapes: list[tuple[int, int, int]] | None = None, audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, + pad_seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = self._build_denoise_inputs( task=task, @@ -2131,6 +2198,7 @@ class MiniMaxH3Pipeline( visual_condition_shapes=visual_condition_shapes, audio_condition_lengths=audio_condition_lengths, keyframe_frame_indices=keyframe_frame_indices, + pad_seq_len=pad_seq_len, ) branch = inputs["branch"] transformer = self._transformer_for_task(task) @@ -2253,6 +2321,7 @@ class MiniMaxH3Pipeline( video = video[..., :height, :width].contiguous() with self._component_on_device(self.audio_vae): audio = self.audio_vae.decode_latent(audio_latent) + audio = self._offload_model_cpu_stage_output(audio) return video, audio @staticmethod @@ -2561,6 +2630,7 @@ class MiniMaxH3Pipeline( base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) video_shift = float(extra.get("flow_shift", self.default_video_shift)) audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) + pad_seq_len = _resolve_pad_seq_len(extra.get("pad_seq_len")) quality_plan = self._quality_policy.resolve( quality=quality, num_inference_steps=num_steps, @@ -2587,6 +2657,7 @@ class MiniMaxH3Pipeline( "visual_condition_shapes": visual_shapes, "audio_condition_lengths": audio_lengths, "keyframe_frame_indices": keyframe_frame_indices, + "pad_seq_len": pad_seq_len, "seed": seed, "num_steps": num_steps, "video_shift": video_shift, @@ -2643,7 +2714,13 @@ class MiniMaxH3Pipeline( height=context["height"], width=context["width"], ) - videos.append(_prepare_minimax_h3_video_output(video)) + # Rebind rather than append the expression: the local would + # otherwise keep the decoded frames' device storage alive for the + # rest of the iteration, and the next seed's diffuse() -- the DiT + # reload this release exists to make room for -- runs before the + # loop rebinds it. post_decode() rebinds for the same reason. + video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) + videos.append(video) audios.append(audio) if videos and isinstance(videos[0], bytes): video = videos[0] if len(videos) == 1 else videos @@ -2966,7 +3043,7 @@ class MiniMaxH3Pipeline( height=shape["height"], width=shape["width"], ) - video = _prepare_minimax_h3_video_output(video) + video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index ca8359904..e40251dff 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -75,11 +75,23 @@ def _load_component_config(component_path: str) -> dict[str, Any]: def _load_remote_component( component_path: str, config: dict[str, Any], + *, + trust_remote_code: bool, ) -> nn.Module: auto_map = config.get("auto_map") or {} class_reference = auto_map.get("AutoModel") if not isinstance(class_reference, str): raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel") + if not trust_remote_code: + raise ValueError( + f"Loading {component_path} executes the modeling code shipped with " + f"the checkpoint (auto_map.AutoModel = {class_reference}). Pass " + "--trust-remote-code (or trust_remote_code=True) to allow it." + ) + # ``trust_remote_code`` is checked here rather than forwarded to + # ``get_class_from_dynamic_module``: that helper takes no such argument and + # would silently absorb it into ``**kwargs``, so forwarding it would read + # like a gate while executing the remote code unconditionally. component_cls = get_class_from_dynamic_module( class_reference, component_path, @@ -144,6 +156,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): *, device: torch.device, load_device: torch.device | None = None, + trust_remote_code: bool = False, ) -> None: super().__init__() self._device_target = device @@ -151,6 +164,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self.remote = _load_remote_component( component_path, self.config_dict, + trust_remote_code=trust_remote_code, ) # Match the reference loader contract before installing inference-only # decoder fast paths. Keyframe encoding remains FP32; decoder Linear @@ -176,6 +190,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self.use_slicing = False self.parallel_size = 1 self.device_module = torch.get_device_module() + self._tile_gather_workspace: torch.Tensor | None = None + self._tile_gather_stats = {"hits": 0, "allocs": 0, "workspace_bytes": 0} + self._checkpoint_tile_gather = None + if self._tile_gather_reuse_enabled(): + self._install_persistent_tile_gather() def load_to_device(self) -> None: if self._stager is not None: @@ -241,6 +260,145 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): parallel_module = importlib.import_module(f"{package}.parallel") return parallel_module.get_parallel_state() + def _tile_gather_reuse_enabled(self) -> bool: + """Whether this device needs the tiled-VAE gather address kept stable. + + XPU only. The accumulation this guards against is an XCCL-side + registration that is kept for every distinct receive-buffer address and + never reclaimed; no other backend in tree does that, so everywhere else + the checkpoint's own method is left in place and behaviour is unchanged. + The stable address comes from one grow-only byte workspace, so what the + adapter holds resident is one buffer for the largest geometry it has + seen rather than one per geometry. + """ + + return self._device_target.type == "xpu" + + def _install_persistent_tile_gather(self) -> None: + """Route the checkpoint's tiled-VAE gather through a bounded workspace. + + The checkpoint's ``_all_gather_tiled_results`` allocates its gather + output afresh on every call. Model-level offload calls ``empty_cache()`` + once per request, so the next request's output lands on a new device + address; XCCL registers a non-reclaimable resource per new receive + address, and the registrations accumulate until the card is full. The + replacement below keeps a single byte workspace alive on this adapter + and carves every gather out of its front, so the address the collective + writes to is the same one every request. + + The workspace only ever grows: a geometry that fits in the current + block reuses it as is, and a larger one replaces the block, which + drops the previous allocation before allocating its own. Resident + memory is therefore bounded by one buffer for the largest geometry the + adapter has seen, not by one buffer per distinct geometry. + + The override is bound on the checkpoint *instance*, not its class: the + class is remote code shared by every component loaded from the same + checkpoint, and only this adapter knows the device it runs on. + """ + + self._checkpoint_tile_gather = self.model._all_gather_tiled_results + self.model._all_gather_tiled_results = self._persistent_tile_gather + logger.info( + "[H3_VAE_GATHER] persistent tile gather installed device=%s", + self._device_target.type, + ) + + def _persistent_tile_gather( + self, + tasks: list[torch.Tensor], + num_tiles: int, + ) -> list[torch.Tensor]: + """Equal-shape replacement for the checkpoint's tiled-result gather. + + Contract kept identical to the checkpoint's method: return a list of + ``num_tiles`` tensors in global tile order, and raise on an empty local + share so a rank that owns no tile cannot silently skip the collective. + + Tile ownership is round-robin (``range(sp_rank, num_tiles, sp_size)``), + so every rank can compute every other rank's task count from + ``num_tiles`` and ``sp_size`` alone. That makes the per-rank payloads + equal once the leading task dimension is padded to ``max_tasks``, which + is what lets a single ``all_gather_into_tensor`` into a stable buffer + replace the variable-shape gather. + + The landing buffer is a view onto the front of one byte workspace that + this adapter keeps alive and only ever grows, so its address stays + stable across requests while resident memory stays bounded by the + largest geometry seen rather than growing per geometry. A byte + workspace is dtype-agnostic, so mixed-precision requests share it too. + + Returned tiles are cloned out of the buffer: the buffer is overwritten + by the next call and callers hold the tiles past that point. + """ + + state = self._native_parallel_state() + group = state["sp_process_group"] + sp_size = int(state["sp_size"]) + sp_rank = int(state["sp_rank"]) + + if not tasks: + raise ValueError(f"Found empty tasks on sp rank {sp_rank}") + + max_tasks = -(-num_tiles // sp_size) + if len(tasks) > max_tasks: + raise ValueError( + f"sp rank {sp_rank} holds {len(tasks)} tiles but round-robin " + f"ownership of {num_tiles} tiles across {sp_size} ranks allows " + f"at most {max_tasks}" + ) + if len(tasks) == max_tasks: + stacked = torch.stack(tasks, dim=0) + else: + # Pad the leading (task) dimension only. The padded slots belong to + # ranks whose share is short by construction, and the unpacking loop + # below never reads them back. + stacked = tasks[0].new_empty((max_tasks, *tasks[0].shape)) + torch.stack(tasks, dim=0, out=stacked[: len(tasks)]) + stacked[len(tasks) :].zero_() + + need_bytes = sp_size * stacked.numel() * stacked.element_size() + workspace = self._tile_gather_workspace + if workspace is None or workspace.device != stacked.device or workspace.numel() < need_bytes: + # Drop the previous block before allocating its replacement so the + # two are never resident at once; only one workspace is ever held. + workspace = None + self._tile_gather_workspace = None + workspace = torch.empty(need_bytes, dtype=torch.uint8, device=stacked.device) + self._tile_gather_workspace = workspace + self._tile_gather_stats["allocs"] += 1 + self._tile_gather_stats["workspace_bytes"] = need_bytes + reuse = "alloc" + else: + self._tile_gather_stats["hits"] += 1 + reuse = "hit" + buffer = workspace[:need_bytes].view(stacked.dtype).view(sp_size, *stacked.shape) + # Quantities, not just presence: a line that only says "installed" + # cannot distinguish a buffer that is being reused from one that is + # reallocated every request, which is the whole failure being fixed. + # Debug level, because this fires once per decoder tile batch (12 times + # per request on the canonical 1344x768 geometry) and the one-shot + # "installed" line above is what an operator needs at info level. + logger.debug( + "[H3_VAE_GATHER] reuse=%s shape=%s/%s need_mib=%.2f workspace_mib=%.2f buf_ptr=0x%x hits=%d allocs=%d", + reuse, + tuple(stacked.shape), + stacked.dtype, + need_bytes / (1024.0 * 1024.0), + self._tile_gather_stats["workspace_bytes"] / (1024.0 * 1024.0), + workspace.data_ptr(), + self._tile_gather_stats["hits"], + self._tile_gather_stats["allocs"], + ) + dist.all_gather_into_tensor(buffer, stacked, group=group) + + results: list[torch.Tensor] = [None] * num_tiles # type: ignore[list-item] + for rank in range(sp_size): + num_rank_tasks = -(-(num_tiles - rank) // sp_size) + for k in range(num_rank_tasks): + results[k * sp_size + rank] = buffer[rank][k].clone() + return results + def _decoder_tile_count(self, latent: torch.Tensor) -> int: """Number of decoder tiles the checkpoint will split ``latent`` into. @@ -464,6 +622,7 @@ class MiniMaxH3AudioVAE(nn.Module): *, device: torch.device, load_device: torch.device | None = None, + trust_remote_code: bool = False, ) -> None: super().__init__() self._device_target = device @@ -471,6 +630,7 @@ class MiniMaxH3AudioVAE(nn.Module): self.remote = _load_remote_component( component_path, self.config_dict, + trust_remote_code=trust_remote_code, ) # The checkpoint's audio VAE contract is FP32 for both reference # encoding and waveform decoding. diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index f21691eab..9c4a6eed6 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -189,6 +189,16 @@ def default_image_to_image_prompt( return result +def _always_preserve_reference_image_size( + *, + model: str | None, + revision: str | None = None, +) -> bool: + """Declare that a pipeline owns reference-image geometry.""" + del model, revision + return True + + _EXTRA_SPECS: dict[str, dict[str, Any]] = { "BagelPipeline": { "extra_body_params": BAGEL_EXTRA_BODY_PARAMS, @@ -230,6 +240,15 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { # Shared T2I/I2V envelopes select the output modality. LingBot's # pipeline owns model-specific validation and normalization. }, + **{ + model_class_name: { + "reference_image_size_resolver": _always_preserve_reference_image_size, + } + for model_class_name in ( + "MiniMaxH3Pipeline", + "MiniMaxH3ModularPipeline", + ) + }, **{ model_class_name: { "extra_body_params": LTX_EXTRA_BODY_PARAMS, From 5f25d9865281029db95c1778ead2c8ad666ac643 Mon Sep 17 00:00:00 2001 From: Pujitha Paladugu Date: Mon, 14 Sep 2026 10:02:39 -0700 Subject: [PATCH 064/209] [Bugfix] Allow Fish Speech phoneme-control tokens through normalize_fish_speech_text (#7043) Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Co-authored-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> --- pyproject.toml | 3 +++ .../models/fish_speech/test_prompt_utils.py | 23 +++++++++++++++++++ .../models/fish_speech/prompt_utils.py | 6 ++++- 3 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 tests/model_executor/models/fish_speech/test_prompt_utils.py diff --git a/pyproject.toml b/pyproject.toml index 2584c8c7a..fc1d5573b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -376,6 +376,9 @@ ba = "ba" SER = "SER" # NumPy ndarray.flags.writeable (canonical flag name, not "writable"). writeable = "writeable" +# ARPAbet phoneme symbol (Fish Speech phoneme-control token example, e.g. +# "IY0"), not a typo of "IT". +IY = "IY" # "convnet" abbreviates convolutional network. It appears in the vendored Fish # Speech DAC codec identifiers (convnet_type, transconvnet_type) and is not a # typo of "convent". diff --git a/tests/model_executor/models/fish_speech/test_prompt_utils.py b/tests/model_executor/models/fish_speech/test_prompt_utils.py new file mode 100644 index 000000000..98700da26 --- /dev/null +++ b/tests/model_executor/models/fish_speech/test_prompt_utils.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.model_executor.models.fish_speech.prompt_utils import normalize_fish_speech_text + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_normalize_fish_speech_text_allows_phoneme_control_tokens(): + text = "Deploy with <|phoneme_start|>K UW2 B ER0 N EH1 T IY0 Z<|phoneme_end|>." + + assert normalize_fish_speech_text(text) == text + + +def test_normalize_fish_speech_text_rejects_unknown_control_token(): + with pytest.raises(ValueError, match=r"unsupported control token\(s\): <\|foo\|>"): + normalize_fish_speech_text("<|foo|>hello") + + +def test_normalize_fish_speech_text_normalizes_legacy_speaker_tag(): + assert normalize_fish_speech_text("hello") == "<|speaker:0|>hello" diff --git a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py index 3918a6f94..61196f18f 100644 --- a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py +++ b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Shared Fish Speech prompt construction helpers.""" from __future__ import annotations @@ -11,6 +14,7 @@ FISH_CLONE_SYSTEM_PROMPT_SUFFIX = "\n\nSpeech:\n" _LEGACY_SPEAKER_TAG_PATTERN = re.compile(r"") _CANONICAL_SPEAKER_TAG_PATTERN = re.compile(r"<\|speaker:\d+\|>") +_PHONEME_TAG_PATTERN = re.compile(r"<\|phoneme_(?:start|end)\|>") _CONTROL_TOKEN_PATTERN = re.compile(r"<\|[^>]+\|>") @@ -21,7 +25,7 @@ def normalize_fish_speech_text(text: str, *, add_default_speaker: bool = False) disallowed_tokens = [ token for token in _CONTROL_TOKEN_PATTERN.findall(normalized) - if not _CANONICAL_SPEAKER_TAG_PATTERN.fullmatch(token) + if not (_CANONICAL_SPEAKER_TAG_PATTERN.fullmatch(token) or _PHONEME_TAG_PATTERN.fullmatch(token)) ] if disallowed_tokens: disallowed_list = ", ".join(sorted(set(disallowed_tokens))) From 731e36fb5786e61fd569ba1979c6e3fdfe0db6ea Mon Sep 17 00:00:00 2001 From: yadong <76457085+yadongtan@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:43:02 +0800 Subject: [PATCH 065/209] [Bugfix][Hardware][Ascend] Fix Qwen3-TTS code predictor dtype and nested graph replay (#6639) Signed-off-by: tanyadong <2209931449@qq.com> Signed-off-by: tanyadong Co-authored-by: tanyadong <2209931449@qq.com> --- .../qwen3_tts/test_code_predictor_dtype.py | 108 ++++++++++++++++++ .../models/common/qwen3_code_predictor.py | 14 ++- 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 438a5ec1d..1998a4dd4 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -572,6 +572,53 @@ class TestCodePredictorPerRowGenerators: class TestCodePredictorModelDtype: """Test the inner model forward with different dtypes.""" + @pytest.mark.parametrize( + ("device_type", "expect_fp32"), + [("cpu", False), ("npu", False), ("xpu", True), ("musa", True), ("cuda", True)], + ) + def test_fp32_fallback_excludes_cpu_and_npu( + self, + mocker: MockerFixture, + loaded_target_classes, + device_type: str, + expect_fp32: bool, + ) -> None: + """The stability fallback preserves accelerator behavior except on NPU.""" + _, _, _, code_predictor_model, _ = loaded_target_classes + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + + inputs = mocker.MagicMock() + inputs.dtype = torch.float16 + inputs.device.type = device_type + fp32_inputs = mocker.MagicMock() + fp32_inputs.device.type = device_type + inputs.float.return_value = fp32_inputs + model = object.__new__(code_predictor_model) + torch.nn.Module.__init__(model) + model.rotary_emb = mocker.Mock(return_value=None) + model.layers = [] + model.norm = mocker.Mock(side_effect=lambda hidden_states: hidden_states) + autocast = mocker.patch.object( + common_mod.torch.amp, + "autocast", + return_value=mocker.MagicMock( + __enter__=mocker.Mock(), + __exit__=mocker.Mock(return_value=False), + ), + ) + + code_predictor_model.forward(model, inputs, mocker.MagicMock()) + + if expect_fp32: + inputs.float.assert_called_once_with() + else: + inputs.float.assert_not_called() + autocast.assert_called_once_with( + device_type, + enabled=expect_fp32, + dtype=torch.float32, + ) + def test_model_forward_float16(self, loaded_target_classes) -> None: """Inner model forward should work in float16.""" _, _, _, code_predictor_model, _ = loaded_target_classes @@ -601,6 +648,67 @@ class TestCodePredictorModelDtype: assert output.shape == (bsz, seq_len, 32) +class TestCodePredictorGraphReplay: + """Test nested device graph handling for the shared code predictor.""" + + @pytest.mark.parametrize("is_capturing", [False, True]) + def test_npu_outer_capture_skips_inner_graph_replay( + self, + mocker: MockerFixture, + loaded_target_classes, + is_capturing: bool, + ) -> None: + _, _, code_predictor_wrapper, _, code_predictor_wrapper_config = loaded_target_classes + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + + predictor = object.__new__(code_predictor_wrapper) + torch.nn.Module.__init__(predictor) + predictor._num_groups = 3 + predictor._model_dtype = torch.float32 + predictor._setup_compile = mocker.Mock() + predictor._padded_bsz = mocker.Mock(side_effect=lambda bsz: bsz) + predictor._ensure_buffers = mocker.Mock() + predictor._proj_buf = torch.zeros(1, 4, 4) + predictor.small_to_mtp_projection = torch.nn.Identity() + regular_output = torch.zeros(1, 4, 4) + predictor._compiled_model_fwd = mocker.Mock(return_value=regular_output) + predictor._lm_heads_list = [ + torch.nn.Linear(4, 8, bias=False), + torch.nn.Linear(4, 8, bias=False), + ] + predictor._codec_embeds_list = [torch.nn.Embedding(8, 4)] + predictor._wrapper_config = code_predictor_wrapper_config( + sampling_mode="per_call", + return_proj_buf=False, + ) + predictor._prefix_graphs_enabled = False + predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)} + graph = mocker.Mock() + graph_output = torch.zeros(1, 4, 4) + predictor._device_graphs = {1: (graph, graph_output)} + + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=True) + npu_mock = mocker.MagicMock() + npu_mock.is_current_stream_capturing.return_value = is_capturing + mocker.patch.object(common_mod.torch, "npu", npu_mock, create=True) + + result = predictor( + layer0_code=torch.zeros(1, dtype=torch.long), + layer0_embed=torch.zeros(1, 4), + last_talker_hidden=torch.zeros(1, 4), + do_sample=False, + ) + + assert result.shape == (1, 3) + npu_mock.is_current_stream_capturing.assert_called_once_with() + if is_capturing: + graph.replay.assert_not_called() + assert predictor._compiled_model_fwd.call_count == 2 + else: + assert graph.replay.call_count == 2 + predictor._compiled_model_fwd.assert_not_called() + + class TestCodePredictorWrapperConfig: """Test wrapper configuration for different models.""" diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 60b4d4525..eff3305d5 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -439,10 +439,11 @@ class CodePredictorBaseModel(nn.Module): # native bf16 support (Turing, Volta). The RMSNorm and RoPE # layers already upcast internally; this extends the same # treatment to attention and MLP. - # autocast to float32 is unsupported on CPU; skip fp32 upcast there - # (CPU uses full-precision intermediates internally). + # CPU autocast to float32 is unsupported, and Ascend fused operators + # require matching activation and weight dtypes. Preserve the existing + # stability fallback on the other accelerator backends. input_dtype = inputs_embeds.dtype - use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type != "cpu" + use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type not in ("cpu", "npu") if use_fp32: inputs_embeds = inputs_embeds.float() hidden_states = inputs_embeds @@ -1035,6 +1036,9 @@ class CodePredictorWrapper(nn.Module): model_fwd = self._compiled_model_fwd lm_heads = self._lm_heads_list codec_embeds = self._codec_embeds_list + # torch-npu cannot replay an inner NPUGraph while the outer talker_mtp + # graph is being captured. Capture state is constant for this forward. + is_npu_capturing = current_omni_platform.is_npu() and torch.npu.is_current_stream_capturing() # Zero the padded region of the buffer proj_buf[:padded_bsz].zero_() @@ -1085,7 +1089,9 @@ class CodePredictorWrapper(nn.Module): # Use captured device graph if available, otherwise call compiled fn. device_graph_entry = self._device_graphs.get(graph_key) - if device_graph_entry is not None: + # Let the outer graph record the regular forward during capture; + # normal inference still uses the inner graph replay fast path. + if device_graph_entry is not None and not is_npu_capturing: device_graph_entry[0].replay() hidden_out = device_graph_entry[1] else: From fb23e8807a151d77c8135e36c85e58d4793c52cb Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Mon, 14 Sep 2026 15:15:27 -0400 Subject: [PATCH 066/209] =?UTF-8?q?[BugFix]=20Fix=20importing=20renamed=20?= =?UTF-8?q?module=20vllm.entrypoints.openai.engine.=E2=80=A6=20(#7531)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Nick Cao --- tests/entrypoints/openai_api/test_audio_format.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index c7fc4fbef..2fb76e3f6 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -290,7 +290,7 @@ class TestResolveAudioFormat: @pytest.mark.parametrize("fmt", ["pcm", "wav"]) def test_chat_audio_metadata_survives_response_serialization(stream, fmt): from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest - from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo + from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo from vllm_omni.entrypoints.openai.protocol.chat_completion import ( OmniChatCompletionResponse, From ed3776645d4f310eb5c762864eb68d04f47333d3 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Mon, 14 Sep 2026 16:50:07 -0400 Subject: [PATCH 067/209] [Tests] Skip global GPU cleanup for parallel online diffusion tests (#7534) Signed-off-by: Nick Cao Co-authored-by: Codex --- .../diffusion/test_common_online.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index 9c7ef8a0a..65089e3ef 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/tests/model_tests/diffusion/test_common_online.py @@ -8,6 +8,7 @@ encoding) using tiny models. """ import pytest +from xdist import is_xdist_worker from tests.helpers.runtime import OmniServer, OnlineOmniClient from tests.model_tests.diffusion.case_filtering import get_parametrized_options @@ -30,6 +31,23 @@ from tests.model_tests.diffusion.task_runners import ( pytestmark = [pytest.mark.diffusion, pytest.mark.xdist] +@pytest.fixture(autouse=True) +def _disable_global_gpu_cleanup_for_parallel_workers( + request: pytest.FixtureRequest, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not wait on total GPU usage while sibling xdist workers are active. + + ``cleanup_test_environment`` observes the whole device, not allocations + owned by the current worker. Waiting for the device to fall below its + global threshold therefore turns server teardown into a cross-worker + barrier when online tests run with xdist. ``OmniServer`` still tears down + its own subprocess tree; retain the broader cleanup for non-xdist runs. + """ + if is_xdist_worker(request): + monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None) + + @pytest.mark.parametrize( "model_name,accelerations,supported_tasks,check_multioutput,check_determinism", get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True), From 7ad51f0007f027e182565bd99d61bfddf58d7e38 Mon Sep 17 00:00:00 2001 From: LOGO127 Date: Tue, 15 Sep 2026 06:24:49 +0800 Subject: [PATCH 068/209] [Misc][HiggsAudioV3] Diagnose skipped reference audio substitution (#7098) Signed-off-by: luozijian --- recipes/BosonAI/Higgs-Audio-V3-TTS.md | 25 +-- .../test_ref_audio_diagnostics.py | 168 ++++++++++++++++++ .../higgs_audio_v3/higgs_audio_v3_talker.py | 14 ++ 3 files changed, 197 insertions(+), 10 deletions(-) create mode 100644 tests/model_executor/models/higgs_audio_v3/test_ref_audio_diagnostics.py diff --git a/recipes/BosonAI/Higgs-Audio-V3-TTS.md b/recipes/BosonAI/Higgs-Audio-V3-TTS.md index 0391fe03c..5f023fe3f 100644 --- a/recipes/BosonAI/Higgs-Audio-V3-TTS.md +++ b/recipes/BosonAI/Higgs-Audio-V3-TTS.md @@ -102,23 +102,28 @@ python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client - Key flags: `--trust-remote-code` and `--omni` are required. - Output: 24 kHz mono WAV. - Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity. +- Reference substitution diagnostics: temporarily set `VLLM_LOGGING_LEVEL=DEBUG` + before starting the server to see why reference audio substitution was skipped. + Count-mismatch messages report the batch-local request index, placeholder count, + and reference code row count without logging text or audio content. They do not + change the existing fallback behavior; legacy counts describe the current span. - Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`). - - `max_num_seqs=16` for both stages. - - Stage 0 and Stage 1 default to the same device (`0`) for single-GPU serving. - - Stage 0 intentionally keeps `enforce_eager=true`. This preserves the Higgs-specific local MLP CUDA graph path, which is the current high-throughput default. - - Stage 1 remains `enforce_eager=true` for the codec decoder. + - `max_num_seqs=16` for both stages. + - Stage 0 and Stage 1 default to the same device (`0`) for single-GPU serving. + - Stage 0 intentionally keeps `enforce_eager=true`. This preserves the Higgs-specific local MLP CUDA graph path, which is the current high-throughput default. + - Stage 1 remains `enforce_eager=true` for the codec decoder. - Deploy profiles: - - High throughput: `vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml`. + - High throughput: `vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml`. Use this for medium/high concurrency. The auto-discovered `higgs_multimodal_qwen3.yaml` is kept as a compatibility/default alias for this profile. - - Low latency: `vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml`. + - Low latency: `vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml`. Use this for low-concurrency serving (for example c1-c4). It sets Stage 0 `enforce_eager=false` and enables vLLM `FULL_DECODE_ONLY` CUDA graph through YAML `compilation_config`; no environment variable is required. - - Profile details: `vllm_omni/deploy/README_higgs_audio_v3.md`. + - Profile details: `vllm_omni/deploy/README_higgs_audio_v3.md`. - Performance note: - - Do not switch the auto-discovered default Stage 0 profile to vLLM + - Do not switch the auto-discovered default Stage 0 profile to vLLM `FULL_DECODE_ONLY` CUDA graph without an end-to-end throughput and audio-quality revalidation. On the H20 SeedTTS c16/full-dataset benchmark, the eager Stage 0 path with Higgs local MLP CUDA graph reproduced @@ -127,5 +132,5 @@ python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client capture fix showed lower c1/c4 latency, so the FULL_DECODE profile is kept as an explicit low-latency option rather than the throughput default. - Known limitations: - - Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture). - - Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph. + - Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture). + - Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph. diff --git a/tests/model_executor/models/higgs_audio_v3/test_ref_audio_diagnostics.py b/tests/model_executor/models/higgs_audio_v3/test_ref_audio_diagnostics.py new file mode 100644 index 000000000..c7417584d --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_ref_audio_diagnostics.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reference-substitution diagnostics without loading a model checkpoint.""" + +import pytest +import torch +from pytest_mock import MockerFixture + +from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def talker() -> mod.HiggsAudioV3TalkerForConditionalGeneration: + # Exercise the real method and fused embedding without constructing Qwen3. + instance = mod.HiggsAudioV3TalkerForConditionalGeneration.__new__(mod.HiggsAudioV3TalkerForConditionalGeneration) + torch.nn.Module.__init__(instance) + instance._last_step_query_start_loc = torch.tensor([0, 4]) + instance.multimodal_embedding = mod.HiggsFusedMultiTextEmbedding(num_codebooks=1, vocab_size=32, hidden_size=1) + with torch.no_grad(): + instance.multimodal_embedding.weight.copy_(torch.arange(32).reshape(32, 1)) + return instance + + +@pytest.mark.parametrize( + "ref_positions,mask,expected_positions,expected_codes", + [ + ([], [True, True], 0, 2), + ([1], [True, True], 1, 2), + ([1, 2, 3], [True, False], 3, 1), + ], + ids=["empty-positions", "too-few-positions", "count-mismatch-after-masking"], +) +def test_explicit_mismatch_logs_counts_without_substitution( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, + mocker: MockerFixture, + ref_positions: list[int], + mask: list[bool], + expected_positions: int, + expected_codes: int, +) -> None: + debug = mocker.spy(mod.logger, "debug") + hidden = torch.zeros(4, 1) + info = { + "audio_input_ids": torch.tensor([[5], [7]]), + "audio_input_ids_mask": torch.tensor(mask), + "audio_placeholder_positions": torch.tensor(ref_positions, dtype=torch.long), + } + result = talker._apply_ref_audio_substitution(hidden, torch.ones(4, dtype=torch.long), torch.arange(4), [info]) + + assert result is hidden + assert torch.count_nonzero(hidden) == 0 + debug.assert_called_once() + message, request_index, position_count, code_count = debug.call_args.args + assert "reference audio" in message + assert "placeholder" in message + assert (request_index, position_count, code_count) == (0, expected_positions, expected_codes) + + +@pytest.mark.parametrize("input_ids", [[1, -100, 2, 3], [1, 2, 3, 4]], ids=["one-slot", "no-slots"]) +def test_legacy_short_span_logs_counts_without_substitution( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, + mocker: MockerFixture, + input_ids: list[int], +) -> None: + debug = mocker.spy(mod.logger, "debug") + hidden = torch.zeros(4, 1) + result = talker._apply_ref_audio_substitution( + hidden, torch.tensor(input_ids), torch.arange(4), [{"audio_input_ids": torch.tensor([[5], [7]])}] + ) + + assert result is hidden + debug.assert_called_once() + message, request_index, position_count, code_count = debug.call_args.args + assert "legacy" in message + assert (request_index, position_count, code_count) == (0, input_ids.count(-100), 2) + + +@pytest.mark.parametrize( + "positions,expected", + [([0, 1, 2, 3], [0.0, 5.0, 7.0, 0.0]), ([2], [7.0]), ([0, 3], [0.0, 0.0])], + ids=["full-prefill", "chunk-intersects-reference", "chunk-outside-reference"], +) +def test_valid_explicit_positions_do_not_log_mismatch( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, + mocker: MockerFixture, + positions: list[int], + expected: list[float], +) -> None: + debug = mocker.spy(mod.logger, "debug") + talker._last_step_query_start_loc = torch.tensor([0, len(positions)]) + hidden = torch.zeros(len(positions), 1) + result = talker._apply_ref_audio_substitution( + hidden, + torch.ones(len(positions), dtype=torch.long), + torch.tensor(positions), + [{"audio_input_ids": torch.tensor([[5], [7]]), "audio_placeholder_positions": torch.tensor([1, 2])}], + ) + + assert result[:, 0].tolist() == expected + debug.assert_not_called() + assert torch.count_nonzero(hidden) == 0 + + +def test_matching_mask_filters_codes_and_positions_without_diagnostic( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, mocker: MockerFixture +) -> None: + debug = mocker.spy(mod.logger, "debug") + result = talker._apply_ref_audio_substitution( + torch.zeros(4, 1), + torch.ones(4, dtype=torch.long), + torch.arange(4), + [ + { + "audio_input_ids": torch.tensor([[5], [7]]), + "audio_input_ids_mask": torch.tensor([True, False]), + "audio_placeholder_positions": torch.tensor([1, 2]), + } + ], + ) + assert result[:, 0].tolist() == [0.0, 5.0, 0.0, 0.0] + debug.assert_not_called() + + +@pytest.mark.parametrize("masked", [False, True], ids=["empty-codes", "all-masked"]) +def test_empty_reference_does_not_log_mismatch( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, mocker: MockerFixture, masked: bool +) -> None: + debug = mocker.spy(mod.logger, "debug") + info = { + "audio_input_ids": torch.tensor([[5], [7]]) if masked else torch.empty((0, 1), dtype=torch.long), + "audio_input_ids_mask": torch.tensor([False, False] if masked else [], dtype=torch.bool), + "audio_placeholder_positions": torch.tensor([1, 2] if masked else [], dtype=torch.long), + } + hidden = torch.zeros(4, 1) + result = talker._apply_ref_audio_substitution(hidden, torch.ones(4, dtype=torch.long), torch.arange(4), [info]) + assert result is hidden + debug.assert_not_called() + + +def test_legacy_decode_does_not_log_mismatch( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, mocker: MockerFixture +) -> None: + debug = mocker.spy(mod.logger, "debug") + talker._last_step_query_start_loc = torch.tensor([0, 1]) + hidden = torch.zeros(1, 1) + result = talker._apply_ref_audio_substitution( + hidden, torch.tensor([-100]), torch.tensor([2]), [{"audio_input_ids": torch.tensor([[5], [7]])}] + ) + assert result is hidden + debug.assert_not_called() + + +@pytest.mark.parametrize("input_ids", [[1, -100, -100, 2], [-100, -100, -100, 2]]) +def test_legacy_sufficient_slots_remain_supported( + talker: mod.HiggsAudioV3TalkerForConditionalGeneration, mocker: MockerFixture, input_ids: list[int] +) -> None: + debug = mocker.spy(mod.logger, "debug") + hidden = torch.zeros(4, 1) + result = talker._apply_ref_audio_substitution( + hidden, torch.tensor(input_ids), torch.arange(4), [{"audio_input_ids": torch.tensor([[5], [7]])}] + ) + targets = [i for i, token in enumerate(input_ids) if token == -100][:2] + assert result[targets, 0].tolist() == [5.0, 7.0] + assert torch.count_nonzero(result) == 2 + assert torch.count_nonzero(hidden) == 0 + debug.assert_not_called() diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index 738e469f1..ba9c9eb2d 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -722,6 +722,13 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): if isinstance(placeholder_positions, torch.Tensor): if placeholder_positions.numel() == 0 or placeholder_positions.numel() != codes.shape[0]: + logger.debug( + "Skipping reference audio substitution for request %d: " + "%d explicit placeholder positions do not match %d reference code rows", + i, + placeholder_positions.numel(), + codes.shape[0], + ) continue request_positions = flat_positions[s:e] ref_positions = placeholder_positions.to( @@ -743,6 +750,13 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): placeholders = span_mask.nonzero(as_tuple=True)[0] n_codes = int(codes.shape[0]) if int(placeholders.numel()) < n_codes: + logger.debug( + "Skipping reference audio substitution for request %d: " + "legacy span contains %d placeholders for %d reference code rows", + i, + placeholders.numel(), + n_codes, + ) continue target = placeholders[:n_codes] + s From 9f85b3722e91677615e2fff899ed81856e36cca5 Mon Sep 17 00:00:00 2001 From: Ethan ZHU Date: Tue, 15 Sep 2026 06:25:00 +0800 Subject: [PATCH 069/209] [Model] Stop MiniCPM-o Thinker at TTS boundaries (#7463) Signed-off-by: Zhu Shihao Co-authored-by: Zhu Shihao --- docs/configuration/stage_configs.md | 4 +- tests/config/test_config_factory.py | 47 ++++++++++++++++++- tests/entrypoints/test_omni_entrypoints.py | 24 ++++++---- .../models/minicpmo_4_5/test_llm2tts.py | 13 ++--- .../models/minicpmo_4_5/test_pipeline.py | 4 ++ vllm_omni/config/omni_config.py | 9 ++-- vllm_omni/config/stage_config.py | 21 +++++++-- vllm_omni/entrypoints/omni_base.py | 13 +++-- .../models/minicpmo_4_5/pipeline.py | 8 +++- 9 files changed, 113 insertions(+), 30 deletions(-) diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 855390739..07beef2de 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -38,7 +38,7 @@ Common `StagePipelineConfig` fields include: | `custom_process_input_func` | Processor applied to this stage's incoming payload. | | `custom_process_next_stage_input_func` | Processor used for full-payload handoff to the next stage. | | `async_chunk_process_next_stage_input_func` | Processor used for async chunk handoff. | -| `sampling_constraints` | Model-owned sampling constraints that deploy defaults cannot override. | +| `sampling_constraints` | Model-owned sampling constraints that deploy defaults cannot override. Scalar values replace deploy defaults; required `stop_token_ids` extend and deduplicate them. | To add or change topology, define and register a new pipeline variant. Use deploy YAML only for runtime placement, resource sizing, connectors, and other @@ -97,7 +97,7 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne | `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. | | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_`; values are names registered under top-level `connectors:`. | | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_`; values are names registered under top-level `connectors:`. | -| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). | +| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Merged with pipeline `sampling_constraints`; scalar constraints win, while required `stop_token_ids` are appended and deduplicated. | | `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. | **Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index af898ae50..82c1fb32e 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1332,6 +1332,10 @@ class TestDeployConfigLoading: assert deploy.duplex_session.max_sessions == max_sessions assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] + assert stages[0].yaml_extras["default_sampling_params"]["stop_token_ids"] == [ + 151704, + 151645, + ] def test_load_minicpmo_default_deploy_config(self): deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml")) @@ -3212,7 +3216,7 @@ class TestSentinelDefaultPrecedence: class TestSamplingConstraintsPrecedence: - """Test that pipeline sampling_constraints override deploy defaults.""" + """Test scalar constraint precedence and additive required stop tokens.""" def test_constraints_win(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) @@ -3232,6 +3236,47 @@ class TestSamplingConstraintsPrecedence: # Pipeline says stop_token_ids=[2150] for talker assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150] + def test_required_stop_tokens_extend_yaml_defaults(self, tmp_path): + pipeline = PipelineConfig( + model_type="required_stop_tokens", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + final_output=True, + sampling_constraints={ + "detokenize": True, + "stop_token_ids": [151704, 151645], + }, + ), + ), + ) + deploy_path = tmp_path / "required-stop-tokens.yaml" + deploy_path.write_text( + """ +stages: + - stage_id: 0 + default_sampling_params: + max_tokens: 7 + stop_token_ids: [100, 151704] +""", + encoding="utf-8", + ) + + legacy_stage = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))[0] + structured_stage = VllmOmniConfig.from_pipeline_config( + pipeline, + deploy_config_path=str(deploy_path), + ).stage_by_id(0) + expected = { + "max_tokens": 7, + "detokenize": True, + "stop_token_ids": [100, 151704, 151645], + } + + assert legacy_stage.yaml_extras["default_sampling_params"] == expected + assert structured_stage.model_config.default_sampling_params == expected + class TestPipelineConfigResolvers: @pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)]) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 6fe04fdb7..be7a0ffac 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -200,20 +200,28 @@ def _make_base(): return obj -def test_resolve_sampling_params_list_preserves_stage_constraints(): +def test_resolve_sampling_params_list_merges_required_stop_tokens(): base = _make_base() base.engine.num_stages = 1 - base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])] + required_stop_ids = [151704, 151645] + base.default_sampling_params_list = [ + SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[99, *required_stop_ids]) + ] base.engine.stage_configs = [ StageConfig( stage_id=0, model_stage="dummy-model", - sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + sampling_constraints={"detokenize": False, "stop_token_ids": required_stop_ids}, ).to_omegaconf() ] base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) - assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}] - caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7]) + assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}] + + resolved_defaults = base.resolve_sampling_params_list(None) + + assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645] + + caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704]) resolved = base.resolve_sampling_params_list(caller_params) @@ -221,10 +229,10 @@ def test_resolve_sampling_params_list_preserves_stage_constraints(): assert resolved[0].seed == 1234 assert resolved[0].max_tokens == 7 assert resolved[0].detokenize is False - assert resolved[0].stop_token_ids == [42] - assert 42 in resolved[0]._all_stop_token_ids + assert resolved[0].stop_token_ids == [100, 151704, 151645] + assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids) assert caller_params.detokenize is True - assert caller_params.stop_token_ids == [7] + assert caller_params.stop_token_ids == [100, 151704] @pytest.mark.parametrize("use_defaults", [False, True]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py index 3fd4867cb..d93755fec 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py +++ b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py @@ -9,7 +9,7 @@ Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2t - both inputs missing -> raises - structured model_intermediate_buffer carries the thinker ids and text - scheduler prompt tokens follow the selected TTS region or output tokens - - MiniCPM-o 4.5 TTS region detection on 151703 / 151704 tokens + - MiniCPM-o 4.5 TTS region detection on 151703 / 151704 / 151645 tokens - plain chat without TTS markers conditions on the generated assistant span - prompt arg is normalized to a list and ``multi_modal_data`` is gated by ``requires_multimodal_data`` @@ -260,13 +260,14 @@ class TestTtsRegionDetection: ) return result[0]["model_intermediate_buffer"], hidden - def test_4_5_markers_detected(self) -> None: + @pytest.mark.parametrize("tts_end_id", [151704, 151645]) + def test_4_5_markers_detected(self, tts_end_id: int) -> None: # prompt: [10, 11] - # output: [151703, 30, 31, 151704, 40] - # full sequence: [10, 11, 151703, 30, 31, 151704, 40] + # output: [151703, 30, 31, end, 40] + # full sequence: [10, 11, 151703, 30, 31, end, 40] # 0 1 2 3 4 5 6 - # 4.5 BOS at idx 2 -> slice starts at 3; EOS at idx 5 -> slice ends at 5. - buffer, hidden = self._run([10, 11], [151703, 30, 31, 151704, 40]) + # 4.5 BOS at idx 2 -> slice starts at 3; end at idx 5 -> slice ends at 5. + buffer, hidden = self._run([10, 11], [151703, 30, 31, tts_end_id, 40]) assert buffer["ids"]["tts"] == [30, 31] assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index f8fdfdcd2..e12010ce4 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -91,6 +91,10 @@ class TestPipelineTopology: assert thinker.final_output_type == "text" assert thinker.owns_tokenizer is True assert thinker.requires_multimodal_data is True + assert thinker.sampling_constraints == { + "detokenize": True, + "stop_token_ids": [151704, 151645], + } def test_talker_stage(self, pipeline: PipelineConfig) -> None: talker = pipeline.get_stage(1) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index b58887099..0b8723c0e 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -47,6 +47,7 @@ from vllm_omni.config.stage_config import ( _select_processor_funcs, build_stage_runtime_overrides, load_deploy_config, + merge_sampling_constraints, normalize_pipeline_cli_overrides, reconcile_diffusion_attention_overrides, ) @@ -1347,10 +1348,10 @@ def _stage_sampling_params( stage_deploy: StageDeployConfig | None, topology: StagePipelineConfig, ) -> dict[str, Any] | None: - sampling: dict[str, Any] = {} - if stage_deploy is not None and stage_deploy.default_sampling_params: - sampling.update(_copy_value(stage_deploy.default_sampling_params)) - sampling.update(_copy_value(topology.sampling_constraints)) + sampling = merge_sampling_constraints( + _copy_value(stage_deploy.default_sampling_params) if stage_deploy is not None else None, + _copy_value(topology.sampling_constraints), + ) return sampling or None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 063a95b0e..9c9186c16 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -970,16 +970,29 @@ def _build_engine_args( return engine_args +def merge_sampling_constraints( + sampling_params: Mapping[str, Any] | None, + constraints: Mapping[str, Any], +) -> dict[str, Any]: + """Apply scalar constraints and extend stop tokens without mutating inputs.""" + resolved_constraints = dict(constraints) + if "stop_token_ids" in resolved_constraints: + caller_stop_ids = (sampling_params or {}).get("stop_token_ids") or [] + required_stop_ids = resolved_constraints["stop_token_ids"] or [] + resolved_constraints["stop_token_ids"] = list(dict.fromkeys([*caller_stop_ids, *required_stop_ids])) + return {**(sampling_params or {}), **resolved_constraints} + + def _build_extras( ps: StagePipelineConfig, ds: StageDeployConfig | None, ) -> dict[str, Any]: """Assemble ``yaml_extras`` (sampling + connectors + pipeline extras).""" extras: dict[str, Any] = {} - sampling: dict[str, Any] = {} - if ds is not None and ds.default_sampling_params: - sampling.update(ds.default_sampling_params) - sampling.update(ps.sampling_constraints) + sampling = merge_sampling_constraints( + ds.default_sampling_params if ds is not None else None, + ps.sampling_constraints, + ) if sampling: extras["default_sampling_params"] = sampling if ds is not None and ds.default_pooling_params: diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 899ff116d..ead0e6f1d 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -18,6 +18,7 @@ from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.stage_config import merge_sampling_constraints from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import ( EngineQueueMessage, @@ -372,12 +373,12 @@ class OmniBase(PDDisaggregationMixin): @staticmethod def _apply_sampling_constraints(params: Any, constraints: Mapping[str, Any]) -> Any: - """Rebuild params with pipeline-required settings without mutating caller input.""" + """Apply pipeline requirements, merging required stops with caller stops.""" if not constraints: return params if isinstance(params, Mapping): - return {**params, **constraints} - if is_dataclass(params): + values = dict(params) + elif is_dataclass(params): values = {field.name: getattr(params, field.name) for field in fields(params) if field.init} elif struct_fields := getattr(params, "__struct_fields__", None): values = { @@ -387,7 +388,11 @@ class OmniBase(PDDisaggregationMixin): } else: raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}") - return type(params)(**{**values, **constraints}) + + resolved = merge_sampling_constraints(values, constraints) + if isinstance(params, Mapping): + return resolved + return type(params)(**resolved) def _record_request_failure_once(self, request_id: str, reason: str) -> None: req_state = self.request_states.get(request_id) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 0ca61a346..e48328d0b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -53,7 +53,13 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( owns_tokenizer=True, requires_multimodal_data=True, engine_output_type="latent", - sampling_constraints={"detokenize": True}, + sampling_constraints={ + "detokenize": True, + # The llm2tts bridge discards this boundary and every row + # after it, so stop Stage 0 as soon as either valid + # MiniCPM-o 4.5 turn terminator is sampled. + "stop_token_ids": [151704, 151645], + }, ), StagePipelineConfig( stage_id=1, From c68eadbbe002c422be8be2297c5a68e645652531 Mon Sep 17 00:00:00 2001 From: kunkun <72174834+kunkunblueberry@users.noreply.github.com> Date: Tue, 15 Sep 2026 06:25:08 +0800 Subject: [PATCH 070/209] =?UTF-8?q?[Performance]=20Optimize=20MammothModa2?= =?UTF-8?q?=20AR=20=E2=86=92=20DiT=20hidden-state=20transfer=20(#7102)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: kunkunblueberry <1833921874@qq.com> Signed-off-by: kunkun <72174834+kunkunblueberry@users.noreply.github.com> --- .../omni_connectors/test_basic_connectors.py | 12 ++ .../test_mammoth_moda2.py | 185 ++++++++++++++++++ .../pipeline_mammothmoda2_dit.py | 10 +- .../stage_input_processors/mammoth_moda2.py | 10 +- 4 files changed, 212 insertions(+), 5 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_mammoth_moda2.py diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 787d5594b..2cd0fcea3 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -33,6 +33,18 @@ def test_tensor_serialization(): assert torch.equal(tensor, deserialized) +def test_bfloat16_tensor_serialization(): + """The raw tensor serializer preserves the Mammoth BF16 wire dtype.""" + import torch + + tensor = torch.tensor([[1.0, -2.5], [3.25, 4.5]], dtype=torch.bfloat16) + serialized = OmniSerializer.serialize({"text_prompt_embeds": tensor}) + deserialized = OmniSerializer.deserialize(serialized) + + assert deserialized["text_prompt_embeds"].dtype == torch.bfloat16 + assert torch.equal(deserialized["text_prompt_embeds"], tensor) + + def test_noncontiguous_tensor_serialization(): """Test torch.Tensor serialization for strided views.""" import torch diff --git a/tests/model_executor/stage_input_processors/test_mammoth_moda2.py b/tests/model_executor/stage_input_processors/test_mammoth_moda2.py new file mode 100644 index 000000000..992908c36 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_mammoth_moda2.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU regression tests for MammothModa2's completed-AR to DiT bridge.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( + MammothModa2DiTPipeline, +) +from vllm_omni.engine.serialization import ( + deserialize_additional_information, + serialize_additional_information, +) +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.model_executor.models.mammoth_moda2.pipeline import MAMMOTH_MODA2_PIPELINE +from vllm_omni.model_executor.stage_input_processors.mammoth_moda2 import ar2dit + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_PROMPT_TOKEN_IDS = [7, 8, 9] +_GENERATED_TOKEN_IDS = [101, 102, 103] + + +def _hidden_states(rows: int = 5, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + return torch.arange(rows * 4, dtype=dtype).reshape(rows, 4) + + +def _ar_output(hidden_states: torch.Tensor | None = None) -> SimpleNamespace: + multimodal_output = {} if hidden_states is None else {"latent": hidden_states} + return SimpleNamespace( + request_id="mammoth-request", + prompt_token_ids=list(_PROMPT_TOKEN_IDS), + outputs=[ + SimpleNamespace( + cumulative_token_ids=list(_GENERATED_TOKEN_IDS), + multimodal_output=multimodal_output, + ) + ], + ) + + +def _prompt( + *, + target_h: int | None = None, + target_w: int | None = None, +) -> dict[str, object]: + mm_processor_kwargs: dict[str, int] = {} + if target_h is not None: + mm_processor_kwargs["target_h"] = target_h + if target_w is not None: + mm_processor_kwargs["target_w"] = target_w + return { + "additional_information": { + "image_height": [384], + "image_width": [640], + "text_guidance_scale": [4.5], + "cfg_range": [0.1, 0.9], + "num_inference_steps": [20], + }, + "mm_processor_kwargs": mm_processor_kwargs, + } + + +def _assert_complete_dit_input(dit_inputs: list[dict[str, object]], hidden_states: torch.Tensor) -> None: + assert len(dit_inputs) == 1 + dit_input = dit_inputs[0] + assert dit_input["prompt_token_ids"] == [0] + + additional_information = dit_input["additional_information"] + assert isinstance(additional_information, dict) + full_hidden_states = additional_information["full_hidden_states"] + assert isinstance(full_hidden_states, torch.Tensor) + assert full_hidden_states.dtype == torch.bfloat16 + assert full_hidden_states.is_contiguous() + assert torch.equal(full_hidden_states, hidden_states) + assert additional_information["full_token_ids"] == [7, 8, 9, 101, 102] + assert additional_information["answer_start_index"] == [3] + assert additional_information["image_height"] == [384] + assert additional_information["image_width"] == [640] + assert additional_information["text_guidance_scale"] == [4.5] + assert additional_information["cfg_range"] == [0.1, 0.9] + assert additional_information["num_inference_steps"] == [20] + + +def test_ar2dit_builds_one_complete_dit_input() -> None: + hidden_states = _hidden_states() + + dit_inputs = ar2dit([_ar_output(hidden_states)], _prompt()) + + _assert_complete_dit_input(dit_inputs, hidden_states) + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_ar2dit_preserves_low_precision_through_engine_core_payload(dtype: torch.dtype) -> None: + hidden_states = _hidden_states(dtype=dtype) + dit_input = ar2dit([_ar_output(hidden_states)], _prompt())[0] + + wire_payload = serialize_additional_information(dit_input["additional_information"]) + assert wire_payload is not None + restored = deserialize_additional_information(wire_payload) + restored_hidden_states = restored["full_hidden_states"] + + assert isinstance(restored_hidden_states, torch.Tensor) + assert restored_hidden_states.dtype == dtype + assert restored_hidden_states.is_contiguous() + assert torch.equal(restored_hidden_states, hidden_states) + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_dit_condition_split_preserves_compact_transfer_dtype(dtype: torch.dtype) -> None: + # Construct only the config read by _split_ar_conditions. This keeps the + # test CPU-only and avoids loading the DiT/VAE weights. + pipeline = object.__new__(MammothModa2DiTPipeline) + object.__setattr__( + pipeline, + "config", + SimpleNamespace( + llm_config=SimpleNamespace(gen_vocab_start_index=100), + image_token_id=20, + video_token_id=21, + vision_start_token_id=22, + vision_end_token_id=23, + ), + ) + hidden_states = _hidden_states(dtype=dtype) + + text_cond, image_cond = pipeline._split_ar_conditions( + full_hidden_states=hidden_states, + full_token_ids=[7, 20, 8, 101, 102], + answer_start_index=3, + ) + + assert text_cond.dtype == dtype + assert image_cond.dtype == dtype + assert text_cond.is_contiguous() + assert image_cond.is_contiguous() + assert torch.equal(text_cond, hidden_states[[0, 2]]) + assert torch.equal(image_cond, hidden_states[[3, 4]]) + + +def test_ar2dit_prefers_processor_image_dimensions() -> None: + hidden_states = _hidden_states() + + dit_inputs = ar2dit([_ar_output(hidden_states)], _prompt(target_h=512, target_w=768)) + + additional_information = dit_inputs[0]["additional_information"] + assert additional_information["image_height"] == [512] + assert additional_information["image_width"] == [768] + + +def test_ar2dit_rejects_missing_latent_output() -> None: + with pytest.raises(ValueError, match="missing latent multimodal output"): + ar2dit([_ar_output()], _prompt()) + + +def test_ar2dit_rejects_hidden_token_row_mismatch() -> None: + with pytest.raises(AssertionError, match="Hidden states length mismatch"): + ar2dit([_ar_output(_hidden_states(rows=4))], _prompt()) + + +def test_mammoth_pipeline_uses_standard_completed_ar_forwarding() -> None: + stage0, stage1 = MAMMOTH_MODA2_PIPELINE.stages + + assert stage0.custom_process_next_stage_input_func is None + assert stage1.custom_process_input_func.endswith(".ar2dit") + assert stage1.sync_process_input_func is None + assert stage1.requires_full_payload_input is False + + +def test_stage_client_forwards_completed_ar_output_to_mammoth_adapter() -> None: + hidden_states = _hidden_states() + client = object.__new__(StageEngineCoreClient) + client.custom_process_input_func = ar2dit + client.requires_multimodal_data = False + + dit_inputs = client.process_engine_inputs( + [_ar_output(hidden_states)], + _prompt(), + ) + + _assert_complete_dit_input(dit_inputs, hidden_states) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index 5298d5514..c96e5548c 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations from collections.abc import Iterable @@ -191,8 +193,12 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): text_mask = questions_mask & ~(visual_token_mask | gen_token_mask) image_mask = answers_mask & gen_token_mask - text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous() - image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous() + # Keep the transferred representation compact until the final device/dtype + # conversion in ``forward``. The masks are dtype-independent, and an + # unconditional host-side float32 expansion here would otherwise double both + # the selected-condition staging footprint and BF16 H2D traffic. + text_cond = full_hidden_states[text_mask].contiguous() + image_cond = full_hidden_states[image_mask].contiguous() return text_cond, image_cond @torch.inference_mode() diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py index 7b23f8e49..fab22b739 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage input processor for MammothModa2 (AR -> DiT).""" from collections.abc import Mapping @@ -91,9 +93,11 @@ def ar2dit( # from the model config. Pass through the raw AR hidden states + token ids and # the question/answer boundary so the pipeline can reconstruct the masks. additional_information = { - # float32 so the tensor crosses the stage boundary (the serializer uses - # numpy, which has no bf16); the DiT re-casts to the model dtype. - "full_hidden_states": full_hidden_states.float().contiguous(), + # The EngineCore payload serializer preserves bfloat16 as raw bytes. + # Keep the AR representation compact across the CPU/IPC boundary; + # the DiT selects its conditioning rows and casts them only when it + # consumes them. + "full_hidden_states": full_hidden_states.contiguous(), "full_token_ids": full_token_ids, "answer_start_index": [len(prompt_token_ids)], "image_height": [int(image_height)], From 8e088c7725462634f6b8cf0981d24b20a37bc900 Mon Sep 17 00:00:00 2001 From: yulun <100981785+Big2Wheel@users.noreply.github.com> Date: Tue, 15 Sep 2026 06:25:17 +0800 Subject: [PATCH 071/209] [Hardware][Ascend][Model] Support MOSS-TTS-Nano on Ascend A2 (#7192) Signed-off-by: Big2Wheel <1404207590@qq.com> --- docs/models/supported_models.md | 2 +- recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md | 128 ++++++++++++++++++ recipes/README.md | 1 + .../test_moss_tts_nano_expansion.py | 8 +- .../test_moss_tts_nano_expansion.py | 28 +++- .../models/moss_tts_nano/test_npu_compat.py | 32 +++++ vllm_omni/deploy/moss_tts_nano.yaml | 7 +- .../moss_tts_nano/modeling_moss_tts_nano.py | 40 ++++-- 8 files changed, 224 insertions(+), 22 deletions(-) create mode 100644 recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md create mode 100644 tests/model_executor/models/moss_tts_nano/test_npu_compat.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a60eb6670..f8db9eb92 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -87,7 +87,7 @@ th { | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — | | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — | -| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — | +| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md) | | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — | | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — | | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — | diff --git a/recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md b/recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md new file mode 100644 index 000000000..5d012fbb4 --- /dev/null +++ b/recipes/OpenMOSS/MOSS-TTS-Nano-NPU.md @@ -0,0 +1,128 @@ +# MOSS-TTS-Nano on Ascend NPU + +## Summary + +- Model: `OpenMOSS-Team/MOSS-TTS-Nano` +- Task: multilingual text-to-speech with a reference voice +- Hardware: one Ascend 910B3 64GB NPU +- Output: 48 kHz mono audio +- Serving API: OpenAI-compatible `/v1/audio/speech` + +MOSS-TTS-Nano combines a small autoregressive language model with +`OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano`. Both components run on the same +NPU. The codec is a separate checkpoint and must be available even when the +main model has already been downloaded. + +## Environment + +Use an aligned vLLM, vLLM-Ascend, and vLLM-Omni environment. The following +configuration was validated: + +- Image: `quay.nju.edu.cn/ascend/vllm-omni:v0.28.0` +- Ascend device: 910B3 64GB +- CANN: 9.1.0 +- vLLM: 0.28.0 +- Execution mode: eager + +See the [NPU installation guide](../../docs/getting_started/installation/npu.md) +for the complete container device and driver mounts. + +## Download checkpoints + +```bash +export HF_HOME=/models/huggingface-cache + +hf download OpenMOSS-Team/MOSS-TTS-Nano + +# This dependency is loaded from its model ID by the MOSS-TTS-Nano config. +hf download OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano +``` + +For an air-gapped server, populate `HF_HOME` with the audio-tokenizer snapshot +before startup and set `HF_HUB_OFFLINE=1`. Do not omit the audio tokenizer: +the main model directory does not contain its weights. + +## Start the server + +Expose one physical NPU and keep the stage's logical device as `0`: + +```bash +export ASCEND_RT_VISIBLE_DEVICES=0 +export HF_HOME=/models/huggingface-cache + +vllm serve OpenMOSS-Team/MOSS-TTS-Nano \ + --omni \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 8091 +``` + +`--trust-remote-code` is required because both checkpoints provide custom +Transformers model implementations. The bundled `moss_tts_nano.yaml` deploy +configuration enables eager execution; ACL graph capture is not used for this +model path. + +## Voice-cloning request + +`ref_audio` accepts an HTTP URL or a base64 data URL. A non-streaming request +returns a WAV file: + +```bash +curl http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenMOSS-Team/MOSS-TTS-Nano", + "input": "Hello, this is a voice cloning test on Ascend NPU.", + "stream": false, + "response_format": "wav", + "ref_audio": "https://example.com/reference.wav" + }' \ + --output output.wav +``` + +Streaming raw PCM uses the same endpoint: + +```bash +curl http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenMOSS-Team/MOSS-TTS-Nano", + "input": "你好，这是一段流式语音合成测试。", + "stream": true, + "stream_format": "audio", + "response_format": "pcm", + "ref_audio": "https://example.com/reference.wav" + }' \ + --output output.pcm +``` + +Raw PCM output is signed 16-bit little-endian mono at 48 kHz. + +## Validation + +The Ascend 910B3 validation covered: + +- English non-streaming WAV output; +- Chinese streaming PCM output; +- repeat generation with the same seed producing byte-identical WAV files; +- two concurrent non-streaming requests; +- waveform checks for a non-empty, finite, non-silent signal at 48 kHz. + +The first request includes one-time NPU kernel warmup. In the validated setup, +the first non-streaming request took about 22.7 seconds, a repeated request +took about 7.5 seconds, and streaming produced its first audio chunk in about +0.38 seconds. These figures are functional reference measurements, not a +performance guarantee. + +## Current limitations + +- The model runs in eager mode; ACL graph mode has not been enabled. +- Concurrent HTTP requests are accepted but generated serially. Do not raise + the stage's `max_num_seqs` above `1`: the remote model and audio tokenizer + share RNG and streaming decode state, and interleaved generation can corrupt + audio. +- The remote model's CUDA-only `flash_attention_2` path is replaced with + PyTorch SDPA on Ascend. +- The current non-CUDA model wrapper loads the language model and audio codec + in float32. Quantized checkpoints have not been validated. +- A reference audio clip is required for `voice_clone` mode. diff --git a/recipes/README.md b/recipes/README.md index 86eba6d87..33e29c086 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -74,6 +74,7 @@ recipes/ | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB | | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB | +| [`OpenMOSS/MOSS-TTS-Nano-NPU.md`](./OpenMOSS/MOSS-TTS-Nano-NPU.md) | Online multilingual voice-cloned TTS | 1x Ascend 910B3 64GB | | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB | | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB | | [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk | diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index 41e2229d3..3d9dafaf9 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -121,7 +121,7 @@ def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: @pytest.mark.advanced_model -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "npu": "A2"}) def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None: """English TTS produces non-empty 48 kHz stereo audio.""" req = _build_request("Hello, this is a short voice cloning demo for testing.", ref_audio_path) @@ -133,7 +133,7 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None: @pytest.mark.advanced_model -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "npu": "A2"}) def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None: """Chinese TTS produces non-empty audio.""" req = _build_request("你好，这是语音合成测试。", ref_audio_path) @@ -145,7 +145,7 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None: @pytest.mark.advanced_model -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "npu": "A2"}) def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None: """Same seed produces identical waveforms.""" req = _build_request("Reproducible output test.", ref_audio_path, seed=123) @@ -157,7 +157,7 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> @pytest.mark.advanced_model -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "npu": "A2"}) def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None: """Batch of two requests returns audio for each.""" requests = [ diff --git a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py index 3ea8c91ea..cf623b203 100644 --- a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py @@ -84,7 +84,7 @@ tts_server_params = [ ] -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -106,7 +106,7 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_002(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -138,7 +138,7 @@ def test_text_to_audio_002(omni_server, online_client, ref_audio_data_url) -> No online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_003(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -160,7 +160,7 @@ def test_text_to_audio_003(omni_server, online_client, ref_audio_data_url) -> No online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_004_ref_text_ignored(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -181,3 +181,23 @@ def test_text_to_audio_004_ref_text_ignored(omni_server, online_client, ref_audi } online_client.send_audio_speech_request(request_config) + + +@hardware_test(res={"cuda": "L4", "npu": "A2"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_005_concurrent_requests_are_serialized(omni_server, online_client, ref_audio_data_url) -> None: + """Concurrent clients produce identical audio when generation is serialized.""" + request_config = { + "model": omni_server.model, + "input": "Concurrent clients should receive clean and deterministic speech.", + "stream": False, + "response_format": "wav", + "ref_audio": ref_audio_data_url, + "seed": 42, + } + + responses = online_client.send_audio_speech_request(request_config, request_num=2) + + assert len(responses) == 2 + assert responses[0].audio_bytes + assert responses[0].audio_bytes == responses[1].audio_bytes diff --git a/tests/model_executor/models/moss_tts_nano/test_npu_compat.py b/tests/model_executor/models/moss_tts_nano/test_npu_compat.py new file mode 100644 index 000000000..59882edc1 --- /dev/null +++ b/tests/model_executor/models/moss_tts_nano/test_npu_compat.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts_nano.modeling_moss_tts_nano import ( + _configure_attention_implementation, + _validate_max_num_seqs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_non_cuda_attention_uses_sdpa(): + class FakeModel: + attention_implementation = None + + def _set_attention_implementation(self, implementation): + self.attention_implementation = implementation + + model = FakeModel() + _configure_attention_implementation(model, torch.device("cpu")) + + assert model.attention_implementation == "sdpa" + + +def test_moss_tts_nano_rejects_concurrent_active_sequences(): + _validate_max_num_seqs(1) + + with pytest.raises(ValueError, match="requires max_num_seqs=1"): + _validate_max_num_seqs(2) diff --git a/vllm_omni/deploy/moss_tts_nano.yaml b/vllm_omni/deploy/moss_tts_nano.yaml index e9eab184b..9a7e3fa0b 100644 --- a/vllm_omni/deploy/moss_tts_nano.yaml +++ b/vllm_omni/deploy/moss_tts_nano.yaml @@ -6,8 +6,9 @@ # * enforce_eager=true — the model uses upstream inference_stream() with # trust_remote_code; CUDA graph capture is not wired up. # * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB. -# * max_num_seqs=4 — supports per-request streaming generators stored in -# self._stream_gens. +# * max_num_seqs=1 — the remote model and audio tokenizer keep process-wide +# RNG and streaming decode state. HTTP requests may arrive concurrently, +# but generation must remain serial to avoid cross-request audio corruption. # # MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling # bypasses the dummy-MM pass at profile time. @@ -16,7 +17,7 @@ trust_remote_code: true stages: - stage_id: 0 - max_num_seqs: 4 + max_num_seqs: 1 gpu_memory_utilization: 0.3 enforce_eager: true enable_prefix_caching: false diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index a0b31bdf2..49c254deb 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MOSS-TTS-Nano single-stage model for vLLM-Omni. Runs in a single AR worker stage. The 0.1B AR LM and the @@ -160,6 +160,33 @@ def _pick(info: dict, key: str, default): return val if val is not None else default +def _configure_attention_implementation(lm: nn.Module, device: torch.device) -> None: + if device.type == "cuda": + try: + import flash_attn # noqa: F401 + + lm._set_attention_implementation("flash_attention_2") + logger.info("MOSS-TTS-Nano using flash_attention_2") + except ImportError: + lm._set_attention_implementation("sdpa") + logger.info("MOSS-TTS-Nano using sdpa (flash_attn not installed)") + else: + # The checkpoint defaults to flash_attention_2, whose external + # flash_attn package is CUDA-only. PyTorch SDPA is dispatched to the + # native backend on non-CUDA accelerators such as Ascend NPU. + lm._set_attention_implementation("sdpa") + logger.info("MOSS-TTS-Nano using sdpa on %s", device.type) + + +def _validate_max_num_seqs(max_num_seqs: int) -> None: + if max_num_seqs != 1: + raise ValueError( + "MOSS-TTS-Nano currently requires max_num_seqs=1 because its " + "remote model and audio tokenizer share RNG and streaming decode " + "state across requests. Higher values can corrupt generated audio." + ) + + class MossTTSNanoForGeneration(nn.Module): """Single-stage MOSS-TTS-Nano model with streaming audio output. @@ -179,6 +206,7 @@ class MossTTSNanoForGeneration(nn.Module): self.vllm_config = vllm_config self.config = vllm_config.model_config.hf_config self.model_path: str = vllm_config.model_config.model + _validate_max_num_seqs(vllm_config.scheduler_config.max_num_seqs) # Eager construction (not in load_weights) -- see module docstring. _patch_torchaudio_load() @@ -209,15 +237,7 @@ class MossTTSNanoForGeneration(nn.Module): trust_remote_code=True, torch_dtype=tts_dtype, ) - if device.type == "cuda": - try: - import flash_attn # noqa: F401 - - lm._set_attention_implementation("flash_attention_2") - logger.info("MOSS-TTS-Nano using flash_attention_2") - except ImportError: - lm._set_attention_implementation("sdpa") - logger.info("MOSS-TTS-Nano using sdpa (flash_attn not installed)") + _configure_attention_implementation(lm, device) lm.to(device=device) lm.eval() From 5f50fc3daf4e65ce777c2ecacd974f11d29d6bf4 Mon Sep 17 00:00:00 2001 From: Yash Jain Date: Mon, 14 Sep 2026 23:25:25 +0100 Subject: [PATCH 072/209] [Kernel] Fuse Ming streaming ISTFT overlap-add and normalization (#7338) Signed-off-by: Yash Jain --- benchmarks/kernels/benchmark_ming_istft.py | 218 ++++++++++++++++++ .../models/ming_tts/test_fused_istft.py | 147 ++++++++++++ .../models/common/ming/audio_dsp.py | 22 +- .../models/common/ming/fused_istft.py | 177 ++++++++++++++ 4 files changed, 563 insertions(+), 1 deletion(-) create mode 100644 benchmarks/kernels/benchmark_ming_istft.py create mode 100644 tests/model_executor/models/ming_tts/test_fused_istft.py create mode 100644 vllm_omni/model_executor/models/common/ming/fused_istft.py diff --git a/benchmarks/kernels/benchmark_ming_istft.py b/benchmarks/kernels/benchmark_ming_istft.py new file mode 100644 index 000000000..b67c9bbfc --- /dev/null +++ b/benchmarks/kernels/benchmark_ming_istft.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Measure Ming ISTFT and optionally its pretrained streaming AudioVAE decoder. + +Example (no model weights needed for the operator benchmark): + python benchmarks/kernels/benchmark_ming_istft.py --compile --iterations 100 + +Add --decoder to measure the published decoder with fixed synthetic latents. +This is a decoder-stage benchmark, not text-to-speech serving latency. +""" + +import argparse +import json +import statistics +import time +from collections.abc import Callable +from functools import partial + +import torch + +from vllm_omni.model_executor.models.common.ming.audio_dsp import ISTFT +from vllm_omni.model_executor.models.common.ming.fused_istft import fused_istft + + +def measure(operation: Callable, warmup: int, iterations: int) -> dict: + for _ in range(warmup): + operation() + torch.accelerator.synchronize() + samples = [] + for _ in range(iterations): + start = time.perf_counter() + operation() + torch.accelerator.synchronize() + samples.append((time.perf_counter() - start) * 1000) + torch.accelerator.reset_peak_memory_stats() + before = torch.accelerator.memory_allocated() + operation() + torch.accelerator.synchronize() + return { + "p50_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "stdev_ms": statistics.pstdev(samples), + "extra_peak_mib": (torch.accelerator.max_memory_allocated() - before) / 1024**2, + } + + +def native_post_fft(module, frames, audio, window): + """The native middle-chunk post-FFT sequence, isolated from cuFFT.""" + count = frames.shape[-1] + size = (count - 1) * module.hop_length + module.win_length + + def fold(value): + return torch.nn.functional.fold( + value, output_size=(1, size), kernel_size=(1, module.win_length), stride=(1, module.hop_length) + )[:, 0, 0, :] + + samples = fold(frames * module.window[None, :, None]) + samples, audio = module._buffer_process(samples, audio, 0, streaming=True) + envelope = fold(module.window.square().expand(1, count, -1).transpose(1, 2)) + envelope, window = module._buffer_process(envelope, window, 0, streaming=True) + if not (envelope > 1e-11).all(): + raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.") + return samples / envelope, audio, window + + +def assert_outputs(actual, reference): + maximum = 0.0 + for value, expected in zip(actual, reference, strict=True): + if expected is None: + assert value is None + else: + torch.testing.assert_close(value, expected, atol=5e-7, rtol=1e-5) + maximum = max(maximum, (value - expected).abs().max().item()) + return maximum + + +def benchmark_istft(args): + for dtype in (torch.float32, torch.bfloat16): + module = ISTFT(3528, 882, 3528).to(device="cuda", dtype=dtype) + for count in (16, 100, 400): + spec = torch.randn(1, 1765, count, device="cuda", dtype=torch.complex64) + _, audio, window = module.forward_native(spec, streaming=True) + frames = torch.fft.irfft(spec, 3528, dim=1) + for scope in ("post_fft", "istft"): + if scope == "post_fft": + native = partial(native_post_fft, module, frames, audio, window) + fused = partial(fused_istft, frames, module.window, 882, audio, window, True, False) + else: + native = partial(module.forward_native, spec, audio, window, streaming=True) + fused = partial(module, spec, audio, window, streaming=True) + error = assert_outputs(fused(), native()) + results = {"native": measure(native, args.warmup, args.iterations)} + if args.compile: + compiled = torch.compile(native) + compiled_output, reference = compiled(), native() + # Inductor can change BF16 envelope rounding. Report that + # difference without weakening the fused-kernel parity gate. + compiled_matches = all( + torch.allclose(value, expected, atol=5e-7, rtol=1e-5) + for value, expected in zip(compiled_output, reference, strict=True) + ) + results["compiled"] = measure(compiled, args.warmup, args.iterations) + results["compiled"]["matches_native_tolerance"] = compiled_matches + results["compiled"]["waveform_max_abs_error"] = ( + (compiled_output[0] - reference[0]).abs().max().item() + ) + results["fused"] = measure(fused, args.warmup, args.iterations) + print( + json.dumps( + { + "scope": scope, + "window_dtype": str(dtype), + "frames": count, + "max_abs_error": error, + "results": results, + } + ), + flush=True, + ) + + +def benchmark_decoder(args): + from huggingface_hub import hf_hub_download + from safetensors import safe_open + + from vllm_omni.model_executor.models.common.ming.audio_vae import Decoder + + model = "inclusionAI/Ming-omni-tts-0.5B" + revision = "9154772e7fbc585907b6237e3190790676f28975" + config_path = hf_hub_download(model, "config.json", revision=revision) + with open(config_path) as source: + config = json.load(source)["audio_tokenizer_config"] + kwargs = config["dec_kwargs"] + backbone = kwargs["backbone"] + backbone["_attn_implementation"] = "sdpa" + decoder = Decoder(backbone, kwargs["output_dim"], kwargs["latent_dim"], config["patch_size"]) + checkpoint = hf_hub_download(model, "model.safetensors", revision=revision) + prefix = "audio.decoder." + with safe_open(checkpoint, framework="pt") as weights: + state = {key.removeprefix(prefix): weights.get_tensor(key) for key in weights.keys() if key.startswith(prefix)} + decoder.load_state_dict(state, strict=True) + decoder = decoder.eval().to(device="cuda", dtype=torch.bfloat16) + chunks = [ + torch.randn(1, size, kwargs["latent_dim"], device="cuda", dtype=torch.bfloat16) for size in (4, 25, 25, 5) + ] + istft = decoder.head.istft + fused_forward = istft.forward + + def run(): + cache = state = None + outputs = [] + for index, chunk in enumerate(chunks): + output, state, cache = decoder.low_level_reconstruct( + chunk, past_key_values=cache, use_cache=True, stream_state=state, last_chunk=index == len(chunks) - 1 + ) + outputs.append(output) + return torch.cat(outputs, dim=-1) + + istft.forward = istft.forward_native + expected = run() + istft.forward = fused_forward + actual = run() + torch.testing.assert_close(actual, expected, atol=5e-7, rtol=1e-5) + # Alternate the order to expose drift when the decoder-level gain is small. + results = {"native": [], "fused": []} + for order in (("native", "fused"), ("fused", "native"), ("native", "fused")): + for name in order: + istft.forward = istft.forward_native if name == "native" else fused_forward + results[name].append(measure(run, args.warmup, args.decoder_iterations)) + print( + json.dumps( + { + "scope": "pretrained_decoder", + "model": model, + "revision": revision, + "backend": "sdpa", + "dtype": "bfloat16", + "latent_chunks": [4, 25, 25, 5], + "max_abs_error": (actual - expected).abs().max().item(), + "results": results, + } + ), + flush=True, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--compile", action="store_true") + parser.add_argument("--decoder", action="store_true") + parser.add_argument("--decoder-iterations", type=int, default=30) + args = parser.parse_args() + if args.warmup < 0 or min(args.iterations, args.decoder_iterations) < 1: + parser.error("warmup must be nonnegative and iteration counts must be positive") + torch.manual_seed(0) + print( + json.dumps( + { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "warmup": args.warmup, + "iterations": args.iterations, + "decoder_iterations": args.decoder_iterations, + } + ) + ) + with torch.inference_mode(): + benchmark_istft(args) + if args.decoder: + benchmark_decoder(args) + + +if __name__ == "__main__": + main() diff --git a/tests/model_executor/models/ming_tts/test_fused_istft.py b/tests/model_executor/models/ming_tts/test_fused_istft.py new file mode 100644 index 000000000..fa7d072d7 --- /dev/null +++ b/tests/model_executor/models/ming_tts/test_fused_istft.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from vllm.triton_utils import HAS_TRITON + +from vllm_omni.model_executor.models.common.ming.audio_dsp import ISTFT + +pytestmark = [pytest.mark.core_model, pytest.mark.tts] +cuda = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None or not HAS_TRITON, + reason="NVIDIA CUDA and Triton required", +) + + +def _assert_outputs_close(actual, expected): + for value, reference in zip(actual, expected, strict=True): + if reference is None: + assert value is None + else: + assert value.dtype == reference.dtype + torch.testing.assert_close(value, reference, atol=5e-7, rtol=1e-5) + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize("hop", [4, 320, 882]) +@pytest.mark.parametrize("window_dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("batch,frames", [(1, 2), (1, 25), (3, 7), (2, 100)]) +@pytest.mark.parametrize("mode", ["offline", "first", "middle", "last", "first_last"]) +@torch.inference_mode() +def test_fused_istft_matches_native(hop, window_dtype, batch, frames, mode, monkeypatch): + torch.manual_seed(0) + module = ISTFT(4 * hop, hop, 4 * hop).to(device="cuda", dtype=window_dtype) + # Slice a larger spectrum so irfft must handle non-contiguous inputs. + spec = torch.randn(batch, 2 * hop + 1, 2 * frames, device="cuda", dtype=torch.complex64)[..., ::2] + kwargs = {"streaming": mode != "offline", "last_chunk": mode in ("last", "first_last")} + if mode in ("middle", "last"): + _, audio, window = module.forward_native(spec, streaming=True) + kwargs.update(audio_buffer=audio, window_buffer=window) + saved_audio, saved_window = audio.clone(), window.clone() + expected = module.forward_native(spec, **kwargs) + + def no_fallback(*args, **kwargs): + pytest.fail("supported input fell back to native ISTFT") + + monkeypatch.setattr(module, "forward_native", no_fallback) + _assert_outputs_close(module(spec, **kwargs), expected) + if mode in ("middle", "last"): + torch.testing.assert_close(audio, saved_audio, rtol=0, atol=0) + torch.testing.assert_close(window, saved_window, rtol=0, atol=0) + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize("hop", [320, 882]) +@pytest.mark.parametrize("window_dtype", [torch.float32, torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_fused_istft_streaming_state(hop, window_dtype): + torch.manual_seed(1) + module = ISTFT(4 * hop, hop, 4 * hop).to(device="cuda", dtype=window_dtype) + spec = torch.randn(2, 2 * hop + 1, 40, device="cuda", dtype=torch.complex64) + # Reusing the same module for a new request must not reuse the old tails. + for _ in range(2): + actual_state = reference_state = (None, None) + actual_chunks, reference_chunks = [], [] + for start, end in [(0, 4), (4, 29), (29, 31), (31, 40)]: + chunk = spec[..., start:end] + actual = module(chunk, *actual_state, streaming=True, last_chunk=end == 40) + expected = module.forward_native(chunk, *reference_state, streaming=True, last_chunk=end == 40) + _assert_outputs_close(actual, expected) + actual_state, reference_state = actual[1:], expected[1:] + actual_chunks.append(actual[0]) + reference_chunks.append(expected[0]) + output = torch.cat(actual_chunks, dim=-1) + torch.testing.assert_close(output, torch.cat(reference_chunks, dim=-1), atol=5e-7, rtol=1e-5) + if window_dtype == torch.float32: + # Low-precision window envelopes already round differently between + # streaming and offline native execution; preserve each contract. + torch.testing.assert_close(output, module.forward_native(spec)[0], atol=5e-7, rtol=1e-5) + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize("invalid", [False, True]) +@torch.inference_mode() +def test_fused_istft_silence_and_invalid_envelope(invalid): + module = ISTFT(1280, 320, 1280).cuda() + spec = torch.zeros(2, 641, 25, device="cuda", dtype=torch.complex64) + if invalid: + module.window.zero_() + for forward in (module.forward_native, module.forward): + with pytest.raises(RuntimeError, match="window envelope underflowed"): + forward(spec) + else: + _assert_outputs_close(module(spec), module.forward_native(spec)) + assert torch.count_nonzero(module(spec)[0]) == 0 + + +@pytest.mark.parametrize( + "device", + [pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, cuda])], +) +@pytest.mark.parametrize("padding", ["same", "center"]) +def test_istft_autograd_fallback(padding, device, monkeypatch): + module = ISTFT(32, 8, 32, padding=padding).to(device) + spec = torch.randn(2, 17, 7, device=device, dtype=torch.complex64, requires_grad=True) + native = module.forward_native + calls = [] + + def tracked(*args, **kwargs): + calls.append(True) + return native(*args, **kwargs) + + monkeypatch.setattr(module, "forward_native", tracked) + result = module(spec) + assert calls == [True] + reference = native(spec) + if padding == "same": + _assert_outputs_close(result, reference) + result = result[0] + else: + torch.testing.assert_close(result, reference) + result.square().sum().backward() + assert spec.grad is not None and torch.isfinite(spec.grad).all() + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize( + "n_fft,hop,frames,dtype", [(32, 8, 1, torch.complex64), (30, 10, 7, torch.complex64), (32, 8, 7, torch.complex128)] +) +@torch.inference_mode() +def test_istft_unsupported_cuda_falls_back(n_fft, hop, frames, dtype, monkeypatch): + module = ISTFT(n_fft, hop, n_fft).cuda() + spec = torch.randn(1, n_fft // 2 + 1, frames, device="cuda", dtype=dtype) + native = module.forward_native + calls = [] + + def tracked(*args, **kwargs): + calls.append(True) + return native(*args, **kwargs) + + monkeypatch.setattr(module, "forward_native", tracked) + _assert_outputs_close(module(spec), native(spec)) + assert calls == [True] diff --git a/vllm_omni/model_executor/models/common/ming/audio_dsp.py b/vllm_omni/model_executor/models/common/ming/audio_dsp.py index be8cd489c..d32415352 100644 --- a/vllm_omni/model_executor/models/common/ming/audio_dsp.py +++ b/vllm_omni/model_executor/models/common/ming/audio_dsp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/istft.py import torch @@ -66,6 +66,26 @@ class ISTFT(nn.Module): window_buffer: torch.Tensor | None = None, streaming: bool = False, last_chunk: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if self.padding == "same" and spec.is_cuda and not torch.is_grad_enabled(): + from .fused_istft import fused_istft, fused_istft_supported + + if fused_istft_supported( + spec, self.window, self.n_fft, self.hop_length, audio_buffer, window_buffer, streaming + ): + frames = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") + return fused_istft( + frames, self.window, self.hop_length, audio_buffer, window_buffer, streaming, last_chunk + ) + return self.forward_native(spec, audio_buffer, window_buffer, streaming, last_chunk) + + def forward_native( + self, + spec: torch.Tensor, + audio_buffer: torch.Tensor | None = None, + window_buffer: torch.Tensor | None = None, + streaming: bool = False, + last_chunk: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: """ Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. diff --git a/vllm_omni/model_executor/models/common/ming/fused_istft.py b/vllm_omni/model_executor/models/common/ming/fused_istft.py new file mode 100644 index 000000000..95da3ff70 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/fused_istft.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Post-FFT ISTFT fusion for Ming's four-way overlapping audio frames. + +Each lane owns one output sample and gathers its contributing frames. Audio +and window tails remain unnormalized so the next streaming chunk can add them. +The low-precision window rounding follows PyTorch's square, fold, and add +boundaries; the inverse FFT and the envelope-underflow check remain unchanged. +""" + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +if HAS_TRITON: + + @triton.jit + def _istft_kernel( + frames, + window, + previous_audio, + previous_window, + output, + envelope_output, + next_audio, + next_window, + frame_stride_b, + frame_stride_w, + frame_stride_t, + audio_stride_b, + audio_stride_t, + window_stride_t, + num_frames, + raw_size, + output_start, + output_size, + hop_size: tl.constexpr, + streaming: tl.constexpr, + has_buffer: tl.constexpr, + block_size: tl.constexpr, + ): + batch = tl.program_id(0) + sample = tl.program_id(1) * block_size + tl.arange(0, block_size) + valid = sample < raw_size + last_frame = sample // hop_size + audio = tl.full((block_size,), 0, tl.float32) + envelope = tl.full((block_size,), 0, tl.float32) + window_dtype = window.dtype.element_ty + + # Match CUDA col2im's ascending frame order. Disable FMA at launch to + # preserve the separately rounded window multiplication. + for overlap in tl.static_range(3, -1, -1): + frame = last_frame - overlap + window_index = sample % hop_size + overlap * hop_size + mask = valid & (frame >= 0) & (frame < num_frames) + weight = tl.load(window + window_index, mask=mask, other=0).to(tl.float32) + value = tl.load( + frames + batch * frame_stride_b + window_index * frame_stride_w + frame * frame_stride_t, + mask=mask, + other=0, + ) + audio += value * weight + envelope += (weight * weight).to(window_dtype).to(tl.float32) + + # F.fold accumulates in FP32, then stores in the window's dtype. + envelope = envelope.to(window_dtype).to(tl.float32) + tail_size: tl.constexpr = 3 * hop_size + if has_buffer: + carry_mask = valid & (sample < tail_size) + carry_audio = tl.load( + previous_audio + batch * audio_stride_b + sample * audio_stride_t, mask=carry_mask, other=0 + ) + carry_window = tl.load(previous_window + sample * window_stride_t, mask=carry_mask, other=0) + audio += carry_audio + envelope = (envelope + carry_window.to(tl.float32)).to(window_dtype).to(tl.float32) + + output_index = sample - output_start + emit = valid & (output_index >= 0) & (output_index < output_size) + tl.store(output + batch * output_size + output_index, audio / envelope, mask=emit) + if batch == 0: + tl.store(envelope_output + output_index, envelope, mask=emit) + + if streaming: + tail_index = sample - (raw_size - tail_size) + tail_mask = valid & (tail_index >= 0) + tl.store(next_audio + batch * tail_size + tail_index, audio, mask=tail_mask) + if batch == 0: + tl.store(next_window + tail_index, envelope, mask=tail_mask) + + +def fused_istft_supported( + spec: torch.Tensor, + window: torch.Tensor, + n_fft: int, + hop: int, + audio_buffer: torch.Tensor | None, + window_buffer: torch.Tensor | None, + streaming: bool, +) -> bool: + if not ( + HAS_TRITON + and spec.is_cuda + and torch.version.hip is None + and spec.dtype == torch.complex64 + and spec.ndim == 3 + and spec.shape[0] > 0 + and spec.shape[1] == n_fft // 2 + 1 + and spec.shape[2] >= 2 + and hop > 0 + and n_fft == 4 * hop + and window.shape == (n_fft,) + and window.is_contiguous() + and window.device == spec.device + and window.dtype in (torch.float32, torch.bfloat16, torch.float16) + ): + return False + if not streaming or (audio_buffer is None and window_buffer is None): + return True + return ( + audio_buffer is not None + and window_buffer is not None + and audio_buffer.shape == (spec.shape[0], 3 * hop) + and window_buffer.shape == (1, 3 * hop) + and audio_buffer.device == window_buffer.device == spec.device + and audio_buffer.dtype == torch.float32 + and window_buffer.dtype == window.dtype + ) + + +def fused_istft( + frames: torch.Tensor, + window: torch.Tensor, + hop: int, + audio_buffer: torch.Tensor | None, + window_buffer: torch.Tensor | None, + streaming: bool, + last_chunk: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Fuse supported post-FFT inputs; dispatch is checked by ``ISTFT``.""" + batch, win_length, num_frames = frames.shape + tail_size = win_length - hop + pad = tail_size // 2 + raw_size = (num_frames - 1) * hop + win_length + has_buffer = streaming and audio_buffer is not None + start = 0 if has_buffer else pad + end = raw_size - (tail_size if streaming and not last_chunk else pad) + output_size = end - start + output = frames.new_empty((batch, output_size)) + envelope = window.new_empty((output_size,)) + next_audio = frames.new_empty((batch, tail_size)) if streaming else audio_buffer + next_window = window.new_empty((1, tail_size)) if streaming else window_buffer + + _istft_kernel[(batch, triton.cdiv(raw_size, 256))]( + frames, + window, + audio_buffer, + window_buffer, + output, + envelope, + next_audio, + next_window, + *frames.stride(), + audio_buffer.stride(0) if streaming and audio_buffer is not None else 0, + audio_buffer.stride(1) if streaming and audio_buffer is not None else 0, + window_buffer.stride(1) if streaming and window_buffer is not None else 0, + num_frames, + raw_size, + start, + output_size, + hop_size=hop, + streaming=streaming, + has_buffer=has_buffer, + block_size=256, + enable_fp_fusion=False, + ) + if not (envelope > 1e-11).all(): + raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.") + return output, next_audio, next_window From 1b7e854b054a0c1f34059d906fff6831d27890f5 Mon Sep 17 00:00:00 2001 From: Tianyao Wu <54675599+twu3202@users.noreply.github.com> Date: Tue, 15 Sep 2026 06:25:34 +0800 Subject: [PATCH 073/209] [Bugfix][MiniCPM-o] Attach a unit's whole frame group and size the vision slot budget like Stage0 (#7271) Signed-off-by: Tianyao Wu Co-authored-by: Claude --- examples/online_serving/minicpmo/README.md | 18 ++- tests/engine/duplex/test_duplex_runtime.py | 80 +++++++++++++ .../models/minicpmo_4_5/duplex/test_input.py | 105 ++++++++++++++++++ .../entrypoints/duplex/realtime_input.py | 12 +- .../models/minicpmo_4_5/duplex/input.py | 36 +++--- .../models/minicpmo_4_5/duplex/policy.py | 10 +- .../models/minicpmo_4_5/duplex/runtime.py | 91 +++++++++++++-- 7 files changed, 311 insertions(+), 41 deletions(-) diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index 2836eeb92..67b268a77 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -183,10 +183,20 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \ --output-dir /tmp/minicpmo_realtime_duplex_video_demo ``` -Detail inside a composite is capped by `scale_resolution=448` at -`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for -stacked frames, which the duplex adapter does not implement yet. Reading small -text or digits out of a wide scene is limited by that, not by frame timing. +Stage 0 processes a stacked unit the way official HD mode does +(`max_slice_nums=[2, 1]`): the base frame keeps its own 66-token block plus the +patches `get_sliced_grid` cuts for its size at `scale_resolution=448`, and the +composite adds one block. The wire still rejects a caller-supplied +`max_slice_nums`, so slicing is Stage 0's decision, not the client's. + +That detail costs context. One 960x540 frame is 66 vision tokens per unit and a +stacked pair is 264, on top of the unit's audio, so the Stage 0 prompt grows by +277 tokens per second stacked against 79 with a single frame. Against the +40960-token context this model ships with, a listen-heavy call reaches the limit +after about 2.5 minutes stacked and about 8.5 minutes unstacked (measured on a +264 s 960x540 clip). Duplex sessions cannot yet evict old units, so keep +multi-minute calls on the default `--stack-frames 1`; a stacked session that +outgrows the context ends with a context-length error. ## Open the experimental browser client diff --git a/tests/engine/duplex/test_duplex_runtime.py b/tests/engine/duplex/test_duplex_runtime.py index 9cbf4d6e1..475cf79d0 100644 --- a/tests/engine/duplex/test_duplex_runtime.py +++ b/tests/engine/duplex/test_duplex_runtime.py @@ -23,6 +23,7 @@ from vllm_omni.engine.duplex.runtime import ( ) from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import ( MiniCPMO45DuplexRuntimeExtension, + _duplex_hd_slice_count, build_duplex_data_plane_prompt, duplex_scheduler_token_budget, ) @@ -481,3 +482,82 @@ def test_placeholder_budget_is_planned_inside_omni_engine_boundary(): assert len(prompt["prompt_token_ids"]) == 16 assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16 + + +def _jpeg_b64(width: int, height: int) -> str: + import base64 + from io import BytesIO + + from PIL import Image + + buffer = BytesIO() + Image.new("RGB", (width, height), (32, 64, 96)).save(buffer, format="JPEG") + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +def _pcm_unit_payload(frames: list[str]) -> dict[str, object]: + import base64 + + return { + "audio": base64.b64encode(b"\x00" * (16000 * 4)).decode("ascii"), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + "video_frames": frames, + } + + +def test_duplex_scheduler_token_budget_one_frame_is_one_block(): + audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) + one_frame = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(960, 540)])) + + assert one_frame - audio_only == 66 + + +@pytest.mark.parametrize( + ("base_size", "expected_blocks"), + [ + ((960, 540), 4), # official camera frame: source + 2 HD patches + composite + ((540, 960), 4), # portrait phone capture + ((640, 360), 4), + ((448, 448), 2), # fits one tile: the processor does not slice it + ((320, 240), 2), + ], +) +def test_duplex_scheduler_token_budget_matches_stage0_hd_slicing(base_size, expected_blocks): + """A stacked pair is processed with max_slice_nums=[2, 1]; the budget must + equal the blocks Stage0 builds for that base-frame size.""" + audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) + pair = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(*base_size), _jpeg_b64(448, 448)])) + + assert pair - audio_only == expected_blocks * 66 + + +@pytest.mark.parametrize( + ("size", "expected_patches"), + [ + # multiple = min(ceil(w * h / 448**2), max_slice_nums); <= 1 is never sliced. + ((448, 448), 0), # exactly one tile + ((320, 240), 0), # smaller than one tile + ((447, 449), 0), # just under one tile + ((449, 449), 2), # just over: the processor cuts a 2-tile grid + ((640, 360), 2), + ((960, 540), 2), # official camera frame + ((540, 960), 2), # same frame rotated + ((1920, 1080), 2), # capped by max_slice_nums=2, not by area + ], +) +def test_duplex_hd_slice_count_matches_official_grid(size, expected_patches): + """Lock the port of ``MiniCPMVImageProcessor.get_sliced_grid``. + + Stage0 slices a stacked unit with ``max_slice_nums=[2, 1]``, so the patch + count for the base frame decides the scheduler budget. A frame that fits one + ``scale_resolution=448`` tile is not sliced at all. + """ + assert _duplex_hd_slice_count(size, 2) == expected_patches + + +def test_duplex_scheduler_token_budget_unreadable_base_frame_keeps_hd_fallback(): + audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) + pair = duplex_scheduler_token_budget(_pcm_unit_payload(["not-an-image", _jpeg_b64(448, 448)])) + + assert pair - audio_only == 4 * 66 diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_input.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_input.py index 4d6f3e46a..cf23d3ad1 100644 --- a/tests/model_executor/models/minicpmo_4_5/duplex/test_input.py +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_input.py @@ -177,3 +177,108 @@ def test_pcm_commit_reservation_rollback_restores_residual_audio(): ) assert retried.payload is not None assert base64.b64decode(retried.payload["audio"]) == (base64.b64decode(original["audio"]) + b"\x00" * (8_000 * 4)) + + +def _frame_payload(samples: int, frames: list[str]) -> dict[str, object]: + payload = pcm_payload(samples) + payload["video_frames"] = frames + return payload + + +def test_append_attaches_every_frame_of_the_unit_closing_append(): + """A base frame and its stacked composite arrive on one append and must + enter the same model unit (official ``frame_list``), not one per unit.""" + buffer = MiniCPMO45PcmAppendBuffer() + + emitted = buffer.append(_frame_payload(16_000, ["base-0", "stack-0"]), chunk_period_ms=1_000) + + assert emitted is not None + assert emitted["video_frames"] == ["base-0", "stack-0"] + assert buffer._frame_queue == [] + + +def test_stacked_frames_do_not_accumulate_across_units(): + buffer = MiniCPMO45PcmAppendBuffer() + attached: list[list[str]] = [] + + for unit in range(64): + emitted = buffer.append( + _frame_payload(16_000, [f"base-{unit}", f"stack-{unit}"]), + chunk_period_ms=1_000, + ) + assert emitted is not None + attached.append(list(emitted["video_frames"])) + + assert attached[0] == ["base-0", "stack-0"] + assert attached[-1] == ["base-63", "stack-63"] + assert buffer._frame_queue == [] + + +def test_frames_from_partial_appends_ride_the_unit_they_close(): + """200 ms client chunks: the frames arrive with the 5th chunk of each + second and attach to the unit that chunk completes, one group per unit.""" + buffer = MiniCPMO45PcmAppendBuffer() + + emitted_units: list[dict[str, object]] = [] + for unit in range(3): + for chunk in range(5): + frames = [f"base-{unit}", f"stack-{unit}"] if chunk == 4 else [] + emitted = buffer.append(_frame_payload(3_200, frames), chunk_period_ms=1_000) + if emitted is not None: + emitted_units.append(emitted) + + assert [unit["video_frames"] for unit in emitted_units] == [ + ["base-0", "stack-0"], + ["base-1", "stack-1"], + ["base-2", "stack-2"], + ] + assert buffer._frame_queue == [] + + +def test_single_frame_per_unit_still_attaches_one_frame_per_unit(): + buffer = MiniCPMO45PcmAppendBuffer() + + first = buffer.append(_frame_payload(16_000, ["base-0"]), chunk_period_ms=1_000) + second = buffer.append(_frame_payload(16_000, ["base-1"]), chunk_period_ms=1_000) + + assert first is not None and first["video_frames"] == ["base-0"] + assert second is not None and second["video_frames"] == ["base-1"] + + +def test_frame_groups_stay_queued_when_audio_outruns_units(): + """Two appends land before one unit closes: the later group waits for the + next unit instead of merging into the first.""" + buffer = MiniCPMO45PcmAppendBuffer() + + assert buffer.append(_frame_payload(8_000, ["base-0", "stack-0"]), chunk_period_ms=1_000) is None + first = buffer.append(_frame_payload(16_000, ["base-1", "stack-1"]), chunk_period_ms=1_000) + second = buffer.append(pcm_payload(8_000), chunk_period_ms=1_000) + + assert first is not None and first["video_frames"] == ["base-0", "stack-0"] + assert second is not None and second["video_frames"] == ["base-1", "stack-1"] + assert buffer._frame_queue == [] + + +def test_rollback_restores_frame_groups_in_wire_order(): + buffer = MiniCPMO45PcmAppendBuffer() + + first = buffer.prepare_append( + _frame_payload(16_000, ["base-0", "stack-0"]), + operation_id="append-0", + chunk_period_ms=1_000, + ) + second = buffer.prepare_append( + _frame_payload(16_000, ["base-1", "stack-1"]), + operation_id="append-1", + chunk_period_ms=1_000, + ) + assert first is not None and second is not None + assert first.payload is not None and first.payload["video_frames"] == ["base-0", "stack-0"] + + first.rollback() + + assert buffer._frame_queue == [["base-0", "stack-0"], ["base-1", "stack-1"]] + retried = buffer.flush(chunk_period_ms=1_000) + assert retried is not None + assert retried["video_frames"] == ["base-0", "stack-0"] + assert buffer._frame_queue == [["base-1", "stack-1"]] diff --git a/vllm_omni/entrypoints/duplex/realtime_input.py b/vllm_omni/entrypoints/duplex/realtime_input.py index 7fdadbcb2..386b4cfc3 100644 --- a/vllm_omni/entrypoints/duplex/realtime_input.py +++ b/vllm_omni/entrypoints/duplex/realtime_input.py @@ -1068,14 +1068,14 @@ class RealtimeInputTranslator(RealtimeStateOwner): Wire contract matches the official MiniCPM-o duplex loop: one base base64 JPEG per ~1 s audio chunk, optionally followed by that unit's stacked composite tiling the sub-frames captured inside it (at most 2 - images either way). HD slicing (max_slice_nums > 1) is not implemented - by the duplex adapter yet and is rejected explicitly rather than - silently ignored; note official suggests it for composites - (``max_slice_nums=[2, 1]``), so stacked detail is capped at - scale_resolution until that lands. + images either way). A caller-supplied ``max_slice_nums`` is rejected + rather than silently ignored: slicing is Stage 0's decision here, and + Stage 0 already applies the official HD suggestion for a stacked unit + (``max_slice_nums=[2, 1]``). The wire simply does not let the client + choose it. """ if max_slice_nums not in (None, 1): - return "max_slice_nums > 1 (HD slicing) is not implemented by the duplex Realtime adapter" + return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" if not isinstance(video_frames, list): return "video_frames must be a list of base64-encoded images" frames = [frame for frame in video_frames if frame is not None] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py index f7d86f1d7..1b30ec4c2 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py @@ -116,9 +116,12 @@ class MiniCPMO45PcmAppendBuffer: self._turn_had_speech = False self._reservation_seq = 0 self._reservations: list[MiniCPMO45PcmAppendReservation] = [] - # Omni duplex: queued camera frames (base64 JPEG), consumed FIFO at - # one frame per emitted model unit alongside the unit's audio. - self._frame_queue: list[str] = [] + # Omni duplex: queued camera frames (base64 JPEG), grouped per client + # append and consumed FIFO at one group per emitted model unit. A + # group is the base frame plus its optional stacked composite (the + # wire contract allows at most 2 images per append); both must ride + # the same unit, mirroring official ``frame_list``. + self._frame_queue: list[list[str]] = [] def clear(self) -> None: for reservation in self._reservations: @@ -251,7 +254,9 @@ class MiniCPMO45PcmAppendBuffer: ) frames_in = payload.get("video_frames") if isinstance(frames_in, list): - self._frame_queue.extend(frame for frame in frames_in if isinstance(frame, str) and frame) + frame_group = [frame for frame in frames_in if isinstance(frame, str) and frame] + if frame_group: + self._frame_queue.append(frame_group) self._turn_had_speech = self._turn_had_speech or bool(payload.get("is_speech", False)) if not allow_emit: return None @@ -284,16 +289,17 @@ class MiniCPMO45PcmAppendBuffer: out.pop("video_frames", None) out["audio"] = base64.b64encode(emit_raw).decode("ascii") out["sample_rate_hz"] = sample_rate_hz - # Omni duplex: attach at most one queued camera frame per emitted - # model unit (official cadence: one frame per 1 s chunk). The engine - # budgets 66 scheduler slots per attached frame from this payload. - # Official omni cadence is one frame per ~1 s chunk, and the first - # append consumes extra samples (1035 ms first window), so per-unit - # attachment could outrun the units Stage0 actually builds. Attach at - # most ONE frame per emitted payload; the rest stay queued. + # Omni duplex: attach at most one queued frame *group* per emitted + # model unit (official cadence: one ``frame_list`` per 1 s chunk). The + # engine budgets scheduler slots for every frame in this payload. + # The first append consumes extra samples (1035 ms first window), so + # per-unit attachment could outrun the units Stage0 actually builds. + # Attach ONE group per emitted payload; later groups stay queued. A + # group is what one client append carried (base frame + optional + # stacked composite), so the composite never lags its base frame. attached_frames: list[str] = [] if emit_samples + pad_samples >= min_samples and self._frame_queue: - attached_frames = [self._frame_queue.pop(0)] + attached_frames = list(self._frame_queue.pop(0)) out["video_frames"] = attached_frames out["force_listen"] = any(span.force_listen for span in reserved_spans) out["is_speech"] = any(span.is_speech for span in reserved_spans) @@ -400,9 +406,9 @@ class MiniCPMO45PcmAppendBuffer: restored = b"".join(item._raw for item in rolled_back) self._buffer[:0] = restored self._prepend_spans([span for item in rolled_back for span in item._spans]) - restored_frames = [frame for item in rolled_back for frame in item._video_frames] - if restored_frames: - self._frame_queue[:0] = restored_frames + restored_groups = [list(item._video_frames) for item in rolled_back if item._video_frames] + if restored_groups: + self._frame_queue[:0] = restored_groups self._sample_rate_hz = self._sample_rate_hz or reservation._sample_rate_hz self._turn_had_speech = self._turn_had_speech or any(item._turn_had_speech for item in rolled_back) for item in rolled_back: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py index 7a1779617..9a2643521 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py @@ -26,10 +26,12 @@ class MiniCPMO45DuplexPolicy: SAMPLES_PER_AUDIO_TOKEN = 1600 # Vision framing contract (omni duplex). Official streaming_prefill feeds # each frame as + 64 resampler embeddings + inside the - # unit, ahead of the unit's audio embeddings (max_slice_nums=1 in - # streaming, so exactly one 64-token block per frame). A unit may carry - # its base frame plus an optional stacked composite of the sub-frames - # captured inside that unit (2 blocks). + # unit, ahead of the unit's audio embeddings. A lone frame is one block. + # A unit carrying its base frame plus the stacked composite of the + # sub-frames captured inside it is processed with the official HD + # suggestion (max_slice_nums=[2, 1]): the base frame adds the patches the + # processor cuts for its size (two for a 960x540 capture, none for a frame + # that fits one 448x448 tile) and the composite adds one block. VISION_EMBEDS_PER_FRAME = 64 VISION_TOKENS_PER_FRAME = VISION_EMBEDS_PER_FRAME + 2 # + embeds + DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK = 20 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py index c8e746ce7..d5e98d723 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py @@ -3,6 +3,7 @@ from __future__ import annotations +import math from base64 import b64decode from binascii import Error as BinasciiError from typing import Any, cast @@ -23,17 +24,73 @@ _DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600 # matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME. _DUPLEX_VISION_TOKENS_PER_FRAME = 66 # Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD -# sliced (1 source + 2 patches on 960x540) and the composite is not. +# sliced (1 source + up to 2 patches, e.g. on 960x540) and the composite is +# not. The patch count depends on the frame size (see +# ``_duplex_hd_slice_count``); the constant is the fallback when the frame +# header cannot be read. _DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 +_DUPLEX_HD_MAX_SLICE_NUMS = 2 +_DUPLEX_SCALE_RESOLUTION = 448 + + +def _duplex_frames(payload: object) -> list[str]: + if not isinstance(payload, dict): + return [] + frames = payload.get("video_frames") + if not isinstance(frames, list): + return [] + return [frame for frame in frames if isinstance(frame, str) and frame] def _duplex_frame_count(payload: object) -> int: - if not isinstance(payload, dict): + return len(_duplex_frames(payload)) + + +def _duplex_frame_size(frame_b64: str) -> tuple[int, int] | None: + """Pixel size of a base64 JPEG/PNG frame from its header, or ``None``.""" + from io import BytesIO + + try: + from PIL import Image + + with Image.open(BytesIO(b64decode(frame_b64, validate=True))) as image: + width, height = image.size + except Exception: # noqa: BLE001 - Stage0 rejects bad frames with a reason + return None + if width <= 0 or height <= 0: + return None + return int(width), int(height) + + +def _duplex_hd_slice_count(image_size: tuple[int, int], max_slice_nums: int) -> int: + """Number of HD patches ``MiniCPMVImageProcessor.get_sliced_grid`` adds. + + Port of the official grid selection (``scale_resolution=448``): an image + whose area fits in one 448x448 tile is not sliced at all, otherwise the + grid closest to the image aspect ratio among ``multiple-1 .. multiple+1`` + splits is used. Stage0 runs the same processor, so the scheduler budget + and the worker-built embeddings agree. + """ + width, height = image_size + ratio = width * height / (_DUPLEX_SCALE_RESOLUTION * _DUPLEX_SCALE_RESOLUTION) + multiple = min(math.ceil(ratio), max_slice_nums) + if multiple <= 1: return 0 - frames = payload.get("video_frames") - if not isinstance(frames, list): - return 0 - return sum(1 for frame in frames if isinstance(frame, str) and frame) + log_ratio = math.log(width / height) + best_grid = (1, 1) + min_error = float("inf") + for split_grids_nums in (multiple - 1, multiple, multiple + 1): + if split_grids_nums == 1 or split_grids_nums > max_slice_nums: + continue + for m in range(1, split_grids_nums + 1): + if split_grids_nums % m: + continue + grid = (m, split_grids_nums // m) + error = abs(log_ratio - math.log(grid[0] / grid[1])) + if error < min_error: + best_grid = grid + min_error = error + return best_grid[0] * best_grid[1] def _duplex_vision_tokens(payload: object) -> int: @@ -41,15 +98,25 @@ def _duplex_vision_tokens(payload: object) -> int: Audio is never stacked: a unit still carries one second of soundtrack. ``stack_frames`` only adds a second *image*. Official HD on that pair is - ``[2, 1]``, so the base frame reserves three 66-token blocks and every - extra frame reserves one. + ``max_slice_nums=[2, 1]``: the base frame keeps its source block plus the + HD patches the processor cuts for its size (two for a 960x540 camera + frame, none for anything that fits in one 448x448 tile) and every extra + frame is one block. A single frame is never sliced. The count must match + the embeddings Stage0 builds exactly: surplus slots become pad + embeddings inside the KV. """ - count = _duplex_frame_count(payload) + frames = _duplex_frames(payload) + count = len(frames) if count <= 0: return 0 - if count >= 2: - return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME - return count * _DUPLEX_VISION_TOKENS_PER_FRAME + if count == 1: + return _DUPLEX_VISION_TOKENS_PER_FRAME + base_size = _duplex_frame_size(frames[0]) + if base_size is None: + base_blocks = _DUPLEX_HD_SLICES_PER_BASE_FRAME + else: + base_blocks = 1 + _duplex_hd_slice_count(base_size, _DUPLEX_HD_MAX_SLICE_NUMS) + return (base_blocks + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME def _duplex_pcm_sample_count(payload: object) -> int | None: From a22a7dd68fa508767c5ff6d55a1ae64e96677ff7 Mon Sep 17 00:00:00 2001 From: Aman <67479108+Bezdarnost@users.noreply.github.com> Date: Tue, 15 Sep 2026 04:25:42 +0600 Subject: [PATCH 074/209] [Model] Add VoxCPM2 startup LoRA adapter support (#7108) Signed-off-by: bezdarnost Co-authored-by: bezdarnost --- docs/serving/speech_api.md | 34 ++++ .../online_serving/test_voxcpm2_tts_lora.py | 106 +++++++++++ .../models/voxcpm2/test_startup_lora.py | 177 ++++++++++++++++++ .../model_executor/models/voxcpm2/lora.py | 95 ++++++++++ .../models/voxcpm2/runtime_config.py | 7 +- .../models/voxcpm2/voxcpm2_talker.py | 14 ++ 6 files changed, 432 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/online_serving/test_voxcpm2_tts_lora.py create mode 100644 tests/model_executor/models/voxcpm2/test_startup_lora.py create mode 100644 vllm_omni/model_executor/models/voxcpm2/lora.py diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 38cca54d6..d4fad4206 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -777,6 +777,40 @@ hidden-state dtype before the output projection. Consequently, a float32 stage c | ------- | ------------- | | `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. | +#### Startup LoRA adapter + +To serve a fine-tuned VoxCPM2 voice, set +`voxcpm2_runtime_config.startup_lora_path` in the stage's `hf_overrides`. +Copy `vllm_omni/deploy/voxcpm2.yaml` to a local deploy config and add this key +alongside its existing runtime settings: + +```yaml +# Under stages[0].engine_extras.hf_overrides.voxcpm2_runtime_config: +startup_lora_path: /absolute/path/to/adapter +``` + +Then launch with the modified config: + +```bash +vllm serve openbmb/VoxCPM2 --omni --deploy-config /absolute/path/to/voxcpm2-lora.yaml +``` + +The directory must be accessible to the worker and contain the native VoxCPM2 +training export: `lora_config.json` (with a `lora_config` object containing +`r`, `alpha`, enabled groups, and target module names) and +`lora_weights.safetensors`. PEFT checkpoints and pickle checkpoints are not +accepted. Missing, unexpected, non-finite, or incorrectly shaped adapter +tensors fail model loading rather than silently loading a partial adapter. + +The adapter is merged into the base LM, residual LM, LocDiT, and optional +projection layers selected by its configuration, after base weight loading +and before compilation or CUDA graph capture. All requests use that adapter; +`voice` still selects a reference voice, not a LoRA adapter. Changing adapters +requires restarting the server. Runtime loading/unloading, per-request +multi-LoRA selection, and `load_format=dummy` are not supported by this path. +Weight fusion rounds to the base model's dtype, so numerical and speech-quality +parity should be checked against the upstream adapter on your deployment. + ### MOSS-TTS-Nano | Model | Description | diff --git a/tests/e2e/online_serving/test_voxcpm2_tts_lora.py b/tests/e2e/online_serving/test_voxcpm2_tts_lora.py new file mode 100644 index 000000000..37efbea37 --- /dev/null +++ b/tests/e2e/online_serving/test_voxcpm2_tts_lora.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Serve a native VoxCPM2 adapter through the speech API.""" + +import json +from pathlib import Path + +import pytest +import torch +import yaml +from huggingface_hub import hf_hub_download +from safetensors import safe_open +from safetensors.torch import save_file + +from tests.helpers.fixtures.runtime import omni_fixture_lock +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams, get_model_prefix, iter_omni_server +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.slow, pytest.mark.tts] +MODEL = "openbmb/VoxCPM2" + + +@pytest.fixture(scope="module") +def omni_server(request, run_level, tmp_path_factory): + directory = tmp_path_factory.mktemp("voxcpm2-lora") + model = get_model_prefix() + MODEL + weights = Path(model) / "model.safetensors" + if not weights.is_file(): + weights = Path(hf_hub_download(MODEL, "model.safetensors")) + projections = ["enc_to_lm_proj", "lm_to_dit_proj", "res_to_dit_proj", "fusion_concat_proj"] + generator = torch.Generator().manual_seed(42) + tensors = {} + with safe_open(weights, framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + if not key.endswith(".weight"): + continue + name = key.removesuffix(".weight") + attention = name.startswith(("base_lm.", "residual_lm.", "feat_decoder.estimator.")) and name.rsplit( + ".", 1 + )[-1] in {"q_proj", "k_proj", "v_proj", "o_proj"} + if not attention and name not in projections: + continue + out_features, in_features = checkpoint.get_slice(key).get_shape() + # Small nonzero deltas exercise every group without retraining a voice. + tensors[f"{name}.lora_A"] = torch.randn(2, in_features, generator=generator) * 0.001 + tensors[f"{name}.lora_B"] = torch.randn(out_features, 2, generator=generator) * 0.001 + assert tensors + save_file(tensors, str(directory / "lora_weights.safetensors")) + (directory / "lora_config.json").write_text( + json.dumps({"lora_config": {"r": 2, "alpha": 2, "enable_lm": True, "enable_dit": True, "enable_proj": True}}) + ) + + with open(get_deploy_config_path("voxcpm2.yaml")) as stream: + config = yaml.safe_load(stream) + stage = config["stages"][0] + stage.update( + max_num_seqs=4, + kv_cache_memory_bytes=1024**3, + gpu_memory_utilization=0.65, + max_model_len=2048, + max_num_batched_tokens=2048, + ) + stage["default_sampling_params"]["max_tokens"] = 256 + runtime = stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"] + runtime.update(startup_lora_path=str(directory), unified_decode_graph_max_batch_size=4) + deploy_path = directory / "deploy.yaml" + deploy_path.write_text(yaml.safe_dump(config)) + request.param = OmniServerParams( + model=MODEL, + stage_config_path=str(deploy_path), + server_args=["--trust-remote-code", "--disable-log-stats"], + ) + yield from iter_omni_server(request, run_level, omni_fixture_lock) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_startup_lora_concurrent_speech(omni_server, online_client): + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The weather is nice today, perfect for a walk in the park.", + "voice": "default", + "response_format": "wav", + "stream": False, + "timeout": 300, + "min_audio_bytes": 40_000, + }, + request_num=4, + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_startup_lora_streaming_speech(omni_server, online_client): + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The weather is nice today, perfect for a walk in the park.", + "voice": "default", + "response_format": "wav", + "stream": True, + "stream_format": "audio", + "timeout": 300, + "min_audio_bytes": 40_000, + } + ) diff --git a/tests/model_executor/models/voxcpm2/test_startup_lora.py b/tests/model_executor/models/voxcpm2/test_startup_lora.py new file mode 100644 index 000000000..513cffca4 --- /dev/null +++ b/tests/model_executor/models/voxcpm2/test_startup_lora.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Native VoxCPM2 adapter validation and merge parity, without model weights.""" + +import importlib.util +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn +from torch.nn import functional as F + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def merge_lora(): + # Keep these L1 checks independent of vLLM custom-op/platform imports. + path = Path(__file__).resolve().parents[4] / "vllm_omni/model_executor/models/voxcpm2/lora.py" + spec = importlib.util.spec_from_file_location("voxcpm2_startup_lora", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.merge_voxcpm2_lora + + +@pytest.fixture +def adapter(tmp_path): + def attention(): + return nn.ModuleDict({name: nn.Linear(3, 4, bias=True) for name in ("q_proj", "k_proj", "v_proj", "o_proj")}) + + base_lm = nn.ModuleDict({"layers": nn.ModuleList([attention()])}) + residual_lm = nn.ModuleDict({"layers": nn.ModuleList([attention()])}) + tts = nn.ModuleDict( + { + "feat_decoder": nn.ModuleDict({"estimator": nn.ModuleDict({"decoder": attention()})}), + "fusion_concat_proj": nn.Linear(3, 4), + "stop_head": nn.Linear(3, 2), + } + ) + roots = {"base_lm": base_lm, "residual_lm": residual_lm, "tts": tts} + native = nn.ModuleDict({"base_lm": base_lm, "residual_lm": residual_lm, **dict(tts.items())}) + targets = {name: module for name, module in native.named_modules() if isinstance(module, nn.Linear)} + targets.pop("stop_head") + generator = torch.Generator().manual_seed(42) + tensors = { + f"{name}.lora_{suffix}": torch.randn(shape, generator=generator) + for name in targets + for suffix, shape in (("A", (2, 3)), ("B", (4, 2))) + } + config = { + "r": 2, + "alpha": 3, + "enable_lm": True, + "enable_dit": True, + "enable_proj": True, + "target_proj_modules": ["fusion_concat_proj"], + } + + def save(): + (tmp_path / "lora_config.json").write_text(json.dumps({"lora_config": config})) + save_file(tensors, str(tmp_path / "lora_weights.safetensors")) + + save() + return tmp_path, roots, targets, tensors, config, save + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_merge_matches_native_linear_output(merge_lora, adapter, dtype): + path, roots, targets, tensors, _, _ = adapter + inputs = torch.tensor([[0.5, -0.2, 0.1]], dtype=dtype) + expected = {} + identities = {} + for name, layer in targets.items(): + layer.to(dtype=dtype) + identities[name] = id(layer.weight) + # Native LoRALinear: linear(x, W, bias) + linear(linear(x, A), B) * alpha/r. + expected[name] = F.linear(inputs.float(), layer.weight.float(), layer.bias.float()) + 1.5 * F.linear( + F.linear(inputs.float(), tensors[f"{name}.lora_A"]), tensors[f"{name}.lora_B"] + ) + untouched = roots["tts"].stop_head.weight.detach().clone() + assert merge_lora(str(path), **roots) == 13 + for name, layer in targets.items(): + tolerance = 0.025 if dtype == torch.bfloat16 else 1e-6 + torch.testing.assert_close(layer(inputs).float(), expected[name], atol=tolerance, rtol=tolerance) + assert id(layer.weight) == identities[name] + assert layer.weight.dtype == dtype + assert set(dict(layer.named_parameters())) == {"weight", "bias"} + torch.testing.assert_close(roots["tts"].stop_head.weight, untouched, rtol=0, atol=0) + + +@pytest.mark.parametrize("group", ["lm", "dit", "proj"]) +def test_independently_enabled_groups(merge_lora, adapter, group): + path, roots, targets, tensors, config, save = adapter + for key in ("lm", "dit", "proj"): + config[f"enable_{key}"] = key == group + prefixes = {"lm": ("base_lm.", "residual_lm."), "dit": ("feat_decoder.",), "proj": ("fusion_concat_proj.",)} + for key in list(tensors): + if not key.startswith(prefixes[group]): + del tensors[key] + original = {name: layer.weight.detach().clone() for name, layer in targets.items()} + save() + assert merge_lora(str(path), **roots) == {"lm": 8, "dit": 4, "proj": 1}[group] + for name, layer in targets.items(): + if not (name + ".").startswith(prefixes[group]): + torch.testing.assert_close(layer.weight, original[name], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "corruption", ["missing", "extra", "shape", "rank", "alpha", "nonfinite", "dtype", "target", "flag"] +) +def test_rejects_invalid_checkpoint_before_mutating_weights(merge_lora, adapter, corruption): + path, roots, targets, tensors, config, save = adapter + key = "fusion_concat_proj.lora_B" # Validate even the last component before merging any layer. + if corruption == "missing": + del tensors[key] + elif corruption == "extra": + tensors["audio_vae.lora_A"] = torch.ones(2, 3) + elif corruption == "shape": + tensors[key] = torch.ones(5, 2) + elif corruption == "rank": + config["r"] = 0 + elif corruption == "alpha": + config["alpha"] = float("nan") + elif corruption == "nonfinite": + tensors[key][0, 0] = float("nan") + elif corruption == "dtype": + tensors[key] = tensors[key].long() + elif corruption == "target": + config["target_modules_lm"] = ["does_not_exist"] + elif corruption == "flag": + config["enable_dit"] = "false" + save() + original = {name: layer.weight.detach().clone() for name, layer in targets.items()} + with pytest.raises(ValueError, match="VoxCPM2"): + merge_lora(str(path), **roots) + for name, layer in targets.items(): + torch.testing.assert_close(layer.weight, original[name], rtol=0, atol=0) + + +def test_requires_native_metadata_and_safetensors(merge_lora, adapter): + path, roots, _, _, _, save = adapter + (path / "lora_config.json").write_text(json.dumps({"r": 2, "alpha": 3})) + with pytest.raises(ValueError, match="lora_config object"): + merge_lora(str(path), **roots) + save() + (path / "lora_weights.safetensors").unlink() + with pytest.raises(FileNotFoundError): + merge_lora(str(path), **roots) + + +def test_talker_loads_base_before_merge_and_rejects_reload(adapter): + pytest.importorskip("vllm") + from vllm_omni.model_executor.models.voxcpm2.runtime_config import _VoxCPM2RuntimeConfig + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import VoxCPM2TalkerForConditionalGeneration + + path, roots, targets, tensors, _, _ = adapter + talker = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration) + nn.Module.__init__(talker) + talker.model = roots["base_lm"] + talker.residual_model = roots["residual_lm"] + talker._tts = roots["tts"] + talker._runtime_config = _VoxCPM2RuntimeConfig(startup_lora_path=str(path)) + talker._startup_lora_applied = False + talker._patch_size, talker._feat_dim, talker._side_dtype = 4, 64, torch.float32 + weights = [(f"base_lm.{name}", torch.ones_like(param)) for name, param in talker.model.named_parameters()] + loaded = talker.load_weights(iter(weights)) + name = "base_lm.layers.0.q_proj" + expected = 1 + 1.5 * tensors[f"{name}.lora_B"] @ tensors[f"{name}.lora_A"] + torch.testing.assert_close(targets[name].weight, expected) + assert "model.layers.0.q_proj.weight" in loaded + assert talker._startup_lora_applied + with pytest.raises(ValueError, match="restart the server"): + talker.load_weights(iter(weights)) + torch.testing.assert_close(targets[name].weight, expected) diff --git a/vllm_omni/model_executor/models/voxcpm2/lora.py b/vllm_omni/model_executor/models/voxcpm2/lora.py new file mode 100644 index 000000000..cc3640582 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm2/lora.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Merge a native VoxCPM2 LoRA checkpoint at startup for all requests.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import torch +from safetensors.torch import load_file +from torch import nn + + +def merge_voxcpm2_lora( + adapter_path: str, + *, + base_lm: nn.Module, + residual_lm: nn.Module, + tts: nn.Module, +) -> int: + """Validate and merge one local native safetensors adapter. + + Call once after base weight loading and before fusion or graph capture. + Validate all checkpoint keys and tensors before changing weights. + """ + path = Path(adapter_path) + with (path / "lora_config.json").open(encoding="utf-8") as f: + metadata = json.load(f) + config = metadata.get("lora_config") if isinstance(metadata, dict) else None + if not isinstance(config, dict): + raise ValueError("VoxCPM2 lora_config.json must contain a lora_config object") + rank, alpha = config.get("r"), config.get("alpha") + if type(rank) is not int or rank <= 0: + raise ValueError("VoxCPM2 LoRA r must be a positive integer") + if isinstance(alpha, bool) or not isinstance(alpha, (int, float)) or not math.isfinite(alpha): + raise ValueError("VoxCPM2 LoRA alpha must be finite") + + targets: dict[str, nn.Linear] = {} + groups = ( + ("lm", {"base_lm": base_lm, "residual_lm": residual_lm}, "target_modules_lm"), + ("dit", {"feat_decoder.estimator": tts.feat_decoder.estimator}, "target_modules_dit"), + ("proj", {"": tts}, "target_proj_modules"), + ) + for group, roots, target_key in groups: + enabled = config.get(f"enable_{group}", False) + if type(enabled) is not bool: + raise ValueError(f"VoxCPM2 LoRA enable_{group} must be a boolean") + if not enabled: + continue + defaults = ( + ["enc_to_lm_proj", "lm_to_dit_proj", "res_to_dit_proj", "fusion_concat_proj"] + if group == "proj" + else ["q_proj", "v_proj", "k_proj", "o_proj"] + ) + names = config.get(target_key, defaults) + if not isinstance(names, list) or not names or any(not isinstance(n, str) or not n for n in names): + raise ValueError(f"VoxCPM2 LoRA {target_key} must be a non-empty list of names") + for prefix, root in roots.items(): + matched = set() + for name, module in root.named_modules(): + selector = name if group == "proj" else name.rsplit(".", 1)[-1] + if selector in names and isinstance(module, nn.Linear): + targets[f"{prefix}.{name}" if prefix else name] = module + matched.add(selector) + if missing := set(names) - matched: + raise ValueError(f"VoxCPM2 LoRA targets not found in {prefix or 'tts'}: {sorted(missing)}") + if not targets: + raise ValueError("VoxCPM2 LoRA config enables no target layers") + + state = load_file(str(path / "lora_weights.safetensors"), device="cpu") + expected = {f"{name}.lora_{suffix}" for name in targets for suffix in ("A", "B")} + if set(state) != expected: + raise ValueError( + "VoxCPM2 LoRA checkpoint keys do not match the configured targets: " + f"missing={sorted(expected - set(state))}, unexpected={sorted(set(state) - expected)}" + ) + for name, module in targets.items(): + if not module.weight.is_floating_point() or module.weight.is_meta: + raise ValueError(f"VoxCPM2 LoRA requires materialized floating-point weights: {name}") + for suffix, shape in (("A", (rank, module.in_features)), ("B", (module.out_features, rank))): + key = f"{name}.lora_{suffix}" + tensor = state[key] + if tensor.shape != shape or not tensor.is_floating_point() or not torch.isfinite(tensor).all(): + raise ValueError(f"Invalid VoxCPM2 LoRA tensor {key}: expected finite floating-point shape {shape}") + + with torch.no_grad(): + for name, module in targets.items(): + # Merge one layer at a time in fp32, preserving the base Parameter. + a = state[f"{name}.lora_A"].to(device=module.weight.device, dtype=torch.float32) + b = state[f"{name}.lora_B"].to(device=module.weight.device, dtype=torch.float32) + merged = torch.addmm(module.weight.float(), b, a, alpha=alpha / rank) + module.weight.copy_(merged) + return len(targets) diff --git a/vllm_omni/model_executor/models/voxcpm2/runtime_config.py b/vllm_omni/model_executor/models/voxcpm2/runtime_config.py index 1446729e7..207be70f2 100644 --- a/vllm_omni/model_executor/models/voxcpm2/runtime_config.py +++ b/vllm_omni/model_executor/models/voxcpm2/runtime_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -13,6 +13,7 @@ logger = init_logger(__name__) @dataclasses.dataclass(frozen=True) class _VoxCPM2RuntimeConfig: + startup_lora_path: str | None = None enable_profiling: bool = False enable_nvtx_profile: bool = False enable_loc_dit_layer_nvtx: bool = False @@ -105,6 +106,10 @@ class _VoxCPM2RuntimeConfig: @staticmethod def _coerce_value(key: str, value: Any, default: Any) -> Any: + if key == "startup_lora_path": + if value is not None and (not isinstance(value, str) or not value.strip()): + raise ValueError("VoxCPM2 startup_lora_path must be a non-empty local directory path or null") + return value if isinstance(default, bool): if isinstance(value, str): return value.strip().lower() in ("1", "true", "yes", "on") diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 5febb8912..78c0b26ea 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -46,6 +46,7 @@ from vllm_omni.utils.speaker_cache import ( ) from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest +from .lora import merge_voxcpm2_lora from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM from .runtime_config import _VoxCPM2RuntimeConfig from .voxcpm2_import_utils import import_voxcpm2_core @@ -834,6 +835,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self.vllm_config = vllm_config self.config = vllm_config.model_config.hf_config self._runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(vllm_config) + self._startup_lora_applied = False + if self._runtime_config.startup_lora_path and vllm_config.load_config.load_format == "dummy": + raise ValueError("VoxCPM2 startup LoRA requires real base weights; load_format=dummy is unsupported") global _ENABLE_NVTX_PROFILE _ENABLE_NVTX_PROFILE = self._runtime_config.enable_nvtx_profile @@ -3284,6 +3288,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"base_lm.": "model."}) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + if self._startup_lora_applied: + raise ValueError("Reloading VoxCPM2 weights with a startup LoRA is unsupported; restart the server") + def _base_lm_only(ws): for name, tensor in ws: if name.startswith("base_lm."): @@ -3298,6 +3305,13 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # them as missing from the checkpoint. loaded |= {name for name, _ in self.named_parameters() if name.startswith(("_tts.", "residual_model."))} + if adapter_path := self._runtime_config.startup_lora_path: + merged = merge_voxcpm2_lora( + adapter_path, base_lm=self.model, residual_lm=self.residual_model, tts=self._tts + ) + self._startup_lora_applied = True + logger.info("Merged VoxCPM2 startup LoRA from %s into %d linear layers", adapter_path, merged) + logger.info( "Loaded VoxCPM2 (patch=%d, feat_dim=%d, dtype=%s)", self._patch_size, self._feat_dim, self._side_dtype ) From 27f0e8c1ead8249c81c67bf0eb885f182508b05e Mon Sep 17 00:00:00 2001 From: Zhichao Zhang <60429419+smartDream-chao@users.noreply.github.com> Date: Tue, 15 Sep 2026 06:25:51 +0800 Subject: [PATCH 075/209] [Bugfix][Qwen3-TTS] Gate async-chunk Code2Wav on non_streaming_mode (#4371) (#6898) Signed-off-by: Zhichao Zhang <60429419+smartDream-chao@users.noreply.github.com> --- .../text_to_speech/qwen3_tts/end2end.py | 11 +++ .../openai_api/test_serving_speech.py | 61 ++++++++++++ .../test_qwen3_tts_async_chunk.py | 93 ++++++++++++++++++- .../openai/tts_adapters/qwen3_tts.py | 6 ++ .../stage_input_processors/qwen3_tts.py | 18 +++- .../stage_input_processors/tts_utils.py | 33 ++++++- 6 files changed, 215 insertions(+), 7 deletions(-) diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index 75f562af6..54b55a193 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Offline inference demo for Qwen3 TTS via vLLM Omni. Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base @@ -166,6 +169,9 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult: "language": [language], "speaker": [speaker], "max_new_tokens": [2048], + # Optional Code2Wav emit opt-in under deploy.async_chunk. + # Independent of prompt-mode non_streaming_mode (#4198 / #6898). + "full_utterance_decode": [True], } inputs.append( { @@ -185,6 +191,7 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult: "speaker": [speaker], "instruct": [instruct], "max_new_tokens": [2048], + "full_utterance_decode": [True], } inputs = { "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), @@ -226,6 +233,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult: "instruct": [instruct], "max_new_tokens": [2048], "non_streaming_mode": [True], + "full_utterance_decode": [True], } inputs.append( { @@ -244,6 +252,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult: "instruct": [instruct], "max_new_tokens": [2048], "non_streaming_mode": [True], + "full_utterance_decode": [True], } inputs = { "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), @@ -291,6 +300,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que "language": [language], "x_vector_only_mode": [x_vector_only_mode], "max_new_tokens": [2048], + "full_utterance_decode": [True], } inputs.append( { @@ -307,6 +317,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que "language": [syn_lang_single], "x_vector_only_mode": [x_vector_only_mode], "max_new_tokens": [2048], + "full_utterance_decode": [True], } inputs = { "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 732f86f2f..e7666ee33 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -2241,6 +2241,7 @@ class TestTTSMethods: assert params["task_type"] == ["Base"] assert params["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in params def test_build_tts_params_base_omits_non_streaming_mode_by_default(self, speech_server): """Base task should keep using the model default when no override is sent.""" @@ -2255,6 +2256,7 @@ class TestTTSMethods: assert params["task_type"] == ["Base"] assert "non_streaming_mode" not in params + assert "full_utterance_decode" not in params def test_build_tts_params_explicit_non_streaming_mode_overrides_voicedesign_default(self, speech_server): """Explicit false should not be replaced by the VoiceDesign fallback.""" @@ -2269,6 +2271,38 @@ class TestTTSMethods: assert params["task_type"] == ["VoiceDesign"] assert params["non_streaming_mode"] == [False] + assert "full_utterance_decode" not in params + + def test_build_tts_params_streaming_voicedesign_keeps_prompt_mode_not_full_decode(self, speech_server): + """Streaming VoiceDesign defaults: prompt-mode True, no full_utterance_decode.""" + req = OpenAICreateSpeechRequest( + input="Hello", + task_type="VoiceDesign", + instructions="warm and calm", + stream=True, + response_format="pcm", + ) + + params = speech_server._build_tts_params(req) + + assert params["task_type"] == ["VoiceDesign"] + assert params["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in params + + def test_build_tts_params_streaming_customvoice_explicit_non_streaming_mode_still_windowed(self, speech_server): + """Explicit non_streaming_mode=True must not inject full_utterance_decode.""" + req = OpenAICreateSpeechRequest( + input="Hello", + voice="Vivian", + non_streaming_mode=True, + stream=True, + response_format="pcm", + ) + + params = speech_server._build_tts_params(req) + + assert params["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in params def test_load_supported_speakers(self, mocker: MockerFixture): """Test _load_supported_speakers.""" @@ -5287,9 +5321,34 @@ class TestTTSAsyncOffloading: assert tts_params["task_type"] == ["VoiceDesign"] assert tts_params["non_streaming_mode"] == [False] + assert "full_utterance_decode" not in tts_params prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"] assert prompt["additional_information"] is tts_params assert prompt["additional_information"]["non_streaming_mode"] == [False] + assert "full_utterance_decode" not in prompt["additional_information"] + + def test_prepare_speech_generation_qwen3_streaming_voicedesign_default_params( + self, qwen3_tts_server, mocker: MockerFixture + ): + """Streaming VoiceDesign defaults keep prompt-mode True without full_utterance_decode.""" + qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + + request = OpenAICreateSpeechRequest( + input="hello", + task_type="VoiceDesign", + instructions="warm and calm", + stream=True, + response_format="pcm", + ) + _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + assert tts_params["task_type"] == ["VoiceDesign"] + assert tts_params["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in tts_params + prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"] + assert prompt["additional_information"]["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in prompt["additional_information"] def test_prepare_speech_generation_qwen3_base_non_streaming_mode_true( self, qwen3_tts_server, mocker: MockerFixture @@ -5312,9 +5371,11 @@ class TestTTSAsyncOffloading: assert tts_params["task_type"] == ["Base"] assert tts_params["ref_text"] == ["reference transcript"] assert tts_params["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in tts_params prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"] assert prompt["additional_information"] is tts_params assert prompt["additional_information"]["non_streaming_mode"] == [True] + assert "full_utterance_decode" not in prompt["additional_information"] def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server): """After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC.""" diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 996486d49..4b542bf77 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections import defaultdict from types import SimpleNamespace @@ -32,11 +32,15 @@ _FRAME = [1, 2, 3, 4] _Q = len(_FRAME) -def _req(rid, *, finished, initial_codec_chunk_frames=None): - ai = None +def _req(rid, *, finished, initial_codec_chunk_frames=None, non_streaming_mode=None, full_utterance_decode=None): + entries = {} if initial_codec_chunk_frames is not None: - entry = SimpleNamespace(list_data=[initial_codec_chunk_frames]) - ai = SimpleNamespace(entries={"initial_codec_chunk_frames": entry}) + entries["initial_codec_chunk_frames"] = SimpleNamespace(list_data=[initial_codec_chunk_frames]) + if non_streaming_mode is not None: + entries["non_streaming_mode"] = SimpleNamespace(list_data=[non_streaming_mode]) + if full_utterance_decode is not None: + entries["full_utterance_decode"] = SimpleNamespace(list_data=[full_utterance_decode]) + ai = SimpleNamespace(entries=entries) if entries else None return SimpleNamespace( external_req_id=rid, is_finished=lambda: finished, @@ -123,6 +127,85 @@ def test_flush_on_finish(): assert len(p.codes.audio) == _Q * 24 +def test_full_utterance_decode_defers_until_finished(): + """full_utterance_decode=True should not emit windowed chunks mid-utterance (#4371).""" + tm = _tm() + tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(50)] + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output=None, + request=_req("r", finished=False, full_utterance_decode=True), + is_finished=False, + ) + assert p is None + + +def test_full_utterance_decode_emits_full_utterance_on_finish(): + tm = _tm() + tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(50)] + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output=None, + request=_req("r", finished=True, full_utterance_decode=True), + is_finished=True, + ) + assert p is not None + assert p.meta.finished.item() is True + assert len(p.codes.audio) == _Q * 50 + + +def test_full_utterance_decode_false_keeps_windowed_emit(): + # Dynamic IC defers emit at 25 frames (see _CASES); 41 hits the first + # steady-state window boundary and proves streaming path is unchanged. + tm = _tm() + tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(41)] + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req("r", finished=False, full_utterance_decode=False), + is_finished=False, + ) + assert p is not None + assert len(p.codes.audio) == _Q * 25 + + +def test_non_streaming_mode_alone_does_not_defer_async_chunk(): + """Prompt-mode non_streaming_mode must not gate Code2Wav emits (#6898). + + VoiceDesign serving defaults non_streaming_mode=True even for streaming + responses; those requests still need windowed async-chunk audio for TTFA. + """ + tm = _tm() + tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(41)] + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req("r", finished=False, non_streaming_mode=True), + is_finished=False, + ) + assert p is not None + assert len(p.codes.audio) == _Q * 25 + + +def test_voicedesign_prompt_mode_with_streaming_keeps_windowed_emit(): + """Streaming VoiceDesign-style defaults: prompt-mode True, decode False.""" + tm = _tm() + tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(41)] + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req( + "r", + finished=False, + non_streaming_mode=True, + full_utterance_decode=False, + ), + is_finished=False, + ) + assert p is not None + assert len(p.codes.audio) == _Q * 25 + + _CASES = [ # ── IC boundary rule ────────────────────────────────────────────── # initial_codec_chunk_frames only controls the first emitted chunk. diff --git a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py index 3a2bb9860..07ee715ea 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py @@ -374,6 +374,12 @@ class Qwen3TTSAdapter(ARTTSAdapter): elif params["task_type"][0] == "VoiceDesign": params["non_streaming_mode"] = [True] + # Do not set ``full_utterance_decode`` here. That flag is an explicit + # Code2Wav emit opt-in (offline / advanced callers via + # additional_information) and must stay independent of prompt-mode + # ``non_streaming_mode`` (#4198 / #6898). Online serving keeps the + # default windowed async-chunk path (incl. #5202 incremental decode). + return params def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 423d288fe..431569626 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Stage input processor for Qwen3-TTS: Talker -> Code2Wav.""" import time @@ -24,6 +27,7 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( parse_chunk_ramp, ) from vllm_omni.model_executor.stage_input_processors.tts_utils import ( + extract_full_utterance_decode_from_request, extract_language_from_prompt, extract_language_from_request, extract_speaker_from_prompt, @@ -180,6 +184,15 @@ def talker2code2wav_async_chunk( initial_chunk_size = chunk_size length = len(transfer_manager.code_prompt_token_ids[request_id]) + # Full-utterance Code2Wav is gated by ``full_utterance_decode`` (explicit + # opt-in via additional_information). Do not reuse ``non_streaming_mode``: + # that flag is prompt construction only and VoiceDesign defaults it to True + # even for streaming responses (#4198, #6898 review). When True, defer + # connector chunks until the talker finishes so Code2Wav decodes once. + full_utterance_decode = extract_full_utterance_decode_from_request(request) + if full_utterance_decode is True and not finished: + return None + if length <= 0: if finished: return OmniPayloadStruct( @@ -192,7 +205,10 @@ def talker2code2wav_async_chunk( ) return None - if adaptive_enabled: + if full_utterance_decode is True: + first_chunk = int(transfer_manager.put_req_chunk.get(request_id, 0)) <= 0 + context_length = length + elif adaptive_enabled: _adaptive_states = getattr(transfer_manager, "_adaptive_states", None) if _adaptive_states is None: _adaptive_states = {} diff --git a/vllm_omni/model_executor/stage_input_processors/tts_utils.py b/vllm_omni/model_executor/stage_input_processors/tts_utils.py index 1bc78b4a2..806f5dbcf 100644 --- a/vllm_omni/model_executor/stage_input_processors/tts_utils.py +++ b/vllm_omni/model_executor/stage_input_processors/tts_utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The Qwen team. """Shared TTS utility functions for speaker and language extraction. @@ -103,6 +103,37 @@ def extract_speaker_from_prompt( return None +# ============================================================================= +# Code2Wav emit-mode helpers (independent of prompt construction) +# ============================================================================= + + +def extract_full_utterance_decode_from_request(request: Any) -> bool | None: + """Extract ``full_utterance_decode`` from a request's additional_information. + + This is a *response / Code2Wav emit* intent flag. It must not be confused + with ``non_streaming_mode``, which only controls Qwen3-TTS prompt + construction and is independent of HTTP/WebSocket streaming (#4198). + + Returns: + True/False when explicitly set, otherwise None. + """ + additional_information = getattr(request, "additional_information", None) + if additional_information is None: + return None + entries = getattr(additional_information, "entries", None) + if not isinstance(entries, dict): + return None + entry = entries.get("full_utterance_decode") + if entry is None: + return None + list_data = getattr(entry, "list_data", None) + if isinstance(list_data, list) and list_data: + raw = list_data[0] + return raw if isinstance(raw, bool) else None + return None + + # ============================================================================= # Language helpers # ============================================================================= From f062124ba428586965ec04bce44c315f86561cb4 Mon Sep 17 00:00:00 2001 From: ZacheryAU Date: Tue, 15 Sep 2026 06:25:59 +0800 Subject: [PATCH 076/209] [Frontend][Benchmark] Add duplex performance metrics for OmniInteract and Omni-DuplexEval benchmark (#7242) --- docs/cli/bench/serve.md | 8 +- tests/benchmarks/metrics/test_metrics.py | 34 +-- tests/benchmarks/patch/test_patch.py | 19 +- tests/benchmarks/test_omniinteract.py | 70 ++++- tests/clients/test_duplex_client.py | 157 ++++++++++- .../fullduplex/test_omni_duplex_eval_cli.py | 266 +++++++++++++++++- .../test_minicpmo_4_5_duplex.py | 11 +- .../test_minicpmo_realtime_duplex_drivers.py | 10 +- tests/engine/test_output_processor.py | 129 +++++++++ .../openai_api/test_duplex_handler.py | 91 +++++- .../duplex/omni_duplex_eval_runner.py | 66 ++++- .../benchmarks/duplex_session_metrics.py | 229 +++++++++++++++ vllm_omni/benchmarks/metrics/metrics.py | 8 +- vllm_omni/benchmarks/omniinteract.py | 51 +--- vllm_omni/benchmarks/patch/patch.py | 36 ++- vllm_omni/clients/duplex.py | 258 +++++++++++++++-- .../cli/benchmark/omni_duplex_eval.py | 49 +++- vllm_omni/entrypoints/duplex/protocol.py | 90 +++++- .../entrypoints/duplex/runtime_bridge.py | 25 +- .../entrypoints/duplex/session_runner.py | 2 + vllm_omni/outputs/output_processor.py | 36 ++- 21 files changed, 1484 insertions(+), 161 deletions(-) create mode 100644 vllm_omni/benchmarks/duplex_session_metrics.py diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md index a162c001a..8dc5e1ed8 100644 --- a/docs/cli/bench/serve.md +++ b/docs/cli/bench/serve.md @@ -316,8 +316,12 @@ arrives, checkpointing the response's history position so a later committed user outputs are ineligible and omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video horizon. -TTFT, TTFP, and RTF start at client receipt of `response.created`. TPOT/ITL use engine stage-0 timing; ITL is emitted only when -every token interval is present. +Per-response TTFT and TTFP start when the server begins executing the native model-turn request that owns the response. RTF +continues to use client receipt of `response.created` through the last audio packet, divided by emitted audio duration. Global +TTFT, TTFP, and RTF cover the complete input-stream window. TPOT/ITL use engine stage-0 timing; ITL is emitted only when every +token interval is present within a continuous generation segment. Model-unit pacing and gaps between generation segments are +excluded from TPOT/ITL. Raw request metrics retain `response_created_to_first_text_ms` and +`response_created_to_first_audio_ms` as client-envelope diagnostics. The checked-in local performance configuration measures four deterministic cases from each OmniInteract subset (12 videos total), with no benchmark warmups and a maximum concurrency of two. Each subset also sends one readiness request before its diff --git a/tests/benchmarks/metrics/test_metrics.py b/tests/benchmarks/metrics/test_metrics.py index d7dd6c14a..0aae841af 100644 --- a/tests/benchmarks/metrics/test_metrics.py +++ b/tests/benchmarks/metrics/test_metrics.py @@ -150,16 +150,16 @@ def test_unmeasured_duplex_latency_does_not_add_zero_samples(): measured.audio_ttfp = 0.2 measured.audio_rtf = 0.5 measured.duplex_session_metrics = { - "mean_ttft_ms": 100.0, - "mean_ttfp_ms": 200.0, - "mean_rtf": 0.5, + "ttft_ms": 100.0, + "ttfp_ms": 200.0, + "rtf": 0.5, } listen_only = _make_output(100, output_tokens=0) listen_only.ttft = listen_only.audio_ttfp = listen_only.audio_rtf = 0.0 listen_only.duplex_session_metrics = { - "mean_ttft_ms": None, - "mean_ttfp_ms": None, - "mean_rtf": None, + "ttft_ms": None, + "ttfp_ms": None, + "rtf": None, } metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0}) @@ -174,9 +174,9 @@ def test_unmeasured_duplex_latency_does_not_add_zero_samples(): def test_all_unmeasured_duplex_latency_is_not_reported_as_zero(): listen_only = _make_output(100, output_tokens=0) listen_only.duplex_session_metrics = { - "mean_ttft_ms": None, - "mean_ttfp_ms": None, - "mean_rtf": None, + "ttft_ms": None, + "ttfp_ms": None, + "rtf": None, } metrics = _calculate_test_metrics( @@ -196,13 +196,13 @@ def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput(): missing_tpot.itl = [] missing_tpot.text_latency = missing_tpot.ttft = 0.1 missing_tpot.tpot_measured = False - missing_tpot.duplex_session_metrics = {"mean_ttft_ms": 100.0} + missing_tpot.duplex_session_metrics = {"ttft_ms": 100.0} slow_ttft = _make_output(100, output_tokens=5) slow_ttft.ttft = 1.0 slow_ttft.text_latency = 1.4 slow_ttft.itl = [0.1] * 4 - slow_ttft.duplex_session_metrics = {"mean_ttft_ms": 1000.0} + slow_ttft.duplex_session_metrics = {"ttft_ms": 1000.0} metrics = _calculate_test_metrics( [missing_tpot, slow_ttft], @@ -219,7 +219,7 @@ def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero(): output.itl = [] output.text_latency = output.ttft output.tpot_measured = False - output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + output.duplex_session_metrics = {"ttft_ms": 100.0} metrics = _calculate_test_metrics([output], {"tpot": float("inf")}) @@ -236,7 +236,7 @@ def test_duplex_response_timings_do_not_build_a_session_token_timeline(): {"response_id": "r1", "stage0_tokens": {"itls_ms": [100.0, 100.0]}}, {"response_id": "r2", "stage0_tokens": {"itls_ms": [100.0, 100.0]}}, ] - output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + output.duplex_session_metrics = {"ttft_ms": 100.0} metrics = _calculate_test_metrics([output]) @@ -252,7 +252,7 @@ def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback(): output.itl = [] output.text_latency = output.ttft output.tpot_measured = False - output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + output.duplex_session_metrics = {"ttft_ms": 100.0} def tokenizer(text, *, add_special_tokens): assert text == output.generated_text @@ -282,7 +282,7 @@ def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback(): def test_zero_itl_does_not_create_zero_tpot(): output = _make_output(100, output_tokens=3) output.itl = [0.0, 0.0] - output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + output.duplex_session_metrics = {"ttft_ms": 100.0} metrics = _calculate_test_metrics([output], {"tpot": 1.0}) @@ -357,9 +357,9 @@ def test_single_token_responses_do_not_report_zero_tpot(capsys): def test_duplex_goodput_does_not_pair_measurements_from_different_requests(): text_only, audio_only = _make_output(100), _make_output(100) text_only.ttft, text_only.audio_ttfp = 0.1, 0.0 - text_only.duplex_session_metrics = {"mean_ttft_ms": 100.0, "mean_ttfp_ms": None, "mean_rtf": None} + text_only.duplex_session_metrics = {"ttft_ms": 100.0, "ttfp_ms": None, "rtf": None} audio_only.ttft, audio_only.audio_ttfp = 0.0, 0.2 - audio_only.duplex_session_metrics = {"mean_ttft_ms": None, "mean_ttfp_ms": 200.0, "mean_rtf": None} + audio_only.duplex_session_metrics = {"ttft_ms": None, "ttfp_ms": 200.0, "rtf": None} metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0}) diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index 6dd35e677..60e071ad5 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -287,20 +287,25 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) "ttft": "conversation.item.create client send to first non-empty text delta", "ttfp": "conversation.item.create client send to first audio packet", "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", + "tpot": "Stage-0 engine mean time per output token", }, "ttft_ms": pytest.approx(20.0, abs=2.0), + "tpot_ms": 10.0, "ttfp_ms": pytest.approx(30.0, abs=2.0), "rtf": pytest.approx(0.3, abs=0.03), "audio_generation_ms": pytest.approx(30.0, abs=2.0), "audio_duration_ms": 100.0, } - assert output.duplex_session_metrics == { - "session_id": session_id, - "audio_turn_count": 4, - "mean_ttft_ms": pytest.approx(20.0, abs=2.0), - "mean_ttfp_ms": pytest.approx(30.0, abs=2.0), - "mean_rtf": pytest.approx(0.3, abs=0.03), - } + session = output.duplex_session_metrics + assert session["session_id"] == session_id + assert session["audio_turn_count"] == 4 + assert session["ttft_ms"]["count"] == 4 + assert session["ttft_ms"]["mean"] == pytest.approx(20.0, abs=2.0) + assert session["tpot_ms"] == {"count": 4, "mean": 10.0, "p50": 10.0, "p99": 10.0} + assert session["ttfp_ms"]["count"] == 4 + assert session["ttfp_ms"]["mean"] == pytest.approx(30.0, abs=2.0) + assert session["rtf"]["count"] == 4 + assert session["rtf"]["mean"] == pytest.approx(0.3, abs=0.03) def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete(): diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py index b9dd10870..73cedfb32 100644 --- a/tests/benchmarks/test_omniinteract.py +++ b/tests/benchmarks/test_omniinteract.py @@ -148,6 +148,72 @@ def _successful_output( return case, collector, result +def test_response_metrics_include_engine_tpot_and_stream_window(): + first_audio = _audio("r1") + first_audio["metadata"] = { + "audio_duration_ms": 100, + "vllm_omni": { + "stage_metrics": { + "0": { + "num_tokens_out": 3, + "vllm_tpot_ms": 10.0, + "vllm_itls_ms": [10.0, 10.0], + } + } + }, + } + second_audio = _audio("r2") + second_audio["metadata"] = { + "audio_duration_ms": 200, + "vllm_omni": { + "stage_metrics": { + "0": { + "num_tokens_out": 2, + "vllm_tpot_ms": 20.0, + "vllm_itls_ms": [20.0], + } + } + }, + } + collector = _collector( + (_created("r1"), 10.1), + (_text("r1", "first"), 10.2), + (first_audio, 10.3), + (_done("r1"), 10.4), + (_created("r2"), 12.1), + (_text("r2", "second"), 12.2), + (second_audio, 12.5), + (_done("r2"), 12.6), + ) + result = oi.OmniInteractCaseResult("1q1a", "video.mp4", "", session_id="session") + + oi._populate_response_metrics(result, collector, stream_start=10.0) + + assert result.output_tokens == 5 + assert [metric["tpot_ms"] for metric in result.duplex_request_metrics] == [10.0, 20.0] + assert result.duplex_session_metrics == { + "session_id": "session", + "audio_turn_count": 2, + "ttft_ms": {"count": 2, "mean": 100.0, "p50": 100.0, "p99": 100.0}, + "tpot_ms": {"count": 2, "mean": 15.0, "p50": 10.0, "p99": 20.0}, + "ttfp_ms": {"count": 2, "mean": 300.0, "p50": 200.0, "p99": 400.0}, + "rtf": {"count": 2, "mean": 2.0, "p50": 2.0, "p99": 2.0}, + "stream_ttft_ms": 200.0, + "stream_ttfp_ms": 300.0, + "stream_rtf": 8.333333, + "stream_audio_generation_ms": 2500.0, + "stream_audio_duration_ms": 300.0, + "stream_measurement_origin": { + "ttft": "input stream start to first non-empty text delta", + "ttfp": "input stream start to first audio packet", + "rtf": ( + "input stream start-to-last-audio receive time divided by total emitted audio duration; " + "includes concurrent realtime input" + ), + }, + } + + def _write_success( root: Path, case: data.OmniInteractCase, @@ -645,7 +711,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa transcript="timing metadata missing", output_tokens=0, duplex_request_metrics=[{"request_metrics": {"ttft_ms": 1.0}}], - duplex_session_metrics={"mean_ttft_ms": 10.0}, + duplex_session_metrics={"ttft_ms": 10.0}, ) monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run) @@ -713,7 +779,7 @@ async def test_adapter_reports_exact_or_weighted_token_timing( success=True, output_tokens=5, duplex_request_metrics=request_metrics, - duplex_session_metrics={"mean_ttft_ms": 100.0}, + duplex_session_metrics={"ttft_ms": 100.0}, ) monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run) diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index ce1795e0f..8a74a2087 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -736,11 +736,162 @@ def test_event_collector_reports_engine_token_and_audio_intervals(): "measurement_origin": { "ttft": "input_audio_buffer.commit client send to first non-empty text delta", "ttfp": "input_audio_buffer.commit client send to first audio packet", + "tpot": "Stage-0 engine mean time per output token", }, "ttft_ms": 250.0, "ttfp_ms": 300.0, "audio_generation_ms": 460.0, "audio_duration_ms": 240.0, + "tpot_ms": 15.0, + } + + +def test_session_summary_omits_unmeasured_zero_tpot(): + summary = summarize_session_request_metrics( + [ + {"ttft_ms": 100.0, "tpot_ms": 0.0}, + {"ttft_ms": 200.0, "tpot_ms": None}, + ], + session_id="session", + ) + + assert summary["ttft_ms"] == { + "count": 2, + "mean": 150.0, + "p50": 100.0, + "p99": 200.0, + } + assert "tpot_ms" not in summary + + +def test_event_collector_reports_global_metrics_across_responses(): + collector = EventCollector() + for response_id, created_at_s, text_at_s, audio_times, durations in ( + ("resp-a", 10.1, 10.2, (10.3, 10.4), (100, 200)), + ("resp-b", 12.1, 12.2, (12.3, 12.5), (150, 300)), + ): + collector.add( + {"type": "response.created", "response": {"id": response_id}}, + received_at_s=created_at_s, + ) + collector.add( + { + "type": "response.audio_transcript.delta", + "response_id": response_id, + "delta": response_id, + }, + received_at_s=text_at_s, + ) + for received_at_s, duration_ms in zip(audio_times, durations, strict=True): + collector.add( + { + "type": "response.audio.delta", + "response_id": response_id, + "delta": base64.b64encode(b"audio").decode("ascii"), + "sample_rate_hz": 16_000, + "metadata": {"audio_duration_ms": duration_ms}, + }, + received_at_s=received_at_s, + ) + + metrics = collector.global_timing_summary( + after_s=10.0, + window_started_at_s=10.0, + response_ids=["resp-a", "resp-b"], + measurement_origin={ + "ttft": "input stream start to first non-empty text delta", + "ttfp": "input stream start to first audio packet", + "rtf": "input stream start-to-last-audio receive time divided by total emitted audio duration", + }, + ) + + assert metrics == { + "source": "client_monotonic_receive", + "response_ids": ["resp-a", "resp-b"], + "measurement_origin": { + "ttft": "input stream start to first non-empty text delta", + "ttfp": "input stream start to first audio packet", + "rtf": "input stream start-to-last-audio receive time divided by total emitted audio duration", + }, + "ttft_ms": 200.0, + "ttfp_ms": 300.0, + "audio_generation_ms": 2500.0, + "audio_duration_ms": 500.0, + } + + +def test_event_collector_prefers_server_request_start_for_ttf(): + collector = EventCollector() + collector.add( + { + "type": "response.created", + "response": { + "id": "resp-a", + "metadata": { + "duplex_event": { + "response_request_metrics": { + "source": "server_monotonic_request_start", + "measurement_origin": { + "ttft": "native request start to first text output", + "ttfp": "native request start to first audio output", + }, + "ttft_ms": 125.0, + "ttfp_ms": 175.0, + } + } + }, + }, + }, + received_at_s=10.0, + ) + stage_metrics = { + "0": { + "num_tokens_out": 4, + "vllm_ttft_ms": 120.0, + "vllm_tpot_ms": 15.0, + "vllm_itl_ms": 14.0, + "vllm_itls_ms": [10.0, 14.0, 18.0], + } + } + for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): + collector.add( + { + "type": "response.audio.delta", + "response_id": "resp-a", + "delta": base64.b64encode(b"audio").decode("ascii"), + "sample_rate_hz": 16_000, + "metadata": { + "audio_duration_ms": cumulative_audio_ms, + "vllm_omni": {"stage_metrics": stage_metrics}, + }, + }, + received_at_s=received_at_s, + ) + collector.add( + {"type": "response.audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, + received_at_s=10.15, + ) + + timing = collector.timing_summary( + after_s=10.0, + input_committed_at_s=9.9, + response_id="resp-a", + ) + + assert timing["request_metrics"] == { + "source": "server_request_start_and_client_receive", + "measurement_origin": { + "ttft": "native request start to first text output", + "ttfp": "native request start to first audio output", + "tpot": "Stage-0 engine mean time per output token", + }, + "ttft_ms": 125.0, + "ttfp_ms": 175.0, + "audio_generation_ms": 460.0, + "audio_duration_ms": 240.0, + "response_created_to_first_text_ms": 150.0, + "response_created_to_first_audio_ms": 200.0, + "tpot_ms": 15.0, } @@ -806,9 +957,9 @@ def test_summarize_session_request_metrics_averages_audio_turns(): assert summary == { "session_id": "sess-1", "audio_turn_count": 2, - "mean_ttft_ms": 200.0, - "mean_ttfp_ms": 300.0, - "mean_rtf": 0.6, + "ttft_ms": {"count": 2, "mean": 200.0, "p50": 100.0, "p99": 300.0}, + "ttfp_ms": {"count": 2, "mean": 300.0, "p50": 200.0, "p99": 400.0}, + "rtf": {"count": 2, "mean": 0.6, "p50": 0.5, "p99": 0.7}, } diff --git a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py index a4d7faf9b..7f4965768 100644 --- a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py +++ b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py @@ -15,6 +15,15 @@ import pytest from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge +from vllm_omni.benchmarks.duplex.omni_duplex_eval_runner import GenerateSampleResult +from vllm_omni.benchmarks.duplex_session_metrics import ( + DUPLEX_METRICS_FILENAME, + build_duplex_metrics_report, + collect_duplex_session_metrics, + merge_duplex_metrics_report, + read_duplex_metrics_report, +) +from vllm_omni.clients.duplex import EventCollector from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand @@ -44,7 +53,25 @@ def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytes output.parent.mkdir(parents=True, exist_ok=True) output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8") output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8") - return output + return GenerateSampleResult( + output=output, + request_metrics=[ + { + "sample_id": sample.id, + "split": sample.split, + "ttft_ms": 10.0, + "ttfp_ms": 20.0, + "rtf": 0.5, + } + ], + session_metrics={ + "sample_id": sample.id, + "split": sample.split, + "stream_ttft_ms": 10.0, + "stream_ttfp_ms": 20.0, + "stream_rtf": 0.5, + }, + ) class FakeJudge: def __init__(self, *args, **kwargs): @@ -75,6 +102,19 @@ def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytes ] ) OmniDuplexEvalSubcommand.cmd(generate) + duplex_metrics = json.loads((response_root / DUPLEX_METRICS_FILENAME).read_text(encoding="utf-8")) + assert duplex_metrics["duplex_request_metrics"] == [ + { + "sample_id": "sample-1", + "split": "PR_correction", + "ttft_ms": 10.0, + "ttfp_ms": 20.0, + "rtf": 0.5, + } + ] + assert duplex_metrics["duplex_stream_ttft_ms"]["mean"] == 10.0 + assert duplex_metrics["duplex_stream_ttfp_ms"]["mean"] == 20.0 + assert duplex_metrics["duplex_stream_rtf"]["mean"] == 0.5 evaluate = parser.parse_args( [ "evaluate", @@ -144,7 +184,7 @@ async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path async with websockets.serve(handler, "127.0.0.1", 0) as server: port = server.sockets[0].getsockname()[1] - output = await runner.generate_sample( + result = await runner.generate_sample( sample, url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1", model="mock", @@ -152,13 +192,231 @@ async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path output_root=tmp_path / "responses", ) - assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}] - meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8")) + assert json.loads(result.output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}] + meta = json.loads(result.output.with_name("sample.meta.json").read_text(encoding="utf-8")) assert meta["response_done"] is True assert meta["drain_timeout"] is None + assert "duplex_request_metrics" not in meta + assert result.request_metrics + assert result.request_metrics[0]["sample_id"] == "sample" + assert result.session_metrics["stream_ttfp_ms"] is not None assert any(event["type"] == "playback.ack" for event in received) +def _pr_manifest(tmp_path: Path, sample_ids: tuple[str, ...]) -> Path: + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps( + [ + { + "id": sample_id, + "split": "PR_correction", + "question_text": "What changed?", + "answer1": "The object moved.", + } + for sample_id in sample_ids + ] + ), + encoding="utf-8", + ) + return manifest + + +def _run_generate_cli(manifest: Path, response_root: Path) -> None: + parser = argparse.ArgumentParser() + OmniDuplexEvalSubcommand.add_cli_args(parser) + OmniDuplexEvalSubcommand.cmd( + parser.parse_args( + [ + "generate", + "--dataset", + str(manifest), + "--family", + "pr", + "--model", + "mock", + "--ref-audio", + str(manifest), + "--response-root", + str(response_root), + ] + ) + ) + + +def _metrics_for(sample_id: str, *, ttft_ms: float, ttfp_ms: float, rtf: float) -> GenerateSampleResult: + split = "PR_correction" + return GenerateSampleResult( + output=Path(split) / f"{sample_id}.json", + request_metrics=[{"sample_id": sample_id, "split": split, "ttft_ms": ttft_ms, "ttfp_ms": ttfp_ms, "rtf": rtf}], + session_metrics={ + "sample_id": sample_id, + "split": split, + "stream_ttft_ms": ttft_ms, + "stream_ttfp_ms": ttfp_ms, + "stream_rtf": rtf, + }, + ) + + +def _fake_generate(results: dict[str, GenerateSampleResult | None]): + async def generate( + sample: DuplexSample, + *, + output_root: str | Path, + url: str, + model: str, + ref_audio: str | Path, + fps: float = 1.0, + mix: str = "question", + pace: str = "realtime", + clock: str = "media", + overwrite: bool = False, + unit_ms: int = 1000, + ) -> GenerateSampleResult: + _ = (url, model, ref_audio, fps, mix, pace, clock, overwrite, unit_ms) + result = results[sample.id] + if result is None: + return GenerateSampleResult(output=Path(output_root) / sample.split / f"{sample.id}.json") + return result + + return generate + + +def test_merge_duplex_metrics_report_replaces_only_incoming_sample_keys(): + existing = build_duplex_metrics_report( + request_metrics=[ + {"sample_id": "kept", "split": "PR_correction", "ttft_ms": 11.0}, + {"sample_id": "replaced", "split": "PR_correction", "ttft_ms": 12.0}, + ], + session_metrics=[ + {"sample_id": "kept", "split": "PR_correction", "stream_ttft_ms": 11.0}, + {"sample_id": "replaced", "split": "PR_correction", "stream_ttft_ms": 12.0}, + ], + ) + merged = merge_duplex_metrics_report( + existing, + request_metrics=[{"sample_id": "replaced", "split": "PR_correction", "ttft_ms": 99.0}], + session_metrics=[{"sample_id": "replaced", "split": "PR_correction", "stream_ttft_ms": 99.0}], + ) + assert [row["sample_id"] for row in merged["duplex_request_metrics"]] == ["kept", "replaced"] + assert merged["duplex_request_metrics"][1]["ttft_ms"] == 99.0 + assert merged["duplex_stream_ttft_ms"]["mean"] == 55.0 + + +def test_read_duplex_metrics_report_returns_none_when_missing(tmp_path: Path): + assert read_duplex_metrics_report(tmp_path / DUPLEX_METRICS_FILENAME) is None + + +def test_read_duplex_metrics_report_rejects_invalid_json(tmp_path: Path): + path = tmp_path / DUPLEX_METRICS_FILENAME + path.write_text("{not-json", encoding="utf-8") + with pytest.raises(json.JSONDecodeError): + read_duplex_metrics_report(path) + + +def test_merge_duplex_metrics_report_without_existing_is_passthrough(): + merged = merge_duplex_metrics_report( + None, + request_metrics=[{"sample_id": "a", "split": "PR_correction", "ttft_ms": 1.0}], + session_metrics=[{"sample_id": "a", "split": "PR_correction", "stream_ttft_ms": 1.0}], + ) + assert merged["duplex_request_metrics"][0]["ttft_ms"] == 1.0 + assert merged["duplex_stream_ttft_ms"]["mean"] == 1.0 + + +def test_merge_duplex_metrics_report_rejects_non_object_rows(): + with pytest.raises(ValueError, match="duplex_request_metrics entries must be JSON objects"): + merge_duplex_metrics_report( + {"duplex_request_metrics": ["bad"]}, + request_metrics=[], + session_metrics=[], + ) + + +def test_cli_generate_all_skipped_keeps_existing_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + manifest = _pr_manifest(tmp_path, ("sample-1",)) + response_root = tmp_path / "responses" + monkeypatch.setattr( + cli, + "generate_sample", + _fake_generate({"sample-1": _metrics_for("sample-1", ttft_ms=10.0, ttfp_ms=20.0, rtf=0.5)}), + ) + _run_generate_cli(manifest, response_root) + monkeypatch.setattr(cli, "generate_sample", _fake_generate({"sample-1": None})) + _run_generate_cli(manifest, response_root) + duplex_metrics = json.loads((response_root / DUPLEX_METRICS_FILENAME).read_text(encoding="utf-8")) + assert duplex_metrics["duplex_request_metrics"][0]["ttft_ms"] == 10.0 + assert duplex_metrics["duplex_stream_ttft_ms"]["mean"] == 10.0 + assert duplex_metrics["duplex_stream_ttfp_ms"]["mean"] == 20.0 + assert duplex_metrics["duplex_stream_rtf"]["mean"] == 0.5 + + +def test_cli_generate_partial_resume_merges_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + manifest = _pr_manifest(tmp_path, ("sample-1", "sample-2")) + response_root = tmp_path / "responses" + monkeypatch.setattr( + cli, + "generate_sample", + _fake_generate( + { + "sample-1": _metrics_for("sample-1", ttft_ms=10.0, ttfp_ms=20.0, rtf=0.5), + "sample-2": _metrics_for("sample-2", ttft_ms=30.0, ttfp_ms=40.0, rtf=0.25), + } + ), + ) + _run_generate_cli(manifest, response_root) + monkeypatch.setattr( + cli, + "generate_sample", + _fake_generate( + { + "sample-1": None, + "sample-2": _metrics_for("sample-2", ttft_ms=90.0, ttfp_ms=80.0, rtf=1.5), + } + ), + ) + _run_generate_cli(manifest, response_root) + duplex_metrics = json.loads((response_root / DUPLEX_METRICS_FILENAME).read_text(encoding="utf-8")) + by_id = {row["sample_id"]: row for row in duplex_metrics["duplex_request_metrics"]} + assert by_id["sample-1"]["ttft_ms"] == 10.0 + assert by_id["sample-2"]["ttft_ms"] == 90.0 + assert duplex_metrics["duplex_stream_ttft_ms"]["mean"] == 50.0 + assert duplex_metrics["duplex_stream_ttfp_ms"]["mean"] == 50.0 + assert duplex_metrics["duplex_stream_rtf"]["mean"] == 1.0 + + +def test_collect_duplex_session_metrics_matches_omniinteract_window(): + collector = EventCollector() + audio = { + "type": "response.audio.delta", + "response_id": "r1", + "format": "pcm16", + "delta": base64.b64encode(bytes((1, 0)) * 2400).decode(), + "sample_rate_hz": 24_000, + "metadata": {"audio_duration_ms": 100}, + } + collector.add({"type": "response.created", "response": {"id": "r1"}}, received_at_s=10.1) + collector.add( + {"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."}, + received_at_s=10.2, + ) + collector.add(audio, received_at_s=10.3) + collector.add({"type": "response.done", "response": {"id": "r1"}}, received_at_s=10.4) + bundle = collect_duplex_session_metrics(collector, stream_start=10.0, session_id="sample") + assert bundle.request_metrics[0]["ttft_ms"] == 100.0 + assert bundle.request_metrics[0]["ttfp_ms"] == 200.0 + assert bundle.request_metrics[0]["rtf"] == 2.0 + assert bundle.session_metrics["stream_ttft_ms"] == 200.0 + assert bundle.session_metrics["stream_ttfp_ms"] == 300.0 + report = build_duplex_metrics_report( + request_metrics=bundle.request_metrics, + session_metrics=[bundle.session_metrics], + ) + assert report["duplex_stream_ttft_ms"]["mean"] == 200.0 + assert report["duplex_stream_ttfp_ms"]["mean"] == 300.0 + + def test_judge_exercises_openai_http_schema(): requests = [] diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py index a8ed72511..c04d3acec 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py @@ -30,7 +30,7 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import run_multi_session, ) from tests.helpers.mark import hardware_test -from vllm_omni.clients.duplex import build_realtime_url +from vllm_omni.clients.duplex import build_realtime_url, metric_mean from vllm_omni.experimental.fullduplex.video_stacking import concat_frames_b64 pytestmark = pytest.mark.omni @@ -59,9 +59,12 @@ def _assert_session_metrics(metrics: object, *, expected_count: int) -> None: assert isinstance(metrics, dict) assert isinstance(metrics["session_id"], str) assert metrics["audio_turn_count"] == expected_count - assert metrics["mean_ttft_ms"] is not None and metrics["mean_ttft_ms"] >= 0 - assert metrics["mean_ttfp_ms"] is not None and metrics["mean_ttfp_ms"] >= 0 - assert metrics["mean_rtf"] is not None and metrics["mean_rtf"] >= 0 + ttft_ms = metric_mean(metrics["ttft_ms"]) + ttfp_ms = metric_mean(metrics["ttfp_ms"]) + rtf = metric_mean(metrics["rtf"]) + assert ttft_ms is not None and ttft_ms >= 0 + assert ttfp_ms is not None and ttfp_ms >= 0 + assert rtf is not None and rtf >= 0 async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: float) -> list[dict[str, object]]: diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index d45c535db..ecd885dd6 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -1209,7 +1209,9 @@ def test_realtime_duplex_demo_partitions_timing_by_response_identity(): "ttft": "input_audio_buffer.commit client send to first non-empty text delta", "ttfp": "input_audio_buffer.commit client send to first audio packet", "rtf": "commit-to-last-audio receive time divided by emitted audio duration", + "tpot": "Stage-0 engine mean time per output token", }, + "tpot_ms": None, }, { "session_id": "seed-tts-session", @@ -1225,15 +1227,17 @@ def test_realtime_duplex_demo_partitions_timing_by_response_identity(): "ttft": "input_audio_buffer.commit client send to first non-empty text delta", "ttfp": "input_audio_buffer.commit client send to first audio packet", "rtf": "commit-to-last-audio receive time divided by emitted audio duration", + "tpot": "Stage-0 engine mean time per output token", }, + "tpot_ms": None, }, ] assert session_metrics == { "session_id": "seed-tts-session", "audio_turn_count": 2, - "mean_ttft_ms": 150.0, - "mean_ttfp_ms": 250.0, - "mean_rtf": 3.125, + "ttft_ms": {"count": 2, "mean": 150.0, "p50": 150.0, "p99": 150.0}, + "ttfp_ms": {"count": 2, "mean": 250.0, "p50": 200.0, "p99": 300.0}, + "rtf": {"count": 2, "mean": 3.125, "p50": 2.5, "p99": 3.75}, } diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 62ed8334b..8618061d8 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard.""" +from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import MagicMock @@ -38,6 +39,21 @@ _DETOK = MagicMock( ) _LOGPROBS = MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None) + +@dataclass +class _StreamingUpdate: + final: bool + prompt: object | None + prompt_token_ids: list[int] + arrival_time: float + + +@dataclass +class _FinishedRequestMetric: + request_id: str + mean_time_per_output_token: float + + _DEFAULT_STATE_KWARGS = dict( request_id="r", external_req_id="r", @@ -130,6 +146,119 @@ def test_native_text_metrics_include_segment_generation_token_count(monkeypatch) assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 27 +def test_native_text_metrics_exclude_cross_segment_gap(monkeypatch): + monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None) + processor = object.__new__(MultimodalOutputProcessor) + processor._native_text_metrics_by_request = {} + processor.lora_states = {} + state = _make_state(RequestOutputKind.DELTA) + iteration_stats = MagicMock() + + def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args): + native_stats.num_generation_tokens = 2 if timestamp == 10.03 else 1 + native_stats.first_token_ts = 10.02 if timestamp == 10.03 else timestamp + native_stats.last_token_ts = timestamp + native_stats.first_token_latency = 0.01 + + iteration_stats.update_from_output.side_effect = update_segment + processor._update_stats_from_output( + state, + MagicMock(), + 10.0, + iteration_stats, + ) + first_segment = processor.pop_native_text_metrics("r") + assert first_segment["vllm_tpot_ms"] == 0.0 + assert first_segment["vllm_itls_ms"] == [] + + state.apply_streaming_update( + _StreamingUpdate( + final=False, + prompt=None, + prompt_token_ids=[], + arrival_time=10.01, + ) + ) + state.is_prefilling = True + processor._update_stats_from_output( + state, + MagicMock(), + 10.02, + iteration_stats, + ) + state.is_prefilling = False + processor._update_stats_from_output( + state, + MagicMock(), + 10.03, + iteration_stats, + ) + second_segment = processor.pop_native_text_metrics("r") + + assert second_segment["num_generation_tokens"] == 2 + assert second_segment["vllm_itls_ms"] == pytest.approx([10.0]) + assert second_segment["vllm_tpot_ms"] == pytest.approx(10.0) + + +def test_native_text_tpot_weights_multi_token_engine_output(monkeypatch): + monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None) + processor = object.__new__(MultimodalOutputProcessor) + processor._native_text_metrics_by_request = {} + processor.lora_states = {} + state = _make_state(RequestOutputKind.DELTA) + iteration_stats = MagicMock() + + def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args): + native_stats.num_generation_tokens = 1 if timestamp == 10.0 else 4 + native_stats.first_token_ts = 10.0 + native_stats.last_token_ts = timestamp + native_stats.first_token_latency = 0.01 + + iteration_stats.update_from_output.side_effect = update_segment + state.is_prefilling = True + processor._update_stats_from_output(state, MagicMock(), 10.0, iteration_stats) + state.is_prefilling = False + processor._update_stats_from_output(state, MagicMock(), 10.03, iteration_stats) + record = processor.pop_native_text_metrics("r") + + assert record["num_generation_tokens"] == 4 + assert record["vllm_itls_ms"] == pytest.approx([30.0]) + assert record["vllm_itl_ms"] == pytest.approx(30.0) + assert record["vllm_tpot_ms"] == pytest.approx(10.0) + assert "_tpot_elapsed_ms" not in record + assert "_tpot_intervals" not in record + + +@pytest.mark.parametrize( + ("finished_tpot_s", "expected_tpot_ms"), + [(0.0, 19.0), (0.023, 23.0)], +) +def test_native_text_tpot_only_accepts_positive_finished_metric( + monkeypatch, + finished_tpot_s, + expected_tpot_ms, +): + monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_finished", lambda *args, **kwargs: None) + processor = object.__new__(MultimodalOutputProcessor) + processor._native_text_metrics_by_request = {"r": {"vllm_tpot_ms": 19.0}} + state = _make_state(RequestOutputKind.DELTA) + iteration_stats = MagicMock() + iteration_stats.finished_requests = [ + _FinishedRequestMetric( + request_id="r", + mean_time_per_output_token=finished_tpot_s, + ) + ] + + processor._update_stats_from_finished( + state, + FinishReason.STOP, + iteration_stats, + ) + + assert processor.pop_native_text_metrics("r")["vllm_tpot_ms"] == expected_tpot_ms + + def test_delta_drains_output_modality_per_step(): """DELTA drains the mm_type key (output modality) but preserves hidden-state keys.""" s = _make_state(RequestOutputKind.DELTA) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index b4a943873..0b64e7f14 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -9007,19 +9007,82 @@ async def test_minicpmo_native_duplex_continuous_speak_reuses_active_response_un assert session.turn_id == 1 +def test_response_timing_binds_latest_request_start_for_model_turn(): + session = DuplexSession( + session_id="sid-response-request-timing", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.mark_model_turn_request_started(0, 10.0) + session.mark_model_turn_request_started(0, 11.0) + session.begin_response(turn_id=0) + + assert session.mark_response_first_outputs( + observed_at_s=11.2, + has_text=True, + has_audio=False, + )["ttft_ms"] == pytest.approx(200.0) + + session.mark_model_turn_request_started(0, 12.0) + assert session.mark_response_first_outputs( + observed_at_s=12.3, + has_text=False, + has_audio=True, + )["ttfp_ms"] == pytest.approx(1300.0) + + +def test_response_tpot_fallback_ignores_single_token_segment(): + session = DuplexSession( + session_id="sid-response-tpot", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.begin_response(turn_id=0) + + first_metrics = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 1, "vllm_tpot_ms": 900.0}}) + assert "vllm_tpot_ms" not in first_metrics["0"] + + second_metrics = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 2, "vllm_tpot_ms": 15.0}}) + assert second_metrics["0"]["num_tokens_out"] == 3 + assert second_metrics["0"]["vllm_tpot_ms"] == 15.0 + + +def test_response_tpot_keeps_token_weighted_value_when_itls_exist(): + session = DuplexSession( + session_id="sid-response-tpot-multi-token", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.begin_response(turn_id=0) + metrics = session.accumulate_response_stage_metrics( + { + "0": { + "num_tokens_out": 4, + "vllm_tpot_ms": 10.0, + "vllm_itls_ms": [30.0], + } + } + ) + assert metrics["0"]["vllm_itls_ms"] == [30.0] + assert metrics["0"]["vllm_itl_ms"] == 30.0 + assert metrics["0"]["vllm_tpot_ms"] == 10.0 + + @pytest.mark.asyncio -async def test_continuous_response_metrics_accumulate_only_owned_model_units(): +async def test_continuous_response_metrics_accumulate_only_owned_model_units(monkeypatch): handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) session = DuplexSession( session_id="sid-response-metrics", config=DuplexSessionConfig(extra_body={"auto_response": True}), ) + session.mark_model_turn_request_started(0, 10.0) + monkeypatch.setattr( + "vllm_omni.entrypoints.duplex.runtime_bridge.time.monotonic", + lambda: 10.25, + ) sent: list[dict[str, Any]] = [] async def send_json(payload: dict[str, Any]) -> None: sent.append(payload) - async def emit(*, text: str, audio: str, tokens: int, ttft_ms: float, itls: list[float]) -> None: + async def emit(*, text: str, audio: str, tokens: int, ttft_ms: float, itls: list[float], tpot_ms: float) -> None: await handler._send_one_native_duplex_event( send_json, { @@ -9038,7 +9101,7 @@ async def test_continuous_response_metrics_accumulate_only_owned_model_units(): "num_tokens_out": tokens, "stage_gen_time_ms": 80.0, "vllm_ttft_ms": ttft_ms, - "vllm_tpot_ms": 9.0, + "vllm_tpot_ms": tpot_ms, "vllm_itls_ms": itls, } }, @@ -9046,21 +9109,32 @@ async def test_continuous_response_metrics_accumulate_only_owned_model_units(): session=session, ) - await emit(text="first", audio="audio-a", tokens=4, ttft_ms=50.0, itls=[10.0, 11.0]) - await emit(text="second", audio="audio-b", tokens=6, ttft_ms=70.0, itls=[12.0, 13.0, 14.0]) + await emit(text="first", audio="audio-a", tokens=3, ttft_ms=50.0, itls=[10.0, 11.0], tpot_ms=10.5) + await emit(text="second", audio="audio-b", tokens=4, ttft_ms=70.0, itls=[12.0, 13.0, 14.0], tpot_ms=13.0) deltas = [payload for payload in sent if payload.get("type") == "response.output_audio.delta"] assert len(deltas) == 2 first_metrics = deltas[0]["vllm_omni"]["stage_metrics"]["0"] second_metrics = deltas[1]["vllm_omni"]["stage_metrics"]["0"] - assert first_metrics["num_tokens_out"] == 4 - assert second_metrics["num_tokens_out"] == 10 + assert first_metrics["num_tokens_out"] == 3 + assert first_metrics["vllm_tpot_ms"] == pytest.approx(10.5) + assert deltas[0]["vllm_omni"]["response_request_metrics"] == { + "source": "server_monotonic_request_start", + "measurement_origin": { + "ttft": "native model-turn request execution start to first non-empty text output", + "ttfp": "native model-turn request execution start to first audio output", + }, + "ttft_ms": 250.0, + "ttfp_ms": 250.0, + } + assert second_metrics["num_tokens_out"] == 7 assert second_metrics["vllm_ttft_ms"] == 50.0 assert second_metrics["vllm_itls_ms"] == [10.0, 11.0, 12.0, 13.0, 14.0] + assert second_metrics["vllm_tpot_ms"] == pytest.approx(12.0) first_response_id = deltas[0]["response_id"] session.end_response() - await emit(text="third", audio="audio-c", tokens=3, ttft_ms=90.0, itls=[15.0, 16.0]) + await emit(text="third", audio="audio-c", tokens=3, ttft_ms=90.0, itls=[15.0, 16.0], tpot_ms=15.5) third_delta = [payload for payload in sent if payload.get("type") == "response.output_audio.delta"][-1] third_metrics = third_delta["vllm_omni"]["stage_metrics"]["0"] @@ -9068,6 +9142,7 @@ async def test_continuous_response_metrics_accumulate_only_owned_model_units(): assert third_metrics["num_tokens_out"] == 3 assert third_metrics["vllm_ttft_ms"] == 90.0 assert third_metrics["vllm_itls_ms"] == [15.0, 16.0] + assert third_metrics["vllm_tpot_ms"] == pytest.approx(15.5) @pytest.mark.asyncio diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py index d5a28aa06..db500ebf0 100644 --- a/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py @@ -8,24 +8,58 @@ from __future__ import annotations import asyncio import hashlib import json +import time +from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING import pybase64 as base64 -from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, build_realtime_url, wait_for +from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample +from vllm_omni.experimental.fullduplex.client import ( + RealtimeDuplexClient, + RealtimeEventCollector, + build_realtime_url, + wait_for, +) + +if TYPE_CHECKING: + from vllm_omni.clients.duplex import EventCollector from .omni_duplex_eval_clock import extract_timed_sentences from .omni_duplex_eval_media import iter_av_units, iter_jpegs, materialize_media, read_audio_pcm16, video_duration +@dataclass(frozen=True) +class GenerateSampleResult: + """Timed-sentence path plus duplex metrics for one generate sample.""" + + output: Path + request_metrics: list[dict[str, object]] = field(default_factory=list) + session_metrics: dict[str, object] = field(default_factory=dict) + + def _ref_audio(path: str | Path) -> str: value = Path(path).expanduser().read_bytes() return "data:audio/wav;base64," + base64.b64encode(value).decode("ascii") +def _event_collector_from_realtime(source: RealtimeEventCollector) -> EventCollector: + """Replay experimental collector events onto the public EventCollector.""" + from vllm_omni.clients.duplex import EventCollector + + collector = EventCollector() + for event, received_at_s in zip(source.events, source.event_received_at_s, strict=True): + collector.add(event, received_at_s=received_at_s) + return collector + + +def _with_sample_identity(payload: dict[str, object], sample: DuplexSample) -> dict[str, object]: + return {"sample_id": sample.id, "split": sample.split, **payload} + + async def generate_sample( - sample: Any, + sample: DuplexSample, *, url: str, model: str, @@ -37,11 +71,17 @@ async def generate_sample( clock: str = "media", overwrite: bool = False, unit_ms: int = 1000, -) -> Path: +) -> GenerateSampleResult: + """Generate one sample's timed sentences and collect duplex session metrics. + + Timed sentences and ``*.meta.json`` stay the quality-eval artifacts. Duplex + request/session metrics are returned for the CLI to write + ``duplex_metrics.json``; they are not written into meta. + """ output = Path(output_root) / sample.split / f"{sample.id}.json" meta_path = output.with_name(output.stem + ".meta.json") if output.exists() and not overwrite: - return output + return GenerateSampleResult(output=output) if mix != "question": raise NotImplementedError("v1 supports mix=question; soundtrack mixing is reserved for P1") media_dir = output.parent / ".media" @@ -57,10 +97,12 @@ async def generate_sample( response_done = False drain_timeout = None close_timeout = None + stream_start: float | None = None async with client: await client.configure(model, ref_audio=_ref_audio(ref_audio), instructions="Streaming Omni Conversation.") ack_task = asyncio.create_task(_ack_playback(client)) try: + stream_start = time.monotonic() await client.stream_av_units(iter_av_units(pcm, frames, unit_ms=unit_ms), realtime=realtime) await client.commit() try: @@ -84,6 +126,7 @@ async def generate_sample( except TimeoutError as exc: close_timeout = str(exc) events = list(client.events.events) + collector = _event_collector_from_realtime(client.events) timed = [sentence.as_dict() for sentence in extract_timed_sentences(events, clock=clock)] output.parent.mkdir(parents=True, exist_ok=True) output.write_text(json.dumps(timed, ensure_ascii=False, indent=2), encoding="utf-8") @@ -102,7 +145,18 @@ async def generate_sample( "ref_audio_sha256": hashlib.sha256(Path(ref_audio).read_bytes()).hexdigest(), } meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") - return output + from vllm_omni.benchmarks.duplex_session_metrics import collect_duplex_session_metrics + + bundle = collect_duplex_session_metrics( + collector, + stream_start=stream_start if stream_start is not None else 0.0, + session_id=sample.id, + ) + return GenerateSampleResult( + output=output, + request_metrics=[_with_sample_identity(metric, sample) for metric in bundle.request_metrics], + session_metrics=_with_sample_identity(bundle.session_metrics, sample), + ) async def _ack_playback(client: RealtimeDuplexClient) -> None: diff --git a/vllm_omni/benchmarks/duplex_session_metrics.py b/vllm_omni/benchmarks/duplex_session_metrics.py new file mode 100644 index 000000000..ba3bab5ac --- /dev/null +++ b/vllm_omni/benchmarks/duplex_session_metrics.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Assemble duplex serve metrics from a public EventCollector. + +This module does not instrument the server. It reads the same +``EventCollector.timing_summary`` / ``global_timing_summary`` surfaces that +OmniInteract already uses, derives RTF with ``compute_audio_rtf``, and folds +per-session ``stream_*`` values into a run-level report. +""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from vllm_omni.metrics.definitions import compute_audio_rtf + +if TYPE_CHECKING: + from vllm_omni.clients.duplex import EventCollector + +DUPLEX_METRICS_FILENAME = "duplex_metrics.json" + +_REQUEST_MEASUREMENT_ORIGIN = { + "tpot": "Stage-0 engine mean time per output token", + "rtf": "response.created client receive to last audio packet divided by emitted audio duration", +} +_STREAM_MEASUREMENT_ORIGIN = { + "ttft": "input stream start to first non-empty text delta", + "ttfp": "input stream start to first audio packet", + "rtf": ( + "input stream start-to-last-audio receive time divided by total emitted audio duration; " + "includes concurrent realtime input" + ), +} + + +@dataclass(frozen=True) +class DuplexSessionMetricBundle: + """One session's request rows plus the session/stream summary.""" + + request_metrics: list[dict[str, object]] + session_metrics: dict[str, object] + output_tokens: int + + +def audio_rtf_from_raw_metric(raw_metric: Mapping[str, object]) -> float | None: + """Derive audio RTF from client-reported generation and duration milliseconds.""" + generation_ms = raw_metric.get("audio_generation_ms") + duration_ms = raw_metric.get("audio_duration_ms") + if not isinstance(generation_ms, int | float) or not isinstance(duration_ms, int | float) or duration_ms <= 0: + return None + return round(compute_audio_rtf(float(generation_ms) / 1000.0, float(duration_ms) / 1000.0), 6) + + +def collect_duplex_session_metrics( + collector: EventCollector, + *, + stream_start: float, + session_id: str | None, +) -> DuplexSessionMetricBundle: + """Build per-response and session/stream metrics for one duplex session. + + ``stream_start`` is the client monotonic time just before input media is + pushed. Per-response TTFT/TTFP prefer server ``response_request_metrics`` + when present; RTF and the session ``stream_*`` window stay client-receive. + """ + from vllm_omni.clients.duplex import summarize_session_request_metrics + + request_metrics: list[dict[str, object]] = [] + output_tokens = 0 + for request_index, response_id in enumerate(collector.response_ids): + timing = collector.timing_summary( + after_s=stream_start, + input_committed_at_s=None, + response_id=response_id, + measurement_origin=_REQUEST_MEASUREMENT_ORIGIN, + ) + raw_metric = timing.get("request_metrics") + stage0 = timing.get("stage0_tokens") + metric: dict[str, object] = { + "session_id": session_id, + "request_index": request_index, + "response_id": response_id, + } + if isinstance(raw_metric, dict): + metric.update(raw_metric) + metric["rtf"] = audio_rtf_from_raw_metric(raw_metric) + if isinstance(stage0, dict): + metric["stage0_tokens"] = dict(stage0) + output_tokens += int(stage0.get("output_token_count") or 0) + if isinstance(raw_metric, dict) or isinstance(stage0, dict): + request_metrics.append(metric) + session_metrics = summarize_session_request_metrics( + request_metrics, + session_id=session_id, + ) + stream_metrics = collector.global_timing_summary( + after_s=stream_start, + window_started_at_s=stream_start, + response_ids=list(collector.response_ids), + measurement_origin=_STREAM_MEASUREMENT_ORIGIN, + ) + if stream_metrics: + session_metrics.update( + { + "stream_ttft_ms": stream_metrics.get("ttft_ms"), + "stream_ttfp_ms": stream_metrics.get("ttfp_ms"), + "stream_rtf": audio_rtf_from_raw_metric(stream_metrics), + "stream_audio_generation_ms": stream_metrics.get("audio_generation_ms"), + "stream_audio_duration_ms": stream_metrics.get("audio_duration_ms"), + "stream_measurement_origin": stream_metrics.get("measurement_origin"), + } + ) + return DuplexSessionMetricBundle( + request_metrics=request_metrics, + session_metrics=session_metrics, + output_tokens=output_tokens, + ) + + +def duplex_stream_metrics(session_metrics: Sequence[Mapping[str, object]]) -> dict[str, object]: + """Summarize finite, non-negative per-session ``stream_*`` values.""" + from vllm_omni.clients.duplex import distribution_summary + + result: dict[str, object] = {} + for session_key, result_key, digits in ( + ("stream_ttft_ms", "duplex_stream_ttft_ms", 3), + ("stream_ttfp_ms", "duplex_stream_ttfp_ms", 3), + ("stream_rtf", "duplex_stream_rtf", 6), + ): + values = [ + float(value) + for metric in session_metrics + if isinstance((value := metric.get(session_key)), int | float) + and not isinstance(value, bool) + and math.isfinite(value) + and value >= 0 + ] + summary = distribution_summary(values, digits=digits) + if summary is not None: + result[result_key] = summary + return result + + +def build_duplex_metrics_report( + *, + request_metrics: Sequence[Mapping[str, object]], + session_metrics: Sequence[Mapping[str, object]], +) -> dict[str, object]: + """Flatten this generate run into the ``duplex_metrics.json`` payload.""" + report: dict[str, object] = { + "duplex_request_metrics": [dict(metric) for metric in request_metrics], + "duplex_session_metrics": [dict(metric) for metric in session_metrics], + } + report.update(duplex_stream_metrics(session_metrics)) + return report + + +def sample_metric_key(metric: Mapping[str, object]) -> tuple[str, str] | None: + """Identity used to merge resume rows: ``(split, sample_id)``.""" + split = metric.get("split") + sample_id = metric.get("sample_id") + if isinstance(split, str) and split and isinstance(sample_id, str) and sample_id: + return (split, sample_id) + return None + + +def merge_metric_rows( + existing: Sequence[Mapping[str, object]], + incoming: Sequence[Mapping[str, object]], +) -> list[dict[str, object]]: + """Replace existing rows that share an incoming ``(split, sample_id)``. + + Incoming rows without a key are appended. Existing rows without a key are + kept because they cannot be matched. + """ + incoming_keys = {key for metric in incoming if (key := sample_metric_key(metric)) is not None} + merged = [dict(metric) for metric in existing if (key := sample_metric_key(metric)) not in incoming_keys] + merged.extend(dict(metric) for metric in incoming) + return merged + + +def read_duplex_metrics_report(path: Path) -> dict[str, object] | None: + """Load an on-disk report. Missing files return ``None``; corrupt JSON raises.""" + if not path.is_file(): + return None + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError(f"{path} must contain a JSON object") + return payload + + +def _metric_rows(value: object, *, field_name: str) -> list[Mapping[str, object]]: + if value is None: + return [] + if not isinstance(value, list): + raise ValueError(f"{field_name} must be a JSON list") + rows: list[Mapping[str, object]] = [] + for item in value: + if not isinstance(item, dict): + raise ValueError(f"{field_name} entries must be JSON objects") + rows.append(item) + return rows + + +def merge_duplex_metrics_report( + existing: Mapping[str, object] | None, + *, + request_metrics: Sequence[Mapping[str, object]], + session_metrics: Sequence[Mapping[str, object]], +) -> dict[str, object]: + """Merge this run into a previous ``duplex_metrics.json`` payload.""" + previous_requests = _metric_rows( + None if existing is None else existing.get("duplex_request_metrics"), + field_name="duplex_request_metrics", + ) + previous_sessions = _metric_rows( + None if existing is None else existing.get("duplex_session_metrics"), + field_name="duplex_session_metrics", + ) + return build_duplex_metrics_report( + request_metrics=merge_metric_rows(previous_requests, request_metrics), + session_metrics=merge_metric_rows(previous_sessions, session_metrics), + ) diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 52eec86b9..759e60667 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -916,11 +916,9 @@ def calculate_metrics( all_tpots.append(tpot) itls += outputs[i].itl session_metrics = getattr(outputs[i], "duplex_session_metrics", None) - ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None - audio_ttfp_measured = ( - not isinstance(session_metrics, dict) or session_metrics.get("mean_ttfp_ms") is not None - ) - audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_rtf") is not None + ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("ttft_ms") is not None + audio_ttfp_measured = not isinstance(session_metrics, dict) or session_metrics.get("ttfp_ms") is not None + audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("rtf") is not None if ttft_measured: ttfts.append(outputs[i].ttft) if audio_ttfp_measured: diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py index 6b861d339..83e224f34 100644 --- a/vllm_omni/benchmarks/omniinteract.py +++ b/vllm_omni/benchmarks/omniinteract.py @@ -946,13 +946,9 @@ class _RealtimeSession: def _audio_rtf_from_raw_metric(raw_metric: dict[str, object]) -> float | None: """Derive the audio RTF from a client-reported raw request metric.""" - from vllm_omni.metrics.definitions import compute_audio_rtf + from vllm_omni.benchmarks.duplex_session_metrics import audio_rtf_from_raw_metric - generation_ms = raw_metric.get("audio_generation_ms") - duration_ms = raw_metric.get("audio_duration_ms") - if not isinstance(generation_ms, int | float) or not isinstance(duration_ms, int | float) or duration_ms <= 0: - return None - return round(compute_audio_rtf(float(generation_ms) / 1000.0, float(duration_ms) / 1000.0), 6) + return audio_rtf_from_raw_metric(raw_metric) def _populate_response_metrics( @@ -961,45 +957,16 @@ def _populate_response_metrics( *, stream_start: float, ) -> None: - measurement_origin = { - "ttft": "response.created client receive to first non-empty text delta", - "ttfp": "response.created client receive to first audio packet", - "rtf": "response.created client receive to last audio packet divided by emitted audio duration", - } - request_metrics: list[dict[str, object]] = [] - output_tokens = 0 - for request_index, response_id in enumerate(collector.response_ids): - timing = collector.timing_summary( - after_s=stream_start, - input_committed_at_s=None, - response_id=response_id, - measurement_origin=measurement_origin, - ) - raw_metric = timing.get("request_metrics") - stage0 = timing.get("stage0_tokens") - metric = { - "session_id": result.session_id, - "request_index": request_index, - "response_id": response_id, - **(raw_metric if isinstance(raw_metric, dict) else {}), - } - if isinstance(raw_metric, dict): - # The client reports raw data only; the RTF is derived here with - # the canonical server-side metric definition. - metric["rtf"] = _audio_rtf_from_raw_metric(raw_metric) - if isinstance(stage0, dict): - metric["stage0_tokens"] = dict(stage0) - output_tokens += int(stage0.get("output_token_count") or 0) - if isinstance(raw_metric, dict) or isinstance(stage0, dict): - request_metrics.append(metric) - result.output_tokens = output_tokens - result.duplex_request_metrics = request_metrics - from vllm_omni.clients.duplex import summarize_session_request_metrics + from vllm_omni.benchmarks.duplex_session_metrics import collect_duplex_session_metrics - result.duplex_session_metrics = summarize_session_request_metrics( - request_metrics, + bundle = collect_duplex_session_metrics( + collector, + stream_start=stream_start, session_id=result.session_id, ) + result.output_tokens = bundle.output_tokens + result.duplex_request_metrics = bundle.request_metrics + result.duplex_session_metrics = bundle.session_metrics async def run_omniinteract_case( diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 3313e691f..a230b9c2e 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -2272,9 +2272,11 @@ async def _async_request_omniinteract( output.audio_duration = case_result.audio_bytes / (24_000 * 2) output.audio_frames = case_result.audio_bytes // 2 session_metrics = case_result.duplex_session_metrics - output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0 - output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0 - output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0) + from vllm_omni.clients.duplex import metric_mean + + output.ttft = (metric_mean(session_metrics.get("ttft_ms")) or 0.0) / 1000.0 + output.audio_ttfp = (metric_mean(session_metrics.get("ttfp_ms")) or 0.0) / 1000.0 + output.audio_rtf = metric_mean(session_metrics.get("rtf")) or 0.0 token_timing_measured = _apply_stage0_token_timings( output, [request_metric.get("stage0_tokens") for request_metric in case_result.duplex_request_metrics], @@ -2524,9 +2526,11 @@ async def async_request_openai_realtime_duplex( await client.close_session(timeout_s=30.0) output.generated_text = " ".join(filter(None, turn_transcripts)) - output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0 - output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0 - output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0) + from vllm_omni.clients.duplex import metric_mean + + output.ttft = (metric_mean(session_metrics.get("ttft_ms")) or 0.0) / 1000.0 + output.audio_ttfp = (metric_mean(session_metrics.get("ttfp_ms")) or 0.0) / 1000.0 + output.audio_rtf = metric_mean(session_metrics.get("rtf")) or 0.0 output.audio_duration = ( sum(float(metric.get("audio_duration_ms") or 0.0) for metric in turn_metrics) / 1000.0 ) @@ -2926,7 +2930,7 @@ async def benchmark( def measured_ttft(output: RequestFuncOutput) -> float | None: session_metrics = getattr(output, "duplex_session_metrics", None) - if isinstance(session_metrics, dict) and session_metrics.get("mean_ttft_ms") is None: + if isinstance(session_metrics, dict) and session_metrics.get("ttft_ms") is None: return None return output.ttft @@ -2998,6 +3002,24 @@ async def benchmark( ] if duplex_session_metrics: result["duplex_session_metrics"] = duplex_session_metrics + from vllm_omni.clients.duplex import distribution_summary + + for session_key, result_key, digits in ( + ("stream_ttft_ms", "duplex_stream_ttft_ms", 3), + ("stream_ttfp_ms", "duplex_stream_ttfp_ms", 3), + ("stream_rtf", "duplex_stream_rtf", 6), + ): + values = [ + float(value) + for metric in duplex_session_metrics + if isinstance((value := metric.get(session_key)), int | float) + and not isinstance(value, bool) + and np.isfinite(value) + and value >= 0 + ] + summary = distribution_summary(values, digits=digits) + if summary is not None: + result[result_key] = summary if omniinteract_summary is not None: result["omniinteract"] = omniinteract_summary diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index 15a191cc1..0504c241b 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -35,7 +35,7 @@ import math import random import time import wave -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Protocol @@ -80,6 +80,8 @@ __all__ = [ "image_data_url", "read_pcm16_wav", "reference_audio_data_url", + "distribution_summary", + "metric_mean", "summarize_session_request_metrics", "wait_for_condition", "write_pcm16_wav", @@ -1248,6 +1250,15 @@ def _rounded_ms(value: float) -> float: return round(float(value), 3) +def _finite_number(value: object, *, nonnegative: bool = False) -> float | None: + if not isinstance(value, int | float) or isinstance(value, bool) or not math.isfinite(float(value)): + return None + number = float(value) + if nonnegative and number < 0: + return None + return number + + def _interval_summary(values: list[float]) -> dict[str, float | int]: clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0) if not clean: @@ -1266,6 +1277,36 @@ def _interval_summary(values: list[float]) -> dict[str, float | int]: } +def distribution_summary(values: Sequence[float], *, digits: int = 3) -> dict[str, float | int] | None: + """Summarize values as ``{count, mean, p50, p99}`` for duplex report fields.""" + clean = sorted(float(value) for value in values if math.isfinite(float(value))) + if not clean: + return None + + def nearest_rank(percentile: float) -> float: + index = max(0, math.ceil(percentile * len(clean)) - 1) + return clean[min(index, len(clean) - 1)] + + return { + "count": len(clean), + "mean": round(sum(clean) / len(clean), digits), + "p50": round(nearest_rank(0.50), digits), + "p99": round(nearest_rank(0.99), digits), + } + + +def metric_mean(value: object) -> float | None: + """Read a scalar mean, or the ``mean`` field of a distribution summary.""" + if isinstance(value, Mapping): + nested = value.get("mean") + if isinstance(nested, int | float) and not isinstance(nested, bool) and math.isfinite(float(nested)): + return float(nested) + return None + if isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(float(value)): + return float(value) + return None + + def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None: candidates: list[object] = [event.get("vllm_omni")] metadata = event.get("metadata") @@ -1285,6 +1326,34 @@ def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None: return None +def _event_response_request_metrics(event: dict[str, object]) -> dict[str, object] | None: + direct_metrics = event.get("response_request_metrics") + if isinstance(direct_metrics, dict): + return direct_metrics + candidates: list[object] = [event.get("vllm_omni")] + metadata = event.get("metadata") + if isinstance(metadata, dict): + candidates.extend((metadata, metadata.get("vllm_omni"), metadata.get("duplex_event"))) + response = event.get("response") + if isinstance(response, dict): + response_metadata = response.get("metadata") + if isinstance(response_metadata, dict): + candidates.extend( + ( + response_metadata, + response_metadata.get("vllm_omni"), + response_metadata.get("duplex_event"), + ) + ) + for candidate in candidates: + if not isinstance(candidate, dict): + continue + request_metrics = candidate.get("response_request_metrics") + if isinstance(request_metrics, dict): + return request_metrics + return None + + class EventCollector: """Accumulate events for assertions and latency summaries (tests/benchmarks). @@ -1379,6 +1448,85 @@ class EventCollector: return received_at_s return None + def global_timing_summary( + self, + *, + after_s: float, + window_started_at_s: float, + response_ids: list[str], + measurement_origin: dict[str, str], + ) -> dict[str, object]: + """Summarize one client-observed timing window across responses. + + Returns raw measurements only. Derived metrics such as RTF are computed + by the caller — e.g. with ``vllm_omni.metrics.definitions.compute_audio_rtf``. + """ + selected_response_ids = set(response_ids) + first_text_received_at_s: float | None = None + audio_received_at_s: list[float] = [] + response_audio_duration_ms: dict[str, float] = {} + + for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True): + if received_at_s < after_s: + continue + response_id = self.response_id(event) + if response_id is None or response_id not in selected_response_ids: + continue + if ( + event.get("type") + in { + "response.audio_transcript.delta", + "response.output_text.delta", + "response.text.delta", + } + and isinstance(event.get("delta"), str) + and bool(event["delta"]) + and first_text_received_at_s is None + ): + first_text_received_at_s = received_at_s + if event.get("type") not in _AUDIO_DELTA_TYPES: + continue + delta = event.get("delta") or event.get("audio") + if not isinstance(delta, str) or not delta: + continue + audio_received_at_s.append(received_at_s) + metadata = event.get("metadata") + duration_ms = metadata.get("audio_duration_ms") if isinstance(metadata, dict) else None + parsed_duration_ms = _finite_number(duration_ms, nonnegative=True) + if parsed_duration_ms is not None: + response_audio_duration_ms[response_id] = max( + response_audio_duration_ms.get(response_id, 0.0), + parsed_duration_ms, + ) + + if not audio_received_at_s: + return {} + + audio_duration_ms = sum( + response_audio_duration_ms.get( + response_id, + self._format_for_bytes().duration_ms(len(self.audio_bytes(response_id))), + ) + for response_id in response_ids + ) + audio_generation_ms = max( + 0.0, + (audio_received_at_s[-1] - window_started_at_s) * 1000.0, + ) + return { + "source": "client_monotonic_receive", + "response_ids": list(response_ids), + "measurement_origin": dict(measurement_origin), + "ttft_ms": ( + _rounded_ms((first_text_received_at_s - window_started_at_s) * 1000.0) + if first_text_received_at_s is not None + else None + ), + "ttfp_ms": _rounded_ms((audio_received_at_s[0] - window_started_at_s) * 1000.0), + "audio_generation_ms": _rounded_ms(audio_generation_ms), + "audio_duration_ms": _rounded_ms(audio_duration_ms), + } + def timing_summary( self, *, @@ -1389,6 +1537,7 @@ class EventCollector: ) -> dict[str, object]: """Summarize engine token metrics and client-observed audio cadence.""" stage0_metrics: dict[str, object] | None = None + response_request_metrics: dict[str, object] = {} response_created_at_s: float | None = None first_text_received_at_s: float | None = None audio_received_at_s: list[float] = [] @@ -1419,6 +1568,10 @@ class EventCollector: if isinstance(stage0, dict): stage0_metrics = stage0 + event_request_metrics = _event_response_request_metrics(event) + if event_request_metrics is not None: + response_request_metrics.update(event_request_metrics) + if event.get("type") not in _AUDIO_DELTA_TYPES: continue delta = event.get("delta") or event.get("audio") @@ -1478,6 +1631,31 @@ class EventCollector: } request_started_at_s = input_committed_at_s if input_committed_at_s is not None else response_created_at_s if request_started_at_s is not None: + default_measurement_origin = ( + { + "ttft": "input_audio_buffer.commit client send to first non-empty text delta", + "ttfp": "input_audio_buffer.commit client send to first audio packet", + } + if input_committed_at_s is not None + else { + "ttft": "response.created client receive to first non-empty text delta", + "ttfp": "response.created client receive to first audio packet", + } + ) + resolved_measurement_origin = { + **default_measurement_origin, + **(measurement_origin or {}), + } + if stage0_metrics is not None: + resolved_measurement_origin["tpot"] = "Stage-0 engine mean time per output token" + server_ttft_ms = _finite_number(response_request_metrics.get("ttft_ms"), nonnegative=True) + server_ttfp_ms = _finite_number(response_request_metrics.get("ttfp_ms"), nonnegative=True) + server_origins = response_request_metrics.get("measurement_origin") + if isinstance(server_origins, dict): + if server_ttft_ms is not None and isinstance(server_origins.get("ttft"), str): + resolved_measurement_origin["ttft"] = server_origins["ttft"] + if server_ttfp_ms is not None and isinstance(server_origins.get("ttfp"), str): + resolved_measurement_origin["ttfp"] = server_origins["ttfp"] audio_duration_ms = ( max(cumulative_audio_ms) if cumulative_audio_ms @@ -1493,21 +1671,49 @@ class EventCollector: # compute_audio_rtf — so metric definitions stay out of this # dependency-free client module. result["request_metrics"] = { - "source": "client_monotonic_receive", - "measurement_origin": measurement_origin - or { - "ttft": "input_audio_buffer.commit client send to first non-empty text delta", - "ttfp": "input_audio_buffer.commit client send to first audio packet", - }, + "source": ( + "server_request_start_and_client_receive" + if server_ttft_ms is not None or server_ttfp_ms is not None + else "client_monotonic_receive" + ), + "measurement_origin": resolved_measurement_origin, "ttft_ms": ( - _rounded_ms((first_text_received_at_s - request_started_at_s) * 1000.0) - if first_text_received_at_s is not None - else None + server_ttft_ms + if server_ttft_ms is not None + else ( + _rounded_ms((first_text_received_at_s - request_started_at_s) * 1000.0) + if first_text_received_at_s is not None + else None + ) + ), + "ttfp_ms": ( + server_ttfp_ms + if server_ttfp_ms is not None + else _rounded_ms((audio_received_at_s[0] - request_started_at_s) * 1000.0) ), - "ttfp_ms": _rounded_ms((audio_received_at_s[0] - request_started_at_s) * 1000.0), "audio_generation_ms": _rounded_ms(audio_generation_ms), "audio_duration_ms": _rounded_ms(audio_duration_ms), } + if server_ttft_ms is not None or server_ttfp_ms is not None: + request_metrics = result["request_metrics"] + assert isinstance(request_metrics, dict) + request_metrics["response_created_to_first_text_ms"] = ( + _rounded_ms((first_text_received_at_s - response_created_at_s) * 1000.0) + if first_text_received_at_s is not None and response_created_at_s is not None + else None + ) + request_metrics["response_created_to_first_audio_ms"] = ( + _rounded_ms((audio_received_at_s[0] - response_created_at_s) * 1000.0) + if response_created_at_s is not None + else None + ) + if stage0_metrics is not None: + request_metrics = result["request_metrics"] + assert isinstance(request_metrics, dict) + request_metrics["tpot_ms"] = _finite_number( + stage0_metrics.get("vllm_tpot_ms"), + nonnegative=True, + ) return result def _format_for_bytes(self) -> AudioFormat: @@ -1640,30 +1846,38 @@ def summarize_session_request_metrics( *, session_id: str | None, ) -> dict[str, object]: - """Average client-observed metrics across turns that emitted audio. + """Summarize client- and engine-observed metrics across audio responses. ``request_metrics`` entries are caller-assembled dicts; keys that are absent or non-numeric in an entry are simply skipped. ``rtf`` is not produced by :meth:`EventCollector.timing_summary` (which reports raw data - only) — callers that want ``mean_rtf`` add an ``rtf`` value per turn, - e.g. via ``vllm_omni.metrics.definitions.compute_audio_rtf``. + only) — callers that want session ``rtf`` add an ``rtf`` value per turn, + e.g. via ``vllm_omni.metrics.definitions.compute_audio_rtf``. Zero or + missing ``tpot_ms`` values are omitted from session ``tpot_ms``. + + Aggregatable fields are nested as ``{count, mean, p50, p99}``. """ - def mean(metric: str, *, digits: int = 3) -> float | None: - values = [ + def values(metric: str, *, positive: bool = False) -> list[float]: + return [ float(request[metric]) for request in request_metrics - if isinstance(request.get(metric), int | float) and math.isfinite(float(request[metric])) + if isinstance(request.get(metric), int | float) + and not isinstance(request.get(metric), bool) + and math.isfinite(float(request[metric])) + and (not positive or float(request[metric]) > 0) ] - return round(sum(values) / len(values), digits) if values else None - return { + summary: dict[str, object] = { "session_id": session_id, "audio_turn_count": len(request_metrics), - "mean_ttft_ms": mean("ttft_ms"), - "mean_ttfp_ms": mean("ttfp_ms"), - "mean_rtf": mean("rtf", digits=6), + "ttft_ms": distribution_summary(values("ttft_ms")), + "ttfp_ms": distribution_summary(values("ttfp_ms")), + "rtf": distribution_summary(values("rtf"), digits=6), } + if (tpot := distribution_summary(values("tpot_ms", positive=True))) is not None: + summary["tpot_ms"] = tpot + return summary async def acknowledge_collected_playback(client: DuplexClient, collector: EventCollector) -> None: diff --git a/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py b/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py index a43a5ec71..748100c5d 100644 --- a/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py +++ b/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py @@ -1,16 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CLI for Omni-DuplexEval generate / evaluate / summarize. + +``generate`` still writes per-sample timed sentences and ``*.meta.json``. +Duplex performance metrics are written to ``/duplex_metrics.json``. +Skipped samples keep previously recorded rows; incoming ``(split, sample_id)`` +rows replace the matching ones, then ``duplex_stream_*`` is recomputed. +""" + import argparse import asyncio import concurrent.futures import json from pathlib import Path -from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DEFAULT_DATASET, load_samples +from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DEFAULT_DATASET, DuplexSample, load_samples from vllm_omni.benchmarks.duplex.omni_duplex_eval_eval import evaluate_sample, summarize_scores from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge -from vllm_omni.benchmarks.duplex.omni_duplex_eval_runner import generate_sample +from vllm_omni.benchmarks.duplex.omni_duplex_eval_runner import GenerateSampleResult, generate_sample +from vllm_omni.benchmarks.duplex_session_metrics import ( + DUPLEX_METRICS_FILENAME, + merge_duplex_metrics_report, + read_duplex_metrics_report, +) from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase @@ -54,6 +67,26 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None: summarize.add_argument("--score-root", required=True) +def _write_duplex_metrics(response_root: str | Path, results: list[GenerateSampleResult]) -> Path: + request_metrics = [metric for result in results for metric in result.request_metrics] + session_metrics = [result.session_metrics for result in results if result.session_metrics] + path = Path(response_root) / DUPLEX_METRICS_FILENAME + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + merge_duplex_metrics_report( + read_duplex_metrics_report(path), + request_metrics=request_metrics, + session_metrics=session_metrics, + ), + ensure_ascii=False, + indent=2, + ), + encoding="utf-8", + ) + return path + + def run(args: argparse.Namespace) -> int: if args.action == "summarize": print(json.dumps(summarize_scores(args.score_root), ensure_ascii=False, indent=2)) @@ -71,12 +104,12 @@ def run(args: argparse.Namespace) -> int: if args.concurrency < 1: raise ValueError("--concurrency must be at least 1") - async def generate() -> None: + async def generate() -> list[GenerateSampleResult]: semaphore = asyncio.Semaphore(args.concurrency) - async def generate_one(sample) -> None: + async def generate_one(sample: DuplexSample) -> GenerateSampleResult: async with semaphore: - await generate_sample( + return await generate_sample( sample, url=args.url, model=args.model, @@ -89,16 +122,16 @@ def run(args: argparse.Namespace) -> int: overwrite=args.overwrite, ) - await asyncio.gather(*(generate_one(sample) for sample in samples)) + return list(await asyncio.gather(*(generate_one(sample) for sample in samples))) - asyncio.run(generate()) + _write_duplex_metrics(args.response_root, asyncio.run(generate())) return 0 if args.eval_workers < 1: raise ValueError("--eval-workers must be at least 1") judge = DuplexJudge(args.judge_base_url, args.judge_model, api_key=args.judge_api_key) - def evaluate_one(sample) -> None: + def evaluate_one(sample: DuplexSample) -> None: response_path = Path(args.response_root) / sample.split / f"{sample.id}.json" score_path = Path(args.score_root) / sample.split / f"{sample.id}.json" if score_path.exists() and not args.overwrite: diff --git a/vllm_omni/entrypoints/duplex/protocol.py b/vllm_omni/entrypoints/duplex/protocol.py index d28a7b009..042180838 100644 --- a/vllm_omni/entrypoints/duplex/protocol.py +++ b/vllm_omni/entrypoints/duplex/protocol.py @@ -413,6 +413,10 @@ class ResponseState: active_request_id: str | None = None active_response_id: str | None = None active_response_turn_id: int | None = None + request_started_at_s_by_turn: dict[int, float] = field(default_factory=dict) + active_response_request_started_at_s: float | None = None + active_response_ttft_ms: float | None = None + active_response_ttfp_ms: float | None = None active_response_input_commit_seq: int | None = None active_response_awaits_input_commit: bool = False last_response_id: str | None = None @@ -656,6 +660,45 @@ class DuplexSession: def bind_response_turn(self, turn_id: int | None) -> None: self._response.active_response_turn_id = turn_id + def mark_model_turn_request_started(self, turn_id: int, started_at_s: float) -> None: + """Record the latest native request start that can own one model turn.""" + turn_id = int(turn_id) + started_at_s = float(started_at_s) + if self.active_response_turn_id == turn_id: + if self._response.active_response_request_started_at_s is None: + self._response.active_response_request_started_at_s = started_at_s + return + self._response.request_started_at_s_by_turn[turn_id] = started_at_s + + def mark_response_first_outputs( + self, + *, + observed_at_s: float, + has_text: bool, + has_audio: bool, + ) -> dict[str, object]: + """Return server-monotonic TTF metrics observed for the active response.""" + started_at_s = self._response.active_response_request_started_at_s + if started_at_s is None: + return {} + elapsed_ms = max(0.0, (float(observed_at_s) - started_at_s) * 1000.0) + if has_text and self._response.active_response_ttft_ms is None: + self._response.active_response_ttft_ms = elapsed_ms + if has_audio and self._response.active_response_ttfp_ms is None: + self._response.active_response_ttfp_ms = elapsed_ms + metrics: dict[str, object] = { + "source": "server_monotonic_request_start", + "measurement_origin": { + "ttft": "native model-turn request execution start to first non-empty text output", + "ttfp": "native model-turn request execution start to first audio output", + }, + } + if self._response.active_response_ttft_ms is not None: + metrics["ttft_ms"] = self._response.active_response_ttft_ms + if self._response.active_response_ttfp_ms is not None: + metrics["ttfp_ms"] = self._response.active_response_ttfp_ms + return metrics + def active_response_accepts_model_turn(self, turn_id: int | None) -> bool: if self._response.active_response_id is None: return False @@ -940,6 +983,11 @@ class DuplexSession: def complete_model_turn(self, turn_id: int) -> None: """Advance the model-owned output identity after its terminal signal.""" completed_turn_id = int(turn_id) + self._response.request_started_at_s_by_turn = { + pending_turn_id: started_at_s + for pending_turn_id, started_at_s in self._response.request_started_at_s_by_turn.items() + if pending_turn_id > completed_turn_id + } if completed_turn_id >= self.turn_id: self.turn_id = completed_turn_id + 1 @@ -976,6 +1024,12 @@ class DuplexSession: response_id = f"resp-{self.session_id}-{self.epoch}-{uuid4().hex[:8]}" self._response.active_response_id = response_id self._response.active_response_turn_id = self.turn_id if turn_id is None else int(turn_id) + self._response.active_response_request_started_at_s = self._response.request_started_at_s_by_turn.pop( + self._response.active_response_turn_id, + None, + ) + self._response.active_response_ttft_ms = None + self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = self.input_commit_seq self._response.active_response_awaits_input_commit = self.turn_state == DuplexTurnState.USER_SPEAKING self._response.last_response_id = response_id @@ -1057,17 +1111,22 @@ class DuplexSession: tpot_ms = raw_values.get("vllm_tpot_ms") token_count = raw_values.get("num_tokens_out") - if isinstance(tpot_ms, int | float) and tpot_ms > 0: - weight = max(int(token_count) - 1, 1) if isinstance(token_count, int | float) else 1 - self._response.stage_metric_tpot_weighted_ms[stage_id] = ( - self._response.stage_metric_tpot_weighted_ms.get(stage_id, 0.0) + float(tpot_ms) * weight - ) - self._response.stage_metric_tpot_weight[stage_id] = ( - self._response.stage_metric_tpot_weight.get(stage_id, 0) + weight - ) - current["vllm_tpot_ms"] = self._response.stage_metric_tpot_weighted_ms[stage_id] / float( - self._response.stage_metric_tpot_weight[stage_id] + if isinstance(tpot_ms, int | float) and not isinstance(tpot_ms, bool) and tpot_ms > 0: + weight = ( + max(int(token_count) - 1, 0) + if isinstance(token_count, int | float) and not isinstance(token_count, bool) + else 1 ) + if weight > 0: + self._response.stage_metric_tpot_weighted_ms[stage_id] = ( + self._response.stage_metric_tpot_weighted_ms.get(stage_id, 0.0) + float(tpot_ms) * weight + ) + self._response.stage_metric_tpot_weight[stage_id] = ( + self._response.stage_metric_tpot_weight.get(stage_id, 0) + weight + ) + current["vllm_tpot_ms"] = self._response.stage_metric_tpot_weighted_ms[stage_id] / float( + self._response.stage_metric_tpot_weight[stage_id] + ) for name, value in raw_values.items(): if name not in handled_fields: @@ -1272,6 +1331,9 @@ class DuplexSession: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None + self._response.active_response_request_started_at_s = None + self._response.active_response_ttft_ms = None + self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self._clear_response_metrics() @@ -1589,6 +1651,10 @@ class DuplexSession: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None + self._response.request_started_at_s_by_turn.clear() + self._response.active_response_request_started_at_s = None + self._response.active_response_ttft_ms = None + self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self._clear_response_metrics() @@ -1604,6 +1670,10 @@ class DuplexSession: self.state = DuplexSessionState.CLOSED self.turn_state = DuplexTurnState.IDLE self._response.active_response_turn_id = None + self._response.request_started_at_s_by_turn.clear() + self._response.active_response_request_started_at_s = None + self._response.active_response_ttft_ms = None + self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self.clear_server_vad_audio() diff --git a/vllm_omni/entrypoints/duplex/runtime_bridge.py b/vllm_omni/entrypoints/duplex/runtime_bridge.py index bf839523c..96a82c061 100644 --- a/vllm_omni/entrypoints/duplex/runtime_bridge.py +++ b/vllm_omni/entrypoints/duplex/runtime_bridge.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import base64 import inspect +import time from collections.abc import Mapping from dataclasses import asdict, is_dataclass from enum import Enum @@ -1113,13 +1114,20 @@ class NativeRuntimeBridgeMixin: if response_id is None: response_id = session.begin_response(turn_id=model_turn_id) response_created = True - await send_json( - self._response_created_payload( - session, - response_id, - epoch=session.epoch, - ) + response_request_metrics = session.mark_response_first_outputs( + observed_at_s=time.monotonic(), + has_text=has_text, + has_audio=has_audio, + ) + if response_created: + response_created_payload = self._response_created_payload( + session, + response_id, + epoch=session.epoch, ) + if response_request_metrics: + response_created_payload["response_request_metrics"] = response_request_metrics + await send_json(response_created_payload) response_stage_metrics = session.accumulate_response_stage_metrics( native_result.get("stage_metrics") if isinstance(native_result.get("stage_metrics"), Mapping) else None ) @@ -1137,6 +1145,7 @@ class NativeRuntimeBridgeMixin: speak_payload, native_result, stage_metrics=response_stage_metrics, + response_request_metrics=response_request_metrics, ) await send_json(speak_payload) previous_sent_ms = session.playback.sent_ms @@ -1205,6 +1214,7 @@ class NativeRuntimeBridgeMixin: payload, native_result, stage_metrics=response_stage_metrics, + response_request_metrics=response_request_metrics, ) await send_json(payload) if ( @@ -1365,6 +1375,7 @@ class NativeRuntimeBridgeMixin: native_result: dict[str, object], *, stage_metrics: Mapping[str, object] | None = None, + response_request_metrics: Mapping[str, object] | None = None, ) -> None: metadata: dict[str, object] = {} runtime_impl = native_result.get("runtime_impl") @@ -1390,6 +1401,8 @@ class NativeRuntimeBridgeMixin: for stage_id, values in effective_stage_metrics.items() if isinstance(values, Mapping) } + if response_request_metrics: + metadata["response_request_metrics"] = dict(response_request_metrics) if metadata: payload["vllm_omni"] = metadata diff --git a/vllm_omni/entrypoints/duplex/session_runner.py b/vllm_omni/entrypoints/duplex/session_runner.py index 8ee5060b0..90251e998 100644 --- a/vllm_omni/entrypoints/duplex/session_runner.py +++ b/vllm_omni/entrypoints/duplex/session_runner.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import binascii import json +import time import uuid from collections.abc import Callable from contextlib import suppress @@ -494,6 +495,7 @@ class DuplexSessionRunnerMixin: async def _run() -> bool: nonlocal runtime_closed try: + session.mark_model_turn_request_started(append_turn_id, time.monotonic()) append_ok, emitted_response = await self._append_runtime_input( session, payload, diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 13844518e..bd60e4793 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -59,6 +59,10 @@ def _modality_to_type_string(modality: OutputModality) -> str: return "text" +_TPOT_ELAPSED_MS = "_tpot_elapsed_ms" +_TPOT_INTERVALS = "_tpot_intervals" + + def _mean_time_per_output_token_ms(stats: RequestStateStats) -> float: output_intervals = stats.num_generation_tokens - 1 if output_intervals <= 0: @@ -67,6 +71,17 @@ def _mean_time_per_output_token_ms(stats: RequestStateStats) -> float: return decode_time_s * 1000.0 / float(output_intervals) +def _accumulate_segment_tpot(record: dict[str, object], *, elapsed_ms: float, new_tokens: int) -> None: + """Weight one decode step by its token count. ITL stays one sample per step.""" + if new_tokens <= 0: + return + total_elapsed_ms = float(record.get(_TPOT_ELAPSED_MS) or 0.0) + elapsed_ms + total_intervals = int(record.get(_TPOT_INTERVALS) or 0) + new_tokens + record[_TPOT_ELAPSED_MS] = total_elapsed_ms + record[_TPOT_INTERVALS] = total_intervals + record["vllm_tpot_ms"] = total_elapsed_ms / float(total_intervals) + + class OmniRequestState(RequestState): """Request state for omni models, tracking multimodal outputs. @@ -430,7 +445,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): ) def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: - return self._native_text_metrics_by_request.pop(request_id, {}) + record = self._native_text_metrics_by_request.pop(request_id, {}) + record.pop(_TPOT_ELAPSED_MS, None) + record.pop(_TPOT_INTERVALS, None) + return record def abort_requests(self, request_ids, internal: bool) -> list[str]: aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal) @@ -770,6 +788,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): was_prefilling = req_state.is_prefilling native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 + previous_num_generation_tokens = int(native_stats.num_generation_tokens) if native_stats is not None else 0 # NOTE: We pass ``None`` for *iteration_stats* to the parent so that # the upstream's ``_update_stats_from_output`` logs stats via its own @@ -808,7 +827,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): itls_ms = record.setdefault("vllm_itls_ms", []) itls_ms.append(itl_ms) record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) - record["vllm_tpot_ms"] = _mean_time_per_output_token_ms(native_stats) + _accumulate_segment_tpot( + record, + elapsed_ms=itl_ms, + new_tokens=max(int(native_stats.num_generation_tokens) - previous_num_generation_tokens, 0), + ) + elif not record["vllm_itls_ms"]: + record["vllm_tpot_ms"] = _mean_time_per_output_token_ms(native_stats) def _update_stats_from_finished( self, @@ -836,6 +861,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): finished_request = iteration_stats.finished_requests[-1] if finished_request.request_id != req_state.external_req_id: return - self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = ( - float(finished_request.mean_time_per_output_token) * 1000.0 - ) + finished_tpot_ms = float(finished_request.mean_time_per_output_token) * 1000.0 + record = self._native_text_metric_record(req_state.external_req_id) + if finished_tpot_ms > 0 and not record.get("vllm_itls_ms"): + record["vllm_tpot_ms"] = finished_tpot_ms From 01f4479138c308a9ccf0ac84403b8b27fab0d3ae Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Mon, 14 Sep 2026 16:33:38 -0600 Subject: [PATCH 077/209] [Bugfix] Don't shutdown Engine on py_generator=True (#6334) Signed-off-by: Alex Brooks --- tests/entrypoints/test_omni_entrypoints.py | 61 ++++++++++++++++++++-- vllm_omni/entrypoints/omni.py | 17 ++++-- 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index be7a0ffac..95059a06e 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -3,6 +3,7 @@ from __future__ import annotations +import gc import queue from collections.abc import Callable from types import SimpleNamespace @@ -766,7 +767,7 @@ def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkey f"{engine.submitted[1]['request_id']}-stage0-0", f"{engine.submitted[1]['request_id']}-stage2-final", ] - assert engine.shutdown_called is True + assert engine.shutdown_called is not True def test_omni_generate_returns_list_when_not_using_generator(monkeypatch: pytest.MonkeyPatch): @@ -812,7 +813,7 @@ def test_omni_generate_diffusion_only_yields_single_image_per_request(monkeypatc [f"{engine.submitted[0]['request_id']}-image"], [f"{engine.submitted[1]['request_id']}-image"], ] - assert engine.shutdown_called is True + assert engine.shutdown_called is not True def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request( @@ -845,7 +846,7 @@ def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request( [f"{engine.submitted[1]['request_id']}-image"], ] assert engine.submitted[0]["sampling_params_list"][0].output_kind == RequestOutputKind.FINAL_ONLY - assert engine.shutdown_called is True + assert engine.shutdown_called is not True def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): @@ -1326,3 +1327,57 @@ def test_omni_errored_property_dead_stage(monkeypatch: pytest.MonkeyPatch): assert app.errored is False finally: app.shutdown() + + +def test_omni_pygenerator_does_not_kill_engine(monkeypatch: pytest.MonkeyPatch): + engine = FakeAsyncOmniEngine( + stage_metadata=THREE_STAGE_META, + on_add_request=_enqueue_async_three_stage_outputs, + ) + + _patch_engine(monkeypatch, engine) + + app = Omni("dummy-model") + # Evaluating the generator should not shut down the engine + list(app.generate(["hello"], py_generator=True)) + assert not engine.shutdown_called + + +def test_omni_generator_close_cleans_up(monkeypatch: pytest.MonkeyPatch): + """Ensure that a closed generator cleans things up properly.""" + engine = FakeAsyncOmniEngine( + stage_metadata=THREE_STAGE_META, + on_add_request=_enqueue_async_three_stage_outputs, + ) + + _patch_engine(monkeypatch, engine) + + app = Omni("dummy-model") + + # Create a generator and start to evaluate it to make sure the request isn't aborted yet + my_gen = app.generate(["hello"], py_generator=True) + next(my_gen) + request_id = engine.submitted[0]["request_id"] + assert engine.aborted == [] + assert request_id in app.request_states + + # Close it and make sure the sure it's aborted, but without killing engine + my_gen.close() + assert engine.aborted == [[request_id]] + assert request_id not in app.request_states + assert not engine.shutdown_called + + +def test_del_shutsdown_engine(monkeypatch: pytest.MonkeyPatch): + engine = FakeAsyncOmniEngine( + stage_metadata=THREE_STAGE_META, + on_add_request=_enqueue_async_three_stage_outputs, + ) + + _patch_engine(monkeypatch, engine) + + app = Omni("dummy-model") + assert not engine.shutdown_called + del app + gc.collect() + assert engine.shutdown_called diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 28bbe9219..cf4da3b7a 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import copy @@ -99,11 +102,11 @@ class Omni(OmniBase): sampling_params_list: Sequence[OmniSamplingParams], use_tqdm: bool | Callable[..., tqdm] = True, ) -> Generator[OmniRequestOutput, None, None]: - gen = self._run_generation(prompts, sampling_params_list, use_tqdm) - try: - yield from gen - finally: - self.close() + yield from self._run_generation( + prompts, + sampling_params_list, + use_tqdm, + ) def _run_generation( self, @@ -198,6 +201,10 @@ class Omni(OmniBase): if pbar is not None: pbar.update(1) self._log_summary_and_cleanup(req_id) + except GeneratorExit: + if "active_reqs" in locals() and active_reqs: + self.abort(list(active_reqs)) + raise except Exception: if "active_reqs" in locals() and active_reqs: for req_id in active_reqs: From 5385dda4feeb417bd94a29a5e62ec904b757a0ef Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:19:21 -0700 Subject: [PATCH 078/209] [Bugfix] Use current OpenAI audio event names in duplex client metrics (#7541) Signed-off-by: Yueqian Lin --- tests/clients/test_duplex_client.py | 8 ++++---- .../e2e/features/fullduplex/test_omni_duplex_eval_cli.py | 2 +- vllm_omni/clients/duplex.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index 8a74a2087..ab2a38951 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -776,7 +776,7 @@ def test_event_collector_reports_global_metrics_across_responses(): ) collector.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": response_id, }, @@ -785,7 +785,7 @@ def test_event_collector_reports_global_metrics_across_responses(): for received_at_s, duration_ms in zip(audio_times, durations, strict=True): collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": base64.b64encode(b"audio").decode("ascii"), "sample_rate_hz": 16_000, @@ -856,7 +856,7 @@ def test_event_collector_prefers_server_request_start_for_ttf(): for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "sample_rate_hz": 16_000, @@ -868,7 +868,7 @@ def test_event_collector_prefers_server_request_start_for_ttf(): received_at_s=received_at_s, ) collector.add( - {"type": "response.audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, + {"type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, received_at_s=10.15, ) diff --git a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py index 7f4965768..b1b02ef65 100644 --- a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py +++ b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py @@ -389,7 +389,7 @@ def test_cli_generate_partial_resume_merges_metrics(tmp_path: Path, monkeypatch: def test_collect_duplex_session_metrics_matches_omniinteract_window(): collector = EventCollector() audio = { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "r1", "format": "pcm16", "delta": base64.b64encode(bytes((1, 0)) * 2400).decode(), diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index 0504c241b..a1d9fc1df 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -1475,7 +1475,7 @@ class EventCollector: if ( event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } From 85da94339130337b49e610d08276aedcf84f6b96 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Tue, 15 Sep 2026 08:40:04 +0800 Subject: [PATCH 079/209] [Refactor][Diffusion] Complete generic backend plan consumption (#7313) Signed-off-by: specture724 Co-authored-by: Claude Opus 5 --- docs/design/feature/offloader/README.md | 33 +++- .../models/flux2/test_flux2_transformer_tp.py | 7 +- .../offloader/test_backend_plan_contract.py | 163 ++++++++++++++++++ .../offloader/test_layerwise_backend.py | 23 ++- .../diffusion/offloader/layerwise_backend.py | 44 +---- .../diffusion/offloader/plan_resolver.py | 10 ++ .../diffusion/offloader/sequential_backend.py | 4 +- 7 files changed, 228 insertions(+), 56 deletions(-) create mode 100644 tests/diffusion/offloader/test_backend_plan_contract.py diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index 96cba8f60..0ec81f3b5 100644 --- a/docs/design/feature/offloader/README.md +++ b/docs/design/feature/offloader/README.md @@ -41,7 +41,8 @@ Every backend implements `OffloadBackend`: - `is_enabled()` reports lifecycle state. `disable()` does not promise to restore the pipeline's original device -placement. The caller owns any subsequent rematerialization. +placement. It must restore usable parameter and buffer storage before releasing +backend-owned host weights. The caller owns subsequent device placement. An enable failure must remove every partially installed hook. Rank-local backends restore ordinary tensors so the pipeline can be retried. A failed @@ -54,9 +55,19 @@ must use distinct hook names and remove only hooks they own. ## Discovery and topology `resolve_offload_plan(pipeline, config)` is the single topology entry point. -Every backend calls it once when `enable()` starts and then consumes the -resolved components, block stacks, resident heads, and staged-residency flags -it returns. Backends do not re-read model declarations. +The generic model-level and ordinary layerwise paths call it once when +`enable()` starts and consume its component selection and block stacks. +Ordinary layerwise offload also respects a resolved `skip_reason`, preserving +the legacy missing-DiT no-op without placing auxiliary components. Non-block +DiT and encoder state is placed by resolved block tensor identity, so the +backend does not rediscover block containers by attribute name. + +Two compatibility paths remain outside this completed generic cutover: +model-level pipelines implementing `SupportsModelCpuOffload` still own their +custom phase lifecycle (RFC #6648 J4), and distributed layerwise offload still +reads declarations for nested-submodule staging and loader-owned mmap paths +(J3). The custom model-level protocol is delegated to before generic plan +resolution; it must not be subjected to the generic encoder/DiT swap checks. Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: @@ -102,3 +113,17 @@ remains responsible for transfer, synchronization, and storage ownership. The diffusion [Offloader module design](../../module/diffusion/offloader.md) describes how these feature contracts fit into the larger diffusion runtime. + +## Backend acceptance tests + +`tests/diffusion/offloader/test_backend_plan_contract.py` covers the two +behaviors the generic cutover changed. A DiT attribute aliasing a streamed +block must keep the ring's host residency, so placement cannot follow attribute +names; the CUDA case fails against the previous name-based placement. Both +backends must then run with their selector helpers and the model declaration +made unreadable, proving execution consumes the resolved plan alone. + +Selection errors, rollback, residency, and enable/disable cycles stay in +`test_plan_resolver.py`, `test_layerwise_backend.py` and +`test_sequential_backend.py`. Model-specific phases and executed tied aliases +remain the J4 contract suite. diff --git a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py index c935871a0..489a948d5 100644 --- a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py +++ b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest import torch from pytest_mock import MockerFixture @@ -265,7 +268,7 @@ class TestFlux2TransformerLayerwiseOffload: ] def test_get_blocks_from_dit(self): - from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend + from vllm_omni.diffusion.offloader.block_discovery import get_blocks_from_dit # Block discovery only needs the ModuleList structure, not production dims. model = Flux2Transformer2DModel( @@ -275,6 +278,6 @@ class TestFlux2TransformerLayerwiseOffload: attention_head_dim=4, joint_attention_dim=16, ) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == ["transformer_blocks", "single_transformer_blocks"] assert len(blocks) == 4 diff --git a/tests/diffusion/offloader/test_backend_plan_contract.py b/tests/diffusion/offloader/test_backend_plan_contract.py new file mode 100644 index 000000000..3c406f5d5 --- /dev/null +++ b/tests/diffusion/offloader/test_backend_plan_contract.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Executable contracts for the two behaviors the generic cutover changed. + +Placement now follows resolved block identity, and the backends run on the +resolved plan alone. Declaration parsing, selection errors, rollback, +residency and enable/disable cycles are covered by ``test_plan_resolver.py``, +``test_layerwise_backend.py`` and ``test_sequential_backend.py``. +""" + +from typing import ClassVar + +import pytest +import torch +from torch import nn + +from tests.diffusion.offloader.helpers import patch_offload_runtime +from vllm_omni.diffusion.offloader import layerwise_backend, sequential_backend +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan +from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model] + +STRATEGIES = [OffloadStrategy.MODEL_LEVEL, OffloadStrategy.LAYER_WISE] +CPU = torch.device("cpu") + + +class _Block(nn.Linear): + def __init__(self): + super().__init__(4, 4) + self.register_buffer("scale", torch.tensor(0.5)) + + def forward(self, x): + return super().forward(x).tanh() * self.scale + + +class _Stack(nn.Module): + """Two block containers plus non-block state that must stay resident.""" + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([_Block() for _ in range(3)]) + self.tail = nn.ModuleList([_Block() for _ in range(2)]) + self.bias = nn.Parameter(torch.ones(4)) + self.register_buffer("scale", torch.tensor(0.25)) + self.proj = nn.Linear(4, 4) + + def forward(self, x): + for block in (*self.blocks, *self.tail): + x = block(x) + return self.proj(x) * self.scale + self.bias + + +class _Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = ["resident"] + _offload_plan = OffloadPlan( + block_attrs={"transformer": ("blocks", "tail")}, + encoder_block_attrs={"text_encoder": ("blocks", "tail")}, + ) + + def __init__(self): + super().__init__() + self.transformer = _Stack() + self.text_encoder = _Stack() + self.image_encoder = nn.Linear(4, 4) + self.vae = nn.Linear(4, 4) + self.resident = nn.Linear(4, 4) + + def forward(self, x): + return self.resident(self.vae(self.transformer(self.text_encoder(self.image_encoder(x))))) + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def execution_device(request, monkeypatch): + if request.param == "cuda": + if not torch.cuda.is_available(): + pytest.skip("CUDA required for real transfer/residency coverage") + return torch.device("cuda:0") + patch_offload_runtime(monkeypatch, current_omni_platform, synchronize=True) + monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda: 0) + return CPU + + +def _config(strategy, components): + return OffloadConfig(strategy=strategy, components=components, pin_cpu_memory=False) + + +def _backend(config, device): + kind = ModelLevelOffloadBackend if config.strategy is OffloadStrategy.MODEL_LEVEL else LayerWiseOffloadBackend + return kind(config, device) + + +def _assert_no_offload_hooks(pipeline): + for module in pipeline.modules(): + registry = getattr(module, "_hook_registry", None) + if registry is not None: + assert registry.get_hook("sequential_offload") is None + assert registry.get_hook("layerwise_offload") is None + assert not getattr(module, "_omni_layerwise_enabled", False) + + +@torch.inference_mode() +def test_alias_of_a_streamed_block_is_not_placed(execution_device): + """Placement moves DiT state by resolved block identity, not by attribute. + + An attribute aliasing a streamed block must stay with its ring: copying it + to the device makes the ring record the device as that block's home and + strands those weights there after teardown. + """ + pipeline = _Pipeline() + pipeline.transformer.proj.weight = pipeline.transformer.blocks[1].weight + original = pipeline.transformer.proj.weight.clone() + backend = _backend(_config(OffloadStrategy.LAYER_WISE, frozenset({"dit"})), execution_device) + + backend.enable(pipeline) + try: + assert pipeline.transformer.proj.weight is pipeline.transformer.blocks[1].weight + assert pipeline.transformer.proj.weight.device == CPU + # Non-block state around the alias is still placed. + assert pipeline.transformer.proj.bias.device == execution_device + assert pipeline.transformer.bias.device == execution_device + finally: + backend.disable() + + assert pipeline.transformer.proj.weight.device == CPU + torch.testing.assert_close(pipeline.transformer.proj.weight, original) + + +@pytest.mark.parametrize("strategy", STRATEGIES) +@torch.inference_mode() +def test_backends_execute_the_resolved_plan_alone(execution_device, strategy, monkeypatch): + pipeline = _Pipeline().to(execution_device) + x = torch.randn(2, 4, device=execution_device) + expected = pipeline(x) + config = _config(strategy, frozenset({"dit", "text_encoder"})) + resolved = resolve_offload_plan(pipeline, config) + + def forbidden_read(*args, **kwargs): + pytest.fail("Backend reinterpreted topology after plan resolution") + + # Once resolved, neither declarations nor selector helpers are an input to + # backend execution. Transport options stay readable. + monkeypatch.setattr(config, "offloads", forbidden_read) + monkeypatch.setattr(config, "should_offload_encoder", forbidden_read) + monkeypatch.setattr(pipeline, "_offload_plan", None) + backend_module = sequential_backend if strategy is OffloadStrategy.MODEL_LEVEL else layerwise_backend + monkeypatch.setattr(backend_module, "resolve_offload_plan", lambda *_: resolved) + + backend = _backend(config, execution_device) + try: + backend.enable(pipeline) + torch.testing.assert_close(pipeline(x), expected) + finally: + backend.disable() + _assert_no_offload_hooks(pipeline) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 099e93f79..4c965adaa 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -24,6 +24,11 @@ from tests.diffusion.offloader.helpers import ( patch_offload_runtime, ) from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.block_discovery import ( + get_blocks_attr_names, + get_blocks_from_dit, + set_blocks_attr_names, +) from vllm_omni.diffusion.offloader.layerwise_backend import ( LayerWiseOffloadBackend, LayerwiseOffloadHook, @@ -209,21 +214,21 @@ class _NoAttrsModel(nn.Module): class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): model = _SingleBlockModel(num_blocks=3) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == ["blocks"] assert len(blocks) == 3 assert all(isinstance(b, _DummyBlock) for b in blocks) def test_get_blocks_from_dit_multi_block_attrs(self): model = _MultiBlockModel(num_transformer=2, num_single=3) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"} assert len(blocks) == 5 assert all(isinstance(b, _DummyBlock) for b in blocks) def test_get_blocks_from_dit_empty_blocks(self): model = _EmptyBlocksModel() - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == [] assert blocks == [] @@ -233,17 +238,17 @@ class TestGetBlocksFromDit: AttributeError, match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist", ): - LayerWiseOffloadBackend.get_blocks_from_dit(model) + get_blocks_from_dit(model) def test_get_blocks_from_dit_no_attrs_defined(self): model = _NoAttrsModel(num_blocks=3) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == [] assert blocks == [] def test_get_blocks_from_dit_deprecated_single_attr(self): model = _DeprecatedSingleAttrModel(num_blocks=2) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == ["blocks"] assert len(blocks) == 2 @@ -251,17 +256,17 @@ class TestGetBlocksFromDit: class TestGetBlocksAttrNames: def test_get_blocks_attr_names_new_format(self): model = _MultiBlockModel() - attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model) + attrs = get_blocks_attr_names(model) assert attrs == ["transformer_blocks", "single_transformer_blocks"] def test_get_blocks_attr_names_no_attrs(self): model = _NoAttrsModel() - attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model) + attrs = get_blocks_attr_names(model) assert attrs == [] def test_set_blocks_attr_names(self): model = _NoAttrsModel() - LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"]) + set_blocks_attr_names(model, ["new_blocks"]) assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs") assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"] diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index a1e594f8b..e9f8cf5be 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -14,11 +14,6 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, run_cleanup_steps -from .block_discovery import ( - get_blocks_attr_names, - get_blocks_from_dit, - set_blocks_attr_names, -) from .component_utils import ( clear_encoder_layerwise_state, iter_streamable_dits, @@ -26,7 +21,6 @@ from .component_utils import ( prepare_pipeline_components, set_encoder_layerwise_state, ) -from .config import DIT_COMPONENT from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, @@ -433,9 +427,8 @@ class LayerWiseOffloadBackend(OffloadBackend): return resolved = resolve_offload_plan(pipeline, self.config) - if not resolved.dits and self.config.offloads(DIT_COMPONENT): - # An explicit selection already failed while resolving. - logger.warning("No DiT/transformer modules found for selected DiT layerwise offload") + if resolved.skip_reason is not None: + logger.warning("%s", resolved.skip_reason) return def enable_encoder_blocks(component: ResolvedComponent) -> bool: @@ -458,15 +451,6 @@ class LayerWiseOffloadBackend(OffloadBackend): enable_encoder_blocks=enable_encoder_blocks, ) - if not self.config.offloads(DIT_COMPONENT): - self.enabled = bool(self._encoder_modules or self._staged_components) - if not self.enabled: - raise ValueError( - "None of the selected layerwise offload components have " - "a model-declared streamable or on-demand plan" - ) - return - logger.info("Applying layer-wise offloading on %s", [component.path for component in resolved.dits]) # Apply block-wise offloading hook for each of the blocks in DiT model(s) @@ -475,22 +459,9 @@ class LayerWiseOffloadBackend(OffloadBackend): dit_module = component.module blocks = list(stack.blocks) - # Move non-block modules to GPU (they stay resident) - for name, m in dit_module.named_children(): - if name not in stack.attrs: - m.to(self.device) - logger.debug(f"Moved {name} to device {self.device}") - else: - logger.debug(f"Skipped blocks module {name}") - - # Move top-level params/buffers to GPU (dit_module's own, not sub-modules) - for param in dit_module._parameters.values(): - if param is not None: - param.data = param.data.to(self.device, non_blocking=True) - - for buffer in dit_module._buffers.values(): - if buffer is not None: - buffer.data = buffer.data.to(self.device, non_blocking=True) + # Place the remainder by resolved block tensor identity, just as + # for encoders. Attribute aliases must not move streamed weights. + move_non_block_state_to_device(dit_module, (stack.blocks,), self.device) block_hooks = _install_layerwise_hook_group( blocks, @@ -550,8 +521,3 @@ class LayerWiseOffloadBackend(OffloadBackend): def disable(self) -> None: self._disable(restore_weights=True) - - # Compatibility aliases for existing model integrations. - get_blocks_attr_names = staticmethod(get_blocks_attr_names) - set_blocks_attr_names = staticmethod(set_blocks_attr_names) - get_blocks_from_dit = staticmethod(get_blocks_from_dit) diff --git a/vllm_omni/diffusion/offloader/plan_resolver.py b/vllm_omni/diffusion/offloader/plan_resolver.py index b2e951b4e..08527318f 100644 --- a/vllm_omni/diffusion/offloader/plan_resolver.py +++ b/vllm_omni/diffusion/offloader/plan_resolver.py @@ -91,6 +91,11 @@ class ResolvedOffloadPlan: # still interpret topology themselves. modules: PipelineModules declaration: OffloadPlan | None = None + # Ordinary layerwise offload turns a legacy selection it cannot serve into + # a complete no-op, placement included. Distributed layerwise offload warns + # and keeps preparing its other components instead, so this stays scoped to + # the strategy that honors it. + skip_reason: str | None = None @property def components(self) -> tuple[ResolvedComponent, ...]: @@ -304,6 +309,11 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved residents=tuple(residents), modules=modules, declaration=declaration, + skip_reason=( + "No DiT/transformer modules found for selected DiT layerwise offload" + if config.strategy is OffloadStrategy.LAYER_WISE and dit_selected and not dits + else None + ), ) _validate_unique_ownership(resolved) return resolved diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 1187c8cee..b03be3687 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -13,7 +13,6 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload -from .config import DIT_COMPONENT from .plan_resolver import resolve_offload_plan logger = init_logger(__name__) @@ -333,6 +332,7 @@ class ModelLevelOffloadBackend(OffloadBackend): encoders = [component.module for component in resolved.encoders] vaes = [component.module for component in resolved.vaes] residents = [component.module for component in resolved.residents] + selected_dits = [component.module for component in resolved.dits if component.selected] selected_encoders = [component.module for component in resolved.encoders if component.selected] all_modules = [*dits, *encoders, *vaes, *residents] @@ -360,7 +360,7 @@ class ModelLevelOffloadBackend(OffloadBackend): device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, - offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()), + offload_dit_modules=selected_dits, offload_encoder_modules=selected_encoders, ) except BaseException: From 5761ad551695e9c4b161697dd4eecc26ab91cb04 Mon Sep 17 00:00:00 2001 From: Deep Shah Date: Mon, 14 Sep 2026 20:41:03 -0700 Subject: [PATCH 080/209] [Bugfix] Restore the async PP sampled-token broadcast in the AR runner (#7393) Signed-off-by: Deep Shah --- .../examples/online_serving/qwen3_omni.md | 4 +- tests/worker/test_pp_async_token_broadcast.py | 321 ++++++++++++++++++ vllm_omni/worker/gpu_ar_model_runner.py | 10 + 3 files changed, 333 insertions(+), 2 deletions(-) create mode 100644 tests/worker/test_pp_async_token_broadcast.py diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 4e01ef4c2..2c85c33fe 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -23,8 +23,8 @@ some GPUs. Asynchronous chunk streaming is **enabled by default** within the bundled configuration. For Thinker pipeline parallelism, configure `pipeline_parallel_size` and `devices` -on the Thinker stage and set `async_scheduling: false`. Keep Talker and Code2Wav -at `pipeline_parallel_size: 1`. Audio chunk streaming can remain enabled. +on the Thinker stage. Keep Talker and Code2Wav at `pipeline_parallel_size: 1`. The +default asynchronous scheduling and audio chunk streaming can remain enabled. To explicitly utilize a custom deployment YAML, specify the configuration path: ```bash diff --git a/tests/worker/test_pp_async_token_broadcast.py b/tests/worker/test_pp_async_token_broadcast.py new file mode 100644 index 000000000..d5e938681 --- /dev/null +++ b/tests/worker/test_pp_async_token_broadcast.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Async scheduling + pipeline parallelism: the last rank must send what earlier ranks wait for. + +vLLM's ``GPUModelRunner.sample_tokens`` pairs two collectives on the PP device group: +a non-final rank posts an async broadcast *receive* of the newly sampled token ids +(``_pp_receive_prev_sampled_token_ids_to_input_batch``) and the final rank performs +the matching broadcast *send* (``_pp_broadcast_prev_sampled_token_ids``) right after +``_update_states_after_model_execute``. ``GPUARModelRunner.sample_tokens`` overrides the +upstream method and inherited the receive but not the send. With async scheduling and +PP > 1 the first rank never returned from posting that receive: it is the first +collective on the PP group's communicator, NCCL has to create the communicator with every +rank, and the last rank never joined (observed live on 2x A100: rank 0 inside +``torch.distributed.broadcast``, rank 1 waiting for the next intermediate tensors, engine +dead after the RPC timeout). Had the communicator existed, the same missing send would +have blocked the ``_pp_recv_work.wait()`` in ``_prepare_input_ids`` on the next step. + +These tests drive the real Omni override with a recorded ``torch.distributed.broadcast`` +ledger and, separately, a real two-process gloo group so the receive is proven to +complete only because the Omni sampling path sent. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.v1.worker import gpu_model_runner as upstream_runner_module + +from vllm_omni.worker import gpu_ar_model_runner as omni_runner_module +from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel] + + +class _StopAtBookkeepingError(Exception): + """Raised by the bookkeeping double: everything before it is what we assert on.""" + + +def _pp(rank: int, world_size: int, device_group: object = "pp-group") -> SimpleNamespace: + return SimpleNamespace( + rank=rank, + world_size=world_size, + last_rank=world_size - 1, + is_last_rank=rank == world_size - 1, + is_first_rank=rank == 0, + device_group=device_group, + ) + + +def _input_batch(num_reqs: int) -> SimpleNamespace: + return SimpleNamespace( + num_reqs=num_reqs, + req_ids=[f"req-{i}" for i in range(num_reqs)], + num_tokens_no_spec=np.ones(num_reqs, dtype=np.int64), + is_token_ids=np.zeros((num_reqs, 8), dtype=bool), + sampling_metadata=SimpleNamespace(no_penalties=True, logitsprocs=SimpleNamespace(non_argmax_invariant=[])), + vocab_size=32, + prev_sampled_token_ids=None, + prev_req_id_to_index=None, + ) + + +def _last_rank_runner( + *, + sampled: torch.Tensor, + async_scheduling: bool = True, + broadcast_pp_output: bool = False, + all_chunked_prefill: bool = False, + events: list[str], +) -> GPUARModelRunner: + """A final-rank runner double that reaches the sampler and the state update.""" + num_reqs = sampled.shape[0] + hidden = torch.ones(num_reqs, 8) + runner = object.__new__(GPUARModelRunner) + runner.use_async_scheduling = async_scheduling + runner.broadcast_pp_output = broadcast_pp_output + runner.speculative_config = None + runner.device = torch.device("cpu") + runner.input_batch = _input_batch(num_reqs) + runner.discard_request_mask = SimpleNamespace(np=np.full(num_reqs, all_chunked_prefill, dtype=bool)) + runner.execute_model_state = ExecuteModelState( + scheduler_output=SimpleNamespace(total_num_scheduled_tokens=num_reqs), + logits=torch.zeros(num_reqs, 32), + spec_decode_metadata=None, + spec_decode_common_attn_metadata=None, + hidden_states=hidden, + hidden_states_cpu=None, + sample_hidden_states=hidden, + aux_hidden_states=None, + ec_connector_output=None, + cudagraph_stats=None, + multimodal_outputs={}, + slot_mappings=None, + ) + runner._sample = lambda logits, spec: SimpleNamespace(sampled_token_ids=sampled, logprobs_tensors=None) + runner._update_states_after_model_execute = lambda *args: events.append("state-update") + + def stop(*args): + raise _StopAtBookkeepingError + + runner._bookkeeping_sync = stop + return runner + + +def _non_last_rank_runner(num_reqs: int) -> GPUARModelRunner: + runner = object.__new__(GPUARModelRunner) + runner.use_async_scheduling = True + runner.execute_model_state = None + runner.kv_connector_output = None + runner.device = torch.device("cpu") + runner._pp_recv_work = None # as vLLM's __init__ leaves it before any receive is posted + runner.input_batch = _input_batch(num_reqs) + runner.discard_request_mask = SimpleNamespace(np=np.zeros(num_reqs, dtype=bool)) + runner.requests = {rid: SimpleNamespace(output_token_ids=[]) for rid in runner.input_batch.req_ids} + runner.attach_omni_connector_output = lambda output: output + return runner + + +class _Ledger(list): + """Every torch.distributed.broadcast, plus a shared timeline with the state update.""" + + def __init__(self): + super().__init__() + self.timeline: list[str] = [] + + +@pytest.fixture +def ledger(monkeypatch): + calls = _Ledger() + + class Pending: + def wait(self): + if not any(c["kind"] == "send" for c in calls): + raise RuntimeError("receive has no matching broadcast") + + def broadcast(tensor, src, group, async_op=False): + kind = "receive" if async_op else "send" + calls.append( + dict( + kind=kind, + shape=tuple(tensor.shape), + dtype=tensor.dtype, + src=src, + group=group, + value=tensor.clone(), + ) + ) + calls.timeline.append(kind) + return Pending() if async_op else None + + monkeypatch.setattr(torch.distributed, "broadcast", broadcast) + return calls + + +def _use_pp(monkeypatch, pp): + # The guard lives in the Omni override; the sender/receiver bodies live in vLLM. + monkeypatch.setattr(omni_runner_module, "get_pp_group", lambda: pp) + monkeypatch.setattr(upstream_runner_module, "get_pp_group", lambda: pp) + + +def test_non_last_rank_posts_the_receive_before_the_connector_only_return(monkeypatch, ledger): + _use_pp(monkeypatch, _pp(rank=0, world_size=2)) + runner = _non_last_rank_runner(num_reqs=2) + output = runner.sample_tokens(None) + assert output is not None + assert [c["kind"] for c in ledger] == ["receive"] + assert ledger[0]["src"] == 1 and ledger[0]["group"] == "pp-group" + assert ledger[0]["shape"] == (2, 1) and ledger[0]["dtype"] == torch.int32 + assert runner._pp_recv_work is not None + assert runner.input_batch.prev_sampled_token_ids.shape == (2, 1) + assert runner.input_batch.prev_req_id_to_index == {"req-0": 0, "req-1": 1} + + +def test_last_rank_broadcasts_sampled_tokens_after_the_state_update(monkeypatch, ledger): + pp = _pp(rank=1, world_size=2) + _use_pp(monkeypatch, pp) + sampled = torch.tensor([[4], [7]], dtype=torch.int32) + runner = _last_rank_runner(sampled=sampled, events=ledger.timeline) + with pytest.raises(_StopAtBookkeepingError): + runner.sample_tokens(None) + assert ledger.timeline == ["state-update", "send"], "the send must follow the state update, not precede it" + sends = [c for c in ledger if c["kind"] == "send"] + assert len(sends) == 1 + assert sends[0]["src"] == pp.rank == pp.last_rank + assert sends[0]["group"] == "pp-group" + assert sends[0]["shape"] == (2, 1) and sends[0]["dtype"] == torch.int32 + assert torch.equal(sends[0]["value"], sampled) + + +def test_receive_and_send_form_one_matched_pair(monkeypatch, ledger): + """The negative control from the handoff probe, now expected to complete.""" + _use_pp(monkeypatch, _pp(rank=0, world_size=2)) + first = _non_last_rank_runner(num_reqs=1) + first.sample_tokens(None) + _use_pp(monkeypatch, _pp(rank=1, world_size=2)) + last = _last_rank_runner(sampled=torch.tensor([[4]], dtype=torch.int32), events=[]) + with pytest.raises(_StopAtBookkeepingError): + last.sample_tokens(None) + first._pp_recv_work.wait() + assert [c["kind"] for c in ledger] == ["receive", "send"] + assert ledger[0]["shape"] == ledger[1]["shape"] and ledger[0]["dtype"] == ledger[1]["dtype"] + assert ledger[0]["src"] == ledger[1]["src"] and ledger[0]["group"] == ledger[1]["group"] + + +@pytest.mark.parametrize( + ("async_scheduling", "broadcast_pp_output", "world_size", "reason"), + [ + (False, False, 2, "synchronous scheduling feeds tokens through the scheduler"), + (True, True, 2, "external_launcher already broadcast the logits to every rank"), + (True, False, 1, "a single rank has nobody waiting"), + ], +) +def test_last_rank_does_not_send_outside_the_async_pp_case( + monkeypatch, ledger, async_scheduling, broadcast_pp_output, world_size, reason +): + _use_pp(monkeypatch, _pp(rank=world_size - 1, world_size=world_size)) + runner = _last_rank_runner( + sampled=torch.tensor([[4]], dtype=torch.int32), + async_scheduling=async_scheduling, + broadcast_pp_output=broadcast_pp_output, + events=[], + ) + with pytest.raises(_StopAtBookkeepingError): + runner.sample_tokens(None) + assert ledger == [], reason + + +def test_both_sides_skip_an_all_chunked_prefill_step(monkeypatch, ledger): + """Dummy tokens from unfinished prefill chunks are discarded, so neither side talks.""" + _use_pp(monkeypatch, _pp(rank=0, world_size=2)) + first = _non_last_rank_runner(num_reqs=1) + first.discard_request_mask = SimpleNamespace(np=np.ones(1, dtype=bool)) + first.sample_tokens(None) + assert first._pp_recv_work is None + _use_pp(monkeypatch, _pp(rank=1, world_size=2)) + last = _last_rank_runner(sampled=torch.tensor([[4]], dtype=torch.int32), all_chunked_prefill=True, events=[]) + with pytest.raises(_StopAtBookkeepingError): + last.sample_tokens(None) + assert ledger == [] + + +def test_sender_rejects_multi_token_rows(monkeypatch, ledger): + """The receive buffer is [num_reqs, 1]; the upstream sender guards that contract.""" + _use_pp(monkeypatch, _pp(rank=1, world_size=2)) + runner = _last_rank_runner(sampled=torch.tensor([[4, 5]], dtype=torch.int32), events=[]) + with pytest.raises(AssertionError, match="num_reqs, 1"): + runner.sample_tokens(None) + assert ledger == [] + + +# --- real collective ----------------------------------------------------------------- + + +def _gloo_rank(rank: int, world_size: int, rendezvous: str, result_path: str) -> None: + """Rank 0 runs the Omni receive branch; rank 1 runs the Omni sampling path. + + Both use the real methods bound to doubles; the process group is a real gloo group, + so rank 0's wait can only return because rank 1's ``sample_tokens`` broadcast. + """ + import torch.distributed as dist + + dist.init_process_group( + backend="gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + pp = _pp(rank=rank, world_size=world_size, device_group=dist.group.WORLD) + omni_runner_module.get_pp_group = lambda: pp + upstream_runner_module.get_pp_group = lambda: pp + if rank == world_size - 1: + runner = _last_rank_runner(sampled=torch.tensor([[4], [7]], dtype=torch.int32), events=[]) + try: + runner.sample_tokens(None) + except _StopAtBookkeepingError: + pass + received = None + else: + runner = _non_last_rank_runner(num_reqs=2) + runner.sample_tokens(None) + assert runner._pp_recv_work is not None + # A missing send surfaces here as a gloo timeout instead of a silent hang. + runner._pp_recv_work.wait(timeout=timedelta(seconds=30)) + received = runner.input_batch.prev_sampled_token_ids.tolist() + with open(f"{result_path}.{rank}", "w") as handle: + handle.write(repr(received)) + finally: + dist.destroy_process_group() + + +def test_gloo_pp2_receive_completes_only_because_the_omni_sampling_path_sent(tmp_path): + import torch.multiprocessing as mp + + rendezvous = str(tmp_path / "rendezvous") + result_path = str(tmp_path / "received") + context = mp.get_context("spawn") + processes = [ + context.Process(target=_gloo_rank, args=(rank, 2, rendezvous, result_path), daemon=True) for rank in range(2) + ] + try: + for process in processes: + process.start() + for process in processes: + process.join(timeout=240) + finally: + for process in processes: + if process.is_alive(): + process.kill() + process.join(timeout=30) + exit_codes = [process.exitcode for process in processes] + assert exit_codes == [0, 0], f"rank exit codes {exit_codes}: an unmatched receive times out on rank 0" + with open(f"{result_path}.0") as handle: + assert handle.read() == repr([[4], [7]]) + assert os.path.exists(f"{result_path}.1") diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6625effe8..06fdfa18b 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -2088,6 +2088,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex sampler_output = self._sample(logits, spec_decode_metadata) self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) + if self.use_async_scheduling: + pp = get_pp_group() + # Earlier PP ranks posted the matching async receive in the + # connector-only branch above (inherited from vLLM). Without this + # send their next `_prepare_input_ids` waits forever. + # For torchrun external_launcher PP mode with broadcast_pp_output=True, + # PP outputs have been broadcast to all ranks at logits computation, + # so the sampled token ids must not be sent again. + if not self.broadcast_pp_output and pp.world_size > 1 and pp.is_last_rank: + self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids) self._draft_token_ids = None self._draft_token_req_ids = None From 92715f3f7b67d7835c08002309402ba08b7b237e Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:24:34 -0700 Subject: [PATCH 081/209] [CI][ROCm] Add non-blocking engine and model executor GPU coverage (#7397) Signed-off-by: andyluo7 --- .buildkite/amd/test-amd-ready.yml | 33 ++++++++++++ tests/buildkite/test_amd_ready_r2_01.py | 72 +++++++++++++++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 tests/buildkite/test_amd_ready_r2_01.py diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 585f089a2..181a451bb 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -160,6 +160,39 @@ steps: # - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model +- label: "ROCm · Engine & Model Executor GPU Coverage (R2-01)" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 20 + artifact_paths: + - "artifacts/rocm-r2-01/**/*" + commands: + # Native AMD tests run from /app/vllm-omni, while Buildkite uploads + # artifacts relative to its checkout directory. + - export R2_01_ARTIFACT_DIR="$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-01" + - mkdir -p "$$R2_01_ARTIFACT_DIR" + # The CUDA ready job also covers tests/worker/test_batched_omni_output.py. + # Keep the initial ROCm burn-in scoped to the requested directories. + # Verbose pytest output and JUnit retain node IDs without a second + # collection-only pass over the same GPU-backed test tree. + # The CUDA/L4-only CosyVoice3 ragged-flow parity test cold-JITs multiple + # AITER kernels on ROCm and is covered by the CUDA ready lane. + - >- + pytest -s -v -ra + tests/engine/ tests/model_executor/ + -m 'core_model and cuda and not + (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' + -k 'not test_code2wav_streaming_batch_matches_ragged_flow_numerics' + --run-level "core_model" + --durations=0 + --junitxml="$$R2_01_ARTIFACT_DIR/pytest.xml" + 2>&1 | tee "$$R2_01_ARTIFACT_DIR/pytest.log" + - >- + tail -n 100 "$$R2_01_ARTIFACT_DIR/pytest.log" > + "$$R2_01_ARTIFACT_DIR/pytest-summary.txt" + - label: "Engine Test" agent_pool: mi300_1 depends_on: amd-build diff --git a/tests/buildkite/test_amd_ready_r2_01.py b/tests/buildkite/test_amd_ready_r2_01.py new file mode 100644 index 000000000..9d5147982 --- /dev/null +++ b/tests/buildkite/test_amd_ready_r2_01.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path +from shlex import split + +import pytest +import yaml + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml") +R2_01_LABEL = "ROCm · Engine & Model Executor GPU Coverage (R2-01)" +SINGLE_GPU_MARKERS = ( + "core_model and cuda and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)" +) +EXCLUDED_TEST = "test_code2wav_streaming_batch_matches_ragged_flow_numerics" + + +def _find_step(label: str) -> dict: + pipeline = yaml.safe_load(AMD_READY_PIPELINE.read_text(encoding="utf-8")) + + def walk(steps: list[dict]) -> dict | None: + for step in steps: + if step.get("label") == label: + return step + if nested := walk(step.get("steps", [])): + return nested + return None + + step = walk(pipeline.get("steps", [])) + assert step is not None, f"missing AMD pipeline step: {label}" + return step + + +def _assert_single_gpu_selection(command: str) -> None: + argv = split(command) + assert "tests/engine/" in argv + assert "tests/model_executor/" in argv + assert "tests/worker/test_batched_omni_output.py" not in argv + assert argv[argv.index("-m") + 1] == SINGLE_GPU_MARKERS + assert argv[argv.index("-k") + 1] == f"not {EXCLUDED_TEST}" + assert argv[argv.index("--run-level") + 1] == "core_model" + + +def test_r2_01_is_nonblocking_single_gpu_coverage() -> None: + step = _find_step(R2_01_LABEL) + + assert step["agent_pool"] == "mi300_1" + assert step["grade"] == "NonBlocking" + assert step["timeout_in_minutes"] == 20 + assert step["mirror_hardwares"] == ["amdproduction"] + assert step["artifact_paths"] == ["artifacts/rocm-r2-01/**/*"] + + +def test_r2_01_runs_once_fails_closed_and_publishes_results() -> None: + commands = _find_step(R2_01_LABEL)["commands"] + pytest_commands = [command for command in commands if command.lstrip().startswith("pytest ")] + assert len(pytest_commands) == 1 + run_command = pytest_commands[0] + artifact_dir = "$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-01" + + _assert_single_gpu_selection(run_command) + assert "--collect-only" not in "\n".join(commands) + assert "VLLM_CI_ALLOW_NO_TESTS" not in "\n".join(commands) + assert any(artifact_dir in command for command in commands) + assert "-v" in split(run_command) + assert "-ra" in split(run_command) + assert "--durations=0" in split(run_command) + assert "--junitxml=$$R2_01_ARTIFACT_DIR/pytest.xml" in split(run_command) + assert "$$R2_01_ARTIFACT_DIR/pytest.log" in run_command + assert any("$$R2_01_ARTIFACT_DIR/pytest-summary.txt" in command for command in commands) From 4a8ac297751b03fbe21f2d8a4550a027fe02e7cd Mon Sep 17 00:00:00 2001 From: xuexueligao <89690814+xuexueligao@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:46:13 +0800 Subject: [PATCH 082/209] [Frontend] Add MiniMax H3 video upscale workflow (WF-07) (#7472) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: xuezhihao <89690814+xuexueligao@users.noreply.github.com> Signed-off-by: 汪志鹏 Co-authored-by: 汪志鹏 --- apps/ComfyUI-vLLM-Omni/README.md | 4 + .../ComfyUI-vLLM-Omni/docs/wf07-h3-upscale.md | 21 + .../vLLM-Omni MiniMax H3 Video Upscale.json | 724 ++++++++++++++++++ 3 files changed, 749 insertions(+) create mode 100644 apps/ComfyUI-vLLM-Omni/docs/wf07-h3-upscale.md create mode 100644 apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax H3 Video Upscale.json diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md index 0d607b30a..b848a0716 100644 --- a/apps/ComfyUI-vLLM-Omni/README.md +++ b/apps/ComfyUI-vLLM-Omni/README.md @@ -139,6 +139,10 @@ You can configure per-stage sampling parameters for multi-stage models. > > Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect. +#### H3 video upscale (WF-07) + +The **vLLM-Omni MiniMax H3 Video Upscale** template generates video remotely, upscales it with SeedVR2, and saves the original and upscaled videos with the generated audio and FPS. See [workflow setup](docs/wf07-h3-upscale.md). + #### FastH3 text-to-video FastH3 is fused into MiniMax-H3 when the vLLM-Omni server starts; it is not a request-switchable LoRA. Download one adapter variant and start a non-offloaded FL2VA server. For the Dense / Data-Free profile: diff --git a/apps/ComfyUI-vLLM-Omni/docs/wf07-h3-upscale.md b/apps/ComfyUI-vLLM-Omni/docs/wf07-h3-upscale.md new file mode 100644 index 000000000..dfdbd166f --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/docs/wf07-h3-upscale.md @@ -0,0 +1,21 @@ +# WF-07: H3 video upscale with SeedVR2 + +Generate video and audio through vLLM-Omni H3, upscale the decoded frames with SeedVR2, and save both the original and upscaled videos. The generated audio and FPS connect directly to Create Video. + +## Requirements + +- ComfyUI with ComfyUI-vLLM-Omni and [ComfyUI-SeedVR2_VideoUpscaler](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler#-installation) installed. +- An H3 service using the FL2VA checkpoint partition. Follow the [H3 recipe](../../../recipes/MiniMaxAI/MiniMax-H3.md) for server setup. +- `seedvr2_ema_3b_fp16.safetensors` and `ema_vae_fp16.safetensors` from [SeedVR2's ComfyUI weights](https://huggingface.co/numz/SeedVR2_comfyUI/tree/main), placed in `ComfyUI/models/SEEDVR2`. The SeedVR2 nodes can also download them on first use. + +## Use + +Open **vLLM-Omni MiniMax H3 Video Upscale.json** from the example workflows. An API-format version is provided alongside it. Set the Generate Video URL/model and the LoRA path as seen by the H3 server. + +The preset uses the LightX2V v1.0 768p Turbo adapter in Diffusers format: `minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors`, five sigma points, video/audio flow shifts 6/3, and LoRA scale 1. Preload the same adapter on the server using the [Turbo LoRA recipe](../../../recipes/MiniMaxAI/MiniMax-H3.md#turbo-lora). For base H3, disconnect LoRA, set inference steps to 50, and set video flow shift to 12. + +The H3 preset requests 1344 x 768 at 24 FPS for 5.167 seconds, which gives 124 frames. SeedVR2's `resolution=1536` sets the target short edge, producing 2688 x 1536 for this input. Update this value when changing the input size; it is an output resolution, not a fixed 2x multiplier. + +The SeedVR2 preset uses 3B FP16, five-frame batches, one overlapping frame, uniform batch padding, and LAB color correction. The DiT loader uses CPU offload with 32 swapped blocks; the VAE uses tiled encoding and decoding. Select the local GPU in both model loaders. See the [SeedVR2 settings](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler#-usage) when adjusting memory use or batch size. + +Run the workflow. The two Save Video nodes write MP4s under `WF07/generated` and `WF07/upscaled` in ComfyUI's output directory. diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax H3 Video Upscale.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax H3 Video Upscale.json new file mode 100644 index 000000000..e8a21d065 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax H3 Video Upscale.json @@ -0,0 +1,724 @@ +{ + "version": 0.4, + "last_node_id": 12, + "last_link_id": 13, + "nodes": [ + { + "id": 1, + "type": "VLLMOmniGenerateVideo", + "pos": [ + 360, + 20 + ], + "size": [ + 330, + 480 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "frame", + "type": "IMAGE", + "shape": 7, + "link": null + }, + { + "name": "references", + "type": "VIDEO_REFERENCES", + "shape": 7, + "link": null + }, + { + "name": "sampling_params", + "type": "SAMPLING_PARAMS", + "shape": 7, + "link": 11 + }, + { + "name": "lora", + "type": "REMOTE_LORA", + "shape": 7, + "link": 12 + }, + { + "name": "model_params", + "type": "VIDEO_PARAMS", + "shape": 7, + "link": 10 + }, + { + "name": "fast_h3", + "type": "FASTH3_DEPLOYMENT", + "shape": 7, + "link": null + } + ], + "outputs": [ + { + "name": "video", + "type": "VIDEO", + "links": [ + 1, + 9 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniGenerateVideo" + }, + "widgets_values": [ + "http://127.0.0.1:8091/v1", + "MiniMaxAI/MiniMax-H3", + "A small brass bell hanging from a wooden beam in a quiet garden. A gentle breeze swings the bell and it rings clearly. Static camera, natural daylight, continuous motion, synchronized ringing and soft birdsong.", + "", + 1344, + 768, + 24, + 5.167 + ] + }, + { + "id": 2, + "type": "GetVideoComponents", + "pos": [ + 730, + 20 + ], + "size": [ + 250, + 125 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "video", + "type": "VIDEO", + "link": 1 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 2 + ] + }, + { + "name": "audio", + "type": "AUDIO", + "links": [ + 6 + ] + }, + { + "name": "fps", + "type": "FLOAT", + "links": [ + 5 + ] + }, + { + "name": "bit_depth", + "type": "INT", + "links": [ + 7 + ] + } + ], + "properties": { + "Node name for S&R": "GetVideoComponents" + } + }, + { + "id": 3, + "type": "SeedVR2LoadDiTModel", + "pos": [ + 730, + 290 + ], + "size": [ + 340, + 245 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "localized_name": "torch_compile_args", + "name": "torch_compile_args", + "shape": 7, + "type": "TORCH_COMPILE_ARGS", + "link": null + } + ], + "outputs": [ + { + "localized_name": "SEEDVR2_DIT", + "name": "SEEDVR2_DIT", + "type": "SEEDVR2_DIT", + "links": [ + 3 + ] + } + ], + "properties": { + "Node name for S&R": "SeedVR2LoadDiTModel" + }, + "widgets_values": [ + "seedvr2_ema_3b_fp16.safetensors", + "cuda:0", + 32, + false, + "cpu", + false, + "sdpa" + ] + }, + { + "id": 4, + "type": "SeedVR2VideoUpscaler", + "pos": [ + 1120, + 290 + ], + "size": [ + 340, + 415 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "localized_name": "image", + "name": "image", + "type": "IMAGE", + "link": 2 + }, + { + "localized_name": "dit", + "name": "dit", + "type": "SEEDVR2_DIT", + "link": 3 + }, + { + "localized_name": "vae", + "name": "vae", + "type": "SEEDVR2_VAE", + "link": 13 + } + ], + "outputs": [ + { + "localized_name": "IMAGE", + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ] + } + ], + "properties": { + "Node name for S&R": "SeedVR2VideoUpscaler" + }, + "widgets_values": [ + 42, + "fixed", + 1536, + 0, + 5, + true, + "lab", + 1, + 0, + 0.0, + 0.0, + "cpu", + false + ] + }, + { + "id": 5, + "type": "CreateVideo", + "pos": [ + 1490, + 290 + ], + "size": [ + 250, + 175 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 4 + }, + { + "name": "fps", + "type": "FLOAT", + "widget": { + "name": "fps" + }, + "link": 5 + }, + { + "name": "audio", + "type": "AUDIO", + "shape": 7, + "link": 6 + }, + { + "name": "bit_depth", + "type": "INT", + "widget": { + "name": "bit_depth" + }, + "link": 7 + } + ], + "outputs": [ + { + "name": "VIDEO", + "type": "VIDEO", + "links": [ + 8 + ] + } + ], + "properties": { + "Node name for S&R": "CreateVideo" + }, + "widgets_values": [ + 24, + 8 + ] + }, + { + "id": 6, + "type": "SaveVideo", + "pos": [ + 1790, + 290 + ], + "size": [ + 330, + 280 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "video", + "type": "VIDEO", + "link": 8 + }, + { + "name": "filename_prefix", + "type": "STRING", + "widget": { + "name": "filename_prefix" + }, + "link": null + }, + { + "name": "format", + "type": "COMBO", + "widget": { + "name": "format" + }, + "link": null + }, + { + "name": "codec", + "type": "COMFY_DYNAMICCOMBO_V3", + "widget": { + "name": "codec" + }, + "link": null + } + ], + "outputs": [ + { + "name": "video", + "type": "VIDEO", + "links": [] + } + ], + "properties": { + "Node name for S&R": "SaveVideo" + }, + "widgets_values": [ + "WF07/upscaled", + "mp4", + "auto" + ] + }, + { + "id": 7, + "type": "VLLMOmniMiniMaxH3Params", + "pos": [ + 0, + 20 + ], + "size": [ + 320, + 120 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MiniMax-H3 video params", + "type": "VIDEO_PARAMS", + "links": [ + 10 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniMiniMaxH3Params" + }, + "widgets_values": [ + 3.0, + 6.0 + ] + }, + { + "id": 8, + "type": "VLLMOmniDiffusionSampling", + "pos": [ + 0, + 200 + ], + "size": [ + 320, + 280 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "diffusion sampling params", + "type": "SAMPLING_PARAMS", + "links": [ + 11 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniDiffusionSampling" + }, + "widgets_values": [ + 1, + 5, + 1.0, + 1.0, + false, + true, + 42, + "fixed" + ] + }, + { + "id": 9, + "type": "VLLMOmniRemoteLoRA", + "pos": [ + 0, + 530 + ], + "size": [ + 320, + 210 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "lora", + "type": "REMOTE_LORA", + "links": [ + 12 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniRemoteLoRA" + }, + "widgets_values": [ + "models/minimax-h3-turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", + "h3-turbo-v1.0", + 1.0, + 0 + ] + }, + { + "id": 10, + "type": "SaveVideo", + "pos": [ + 1120, + 20 + ], + "size": [ + 330, + 190 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "video", + "type": "VIDEO", + "link": 9 + }, + { + "name": "filename_prefix", + "type": "STRING", + "widget": { + "name": "filename_prefix" + }, + "link": null + }, + { + "name": "format", + "type": "COMBO", + "widget": { + "name": "format" + }, + "link": null + }, + { + "name": "codec", + "type": "COMFY_DYNAMICCOMBO_V3", + "widget": { + "name": "codec" + }, + "link": null + } + ], + "outputs": [ + { + "name": "video", + "type": "VIDEO", + "links": [] + } + ], + "properties": { + "Node name for S&R": "SaveVideo" + }, + "widgets_values": [ + "WF07/generated", + "mp4", + "auto" + ], + "title": "Save generated video before upscale" + }, + { + "id": 11, + "type": "Note", + "pos": [ + 360, + 560 + ], + "size": [ + 330, + 260 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "Node name for S&R": "Note" + }, + "widgets_values": [ + "WF-07: H3 + SeedVR2\n\nInstall ComfyUI-SeedVR2_VideoUpscaler. The DiT and VAE loaders use SeedVR2 3B FP16 and ema_vae_fp16.\n\nSet the H3 service URL/model and the server-side Turbo LoRA path.\n\n1536 is the target short edge for this 1344 x 768 preset (2x). Original audio and FPS connect directly to Create Video.\n\nFor base H3: disconnect LoRA, use 50 steps and video flow shift 12." + ] + }, + { + "id": 12, + "type": "SeedVR2LoadVAEModel", + "pos": [ + 730, + 590 + ], + "size": [ + 340, + 315 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "localized_name": "torch_compile_args", + "name": "torch_compile_args", + "shape": 7, + "type": "TORCH_COMPILE_ARGS", + "link": null + } + ], + "outputs": [ + { + "localized_name": "SEEDVR2_VAE", + "name": "SEEDVR2_VAE", + "type": "SEEDVR2_VAE", + "links": [ + 13 + ] + } + ], + "properties": { + "Node name for S&R": "SeedVR2LoadVAEModel" + }, + "widgets_values": [ + "ema_vae_fp16.safetensors", + "cuda:0", + true, + 1024, + 128, + true, + 768, + 128, + "false", + "cpu", + false + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "VIDEO" + ], + [ + 2, + 2, + 0, + 4, + 0, + "IMAGE" + ], + [ + 3, + 3, + 0, + 4, + 1, + "SEEDVR2_DIT" + ], + [ + 4, + 4, + 0, + 5, + 0, + "IMAGE" + ], + [ + 5, + 2, + 2, + 5, + 1, + "FLOAT" + ], + [ + 6, + 2, + 1, + 5, + 2, + "AUDIO" + ], + [ + 7, + 2, + 3, + 5, + 3, + "INT" + ], + [ + 8, + 5, + 0, + 6, + 0, + "VIDEO" + ], + [ + 9, + 1, + 0, + 10, + 0, + "VIDEO" + ], + [ + 10, + 7, + 0, + 1, + 4, + "VIDEO_PARAMS" + ], + [ + 11, + 8, + 0, + 1, + 2, + "SAMPLING_PARAMS" + ], + [ + 12, + 9, + 0, + 1, + 3, + "REMOTE_LORA" + ], + [ + 13, + 12, + 0, + 4, + 2, + "SEEDVR2_VAE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.52, + "offset": [ + 100, + 100 + ] + } + } +} From a3dee6fdf52267fd56f0237f47b5bbe68dbb13b2 Mon Sep 17 00:00:00 2001 From: wtz2333 <2955110911@qq.com> Date: Tue, 15 Sep 2026 18:49:55 +0800 Subject: [PATCH 083/209] [Lingbot World] Bound image condition storage and extend temporal RoPE (#6838) Signed-off-by: wtz2333 <2955110911@qq.com> --- docs/design/feature/realtime_ar_diffusion.md | 16 +- .../ar_diffusion_streaming_decode.py | 20 +- .../diffusion/lingbot_world_v2.py | 3 - .../diffusion/lingbot_world_v2_realtime.py | 19 +- recipes/Robbyant/LingBot-World-2.0.md | 8 +- .../test_lingbot_world_camera.py | 29 +- .../test_lingbot_world_transformer.py | 52 ++- .../test_pipeline_lingbot_world.py | 347 +++++++++++++++--- .../test_lingbot_world_v2.py | 24 ++ .../diffusion/models/lingbot_world/camera.py | 7 +- .../models/lingbot_world/pipeline.py | 238 ++++++++---- .../models/lingbot_world/transformer.py | 24 +- 12 files changed, 592 insertions(+), 195 deletions(-) diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md index 426ee5afd..2f86744b1 100644 --- a/docs/design/feature/realtime_ar_diffusion.md +++ b/docs/design/feature/realtime_ar_diffusion.md @@ -143,8 +143,20 @@ fraction to select additional resident sessions up to the model-declared cap. All model-owned reservations are deducted before the paged self-attention pool is allocated. -LingBot reports its persistent full-horizon image-condition tensor. DreamZero -reports the measured 603 MiB per-session Wan VAE causal-convolution state. +LingBot advances a session-owned causal Wan VAE encoder for each condition +block: the opening pixel frame is followed by four zero pixel frames per later +latent frame. It retains one current condition block, with separate committed +and in-flight encoder histories so a failed block cannot overwrite committed +context. The state reservation includes both encoder histories, derived from +the encoder's convolution input grids, and the streaming decoder cache. These +caches are released through the same close/reset hooks. Condition encoding +currently requires the unpatched, non-tiled Wan encoder path; the existing +decoder's tiling fallback remains independent. Real-weight GPU parity with a +full-horizon encode still requires validation because the pointwise posterior +projection now runs per block. + +DreamZero reports the measured 603 MiB per-session Wan VAE causal-convolution +state. ## Concurrency and routing limits diff --git a/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py index 2aa326f6d..6c87d4e63 100644 --- a/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py +++ b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py @@ -24,10 +24,9 @@ Examples:: ``--action-script`` is a JSON list with one camera action list per AR block, for example ``[[["w"], ["w"], ["w"]], [["a"], [], []]]`` for a checkpoint that -generates three latent frames per block. The frames per block, the temporal -compression and the number of blocks that fit the condition horizon are read -off the checkpoint rather than assumed, so the script can be repointed with -``--model``. +generates three latent frames per block. The frames per block and temporal +compression are read off the checkpoint rather than assumed, so the script +can be repointed with ``--model``. """ from __future__ import annotations @@ -49,9 +48,6 @@ if TYPE_CHECKING: import numpy as np _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" -# Mirrors _MAX_RAW_FRAMES in vllm_omni/diffusion/models/lingbot_world/pipeline.py: -# the image-condition horizon is the pipeline's constant, not the checkpoint's. -_MAX_RAW_FRAMES = 117 def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: @@ -67,8 +63,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "--action-script", required=True, help=( - "JSON file holding one camera action list per AR block. Both the per-block frame count " - "and the number of blocks that fit are read off the checkpoint, not assumed." + "JSON file holding one camera action list per AR block. " + "The per-block frame count is read off the checkpoint, not assumed." ), ) parser.add_argument("--output-dir", required=True, help="Directory for decoded frames and metadata.") @@ -154,12 +150,6 @@ def _validate_against_checkpoint( f"--action-script block {block_index} holds {len(block)} frame action lists, but this " f"checkpoint generates {frames_per_block} latent frames per AR block." ) - max_chunks = ((_MAX_RAW_FRAMES - 1) // temporal_compression + 1) // frames_per_block - if len(script) > max_chunks: - raise ValueError( - f"--action-script holds {len(script)} blocks, but this checkpoint's " - f"{_MAX_RAW_FRAMES}-frame image-condition horizon fits at most {max_chunks}." - ) return (len(script) * frames_per_block - 1) * temporal_compression + 1 diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py index 4bbea2f4f..0eb2027b1 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2.py @@ -32,7 +32,6 @@ import numpy as np _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" _MAX_SEQUENCE_LENGTH = 512 _MAX_PIXEL_AREA = 480 * 832 -_MAX_RAW_FRAMES = 117 _MAX_SOURCE_CAMERA_FRAMES = 4096 _TEMPORAL_COMPRESSION = 4 _LATENT_FRAMES_PER_BLOCK = 3 @@ -191,8 +190,6 @@ def build_request( raise ValueError("--height * --width must not exceed 480 * 832 pixels.") if args.num_frames <= 0: raise ValueError("--num-frames must be a positive integer.") - if args.num_frames > _MAX_RAW_FRAMES: - raise ValueError("--num-frames must not exceed 117 raw frames.") if args.num_frames > paths.camera_frames: raise ValueError("Camera arrays must contain at least --num-frames frames.") if (args.num_frames - 1) % _TEMPORAL_COMPRESSION: diff --git a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py index d98fe06f9..f0f1043ab 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session. This example exercises the internal session API. It is not a public HTTP or WebSocket protocol. Each JSONL event advances the world by one three-latent -frame AR block and writes the latent plus its identity metadata. The current -scoped implementation supports at most ten ticks per generation epoch because -its image condition is bounded to 117 pixel frames; reset or create a session -to start a new world. +frame AR block and writes the latent plus its identity metadata. Session-owned +causal VAE encoder state produces each condition block without retaining or +repeating a full image-condition horizon. """ from __future__ import annotations @@ -24,8 +23,6 @@ from typing import Any _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" _CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1" _FRAMES_PER_BLOCK = 3 -# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames. -_MAX_REALTIME_TICKS = 10 def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]: @@ -41,7 +38,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument( "--events", required=True, - help="JSONL file with one prompt/action event per AR block; at most 10 events.", + help="JSONL file with one prompt/action event per AR block.", ) parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.") parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.") @@ -84,12 +81,6 @@ def _load_events(path: Path) -> list[dict[str, Any]]: events.append({"event_id": event_id, "prompt": prompt, "frames": frames}) if not events: raise ValueError("events file must contain at least one event.") - if len(events) > _MAX_REALTIME_TICKS: - raise ValueError( - "events file must contain at most " - f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime " - "image-condition horizon is 117 pixel frames." - ) event_ids = [event["event_id"] for event in events] if event_ids != sorted(set(event_ids)): raise ValueError("event_id values must be unique and strictly increasing.") diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md index c38e24fd2..ab8c498fa 100644 --- a/recipes/Robbyant/LingBot-World-2.0.md +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -48,9 +48,11 @@ python examples/offline_inference/diffusion/lingbot_world_v2.py \ ``` The current causal-fast checkpoint uses four DMD sampling steps. Raw frame -counts must be `9 + 12k`, up to 117 frames. Camera files may contain a -longer bounded source trajectory (the official example contains 269 frames); -the runtime consumes only the prefix needed by the request. +counts must be `9 + 12k` for a non-negative integer `k`; for example, 129 frames +produce eleven AR blocks. Both camera files must contain the same number of +source frames, at least as many as requested and at most 4096. The runtime +consumes only the prefix needed by the request. Longer rollouts also require +sufficient device memory. ## Realtime in-process generation (deprecated) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py index 121364c07..9cb67a142 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py @@ -8,6 +8,7 @@ import math import subprocess import sys from pathlib import Path +from typing import TYPE_CHECKING import numpy as np import pytest @@ -20,7 +21,10 @@ _CAMERA_MODULE = importlib.util.module_from_spec(_SPEC) sys.modules[_SPEC.name] = _CAMERA_MODULE _SPEC.loader.exec_module(_CAMERA_MODULE) -CameraTrajectory = _CAMERA_MODULE.CameraTrajectory +if TYPE_CHECKING: + from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory +else: + CameraTrajectory = _CAMERA_MODULE.CameraTrajectory build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory @@ -138,21 +142,20 @@ def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_pa load_camera_trajectory(_trusted(tmp_path)) -def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None: - _write_trajectory( - tmp_path, - _identity_poses(269), - np.arange(269 * 4, dtype=np.float32).reshape(269, 4), - ) +@pytest.mark.parametrize("num_frames", [129, 269, 4096]) +def test_load_camera_trajectory_preserves_all_bounded_source_frames(tmp_path: Path, num_frames: int) -> None: + poses = _identity_poses(num_frames) + poses[:, 0, 3] = np.arange(num_frames) + intrinsics = np.arange(num_frames * 4, dtype=np.float32).reshape(num_frames, 4) + _write_trajectory(tmp_path, poses, intrinsics) trajectory = load_camera_trajectory(_trusted(tmp_path)) - assert trajectory.poses.shape == (117, 4, 4) - assert trajectory.intrinsics.shape == (117, 4) - np.testing.assert_array_equal( - trajectory.intrinsics.numpy(), - np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117], - ) + assert trajectory.poses.shape == (num_frames, 4, 4) + assert trajectory.intrinsics.shape == (num_frames, 4) + assert trajectory.poses.dtype == trajectory.intrinsics.dtype == torch.float32 + np.testing.assert_array_equal(trajectory.poses.numpy(), poses.astype(np.float32)) + np.testing.assert_array_equal(trajectory.intrinsics.numpy(), intrinsics) def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None: diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py index b4bf7805f..96c0b37f1 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py @@ -35,11 +35,13 @@ def _tiny_model( num_layers: int = 2, num_frames_per_block: int = 1, sliding_window_num_frames: int = 3, + attention_head_dim: int = 2, + rope_max_seq_len: int = 16, ): return module.CausalLingBotWorldTransformer3DModel( patch_size=(1, 2, 2), num_attention_heads=2, - attention_head_dim=2, + attention_head_dim=attention_head_dim, in_channels=36, out_channels=2, text_dim=6, @@ -48,7 +50,7 @@ def _tiny_model( num_layers=num_layers, cross_attn_norm=True, eps=1e-6, - rope_max_seq_len=16, + rope_max_seq_len=rope_max_seq_len, sink_size=1, num_frames_per_block=num_frames_per_block, sliding_window_num_frames=sliding_window_num_frames, @@ -121,7 +123,8 @@ def _assert_self_cache_unchanged(cache, snapshot) -> None: @pytest.mark.cpu -def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None: +@pytest.mark.parametrize("num_chunks", [4, 18]) +def test_tiny_transformer_runs_contiguous_chunks_with_explicit_cache_commit_and_camera_path(num_chunks) -> None: torch.manual_seed(7) module = attention_tests._load_module() model = _tiny_model(module).eval() @@ -135,7 +138,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera assert model.patch_embedding.in_channels == 36 assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2 - for start_frame in range(4): + for start_frame in range(num_chunks): hidden_states = torch.randn(1, 36, 1, 4, 4) timestep = torch.tensor([float(start_frame + 1)]) snapshot = _self_cache_snapshot(cache) @@ -169,7 +172,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera expected_token_end = (start_frame + 1) * 2 * 2 assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention) - assert len(outputs) == 4 + assert len(outputs) == num_chunks assert all(len(outputs) == 1 for outputs in cross_key_outputs) assert len(text_projection_outputs) == 1 @@ -991,3 +994,42 @@ def test_sp4_direct_matches_sp1_bf16(tmp_path): @hardware_test(res={"cuda": "L4"}, num_cards=4) def test_tp2_sp2_paged_matches_sp1_bf16(tmp_path): _run(tmp_path, 2, 2, "paged", torch.bfloat16, 1) + + +@pytest.mark.cpu +@pytest.mark.parametrize("start_frame", [0, 1021, 1023, 1024, 21600]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_temporal_rope_matches_absolute_positions_without_growing_cache(start_frame, dtype): + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1, attention_head_dim=128, rope_max_seq_len=1024).to(dtype) + buffers = {name: (value.data_ptr(), value.shape) for name, value in model.named_buffers()} + cosine, sine = model._rotary_embedding( + frames=3, height=2, width=3, start_frame=start_frame, dtype=dtype, device=torch.device("cpu") + ) + # Independent complex-valued reference uses absolute (time, height, width) + # coordinates, including blocks crossing the original precomputed table. + grids = torch.meshgrid( + torch.arange(start_frame, start_frame + 3, dtype=torch.float64), + torch.arange(2, dtype=torch.float64), + torch.arange(3, dtype=torch.float64), + indexing="ij", + ) + reference = [] + for positions, dim in zip(grids, (44, 42, 42), strict=True): + frequencies = 1.0 / 10000 ** (torch.arange(0, dim, 2, dtype=torch.float64) / dim) + phase = positions.flatten()[:, None] * frequencies[None, :] + reference.append(torch.polar(torch.ones_like(phase), phase)) + expected = torch.cat(reference, dim=-1) + torch.testing.assert_close(cosine, expected.real.float().to(dtype), rtol=0, atol=0) + torch.testing.assert_close(sine, expected.imag.float().to(dtype), rtol=0, atol=0) + assert {name: (value.data_ptr(), value.shape) for name, value in model.named_buffers()} == buffers + + +@pytest.mark.cpu +def test_temporal_rope_extension_keeps_spatial_limit(): + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1) + with pytest.raises(ValueError, match="Spatial RoPE"): + model._rotary_embedding( + frames=3, height=17, width=1, start_frame=16, dtype=torch.float32, device=torch.device("cpu") + ) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index 86d8e70b8..fb9e93a5d 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -198,10 +198,43 @@ class _StubCausalDecoder: return x.new_zeros(x.shape[0], 3, num_frames, x.shape[-2] * 8, x.shape[-1] * 8) +class _StubCausalEncoder(nn.Module): + """Minimal Wan cache layout; the real tiny VAE below checks numerical parity.""" + + def __init__(self): + super().__init__() + self.conv_in = nn.Conv3d(3, 16, 1) + self.down_blocks = nn.ModuleList([lingbot_pipeline.WanResample(16, mode="downsample2d") for _ in range(3)]) + self.mid_block = SimpleNamespace(resnets=[]) + self.conv_out = nn.Conv3d(16, 32, 1) + self.inputs: list[torch.Tensor] = [] + + def forward(self, video, *, feat_cache, feat_idx): + self.inputs.append(video.detach().clone()) + first = feat_cache[0] is None + tail = video[:, :, -2:].clone() + if tail.shape[2] == 1 and feat_cache[0] is not None: + tail = torch.cat((feat_cache[0][:, :, -1:], tail), dim=2) + feat_cache[0] = tail + height, width = video.shape[-2] // 8, video.shape[-1] // 8 + previous_frames = 0 if feat_cache[1] is None else feat_cache[1].shape[2] + feat_cache[1] = video.new_zeros(1, 16, min(2, previous_frames + 1), height, width) + feat_idx[0] += 2 + moments = video.new_zeros(1, 32, 1, height, width) + if first: + moments[:, :16] = 2.0 + return moments + + class _StubVAE(_FakePretrained): dtype = torch.float32 def __init__(self, *, streaming: bool = True): + self.encoder = _StubCausalEncoder() + self.quant_conv = nn.Identity() + self._cached_conv_counts = {"encoder": 2} + self._enc_feat_map = ["module-owned"] + self._enc_conv_idx = [71] if streaming: self.decoder = _StubCausalDecoder() # Recorded, not discarded: on the streaming branch this is the only @@ -209,7 +242,7 @@ class _StubVAE(_FakePretrained): # what carries the checkpoint's latent rescale. self.post_quant_inputs: list[torch.Tensor] = [] self.post_quant_conv = self._record_post_quant - self._cached_conv_counts = {"decoder": _StubCausalDecoder.NUM_CONVS, "encoder": 2} + self._cached_conv_counts["decoder"] = _StubCausalDecoder.NUM_CONVS # Module-owned cache the shared VAE keeps for whole-clip decode; # streaming must never write through it. self._feat_map = ["module-owned"] @@ -498,16 +531,14 @@ def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None: assert spec.sink_frames == 3 assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)] assert spec.cross_attention_lengths == {"text": 512} - # The image condition (9,600 B) plus the streaming decoder's own per-session - # temporal cache. The cache dwarfs the condition at any real resolution, and - # the runner subtracts this field from the KV budget before sizing its pools, - # so leaving the cache out of it over-commits device memory. - assert spec.model_owned_state_bytes_per_session == 9_600 + 2_421_248 + # One condition block, committed + pending encoder caches, and decoder state. + # The stub encoder retains 2 RGB frames at 16x16 and 2 feature frames at 2x2. + encoder_bytes = (2 * 3 * 16 * 16 + 2 * 16 * 2 * 2) * 4 + assert spec.model_owned_state_bytes_per_session == 960 + 2 * encoder_bytes + 2_421_248 def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None: - """The declaration scales with output area, and is the condition alone when - the VAE cannot stream.""" + """Admission includes encoder state even when the decoder cannot stream.""" module = _load_pipeline_module() def spec_pipeline(): @@ -526,13 +557,13 @@ def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None: # than the image condition alone could account for. assert widened - small > 2_000_000 - # A VAE without the causal-cache contract cannot stream, so there is no - # decoder cache to declare and the figure falls back to the condition. + # Removing decoder support leaves the condition and both encoder histories. bare = spec_pipeline() - for attribute in ("decoder", "post_quant_conv", "_cached_conv_counts"): + for attribute in ("decoder", "post_quant_conv"): if hasattr(bare.vae, attribute): delattr(bare.vae, attribute) - assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 9_600 + bare.vae._cached_conv_counts.pop("decoder") + assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 2 * 6_656 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None: @@ -1200,20 +1231,19 @@ def test_request_validation_rejects_unsupported_contracts(request_batch, message _pipeline(module)._parse_request(request_batch) -def test_resource_limits_accept_exact_documented_boundaries() -> None: +def test_request_validation_accepts_frames_beyond_previous_limit() -> None: module = _load_pipeline_module() - sampling = _SamplingParams(height=480, width=832, num_frames=117, max_sequence_length=512) + sampling = _SamplingParams(height=480, width=832, num_frames=129, max_sequence_length=512) parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) - assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 117, 512) + assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 129, 512) @pytest.mark.parametrize( ("sampling", "message"), [ (_SamplingParams(height=480, width=848), "pixel area|480.*832"), - (_SamplingParams(num_frames=129), "num_frames.*117"), (_SamplingParams(max_sequence_length=511), "max_sequence_length.*512"), (_SamplingParams(max_sequence_length=513), "max_sequence_length.*512"), (_SamplingParams(max_sequence_length=512.0), "max_sequence_length.*512"), @@ -1550,9 +1580,9 @@ def test_multi_chunk_generation_uses_one_request_local_cache_and_decodes_accumul assert cache_ref() is None -def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room"): +def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room", session_id: str = "world-1"): return ARDiffusionTickRequest( - session_id="world-1", + session_id=session_id, request_id="legacy-lingbot-request", chunk_index=chunk_index, applied_event_ids=(chunk_index,), @@ -1567,6 +1597,16 @@ def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room") ).to_extra_args() +def test_realtime_rejects_pixel_output_before_condition_encoding() -> None: + pipeline = _pipeline(_load_pipeline_module()) + pipeline._ar_diffusion_kv_state = object() + sampling = _SamplingParams(output_type="np", extra_args=_tick_extra_args(chunk_index=0)) + with pytest.raises(ValueError, match="require output_type='latent'"): + pipeline(_request(sampling=sampling)) + assert pipeline._ar_sessions == {} + assert pipeline.vae.encoder.inputs == [] + + def test_typed_ticks_generate_one_global_block_and_return_standard_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1769,7 +1809,7 @@ def test_typed_tick_rejects_non_contiguous_chunk_index( pipeline(_request(sampling=sampling)) -def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon( +def test_typed_ticks_keep_bounded_encoder_state_beyond_ten_chunks( monkeypatch: pytest.MonkeyPatch, ) -> None: module = _load_pipeline_module() @@ -1777,22 +1817,229 @@ def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon( pipeline._ar_height = 16 pipeline._ar_width = 16 pipeline._ar_diffusion_kv_state = object() - pipeline._ar_sessions["world-1"] = module._LingBotARSessionState( - next_chunk_index=10, - ) + generated = [] monkeypatch.setattr( pipeline, - "encode_prompt", - lambda *args, **kwargs: pytest.fail("the realtime horizon must be validated before prompt encoding"), + "_ar_text_caches", + lambda *args, **kwargs: [SimpleNamespace()], ) - sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=10)) - with pytest.raises( - ValueError, - match=r"at most 10 ticks.*chunk_index 0 through 9.*117 pixel frames", - ): + def generate_block(**kwargs): + generated.append( + { + "condition": kwargs["condition"].clone(), + "start_frame": kwargs["start_frame"], + } + ) + return torch.zeros_like(kwargs["condition"][:, :16]) + + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + + for chunk_index in range(11): + sampling = _SamplingParams( + extra_args=_tick_extra_args(chunk_index=chunk_index), + ) pipeline(_request(sampling=sampling)) + assert pipeline.vae.encode_inputs == [] + assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 32 + assert all(torch.count_nonzero(video) == 0 for video in pipeline.vae.encoder.inputs[1:]) + encoder_cache = pipeline._ar_sessions["world-1"].encoder_cache + assert sum(t.numel() * t.element_size() for t in encoder_cache) == pipeline._condition_encoder_cache_bytes() + assert pipeline.vae._enc_feat_map == ["module-owned"] + assert pipeline.vae._enc_conv_idx == [71] + assert pipeline._ar_sessions["world-1"].next_chunk_index == 11 + assert [item["start_frame"] for item in generated] == list(range(0, 33, 3)) + assert torch.count_nonzero(generated[0]["condition"][:, :4]) > 0 + torch.testing.assert_close( + generated[1]["condition"][:, :4], + torch.zeros_like(generated[1]["condition"][:, :4]), + ) + for item in generated[2:]: + torch.testing.assert_close(item["condition"], generated[1]["condition"]) + + +def _tiny_condition_pipeline(monkeypatch: pytest.MonkeyPatch): + from diffusers import AutoencoderKLWan + + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import retrieve_latents + + pipeline = _pipeline(_load_pipeline_module()) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(7) + pipeline.vae = AutoencoderKLWan( + base_dim=4, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ).eval() + monkeypatch.setattr(lingbot_pipeline, "retrieve_latents", retrieve_latents) + monkeypatch.setattr(pipeline, "_ar_text_caches", lambda *args, **kwargs: [SimpleNamespace()]) + return pipeline + + +@pytest.mark.parametrize("mode", ["realtime", "stepwise"]) +def test_condition_continues_real_causal_vae_beyond_initial_horizon( + monkeypatch: pytest.MonkeyPatch, + mode: str, +) -> None: + pipeline = _tiny_condition_pipeline(monkeypatch) + blocks = 12 + inputs = pipeline._parse_request(_request(sampling=_SamplingParams(num_frames=(blocks * 3 - 1) * 4 + 1))) + with torch.inference_mode(): + expected = pipeline._prepare_condition(inputs, dtype=torch.float32) + module_cache = pipeline.vae._enc_feat_map + seen = [] + + def check(condition): + index = len(seen) + torch.testing.assert_close(condition, expected[:, :, index * 3 : (index + 1) * 3], rtol=1e-6, atol=1e-6) + seen.append(condition.clone()) + + if mode == "realtime": + pipeline._ar_diffusion_kv_state = object() + + def generate_block(**kwargs): + check(kwargs["condition"]) + return torch.zeros_like(kwargs["condition"][:, :16]) + + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + session_id = "world-1" + outputs = ( + pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index)))) + for index in range(blocks) + ) + else: + + def probe_step(**kwargs): + if kwargs["step_index"] == 0: + check(kwargs["condition"]) + return torch.ones_like(kwargs["current_latents"]) + + monkeypatch.setattr(pipeline._dmd_blocks, "probe_step", probe_step) + monkeypatch.setattr(pipeline._dmd_blocks, "commit_block_kv", lambda **kwargs: None) + state = _stepwise_state(num_frames=(blocks * 3 - 1) * 4 + 1) + session_id = state.request_id + outputs = _stepwise_chunks(pipeline, state, _FakeARState(session_id)) + + cache_sizes = [] + for _ in outputs: + state = pipeline._ar_sessions.get(session_id) + if state is not None: + for cache in (state.encoder_cache, state.pending_encoder_cache): + if cache is not None: + size = sum(t.numel() * t.element_size() for t in cache if t is not None) + assert size == pipeline._condition_encoder_cache_bytes() + assert all(t.grad_fn is None for t in cache if t is not None) + cache_sizes.append(size) + assert len(seen) == blocks and len(set(cache_sizes)) == 1 + assert pipeline.vae._enc_feat_map is module_cache and all(t is None for t in module_cache) + pipeline.close_ar_diffusion_session(session_id) + assert pipeline._ar_sessions == {} + + +def test_condition_encoder_interleaving_retry_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _tiny_condition_pipeline(monkeypatch) + pipeline._ar_diffusion_kv_state = object() + prompts = { + "a": _prompt(images=Image.new("RGB", (16, 16), (255, 0, 0))), + "b": _prompt(images=Image.new("RGB", (16, 16), (0, 255, 0))), + } + expected = {} + with torch.inference_mode(): + for name, prompt in prompts.items(): + inputs = pipeline._parse_request(_request(prompt=prompt, sampling=_SamplingParams(num_frames=45))) + expected[name] = pipeline._prepare_condition(inputs, dtype=torch.float32) + assert not torch.equal(expected["a"], expected["b"]) + active = None + failure = None + original_encoder = pipeline.vae.encoder.forward + + def encoder(*args, **kwargs): + output = original_encoder(*args, **kwargs) + if failure == "encoder": + raise RuntimeError("injected encoder failure") + return output + + def generate_block(**kwargs): + index = kwargs["start_frame"] + torch.testing.assert_close(kwargs["condition"], expected[active][:, :, index : index + 3], rtol=1e-6, atol=1e-6) + if failure == "dit": + raise RuntimeError("injected DiT failure") + return torch.zeros_like(kwargs["condition"][:, :16]) + + monkeypatch.setattr(pipeline.vae.encoder, "forward", encoder) + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + + def run(name, index): + nonlocal active + active = name + sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=index, session_id=name)) + return pipeline(_request(prompt=prompts[name], sampling=sampling)) + + run("a", 0) + for index, kind in [(1, "encoder"), (2, "dit")]: + state = pipeline._ar_sessions["a"] + committed = state.encoder_cache + saved = [t.clone() if t is not None else None for t in committed] + failure = kind + with pytest.raises(RuntimeError, match="injected"): + run("a", index) + assert state.next_chunk_index == index and state.encoder_cache is committed + assert state.pending_encoder_cache is None + for before, after in zip(saved, committed, strict=True): + if before is not None: + torch.testing.assert_close(before, after, rtol=0, atol=0) + failure = None + run("b", index - 1) + run("a", index) + + old_a = pipeline._ar_sessions["a"] + pipeline.reset_ar_diffusion_session("a") + assert old_a.encoder_cache is None and old_a.pending_encoder_cache is None + assert "a" not in pipeline._ar_sessions and "b" in pipeline._ar_sessions + run("a", 0) + old_b = pipeline._ar_sessions["b"] + pipeline.close_ar_diffusion_session("b") + assert old_b.encoder_cache is None and old_b.pending_encoder_cache is None + assert "a" in pipeline._ar_sessions and "b" not in pipeline._ar_sessions + pipeline.close_ar_diffusion_session("a") + assert pipeline._ar_sessions == {} + + +def test_stepwise_failed_chunk_releases_pending_encoder_on_close(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _pipeline(_load_pipeline_module()) + state = _stepwise_state() + transformer = pipeline.transformer + transformer.raise_on_call = 1 + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + session = pipeline._ar_sessions[state.request_id] + assert session.encoder_cache is None and session.pending_encoder_cache is not None + with pytest.raises(RuntimeError, match="forced transformer failure"): + pipeline.denoise_step(None, states=[state]) + assert session.next_chunk_index == 0 and session.encoder_cache is None + pipeline.close_ar_diffusion_session(state.request_id) + assert session.encoder_cache is None and session.pending_encoder_cache is None + assert state.request_id not in pipeline._ar_sessions + + +@pytest.mark.parametrize("mode", ["realtime", "stepwise"]) +def test_stateful_condition_rejects_tiled_encoder(monkeypatch: pytest.MonkeyPatch, mode: str) -> None: + pipeline = _pipeline(_load_pipeline_module()) + _enable_vae_tiling(pipeline, tile_sample_min=8) + with pytest.raises(ValueError, match="does not support tiled VAE encoding"): + if mode == "realtime": + pipeline._ar_diffusion_kv_state = object() + pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=0)))) + else: + state = _stepwise_state() + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + assert pipeline.vae.encoder.inputs == [] + assert all(s.encoder_cache is None and s.pending_encoder_cache is None for s in pipeline._ar_sessions.values()) + def test_request_cache_is_released_before_vae_decode() -> None: module = _load_pipeline_module() @@ -1820,25 +2067,25 @@ def test_request_cache_is_released_before_vae_decode() -> None: assert result.output.shape[2] == 21 -def test_117_frame_request_generates_ten_complete_latent_blocks() -> None: +def test_129_frame_request_generates_eleven_complete_latent_blocks() -> None: module = _load_pipeline_module() transformer = _RecordingTransformer() pipeline = _pipeline(module, transformer=transformer) trajectory = _CameraTrajectory( - poses=torch.eye(4).repeat(117, 1, 1), - intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(117, 1), + poses=torch.eye(4).repeat(129, 1, 1), + intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(129, 1), ) sampling = _SamplingParams( - num_frames=117, + num_frames=129, output_type="latent", extra_args={"_lingbot_camera_trajectory": trajectory}, ) result = pipeline(_request(sampling=sampling)) - assert result.output.shape == (1, 16, 30, 2, 2) - assert len(transformer.calls) == 50 - assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 30, 3)) + assert result.output.shape == (1, 16, 33, 2, 2) + assert len(transformer.calls) == 55 + assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 33, 3)) def test_request_cache_becomes_unreachable_after_transformer_error() -> None: @@ -2030,13 +2277,14 @@ def test_stepwise_progress_metadata_and_commit_trace() -> None: assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2) -def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: +@pytest.mark.parametrize("num_chunks", [2, 12]) +def test_stepwise_matches_tick_transformer_trace_and_latents(num_chunks) -> None: module = _load_pipeline_module() transformer = _RecordingTransformer() pipeline = _pipeline(module, transformer=transformer) pipeline._ar_height = 16 pipeline._ar_width = 16 - script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",))) + script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",))) * (num_chunks // 2) fake = _FakeARState("world-1") with pipeline.bind_ar_diffusion_state("world-1", fake): @@ -2063,7 +2311,7 @@ def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: transformer.calls.clear() pipeline.close_ar_diffusion_session("world-1") - stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17) + stepwise_state = _stepwise_state(request_id="req-1", num_frames=12 * num_chunks - 3, script=script, seed=17) stepwise_fake = _FakeARState(stepwise_state.request_id) with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake): stepwise_outputs = _run_stepwise(pipeline, stepwise_state) @@ -2072,6 +2320,15 @@ def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls] for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True): torch.testing.assert_close(output.output["payload"]["latents"], tick_latent) + for step_call, tick_call in zip(transformer.calls, tick_calls, strict=True): + torch.testing.assert_close(step_call["hidden_states"], tick_call["hidden_states"]) + torch.testing.assert_close(step_call["camera_hidden_states"], tick_call["camera_hidden_states"]) + assert stepwise_state.extra["condition"].shape[2] == 3 + assert pipeline.vae.encode_inputs == [] + assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (num_chunks * 3 - 1)) * 2 + assert stepwise_fake.commits == ["main"] * num_chunks + assert len(stepwise_outputs) == num_chunks and stepwise_outputs[-1].finished + assert pipeline.vae.decode_inputs == [] def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None: @@ -2387,6 +2644,7 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None: durations = pipeline.stage_durations assert stage in durations, f"streamed decode is unprofiled; stages seen: {sorted(durations)}" assert durations[stage] > 0.0 + assert durations[f"{type(pipeline).__name__}._prepare_condition_chunk"] > 0.0 # The stage reaches the consumer on every streamed chunk, not just at the end. assert all(stage in (output.stage_durations or {}) for output in outputs) @@ -2398,11 +2656,10 @@ def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> Non pipeline = _streaming_pipeline(module) # A latent of 2x2 exceeds a one-latent-cell tile, so vae.decode would tile. _enable_vae_tiling(pipeline, tile_sample_min=8) - state = _stepwise_state(num_frames=21) - state.sampling.output_type = "np" - - with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): - outputs = _run_stepwise(pipeline, state) + outputs = [ + pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), output_type="np", session_id="tiled") + for _ in range(2) + ] # One whole-clip decode per AR block, each seeing only that block's frames. assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2 @@ -2410,7 +2667,7 @@ def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> Non assert pipeline._streaming_decode_states == {} # Nine frames per chunk instead of 9 then 12: the frame loss being removed. assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)] - assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}] + assert outputs == ["frames-1", "frames-2"] def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monkeypatch) -> None: diff --git a/tests/examples/offline_inference/test_lingbot_world_v2.py b/tests/examples/offline_inference/test_lingbot_world_v2.py index 41d2e73b7..d83288853 100644 --- a/tests/examples/offline_inference/test_lingbot_world_v2.py +++ b/tests/examples/offline_inference/test_lingbot_world_v2.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import json from pathlib import Path import pytest +from examples.offline_inference.diffusion import ar_diffusion_streaming_decode as streaming_example from examples.offline_inference.diffusion import lingbot_world_v2 as example pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -22,3 +24,25 @@ def test_offline_ulysses_argument(tmp_path): args.ulysses_degree = 0 with pytest.raises(ValueError, match="--ulysses-degree"): example.build_omni_kwargs(args, paths) + + +@pytest.mark.parametrize(("num_blocks", "num_frames"), [(10, 117), (11, 129)]) +def test_streaming_action_script_accepts_long_rollouts(tmp_path, num_blocks, num_frames): + action_script = tmp_path / "actions.json" + action_script.write_text(json.dumps([[["w"], [], ["a"]] for _ in range(num_blocks)])) + + script = streaming_example._load_action_script(action_script) + + assert len(script) == num_blocks + assert ( + streaming_example._validate_against_checkpoint(script, frames_per_block=3, temporal_compression=4) == num_frames + ) + + +def test_streaming_action_script_rejects_wrong_frame_count_in_eleventh_block(tmp_path): + action_script = tmp_path / "actions.json" + action_script.write_text(json.dumps([[[], [], []] for _ in range(10)] + [[[], []]])) + script = streaming_example._load_action_script(action_script) + + with pytest.raises(ValueError, match="block 10 holds 2 frame action lists"): + streaming_example._validate_against_checkpoint(script, frames_per_block=3, temporal_compression=4) diff --git a/vllm_omni/diffusion/models/lingbot_world/camera.py b/vllm_omni/diffusion/models/lingbot_world/camera.py index 0e78e53e6..504d850c7 100644 --- a/vllm_omni/diffusion/models/lingbot_world/camera.py +++ b/vllm_omni/diffusion/models/lingbot_world/camera.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Camera trajectory loading and geometry for LingBot World v2 conditioning.""" from __future__ import annotations @@ -19,7 +19,6 @@ import torch _REFERENCE_HEIGHT = 480 _REFERENCE_WIDTH = 832 -_MAX_ACTION_FRAMES = 117 _MAX_SOURCE_ACTION_FRAMES = 4096 _MAX_NPY_HEADER_BYTES = 10_000 _DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW @@ -207,9 +206,7 @@ def _load_bounded_npy( array = np.load(snapshot_file, allow_pickle=False) except (EOFError, TypeError, ValueError): raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None - # Official LingBot trajectories may be longer than one request. Only - # materialize the bounded prefix that any supported request can consume. - result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32) + result: np.ndarray = np.asarray(array, dtype=np.float32) return result diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index 6e32844fc..ccaa677e0 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -7,7 +7,7 @@ from __future__ import annotations import math import os from collections.abc import Callable, Iterable, Iterator, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, ClassVar, cast @@ -15,6 +15,7 @@ import numpy as np import PIL.Image import torch import torch.nn.functional as F +from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanAttentionBlock, WanResample, WanResidualBlock from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel @@ -97,7 +98,6 @@ LINGBOT_DMD_TIMESTEPS = (1000, 750, 500, 250) _CAMERA_SPATIAL_FOLD = 8 _MAX_PIXEL_AREA = 480 * 832 _MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096 -_MAX_RAW_FRAMES = 117 _MAX_SEQUENCE_LENGTH = 512 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" @@ -150,12 +150,13 @@ class _LingBotRequestInputs: @dataclass class _LingBotARSessionState: - """Small model-owned state; attention tensors remain runner-owned.""" + """Model-owned state; attention tensors remain runner-owned.""" next_chunk_index: int = 0 prompt: str | None = None generator_state: torch.Tensor | None = None - image_condition: torch.Tensor | None = None + encoder_cache: list[torch.Tensor | None] | None = None + pending_encoder_cache: list[torch.Tensor | None] | None = None camera_tail: CameraTrajectory | None = None camera_pitch: float = 0.0 @@ -577,6 +578,7 @@ class LingBotWorldCausalDMDPipeline( self.setup_diffusion_pipeline_profiler( profiler_targets=[ "vae.encode", + "_prepare_condition_chunk", "vae.decode", # Streaming decode never reaches vae.decode, so it needs its # own stage or a streamed rollout reports no decode time at @@ -609,11 +611,11 @@ class LingBotWorldCausalDMDPipeline( recent_window_frames = total_window_frames - sink_frames if recent_window_frames <= 0: raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.") - horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 + condition_latent_frames = int(self.transformer.config.num_frames_per_block) condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels) condition_bytes_per_session = ( condition_channels - * horizon_latent_frames + * condition_latent_frames * latent_height * latent_width * torch.empty((), dtype=self.transformer.dtype).element_size() @@ -635,10 +637,34 @@ class LingBotWorldCausalDMDPipeline( ), ), model_owned_state_bytes_per_session=( - condition_bytes_per_session + self._streaming_decode_bytes_per_session() + condition_bytes_per_session + + 2 * self._condition_encoder_cache_bytes() # Committed and in-flight histories. + + self._streaming_decode_bytes_per_session() ), ) + def _condition_encoder_cache_bytes(self) -> int: + """Bound one Wan encoder history from its per-convolution input grids.""" + if getattr(self.vae.config, "patch_size", None) is not None: + raise ValueError("Stateful LingBot conditioning requires an unpatched Wan VAE.") + encoder = self.vae.encoder + height, width = self._ar_height, self._ar_width + elements = CACHE_T * encoder.conv_in.in_channels * height * width + for layer in encoder.down_blocks: + if isinstance(layer, WanResidualBlock): + elements += CACHE_T * (layer.conv1.in_channels + layer.conv2.in_channels) * height * width + elif isinstance(layer, WanResample) and layer.mode in ("downsample2d", "downsample3d"): + height, width = height // 2, width // 2 + if layer.mode == "downsample3d": + # Temporal downsampling retains one frame after spatial downsampling. + elements += layer.time_conv.in_channels * height * width + elif not isinstance(layer, WanAttentionBlock): + raise ValueError("Stateful LingBot conditioning requires the Wan encoder cache layout.") + for layer in encoder.mid_block.resnets: + elements += CACHE_T * (layer.conv1.in_channels + layer.conv2.in_channels) * height * width + elements += CACHE_T * encoder.conv_out.in_channels * height * width + return elements * torch.empty((), dtype=self.vae.dtype).element_size() + def _streaming_decode_bytes_per_session(self) -> int: """Resident streaming-decoder bytes one session can hold, for admission. @@ -678,12 +704,18 @@ class LingBotWorldCausalDMDPipeline( finally: self._ar_diffusion_kv_state = None + def _release_condition_encoder_state(self, session_id: str) -> None: + state = self._ar_sessions.pop(session_id, None) + if state is not None: + state.encoder_cache = None + state.pending_encoder_cache = None + def reset_ar_diffusion_session(self, session_id: str) -> None: - self._ar_sessions.pop(session_id, None) + self._release_condition_encoder_state(session_id) self._release_streaming_decode_state(session_id) def close_ar_diffusion_session(self, session_id: str) -> None: - self._ar_sessions.pop(session_id, None) + self._release_condition_encoder_state(session_id) self._release_streaming_decode_state(session_id) def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs: @@ -810,8 +842,6 @@ class LingBotWorldCausalDMDPipeline( num_frames = getattr(sampling, "num_frames", None) if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: raise ValueError(f"num_frames must be a positive integer, got {num_frames!r}.") - if num_frames > _MAX_RAW_FRAMES: - raise ValueError(f"num_frames must not exceed {_MAX_RAW_FRAMES}.") temporal_factor = self.vae_scale_factor_temporal if (num_frames - 1) % temporal_factor: raise ValueError( @@ -914,16 +944,16 @@ class LingBotWorldCausalDMDPipeline( raise RuntimeError( f"vae.encode returned an incompatible image latent shape: got {tuple(latent_condition.shape)}." ) + return self._condition_from_latents(latent_condition, first_frame=True, dtype=dtype) + + def _condition_from_latents( + self, latent_condition: torch.Tensor, *, first_frame: bool, dtype: torch.dtype + ) -> torch.Tensor: latent_mean, latent_std = self._vae_latent_stats(latent_condition) latent_condition = (latent_condition - latent_mean) / latent_std - temporal_mask = latent_condition.new_zeros( - 1, - self.vae_scale_factor_temporal, - inputs.num_latent_frames, - latent_condition.shape[-2], - latent_condition.shape[-1], - ) - temporal_mask[:, :, 0] = 1 + temporal_mask = latent_condition.new_zeros(1, self.vae_scale_factor_temporal, *latent_condition.shape[2:]) + if first_frame: + temporal_mask[:, :, 0] = 1 condition = torch.cat((temporal_mask, latent_condition), dim=1).to(dtype=dtype) if condition.shape[1] != 20: raise RuntimeError( @@ -932,6 +962,78 @@ class LingBotWorldCausalDMDPipeline( ) return condition + @torch.no_grad() + def _prepare_condition_chunk( + self, + inputs: _LingBotRequestInputs, + *, + start_frame: int, + encoder_cache: list[torch.Tensor | None] | None, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: + """Continue the causal encoder without modifying its committed history.""" + if getattr(self.vae.config, "patch_size", None) is not None or self.vae_scale_factor_temporal != 4: + raise ValueError("Stateful LingBot conditioning requires an unpatched Wan VAE with temporal factor 4.") + if getattr(self.vae, "use_tiling", False) and ( + inputs.height > self.vae.tile_sample_min_height or inputs.width > self.vae.tile_sample_min_width + ): + raise ValueError("Stateful LingBot conditioning does not support tiled VAE encoding.") + if (encoder_cache is None) != (start_frame == 0): + raise ValueError("LingBot condition encoder history must match the current temporal position.") + + count = int(self.vae._cached_conv_counts["encoder"]) + # Wan replaces cache entries with new tensors; a separate list keeps the + # committed history intact until this AR block succeeds. + pending_cache = [None] * count if encoder_cache is None else list(encoder_cache) + if len(pending_cache) != count: + raise ValueError("LingBot condition encoder cache has an incompatible number of slots.") + block_frames = int(self.transformer.config.num_frames_per_block) + zeros = torch.zeros( + 1, + 3, + self.vae_scale_factor_temporal, + inputs.height, + inputs.width, + device=self.device, + dtype=self.vae.dtype, + ) + image = ( + self._prepare_image_tensor(inputs.image, height=inputs.height, width=inputs.width) + .unsqueeze(2) + .to(dtype=self.vae.dtype) + if start_frame == 0 + else None + ) + context = getattr(self.vae, "_execution_context", None) + encoded_frames = [] + with context() if callable(context) else nullcontext(): + for index in range(block_frames): + # Match Wan's encode loop: one opening pixel frame, then four + # fresh zero pixel frames per subsequent latent frame. + pixel_frames = image if start_frame == 0 and index == 0 else zeros + encoded_frames.append(self.vae.encoder(pixel_frames, feat_cache=pending_cache, feat_idx=[0])) + # The posterior projection is pointwise in time. GPU bitwise parity + # with the whole-clip projection still requires the real-weight gate. + moments = self.vae.quant_conv(torch.cat(encoded_frames, dim=2)) + latent_condition = moments.chunk(2, dim=1)[0] + expected_shape = ( + 1, + self.transformer.config.out_channels, + block_frames, + inputs.height // self.vae_scale_factor_spatial, + inputs.width // self.vae_scale_factor_spatial, + ) + if latent_condition.shape != expected_shape: + raise RuntimeError( + f"Wan condition encoder returned {tuple(latent_condition.shape)}, expected {expected_shape}." + ) + condition = self._condition_from_latents( + latent_condition, + first_frame=start_frame == 0, + dtype=dtype, + ) + return condition, pending_cache + def _prepare_camera( self, inputs: _LingBotRequestInputs, @@ -1150,6 +1252,11 @@ class LingBotWorldCausalDMDPipeline( raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.") if tick is None and self._ar_diffusion_kv_state is not None: raise ValueError("LingBot ARDiffusionEngine requests must carry ar_diffusion_tick.") + if tick is not None and inputs.output_type != "latent": + raise ValueError( + "LingBot realtime ticks currently require output_type='latent'; " + "stateful streaming VAE decode is a separate integration step." + ) session_state: _LingBotARSessionState | None = None if tick is not None: if tick.prompt is not None and " ".join(tick.prompt.split()) != " ".join(inputs.prompt.split()): @@ -1166,16 +1273,6 @@ class LingBotWorldCausalDMDPipeline( "fixed cache geometry " f"{self._ar_height}x{self._ar_width}." ) - horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 - max_realtime_ticks = horizon_latent_frames // block_frames - if tick.chunk_index >= max_realtime_ticks: - raise ValueError( - "LingBot realtime generation currently supports at most " - f"{max_realtime_ticks} ticks per generation epoch " - f"(chunk_index 0 through {max_realtime_ticks - 1}) because " - f"the image-condition horizon is {_MAX_RAW_FRAMES} pixel " - "frames; reset or create a session to start a new world." - ) session_state = self._ar_sessions.setdefault( tick.session_id, _LingBotARSessionState(), @@ -1200,26 +1297,12 @@ class LingBotWorldCausalDMDPipeline( condition = self._prepare_condition(inputs, dtype=dtype) else: assert session_state is not None - if session_state.image_condition is None: - horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 - session_state.image_condition = self._prepare_condition( - replace( - inputs, - num_frames=_MAX_RAW_FRAMES, - num_latent_frames=horizon_latent_frames, - ), - dtype=dtype, - ) - block_frames = int(self.transformer.config.num_frames_per_block) - condition_start = tick.chunk_index * block_frames - condition_stop = condition_start + block_frames - if condition_stop > session_state.image_condition.shape[2]: - raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") - condition = session_state.image_condition[ - :, - :, - condition_start:condition_stop, - ] + condition, pending_encoder_cache = self._prepare_condition_chunk( + inputs, + start_frame=tick.chunk_index * block_frames, + encoder_cache=session_state.encoder_cache, + dtype=dtype, + ) camera_pitch: float | None = None if inputs.camera_actions is not None: assert session_state is not None @@ -1299,6 +1382,7 @@ class LingBotWorldCausalDMDPipeline( assert session_state is not None session_state.prompt = inputs.prompt session_state.generator_state = inputs.generator.get_state() + session_state.encoder_cache = pending_encoder_cache session_state.camera_tail = camera_tail if camera_pitch is not None: session_state.camera_pitch = camera_pitch @@ -1307,11 +1391,6 @@ class LingBotWorldCausalDMDPipeline( # Phase 4: either expose model-space latents or invert the checkpoint's # latent normalization and decode to pixel-space video. if tick is not None: - if inputs.output_type != "latent": - raise ValueError( - "LingBot realtime ticks currently require output_type='latent'; " - "stateful streaming VAE decode is a separate integration step." - ) output = { "payload": {"latents": generated_latents}, "metadata": {"ar_diffusion": ARDiffusionChunkMetadata.from_tick(tick).to_dict()}, @@ -1460,12 +1539,6 @@ class LingBotWorldCausalDMDPipeline( self._cached_video_processor = processor return processor - def _horizon_latent_frames(self) -> int: - return (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 - - def _max_realtime_chunks(self) -> int: - return self._horizon_latent_frames() // int(self.transformer.config.num_frames_per_block) - def _require_bound_ar_state(self) -> None: if self._ar_diffusion_kv_state is None: raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.") @@ -1492,13 +1565,6 @@ class LingBotWorldCausalDMDPipeline( ) block_frames = int(self.transformer.config.num_frames_per_block) total_chunks = inputs.num_latent_frames // block_frames - max_chunks = self._max_realtime_chunks() - if total_chunks > max_chunks: - raise ValueError( - "LingBot step execution currently supports at most " - f"{max_chunks} chunks because the image-condition horizon is " - f"{_MAX_RAW_FRAMES} pixel frames." - ) if inputs.camera_action_script is not None and len(inputs.camera_action_script) != total_chunks: raise ValueError( "camera_action_script must contain one action list per generated chunk; " @@ -1510,14 +1576,6 @@ class LingBotWorldCausalDMDPipeline( max_sequence_length=inputs.max_sequence_length, dtype=dtype, ) - image_condition = self._prepare_condition( - replace( - inputs, - num_frames=_MAX_RAW_FRAMES, - num_latent_frames=self._horizon_latent_frames(), - ), - dtype=dtype, - ) camera_trajectory_cache = None camera_embedding_cache = None if inputs.camera_trajectory is not None: @@ -1546,7 +1604,6 @@ class LingBotWorldCausalDMDPipeline( state.chunk_num_steps = len(LINGBOT_DMD_TIMESTEPS) state.extra = { "inputs": inputs, - "image_condition": image_condition, "schedule": _build_shifted_flow_schedule(flow_shift=inputs.flow_shift), "dtype": dtype, "block_frames": block_frames, @@ -1557,7 +1614,12 @@ class LingBotWorldCausalDMDPipeline( "camera_embedding_cache": camera_embedding_cache, } self._ar_text_caches(prompt_embeds, invalidate=False) - self._prepare_next_chunk(state) + self._ar_sessions[state.request_id] = _LingBotARSessionState() + try: + self._prepare_next_chunk(state) + except Exception: + self._release_condition_encoder_state(state.request_id) + raise return state def _prepare_next_chunk(self, state: StepRequestState) -> None: @@ -1566,10 +1628,15 @@ class LingBotWorldCausalDMDPipeline( block_frames = int(extra["block_frames"]) start_frame = state.chunk_index * block_frames stop_frame = start_frame + block_frames - image_condition = extra["image_condition"] - if stop_frame > image_condition.shape[2]: - raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") - condition = image_condition[:, :, start_frame:stop_frame] + session_state = self._ar_sessions[state.request_id] + if session_state.next_chunk_index != state.chunk_index: + raise ValueError("LingBot condition chunks must be contiguous within the session.") + condition, pending_encoder_cache = self._prepare_condition_chunk( + inputs, + start_frame=start_frame, + encoder_cache=session_state.encoder_cache, + dtype=extra["dtype"], + ) previous = extra.get("camera_tail") if extra.get("camera_action_script") is not None: chunk_actions = extra["camera_action_script"][state.chunk_index] @@ -1639,6 +1706,7 @@ class LingBotWorldCausalDMDPipeline( ) state.step_in_chunk = 0 state.step_index = 0 + session_state.pending_encoder_cache = pending_encoder_cache def denoise_step( self, @@ -1702,6 +1770,9 @@ class LingBotWorldCausalDMDPipeline( latents = state.latents if latents is None: raise RuntimeError("LingBot step execution requires latents before post_decode.") + session_state = self._ar_sessions[state.request_id] + if session_state.pending_encoder_cache is None: + raise RuntimeError("LingBot condition encoder has no prepared chunk to commit.") self._dmd_blocks.commit_block_kv( latents=latents, condition=extra["condition"], @@ -1737,10 +1808,15 @@ class LingBotWorldCausalDMDPipeline( ).to_dict() }, } + session_state.encoder_cache = session_state.pending_encoder_cache + session_state.pending_encoder_cache = None + session_state.next_chunk_index += 1 state.chunk_index += 1 finished = state.request_denoise_completed if not finished: self._prepare_next_chunk(state) + else: + self._release_condition_encoder_state(state.request_id) return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py index 32931db0c..c7033cee8 100644 --- a/vllm_omni/diffusion/models/lingbot_world/transformer.py +++ b/vllm_omni/diffusion/models/lingbot_world/transformer.py @@ -749,7 +749,7 @@ def _sinusoidal_embedding(dim: int, timestep: torch.Tensor) -> torch.Tensor: return torch.cat((phase.cos(), phase.sin()), dim=1) -def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]: +def _rope_axis(max_seq_len: int, dim: int, *, start: int = 0) -> tuple[torch.Tensor, torch.Tensor]: if dim == 0: empty = torch.empty(max_seq_len, 0, dtype=torch.float32) return empty, empty.clone() @@ -759,7 +759,7 @@ def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]: 10000, torch.arange(0, dim, 2, dtype=torch.float64) / dim, ) - phase = torch.outer(torch.arange(max_seq_len, dtype=torch.float64), frequencies) + phase = torch.outer(torch.arange(start, start + max_seq_len, dtype=torch.float64), frequencies) return phase.cos().float(), phase.sin().float() @@ -1028,23 +1028,29 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): dtype: torch.dtype, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: - if start_frame + frames > self.config.rope_max_seq_len: - raise ValueError("Temporal RoPE positions exceed rope_max_seq_len.") if height > self.config.rope_max_seq_len or width > self.config.rope_max_seq_len: raise ValueError("Spatial RoPE positions exceed rope_max_seq_len.") + if start_frame + frames <= self._rope_temporal_cosine.shape[0]: + temporal_cosine = self._rope_temporal_cosine[start_frame : start_frame + frames] + temporal_sine = self._rope_temporal_sine[start_frame : start_frame + frames] + else: + # Keep absolute positions and checkpoint frequencies, computing only + # this block instead of growing a table for the whole session. + temporal_dim = self._rope_temporal_cosine.shape[1] * 2 + temporal_cosine, temporal_sine = _rope_axis(frames, temporal_dim, start=start_frame) + temporal_cosine = temporal_cosine.to(self._rope_temporal_cosine) + temporal_sine = temporal_sine.to(self._rope_temporal_sine) def expand_axis(table: torch.Tensor, axis: str) -> torch.Tensor: if axis == "temporal": - return ( - table[start_frame : start_frame + frames].view(frames, 1, 1, -1).expand(frames, height, width, -1) - ) + return table.view(frames, 1, 1, -1).expand(frames, height, width, -1) if axis == "height": return table[:height].view(1, height, 1, -1).expand(frames, height, width, -1) return table[:width].view(1, 1, width, -1).expand(frames, height, width, -1) cosine = torch.cat( ( - expand_axis(self._rope_temporal_cosine, "temporal"), + expand_axis(temporal_cosine, "temporal"), expand_axis(self._rope_height_cosine, "height"), expand_axis(self._rope_width_cosine, "width"), ), @@ -1052,7 +1058,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): ) sine = torch.cat( ( - expand_axis(self._rope_temporal_sine, "temporal"), + expand_axis(temporal_sine, "temporal"), expand_axis(self._rope_height_sine, "height"), expand_axis(self._rope_width_sine, "width"), ), From 952022c8ec62aa0b3ad0ed57d89630e54fd041ec Mon Sep 17 00:00:00 2001 From: dongbo910220 <32610838+dongbo910220@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:28:48 +0800 Subject: [PATCH 084/209] [Diffusion][Perf] Add Qwen-Image QK RoPE Triton path (#5931) Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../benchmark_qwen_image_qk_norm_rope.py | 224 ++++++++++++++++++ .../test_qwen_image_fused_qk_norm_rope.py | 190 +++++++++++++++ .../qwen_image/qwen_image_transformer.py | 92 +++++-- 3 files changed, 484 insertions(+), 22 deletions(-) create mode 100644 benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py create mode 100644 tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py diff --git a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py new file mode 100644 index 000000000..ee69f7e59 --- /dev/null +++ b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Benchmark Qwen-Image Q/K RMSNorm + interleaved RoPE. + +Example: + + CUDA_VISIBLE_DEVICES=0 python benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py \ + --seq-len 4096 --warmup 20 --iters 100 --include-compiled --profile +""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, +) +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + _apply_qwen_image_rotary_emb, +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=4096) + parser.add_argument("--heads", type=int, default=24) + parser.add_argument("--kv-heads", type=int, default=24) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16") + parser.add_argument("--packed-qkv-view", action="store_true") + parser.add_argument("--include-compiled", action="store_true") + parser.add_argument("--profile", action="store_true") + return parser.parse_args() + + +def _dtype(name: str) -> torch.dtype: + return { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + }[name] + + +def _make_inputs(args: argparse.Namespace, dtype: torch.dtype, device: torch.device): + torch.manual_seed(2026) + if args.packed_qkv_view: + qkv_dim = (args.heads + args.kv_heads + args.kv_heads) * args.head_dim + qkv = torch.randn(args.batch, args.seq_len, qkv_dim, device=device, dtype=dtype) + q, k, _ = qkv.split( + [ + args.heads * args.head_dim, + args.kv_heads * args.head_dim, + args.kv_heads * args.head_dim, + ], + dim=-1, + ) + q = q.unflatten(-1, (args.heads, args.head_dim)) + k = k.unflatten(-1, (args.kv_heads, args.head_dim)) + else: + q = torch.randn(args.batch, args.seq_len, args.heads, args.head_dim, device=device, dtype=dtype) + k = torch.randn(args.batch, args.seq_len, args.kv_heads, args.head_dim, device=device, dtype=dtype) + + q_weight = torch.randn(args.head_dim, device=device, dtype=dtype) + k_weight = torch.randn(args.head_dim, device=device, dtype=dtype) + angles = torch.randn(args.seq_len, args.head_dim // 2, device=device, dtype=torch.float32) + freqs = torch.polar(torch.ones_like(angles), angles) + return q, k, q_weight, k_weight, freqs + + +def _native( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + q = F.rms_norm(q, (q.shape[-1],), q_weight, eps) + k = F.rms_norm(k, (k.shape[-1],), k_weight, eps) + return ( + _apply_qwen_image_rotary_emb(q, freqs), + _apply_qwen_image_rotary_emb(k, freqs), + ) + + +def _fused( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + batch, seq_len, num_heads, head_dim = q.shape + num_kv_heads = k.shape[2] + rope_table = torch.cat((freqs.real, freqs.imag), dim=-1) + rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, head_dim) + out_q, out_k = fused_qk_norm_rope( + q.reshape(batch * seq_len, num_heads, head_dim), + k.reshape(batch * seq_len, num_kv_heads, head_dim), + q_weight, + k_weight, + rope_table, + eps, + interleaved=True, + ) + return ( + out_q.reshape(batch, seq_len, num_heads, head_dim), + out_k.reshape(batch, seq_len, num_kv_heads, head_dim), + ) + + +def _measure(fn: Callable[[], tuple[torch.Tensor, torch.Tensor]], warmup: int, iters: int) -> list[float]: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples = [] + for _ in range(iters): + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + return samples + + +def _stats(samples: list[float]) -> dict[str, float]: + samples = sorted(samples) + return { + "median_ms": samples[len(samples) // 2], + "p90_ms": samples[int(len(samples) * 0.9)], + "mean_ms": sum(samples) / len(samples), + } + + +def _profile(fn: Callable[[], tuple[torch.Tensor, torch.Tensor]]) -> str: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + fn() + torch.accelerator.synchronize() + return prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=15) + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires CUDA") + device = torch.device("cuda") + dtype = _dtype(args.dtype) + eps = 1e-6 + q, k, q_weight, k_weight, freqs = _make_inputs(args, dtype, device) + + def native(): + return _native(q, k, q_weight, k_weight, freqs, eps) + + def fused(): + return _fused(q, k, q_weight, k_weight, freqs, eps) + + results: dict[str, object] = { + "device": torch.cuda.get_device_name(), + "batch": args.batch, + "seq_len": args.seq_len, + "heads": args.heads, + "kv_heads": args.kv_heads, + "head_dim": args.head_dim, + "dtype": args.dtype, + "packed_qkv_view": args.packed_qkv_view, + "fused_fast_path_supported": _fused_cuda_supported( + q.reshape(args.batch * args.seq_len, args.heads, args.head_dim), + k.reshape(args.batch * args.seq_len, args.kv_heads, args.head_dim), + args.head_dim, + args.head_dim, + interleaved=True, + ), + } + + native_samples = _measure(native, args.warmup, args.iters) + fused_samples = _measure(fused, args.warmup, args.iters) + expected_q, expected_k = native() + actual_q, actual_k = fused() + torch.accelerator.synchronize() + + native_stats = _stats(native_samples) + fused_stats = _stats(fused_samples) + results["native"] = native_stats + results["fused"] = fused_stats + results["speedup_vs_native"] = native_stats["median_ms"] / fused_stats["median_ms"] + results["max_abs_q"] = (actual_q.float() - expected_q.float()).abs().max().item() + results["max_abs_k"] = (actual_k.float() - expected_k.float()).abs().max().item() + results["mean_abs_q"] = (actual_q.float() - expected_q.float()).abs().mean().item() + results["mean_abs_k"] = (actual_k.float() - expected_k.float()).abs().mean().item() + + if args.include_compiled: + compiled_native = torch.compile(native, dynamic=True, fullgraph=True) + compiled_fused = torch.compile(fused, dynamic=True, fullgraph=True) + compiled_native_stats = _stats(_measure(compiled_native, args.warmup, args.iters)) + compiled_fused_stats = _stats(_measure(compiled_fused, args.warmup, args.iters)) + results["compiled_native"] = compiled_native_stats + results["compiled_fused"] = compiled_fused_stats + results["compiled_fused_speedup_vs_compiled_native"] = ( + compiled_native_stats["median_ms"] / compiled_fused_stats["median_ms"] + ) + + print(json.dumps(results, indent=2)) + if args.profile: + print("\n# fused profile") + print(_profile(fused)) + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py new file mode 100644 index 000000000..2f10ea33e --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + _apply_qwen_image_rotary_emb, + _qwen_image_qk_norm_rope, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +BATCH = 2 +Q_HEADS = 24 +K_HEADS = 24 +HEAD_DIM = 128 +EPS = 1e-6 +SEED = 2026 + + +@dataclass(frozen=True) +class QwenImageQKInput: + q: torch.Tensor + k: torch.Tensor + norm_q: nn.RMSNorm + norm_k: nn.RMSNorm + freqs: torch.Tensor + + +def _make_input( + *, + seq_len: int, + dtype: torch.dtype, + device: torch.device, + packed_qkv_view: bool, +) -> QwenImageQKInput: + gen = torch.Generator(device=device) + gen.manual_seed(SEED + seq_len) + + if packed_qkv_view: + qkv_dim = (Q_HEADS + K_HEADS + K_HEADS) * HEAD_DIM + qkv = torch.randn(BATCH, seq_len, qkv_dim, device=device, dtype=dtype, generator=gen) + q, k, _ = qkv.split( + [ + Q_HEADS * HEAD_DIM, + K_HEADS * HEAD_DIM, + K_HEADS * HEAD_DIM, + ], + dim=-1, + ) + q = q.unflatten(-1, (Q_HEADS, HEAD_DIM)) + k = k.unflatten(-1, (K_HEADS, HEAD_DIM)) + else: + q = torch.randn(BATCH, seq_len, Q_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen) + k = torch.randn(BATCH, seq_len, K_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen) + + norm_q = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype) + norm_k = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype) + norm_q.requires_grad_(False) + norm_k.requires_grad_(False) + with torch.no_grad(): + norm_q.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen)) + norm_k.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen)) + angles = torch.randn(seq_len, HEAD_DIM // 2, device=device, dtype=torch.float32, generator=gen) + freqs = torch.polar(torch.ones_like(angles), angles) + return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs) + + +def _reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: + q = data.norm_q(data.q) + k = data.norm_k(data.k) + if data.q.device.type == "cuda": + return ( + _apply_qwen_image_rotary_emb(q, data.freqs), + _apply_qwen_image_rotary_emb(k, data.freqs), + ) + + rope = RotaryEmbedding(is_neox_style=False) + cos = data.freqs.real.to(data.q.dtype) + sin = data.freqs.imag.to(data.q.dtype) + return rope(q, cos, sin), rope(k, cos, sin) + + +def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: + return _qwen_image_qk_norm_rope( + data.q, + data.k, + data.norm_q, + data.norm_k, + data.freqs, + RotaryEmbedding(is_neox_style=False), + EPS, + ) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference(): + data = _make_input( + seq_len=7, + dtype=torch.float32, + device=torch.device("cuda:0"), + packed_qkv_view=True, + ) + + actual_q, actual_k = _run(data) + expected_q, expected_k = _reference(data) + + torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference(): + data = _make_input( + seq_len=7, + dtype=torch.float16, + device=torch.device("cuda:0"), + packed_qkv_view=True, + ) + + actual_q, actual_k = _run(data) + expected_q, expected_k = _reference(data) + + torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("seq_len", [1, 7, 257, 1024]) +@pytest.mark.parametrize("packed_qkv_view", [False, True]) +def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference( + seq_len: int, + packed_qkv_view: bool, +): + data = _make_input( + seq_len=seq_len, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + packed_qkv_view=packed_qkv_view, + ) + + actual_q, actual_k = _run(data) + expected_q, expected_k = _reference(data) + + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture(): + data = _make_input( + seq_len=257, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + packed_qkv_view=True, + ) + rope = RotaryEmbedding(is_neox_style=False) + + def fn( + q: torch.Tensor, + k: torch.Tensor, + freqs: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return _qwen_image_qk_norm_rope( + q, + k, + data.norm_q, + data.norm_k, + freqs, + rope, + EPS, + ) + + compiled_fn = torch.compile(fn, dynamic=True, fullgraph=True) + expected_q, expected_k = fn(data.q, data.k, data.freqs) + actual_q, actual_k = compiled_fn(data.q, data.k, data.freqs) + + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 696728b90..8e4dac3a3 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -45,6 +45,10 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, +) from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -61,6 +65,53 @@ def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch. return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) +def _qwen_image_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + norm_q: nn.Module, + norm_k: nn.Module, + freqs: torch.Tensor, + rope: RotaryEmbedding, + eps: float, + *, + use_fused: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + head_dim = q.shape[-1] + rotary_dim = freqs.shape[-1] * 2 + if use_fused and _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved=True): + batch, seq_len, num_heads, _ = q.shape + num_kv_heads = k.shape[2] + rope_table = torch.cat((freqs.real, freqs.imag), dim=-1) + rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, rotary_dim) + fused_q, fused_k = fused_qk_norm_rope( + q.reshape(batch * seq_len, num_heads, head_dim), + k.reshape(batch * seq_len, num_kv_heads, head_dim), + norm_q.weight, + norm_k.weight, + rope_table, + eps, + interleaved=True, + ) + return ( + fused_q.reshape(batch, seq_len, num_heads, head_dim), + fused_k.reshape(batch, seq_len, num_kv_heads, head_dim), + ) + + q = norm_q(q) + k = norm_k(k) + if q.device.type == "cuda": + return ( + _apply_qwen_image_rotary_emb(q, freqs), + _apply_qwen_image_rotary_emb(k, freqs), + ) + + # Retain the platform-specific kernels on other accelerators, which may + # not support complex tensors. + cos = torch.real(freqs).to(q.dtype) + sin = torch.imag(freqs).to(q.dtype) + return rope(q, cos, sin), rope(k, cos, sin) + + def _normalize_qwen_image_weight_name(name: str) -> str: name = name.removeprefix("transformer.") if ".to_out.0." in name: @@ -648,28 +699,25 @@ class QwenImageCrossAttention(nn.Module): txt_key = txt_key.unflatten(-1, (self.add_kv_num_heads, self.head_dim)) txt_value = txt_value.unflatten(-1, (self.add_kv_num_heads, self.head_dim)) - img_query = self.norm_q(img_query) - img_key = self.norm_k(img_key) - txt_query = self.norm_added_q(txt_query) - txt_key = self.norm_added_k(txt_key) - - if img_query.device.type == "cuda": - img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs) - img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs) - txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs) - txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs) - else: - # Retain the platform-specific kernels on other accelerators, - # which may not support complex tensors. - img_cos = torch.real(vid_freqs).to(img_query.dtype) - img_sin = torch.imag(vid_freqs).to(img_query.dtype) - txt_cos = torch.real(txt_freqs).to(txt_query.dtype) - txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) - - img_query = self.rope(img_query, img_cos, img_sin) - img_key = self.rope(img_key, img_cos, img_sin) - txt_query = self.rope(txt_query, txt_cos, txt_sin) - txt_key = self.rope(txt_key, txt_cos, txt_sin) + img_query, img_key = _qwen_image_qk_norm_rope( + img_query, + img_key, + self.norm_q, + self.norm_k, + vid_freqs, + self.rope, + self.eps, + use_fused=self.qk_norm, + ) + txt_query, txt_key = _qwen_image_qk_norm_rope( + txt_query, + txt_key, + self.norm_added_q, + self.norm_added_k, + txt_freqs, + self.rope, + self.eps, + ) seq_len_txt = encoder_hidden_states.shape[1] joint_query = torch.cat([txt_query, img_query], dim=1) From 78934753ff7e092b24ada0e34062cf3d121efd25 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Tue, 15 Sep 2026 12:36:34 -0400 Subject: [PATCH 085/209] [Tests] Test image edit input limits in tiny model framework (#7434) Signed-off-by: Nick Cao Co-authored-by: Codex --- tests/model_tests/diffusion/task_runners.py | 50 +++++++++++++++++++ .../diffusion/test_common_online.py | 17 +++++++ 2 files changed, 67 insertions(+) diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py index 010d73777..991c74603 100644 --- a/tests/model_tests/diffusion/task_runners.py +++ b/tests/model_tests/diffusion/task_runners.py @@ -15,6 +15,8 @@ import numpy as np import pytest from PIL import Image +from benchmarks.accuracy.common import pil_to_png_bytes +from tests.helpers.assertions import assert_images_generations_response from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, dummy_messages_from_mix_data from tests.model_tests.diffusion.config_types import DiffusionTasks from vllm_omni.entrypoints.omni import Omni @@ -204,6 +206,25 @@ def _run_online_i2i( return client.send_diffusion_request(request_config) +def _build_online_image_edits_request(model: str, num_images: int) -> dict[str, Any]: + """Build a multipart ``/v1/images/edits`` request with repeated image fields.""" + files = [ + ("image", (f"image_{index}.png", pil_to_png_bytes(INPUT_IMAGE), "image/png")) for index in range(num_images) + ] + return { + "data": { + "model": model, + "prompt": PROMPT, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": 2, + "seed": 42, + }, + "files": files, + } + + ### Offline task runners def run_and_validate_text_to_image_request(omni: Omni): """Run and validate a text to image request.""" @@ -324,6 +345,35 @@ def run_and_validate_online_image_to_image_request(server: OmniServer, client: O _validate_images(_get_online_images(_run_online_i2i(server, client))) +def run_and_validate_online_image_edits( + server: OmniServer, + client: OnlineOmniClient, + num_images: int, + max_multimodal_image_inputs: int, +): + """Run and validate one request through ``/v1/images/edits``.""" + request = _build_online_image_edits_request(server.model, num_images) + if num_images <= max_multimodal_image_inputs: + response = client.send_images_edits_http_request(request)[0] + assert response.success, response.error_message + assert isinstance(response.json_body, dict) + assert_images_generations_response( + response.json_body, + {"json": {"n": 1, "size": f"{WIDTH}x{HEIGHT}"}}, + ) + else: + err_message = ( + "Only a single image is supported by this model." + if max_multimodal_image_inputs == 1 + else f"At most {max_multimodal_image_inputs} images are supported" + ) + client.send_images_edits_http_request( + request, + err_code=400, + err_message=err_message, + ) + + def run_and_validate_online_determinism(server: OmniServer, client: OnlineOmniClient, task_type: DiffusionTasks): """Checks for determinism through the server, dispatching by task type. diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index 65089e3ef..ed7f0e2fa 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/tests/model_tests/diffusion/test_common_online.py @@ -20,12 +20,14 @@ from tests.model_tests.diffusion.config_types import ( from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS from tests.model_tests.diffusion.task_runners import ( run_and_validate_online_determinism, + run_and_validate_online_image_edits, run_and_validate_online_image_to_image_request, run_and_validate_online_image_to_video_request, run_and_validate_online_multi_output, run_and_validate_online_text_to_image_request, run_and_validate_online_text_to_video_request, ) +from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata # NOTE : Hardware marks are added dynamically based on test requirements pytestmark = [pytest.mark.diffusion, pytest.mark.xdist] @@ -84,6 +86,21 @@ def test_online_on_supported_tasks( run_and_validate_online_text_to_image_request(server, client) elif task_type == DiffusionTasks.IMAGE_TO_IMAGE: run_and_validate_online_image_to_image_request(server, client) + max_multimodal_image_inputs = ( + get_diffusion_model_metadata(model_name).max_multimodal_image_inputs or 1 + ) + image_counts = [1] + if max_multimodal_image_inputs != 1: + image_counts.append(max_multimodal_image_inputs) + image_counts.append(max_multimodal_image_inputs + 1) + for num_images in image_counts: + with subtests.test(api="/v1/images/edits", num_images=num_images): + run_and_validate_online_image_edits( + server, + client, + num_images=num_images, + max_multimodal_image_inputs=max_multimodal_image_inputs, + ) elif task_type == DiffusionTasks.TEXT_TO_VIDEO: run_and_validate_online_text_to_video_request(server, client) elif task_type == DiffusionTasks.IMAGE_TO_VIDEO: From c17126c21994edc5aaf1fa3a6248d67bd0ab9d29 Mon Sep 17 00:00:00 2001 From: Mengjie Zhao <38148457+mjZhaoElaine@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:56:51 -0700 Subject: [PATCH 086/209] [Model] Add Gepard-1.0 OpenAI /v1/audio/speech serving (#7499) Signed-off-by: Mengjie Zhao Co-authored-by: Cursor --- .buildkite/cuda/test-weekly.yml | 3 +- docs/serving/speech_api.md | 49 +- .../examples/online_serving/text_to_speech.md | 58 +- .../text_to_speech/README.md | 8 +- .../online_serving/text_to_speech/README.md | 43 +- .../text_to_speech/gepard/run_server.sh | 33 ++ .../text_to_speech/gepard/speech_client.py | 116 ++++ tests/e2e/online_serving/test_gepard_tts.py | 524 ++++++++++++++++++ .../openai_api/test_gepard_tts_adapter.py | 276 +++++++++ .../openai_api/test_serving_speech_metrics.py | 1 + .../openai_api/test_tts_adapter.py | 12 + .../openai_api/test_tts_detection.py | 1 + tests/worker/test_gpu_ar_model_runner.py | 1 + .../entrypoints/openai/serving_speech.py | 49 +- .../openai/tts_adapters/__init__.py | 1 + .../entrypoints/openai/tts_adapters/gepard.py | 168 ++++++ .../openai/tts_adapters/moss_tts.py | 11 +- 17 files changed, 1329 insertions(+), 25 deletions(-) create mode 100755 examples/online_serving/text_to_speech/gepard/run_server.sh create mode 100755 examples/online_serving/text_to_speech/gepard/speech_client.py create mode 100644 tests/e2e/online_serving/test_gepard_tts.py create mode 100644 tests/entrypoints/openai_api/test_gepard_tts_adapter.py create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/gepard.py diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index 7809fcbc7..70ab59356 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.buildkite/cuda/test-weekly.yml @@ -191,13 +191,14 @@ steps: # from downgrading transformers; see # examples/offline_inference/text_to_speech/README.md. - label: "TTS · Gepard-1.0 · L4" - timeout_in_minutes: 60 + timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt - uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt - >- pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py + tests/e2e/online_serving/test_gepard_tts.py -m "tts and cuda" --run-level "full_model" mirror_hardwares: l4_1 diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index d4fad4206..b4b2e7f55 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -6,6 +6,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz. - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz. - **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz. +- **Gepard-1.0** (`nineninesix/gepard-1.0`) -- single-stage native-AR TTS with a 22.05 kHz NanoCodec. Zero-shot default voice only. See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, MOSS-TTS-Nano, and Breeze-TTS-2. @@ -40,11 +41,16 @@ vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091 # CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request) vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \ --omni --port 8091 --trust-remote-code + +# Gepard-1.0 (zero-shot default voice; packaged gepard.yaml) +vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \ + --stage-init-timeout 900 \ + --deploy-config vllm_omni/deploy/gepard.yaml ``` ### Generate Speech -**Using curl:** +**Qwen3-TTS CustomVoice, using curl:** ```bash curl -X POST http://localhost:8091/v1/audio/speech \ @@ -56,7 +62,24 @@ curl -X POST http://localhost:8091/v1/audio/speech \ }' --output output.wav ``` -**Using Python:** +**Gepard-1.0, using curl:** + +Gepard accepts `input` (required), `voice` (`"default"`), `response_format`, +`stream` / `stream_format`, `max_new_tokens`, and `seed`. Unsupported fields +include `speed`, `language`, `instructions`, `task_type`, `ref_audio`, +`ref_text`, `extra_params`, and `word_timestamps`. + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Gepard speaking.", + "voice": "default", + "seed": 7 + }' --output output.wav +``` + +**Qwen3-TTS CustomVoice, using Python:** ```python import httpx @@ -761,6 +784,28 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne | ------- | ------------- | | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. | +### Gepard-1.0 + +| Model | Description | +| ----- | ----------- | +| `nineninesix/gepard-1.0` | Zero-shot native-AR TTS. 22.05 kHz mono. `voice` must be omitted or `"default"`. | + +Gepard request fields: + +| Field | Behavior | +| ----- | -------- | +| `input` | Required. Empty/whitespace-only returns 400. | +| `voice` | Omitted or `"default"` only. Other values 400. | +| `response_format` | `wav` default. Non-streaming: `wav`/`pcm`/`flac`/`mp3`. `opus` returns 400 (22.05 kHz is not an Opus sample rate). Streaming: `pcm`/`wav` only. | +| `stream` / `stream_format` | SSE (`speech.audio.*`) and raw `audio` byte streaming. | +| `speed` | Must be `1.0`. Gepard has no native speed control. | +| `max_new_tokens` | Frame budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms). Default 1000 from deploy YAML; adapter bounds 1..4096. | +| `seed` | Optional. Reaches the in-model 32-head sampler. The packaged YAML currently pins `seed: 42`, so serving is deterministic by default until that pin is removed. | +| `extra_params` | Any key returns 400, including `temperature`/`top_p`/`top_k`. | +| Cloning / style fields | `ref_audio`, `ref_text`, `speaker_embedding`, `instructions`, `language`, `task_type`, `word_timestamps`, and similar declared-but-unsupported fields return 400. | + +See the [Gepard section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#gepard-10) for launch commands. Native-AR recompute preemption is a known limitation under concurrency. + ### OmniVoice | Model | Description | diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 2b4f617f4..e9499551f 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -18,6 +18,7 @@ For the full list of supported architectures across all modalities, see | --- | --- | --- | --- | --- | --- | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | +| Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — | | higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — | @@ -92,7 +93,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others). +Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for Gepard, 24 kHz for the others). For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md). @@ -212,6 +213,61 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server --- +## Gepard-1.0 + +Single-stage native AR TTS at 22.05 kHz mono. Zero-shot only: omit `voice` or pass `"default"`. Voice cloning from reference audio is not available yet. + +### Prerequisites + +Same NeMo NanoCodec install as the [offline Gepard section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#gepard-10). On a host whose CUDA toolkit cannot build kernels — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also `export VLLM_USE_FLASHINFER_SAMPLER=0` before launch. That env is not a deploy-YAML field. + +### Launch + +```bash +vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \ + --stage-init-timeout 900 \ + --deploy-config vllm_omni/deploy/gepard.yaml +# or: +./gepard/run_server.sh +``` + +`--stage-init-timeout 900` matches the online e2e fixture and the offline example. Serve defaults to 300s, which is often too short for a cold download of the talker plus NanoCodec. + +The packaged `vllm_omni/deploy/gepard.yaml` must be passed with `--deploy-config`. The checkpoint self-identifies as `qwen3_5_text`, so omitting the YAML launches a diffusion fallback instead of the Gepard pipeline. The YAML sets `async_chunk: false`, `max_num_seqs: 4`, and currently pins `seed: 42`, so serving is deterministic by default until that YAML seed is removed. Pass an explicit per-request `seed` in tests and clients rather than depending on either default. + +### Sending requests + +```bash +python examples/online_serving/text_to_speech/gepard/speech_client.py \ + --text "Hello, this is Gepard speaking." + +python examples/online_serving/text_to_speech/gepard/speech_client.py \ + --text "Hello, this is Gepard speaking." --seed 7 --stream --output output.pcm +``` + +Gepard accepts `input` (required), `voice` (`"default"`), `response_format`, `stream` / `stream_format`, `max_new_tokens`, and `seed`. Unsupported fields include `speed`, `language`, `instructions`, `task_type`, `ref_audio`, `ref_text`, `extra_params`, and `word_timestamps`. + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Gepard speaking.", + "voice": "default", + "seed": 7 + }' --output output.wav +``` + +### Notes + +- Output: 22.05 kHz mono. `max_new_tokens` is a **frame** budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms at 21.5 fps; adapter bounds 1..4096). +- Supported request fields: `input` (required), `voice` (`default` only), `response_format` (`wav` default; `wav/pcm/flac/mp3` non-streaming; `opus` 400 because 22.05 kHz is not an Opus sample rate; streaming `pcm/wav` only), `stream` / `stream_format`, `max_new_tokens`, `seed`. +- Unsupported: `speed`, `extra_params` (including `temperature`/`top_p`/`top_k`), `ref_audio`, `ref_text`, `speaker_embedding`, `task_type`, `instructions`, `language`, and `word_timestamps`. +- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands. +- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI. +- For the request-field table, see the [Speech API Gepard section](https://github.com/vllm-project/vllm-omni/blob/main/docs/serving/speech_api.md#gepard-10). + +--- + ## higgs-audio v2 2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client. diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 954b96bd3..04ff1a97d 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -19,7 +19,7 @@ list of supported architectures across all modalities, see | Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | -| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz | +| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning later) | ✓ (online) | zero-shot | 22.05 kHz | | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz | | Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz | @@ -167,7 +167,7 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ ``` ### Architecture -``` +```text Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz) (Llama-based) (32k vocab) (Flow Matching) ``` @@ -491,8 +491,8 @@ Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Ref ### Notes - Output: 22.05 kHz mono WAV. - First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower. -- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR. -- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it). +- Offline and online: `examples/offline_inference/text_to_speech/gepard/end2end.py` and `/v1/audio/speech` (see the [online Gepard section](../../online_serving/text_to_speech/README.md#gepard-10)). +- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it). Online serving **must** pass `--deploy-config vllm_omni/deploy/gepard.yaml` because the checkpoint self-identifies as `qwen3_5_text`. - Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`. - Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced. - `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed. diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 7ad939415..b749731e0 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -17,6 +17,7 @@ For the full list of supported architectures across all modalities, see | Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | +| Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — | | IndexTTS-2.5 | native `checkpoints/` bundle | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — | @@ -96,12 +97,52 @@ curl -X POST http://localhost:8091/v1/audio/speech \ }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others). +Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for Gepard and IndexTTS-2, and 24 kHz for many others). For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md). --- +## Gepard-1.0 + +Single-stage native AR TTS at 22.05 kHz mono. Zero-shot only: omit `voice` or pass `"default"`. Voice cloning from reference audio is not available yet. + +### Prerequisites + +Same NeMo NanoCodec install as the [offline Gepard section](../../offline_inference/text_to_speech/README.md#gepard-10). On a host whose CUDA toolkit cannot build kernels — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also `export VLLM_USE_FLASHINFER_SAMPLER=0` before launch (same note as offline). That env is not a deploy-YAML field. + +### Launch + +```bash +vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \ + --stage-init-timeout 900 \ + --deploy-config vllm_omni/deploy/gepard.yaml +# or: +./gepard/run_server.sh +``` + +`--stage-init-timeout 900` matches the online e2e fixture and the offline example. Serve defaults to 300s, which is often too short for a cold download of the talker plus NanoCodec. + +The packaged `vllm_omni/deploy/gepard.yaml` must be passed with `--deploy-config`. The checkpoint self-identifies as `qwen3_5_text`, so omitting the YAML launches a diffusion fallback instead of the Gepard pipeline. The YAML sets `async_chunk: false`, `max_num_seqs: 4`, and currently pins `seed: 42`, so serving is deterministic by default until that YAML seed is removed. Pass an explicit per-request `seed` in tests and clients rather than depending on either default. + +### Sending requests + +```bash +python examples/online_serving/text_to_speech/gepard/speech_client.py \ + --text "Hello, this is Gepard speaking." + +python examples/online_serving/text_to_speech/gepard/speech_client.py \ + --text "Hello, this is Gepard speaking." --seed 7 --stream --output output.pcm +``` + +### Notes + +- Output: 22.05 kHz mono. `max_new_tokens` is a **frame** budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms at 21.5 fps; adapter bounds 1..4096). +- Supported request fields: `input` (required), `voice` (`default` only), `response_format` (`wav` default; `wav/pcm/flac/mp3` non-streaming; `opus` 400 because 22.05 kHz is not an Opus sample rate; streaming `pcm/wav` only), `stream` / `stream_format`, `max_new_tokens`, `seed`. +- Unsupported: `speed`, `extra_params` (including `temperature`/`top_p`/`top_k`), `ref_audio`, `ref_text`, `speaker_embedding`, `task_type`, `instructions`, `language`, and `word_timestamps`. +- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands. +- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI. + ## dots.tts Single-stage, text-only TTS at 48 kHz. The current integration supports diff --git a/examples/online_serving/text_to_speech/gepard/run_server.sh b/examples/online_serving/text_to_speech/gepard/run_server.sh new file mode 100755 index 000000000..420591288 --- /dev/null +++ b/examples/online_serving/text_to_speech/gepard/run_server.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Launch vLLM-Omni server for Gepard-1.0 TTS (zero-shot default voice). +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 ./run_server.sh +# +# Packaged deploy config: vllm_omni/deploy/gepard.yaml (async_chunk=false, +# 22.05 kHz mono, default seed 42 until the YAML seed is removed). +# Cold start downloads the talker and NeMo NanoCodec; the serve default of +# --stage-init-timeout 300 is too tight (offline end2end.py uses 900). +# +# On a host whose CUDA toolkit cannot JIT FlashInfer (no nvcc/ninja, or a +# consumer Blackwell sm_120 card), also: +# export VLLM_USE_FLASHINFER_SAMPLER=0 + +set -e + +MODEL="${MODEL:-nineninesix/gepard-1.0}" +PORT="${PORT:-8091}" + +echo "Starting Gepard-1.0 server with model: $MODEL" + +vllm-omni serve "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --trust-remote-code \ + --omni \ + --stage-init-timeout 900 \ + --deploy-config vllm_omni/deploy/gepard.yaml diff --git a/examples/online_serving/text_to_speech/gepard/speech_client.py b/examples/online_serving/text_to_speech/gepard/speech_client.py new file mode 100755 index 000000000..4db4f1c7d --- /dev/null +++ b/examples/online_serving/text_to_speech/gepard/speech_client.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Client for Gepard-1.0 TTS via /v1/audio/speech. + +Gepard is zero-shot: omit ``voice`` or pass ``"default"``. Output is 22.05 kHz +mono. ``seed`` is optional; the packaged deploy YAML currently pins ``seed: 42`` +until that default is removed, so requests without an explicit seed are still +deterministic. + +Examples: + python speech_client.py --text "Hello, this is Gepard speaking." + python speech_client.py --text "Hello, this is Gepard speaking." --seed 7 + python speech_client.py --text "Hello, this is Gepard speaking." --stream --output output.pcm +""" + +from __future__ import annotations + +import argparse + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "nineninesix/gepard-1.0" + + +def run_tts(args) -> None: + payload = { + "model": args.model, + "input": args.text, + "voice": args.voice, + "response_format": args.response_format, + } + if args.seed is not None: + payload["seed"] = args.seed + if args.max_new_tokens is not None: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + payload["stream_format"] = "audio" + if args.response_format not in ("pcm", "wav"): + print(f"Note: streaming requires pcm/wav; overriding --response-format {args.response_format} to 'pcm'") + payload["response_format"] = "pcm" + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print(f"Voice: {args.voice}") + if args.seed is not None: + print(f"Seed: {args.seed}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + if args.stream: + output_path = args.output or ("output.wav" if payload["response_format"] == "wav" else "output.pcm") + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as resp: + if resp.status_code != 200: + print(f"Error: {resp.status_code}") + print(resp.read().decode()) + return + total_bytes = 0 + with open(output_path, "wb") as f: + for chunk in resp.iter_bytes(): + f.write(chunk) + total_bytes += len(chunk) + print(f"Streamed {total_bytes} bytes to: {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + output_path = args.output or "output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Gepard-1.0 TTS client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--voice", default="default", help="Voice name (only 'default' is supported)") + parser.add_argument("--seed", type=int, default=None, help="Sampling seed") + parser.add_argument( + "--max-new-tokens", + type=int, + default=None, + help="Frame budget (1 token = 1 frame ≈ 46.4 ms at 21.5 fps)", + ) + parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm"], + help="Audio format (default: wav). Streaming is pcm/wav only. Opus is not supported at 22.05 kHz.", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/online_serving/test_gepard_tts.py b/tests/e2e/online_serving/test_gepard_tts.py new file mode 100644 index 000000000..e6887bbbc --- /dev/null +++ b/tests/e2e/online_serving/test_gepard_tts.py @@ -0,0 +1,524 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""E2E online serving tests for Gepard-1.0 via ``/v1/audio/speech``. + +Zero-shot default voice. Needs a GPU and the NeMo NanoCodec; the weekly +``TTS · Gepard-1.0 · L4`` step installs NeMo then runs this file. +""" + +from __future__ import annotations + +import array +import os +import statistics +import struct +import time +import wave +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from io import BytesIO + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + +import pytest +import requests +import soundfile as sf + +from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_bytes_to_text +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +pytest.importorskip("nemo.collections.tts.models") + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +MODEL = "nineninesix/gepard-1.0" +SAMPLE_RATE = 22050 +SAMPLES_PER_FRAME = 1024 +PCM16_BYTES_PER_FRAME = SAMPLES_PER_FRAME * 2 +DEFAULT_TIMEOUT_S = 180.0 + +_DISTINGUISHABLE_PROMPTS = { + "He drinks coffee every morning.": "coffee", + "Machine learning is interesting.": "learning", + "Please close the window before leaving.": "window", + "My favorite color is purple.": "purple", +} + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("gepard.yaml"), + server_args=["--trust-remote-code", "--disable-log-stats"], + stage_init_timeout=900, + ), + id="gepard", + ) +] + + +def _base_config(omni_server, text: str, **extra) -> dict: + cfg = { + "model": omni_server.model, + "input": text, + "voice": "default", + "timeout": DEFAULT_TIMEOUT_S, + "seed": 7, + # NanoCodec is quieter than the 24 kHz models the PCM HNR helper was + # calibrated on; keep the catastrophic-failure check, not a quality gate. + "min_hnr_db": -2.0, + "transcript_escalation_model": "large-v3", + } + cfg.update(extra) + return cfg + + +def _wav_sample_rate(wav_bytes: bytes) -> int: + return struct.unpack_from(" int: + data, _sr = sf.read(BytesIO(wav_bytes), dtype="int16") + return int(data.size) * 2 + + +def _pcm16le_to_wav(pcm: bytes, sample_rate: int = SAMPLE_RATE) -> bytes: + """Whisper helpers expect a container; streamed speech is raw s16le.""" + buf = BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(pcm) + return buf.getvalue() + + +def _pcm16_samples(pcm: bytes) -> array.array: + samples = array.array("h") + samples.frombytes(pcm) + return samples + + +def _pcm16_diag(label: str, left: bytes, right: bytes) -> None: + extra_frames = (len(left) - len(right)) // PCM16_BYTES_PER_FRAME + overlap = min(len(left), len(right)) + first_frame_match = ( + overlap >= PCM16_BYTES_PER_FRAME and left[:PCM16_BYTES_PER_FRAME] == right[:PCM16_BYTES_PER_FRAME] + ) + worst = None + corr = None + first_diff_sample = None + if overlap >= 2 and overlap % 2 == 0: + a = _pcm16_samples(left[:overlap]) + b = _pcm16_samples(right[:overlap]) + worst = max(abs(x - y) for x, y in zip(a, b, strict=True)) + for i, (x, y) in enumerate(zip(a, b, strict=True)): + if x != y: + first_diff_sample = i + break + try: + corr = round(statistics.correlation(a, b), 6) + except statistics.StatisticsError: + corr = None + print( + f"[Gepard pcm-diag] {label} left={len(left)} right={len(right)} " + f"extra_frames={extra_frames} first_frame_match={first_frame_match} " + f"first_diff_sample={first_diff_sample} max_abs={worst} corr={corr}" + ) + + +def _scan_logs_for_preemption(omni_server) -> None: + log_paths = getattr(omni_server, "_stage_log_paths", {}) or {} + hits: list[str] = [] + for path in log_paths.values(): + if path is None or not path.exists(): + continue + text = path.read_text(encoding="utf-8", errors="replace").lower() + if "preempt" in text: + hits.append(str(path)) + assert not hits, f"preemption mentioned in stage logs: {hits}" + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +@pytest.mark.parametrize("response_format", ["wav", "pcm"]) +def test_text_to_audio_basic(omni_server, online_client, response_format: str, run_level: str) -> None: + text = "Hello, this is Gepard speaking." + # whisper-small/large-v3 mishear 22.05 kHz WAV of this prompt as + # "Jeb Ard" / "Jeff Bard" (cosine 0.80 / 0.75). Offline uses keyword + # containment; the PCM twin below still goes through the 0.9 ASR gate. + if response_format == "wav": + url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" + payload = { + "model": omni_server.model, + "input": text, + "voice": "default", + "seed": 7, + "stream": False, + "response_format": "wav", + } + r = requests.post(url, json=payload, timeout=DEFAULT_TIMEOUT_S) + r.raise_for_status() + audio = r.content + assert _wav_sample_rate(audio) == SAMPLE_RATE + pcm_len = _wav_pcm_payload_len(audio) + else: + [resp] = online_client.send_audio_speech_request( + _base_config(omni_server, text, stream=False, response_format=response_format) + ) + assert resp.audio_bytes + pcm_len = len(resp.audio_bytes) + assert pcm_len % PCM16_BYTES_PER_FRAME == 0 + if run_level in {"advanced_model", "full_model"}: + duration = pcm_len / 2 / SAMPLE_RATE + assert 0.5 < duration < 30.0, f"implausible duration {duration:.2f}s" + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_seed_determinism(omni_server, online_client) -> None: + text = "Hello, this is Gepard speaking." + [a] = online_client.send_audio_speech_request(_base_config(omni_server, text, seed=7, response_format="pcm")) + [b] = online_client.send_audio_speech_request(_base_config(omni_server, text, seed=7, response_format="pcm")) + [c] = online_client.send_audio_speech_request(_base_config(omni_server, text, seed=11, response_format="pcm")) + assert a.audio_bytes == b.audio_bytes + assert a.audio_bytes != c.audio_bytes + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_streaming_pcm_incremental(omni_server, online_client) -> None: + text = "Hello, this is Gepard speaking." + url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" + payload = { + "model": omni_server.model, + "input": text, + "voice": "default", + "seed": 7, + "stream": True, + "stream_format": "audio", + "response_format": "pcm", + } + start = time.perf_counter() + arrivals: list[tuple[float, int]] = [] + streamed = bytearray() + with requests.post(url, json=payload, stream=True, timeout=DEFAULT_TIMEOUT_S) as resp: + resp.raise_for_status() + content_type = resp.headers.get("content-type") or "" + assert "pcm" in content_type or "octet-stream" in content_type + for chunk in resp.iter_content(chunk_size=None): + if not chunk: + continue + arrivals.append((time.perf_counter(), len(chunk))) + streamed.extend(chunk) + assert len(arrivals) >= 2, f"expected incremental chunks, got {len(arrivals)}" + assert arrivals[1][0] > arrivals[0][0] + ttfa_ms = (arrivals[0][0] - start) * 1000.0 + cadence_ms = [(arrivals[i][0] - arrivals[i - 1][0]) * 1000.0 for i in range(1, len(arrivals))] + print(f"[Gepard streaming] TTFA_ms={ttfa_ms:.1f} chunk_cadence_ms={cadence_ms}") + + [non_stream] = online_client.send_audio_speech_request( + _base_config(omni_server, text, seed=7, stream=False, response_format="pcm") + ) + assert bytes(streamed) == non_stream.audio_bytes + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_seeded_streaming_stops(omni_server) -> None: + url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" + payload = { + "model": omni_server.model, + "input": "Hello, this is Gepard speaking.", + "voice": "default", + "seed": 7, + "stream": True, + "stream_format": "audio", + "response_format": "pcm", + } + with requests.post(url, json=payload, stream=True, timeout=DEFAULT_TIMEOUT_S) as resp: + resp.raise_for_status() + audio = b"".join(chunk for chunk in resp.iter_content(chunk_size=None) if chunk) + duration = len(audio) / 2 / SAMPLE_RATE + assert duration < 30.0, f"seeded stream ran {duration:.1f}s — stop constraint dropped?" + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_sse_stream_emits_delta_then_done(omni_server) -> None: + url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" + payload = { + "model": omni_server.model, + "input": "Hello, this is Gepard speaking.", + "voice": "default", + "seed": 7, + "stream": True, + "stream_format": "sse", + "response_format": "pcm", + } + with requests.post(url, json=payload, stream=True, timeout=DEFAULT_TIMEOUT_S) as resp: + resp.raise_for_status() + assert resp.headers.get("content-type", "").startswith("text/event-stream") + body = b"".join(resp.iter_content(chunk_size=None)).decode("utf-8", errors="replace") + assert "speech.audio.delta" in body + assert "speech.audio.done" in body + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_concurrent_requests_stay_isolated(omni_server, online_client, run_level: str) -> None: + def _one(text: str): + return online_client.send_audio_speech_request(_base_config(omni_server, text, seed=7, response_format="wav"))[ + 0 + ] + + texts = list(_DISTINGUISHABLE_PROMPTS) + with ThreadPoolExecutor(max_workers=4) as pool: + responses = list(pool.map(_one, texts)) + + assert all(r.success for r in responses) + for i in range(len(responses)): + for j in range(i + 1, len(responses)): + assert responses[i].audio_bytes != responses[j].audio_bytes + + if run_level in {"advanced_model", "full_model"}: + for text, resp in zip(texts, responses, strict=True): + keyword = _DISTINGUISHABLE_PROMPTS[text] + transcript = convert_audio_bytes_to_text(resp.audio_bytes, language="en").lower() + assert keyword in transcript, f"expected {keyword!r} in {transcript!r} for {text!r}" + + _scan_logs_for_preemption(omni_server) + + +def _speech_url(omni_server) -> str: + return f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" + + +def _stream_payload(omni_server, text: str, **extra) -> dict: + payload = { + "model": omni_server.model, + "input": text, + "voice": "default", + "seed": 7, + "stream": True, + "stream_format": "audio", + "response_format": "pcm", + } + payload.update(extra) + return payload + + +@dataclass +class _StreamCapture: + pcm: bytes + arrivals: list[tuple[float, int]] + + @property + def duration_s(self) -> float: + return len(self.pcm) / 2 / SAMPLE_RATE + + +def _capture_stream( + url: str, + payload: dict, + timeout_s: float, + *, + sleep_after_first_s: float = 0.0, +) -> _StreamCapture: + start = time.perf_counter() + arrivals: list[tuple[float, int]] = [] + pcm = bytearray() + with requests.post(url, json=payload, stream=True, timeout=timeout_s) as resp: + resp.raise_for_status() + first = True + for chunk in resp.iter_content(chunk_size=None): + if not chunk: + continue + arrivals.append((time.perf_counter() - start, len(chunk))) + pcm.extend(chunk) + if first and sleep_after_first_s > 0: + time.sleep(sleep_after_first_s) + first = False + return _StreamCapture(bytes(pcm), arrivals) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_concurrent_streaming_stays_isolated(omni_server, run_level: str) -> None: + url = _speech_url(omni_server) + texts = list(_DISTINGUISHABLE_PROMPTS) + + def _one(text: str) -> tuple[str, _StreamCapture]: + return text, _capture_stream(url, _stream_payload(omni_server, text), DEFAULT_TIMEOUT_S) + + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(_one, texts)) + + captures = [cap for _, cap in results] + assert all(len(c.arrivals) >= 2 for c in captures) + for i in range(len(captures)): + for j in range(i + 1, len(captures)): + assert captures[i].pcm != captures[j].pcm + + if run_level in {"advanced_model", "full_model"}: + for text, cap in results: + keyword = _DISTINGUISHABLE_PROMPTS[text] + transcript = convert_audio_bytes_to_text(_pcm16le_to_wav(cap.pcm), language="en").lower() + assert keyword in transcript, f"expected {keyword!r} in {transcript!r} for {text!r}" + + _scan_logs_for_preemption(omni_server) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_slow_client_still_completes(omni_server, online_client) -> None: + """Slow reader must finish, and seed=7 must match a prior non-stream. + + Runs before ``test_client_disconnect_releases_capacity`` so an abort + storm cannot contaminate this comparison. Prints overlap diagnostics + before the bit-exact asserts. + """ + text = "Hello, this is Gepard speaking." + [baseline] = online_client.send_audio_speech_request( + _base_config(omni_server, text, seed=7, stream=False, response_format="pcm") + ) + cap = _capture_stream( + url=_speech_url(omni_server), + payload=_stream_payload(omni_server, text), + timeout_s=DEFAULT_TIMEOUT_S, + sleep_after_first_s=3.0, + ) + assert len(cap.arrivals) >= 2 + assert len(cap.pcm) % PCM16_BYTES_PER_FRAME == 0 + assert 0.5 < cap.duration_s < 30.0 + print(f"[Gepard slow-client] chunks={len(cap.arrivals)} duration_s={cap.duration_s:.3f}") + _pcm16_diag("slow-vs-prior-nonstream", cap.pcm, baseline.audio_bytes) + [after] = online_client.send_audio_speech_request( + _base_config(omni_server, text, seed=7, stream=False, response_format="pcm") + ) + _pcm16_diag("after-nonstream-vs-prior", after.audio_bytes, baseline.audio_bytes) + assert cap.pcm == baseline.audio_bytes + assert after.audio_bytes == baseline.audio_bytes + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_client_disconnect_releases_capacity(omni_server, online_client) -> None: + """A follow-up request must complete after in-flight streams are closed. + + Logs PCM diagnostics against a pre-abort ``seed=7`` baseline. Bit-exact + equality after abort is not asserted here: a dying stream can still share + a GPU batch with the next generate after ``abort()`` returns, and Gepard's + in-model stop head / bf16 GEMM are not row-stable in that batch. + """ + url = _speech_url(omni_server) + text = "Hello, this is Gepard speaking." + [baseline] = online_client.send_audio_speech_request( + _base_config(omni_server, text, seed=7, stream=False, response_format="pcm") + ) + long_payload = _stream_payload( + omni_server, + "Please keep talking about the weather, the window, coffee, and purple flowers.", + max_new_tokens=80, + ) + for _ in range(4): + resp = requests.post(url, json=long_payload, stream=True, timeout=DEFAULT_TIMEOUT_S) + resp.raise_for_status() + try: + chunk = next(c for c in resp.iter_content(chunk_size=None) if c) + assert chunk + finally: + resp.close() + + t0 = time.perf_counter() + follow = requests.post( + url, + json={ + "model": omni_server.model, + "input": text, + "voice": "default", + "seed": 7, + "stream": False, + "response_format": "pcm", + }, + timeout=DEFAULT_TIMEOUT_S, + ) + follow.raise_for_status() + elapsed = time.perf_counter() - t0 + pcm_len = len(follow.content) + assert pcm_len % PCM16_BYTES_PER_FRAME == 0 + print(f"[Gepard disconnect] followup_s={elapsed:.2f} pcm_bytes={pcm_len}") + _pcm16_diag("disconnect-followup-vs-prior", follow.content, baseline.audio_bytes) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_mixed_stream_and_nonstream_isolated(omni_server, online_client, run_level: str) -> None: + url = _speech_url(omni_server) + stream_text = "He drinks coffee every morning." + nonstream_text = "My favorite color is purple." + + def _stream(): + return _capture_stream(url, _stream_payload(omni_server, stream_text), DEFAULT_TIMEOUT_S) + + def _nonstream(): + return online_client.send_audio_speech_request( + _base_config(omni_server, nonstream_text, seed=7, response_format="wav") + )[0] + + with ThreadPoolExecutor(max_workers=2) as pool: + stream_future = pool.submit(_stream) + nonstream_future = pool.submit(_nonstream) + streamed = stream_future.result() + nonstream = nonstream_future.result() + + assert streamed.pcm + assert nonstream.success + assert streamed.pcm != nonstream.audio_bytes + if run_level in {"advanced_model", "full_model"}: + stream_tr = convert_audio_bytes_to_text(_pcm16le_to_wav(streamed.pcm), language="en").lower() + nonstream_tr = convert_audio_bytes_to_text(nonstream.audio_bytes, language="en").lower() + assert "coffee" in stream_tr, stream_tr + assert "purple" in nonstream_tr, nonstream_tr + _scan_logs_for_preemption(omni_server) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_max_new_tokens_caps_whole_frames(omni_server) -> None: + frames = 16 + url = _speech_url(omni_server) + r = requests.post( + url, + json={ + "model": omni_server.model, + "input": "Hello, this is Gepard speaking.", + "voice": "default", + "seed": 7, + "stream": False, + "response_format": "pcm", + "max_new_tokens": frames, + }, + timeout=DEFAULT_TIMEOUT_S, + ) + assert r.status_code == 200, r.text + pcm_len = len(r.content) + assert pcm_len % PCM16_BYTES_PER_FRAME == 0 + assert pcm_len == frames * PCM16_BYTES_PER_FRAME, f"got {pcm_len} bytes, want {frames} frames" + duration = pcm_len / 2 / SAMPLE_RATE + assert duration < 1.0, f"budget cap should stop near 16 frames, got {duration:.2f}s" + + over = requests.post( + url, + json={ + "model": omni_server.model, + "input": "Hello, this is Gepard speaking.", + "voice": "default", + "max_new_tokens": 4097, + }, + timeout=DEFAULT_TIMEOUT_S, + ) + assert over.status_code == 400 + assert "max_new_tokens" in over.text diff --git a/tests/entrypoints/openai_api/test_gepard_tts_adapter.py b/tests/entrypoints/openai_api/test_gepard_tts_adapter.py new file mode 100644 index 000000000..f9e1e3dd2 --- /dev/null +++ b/tests/entrypoints/openai_api/test_gepard_tts_adapter.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU tests for the Gepard-1.0 speech adapter. + +Covers detection, the reject matrix, voice default, empty input, extra_params, +max_new_tokens frame bounds, and build() parity with the offline prompt helper. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest +import torch +from pydantic import ValidationError +from pytest_mock import MockerFixture + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.tts_adapters import detect_tts_model_type, resolve_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.gepard import GepardAdapter +from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig +from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODEL = "nineninesix/gepard-1.0" + + +def _adapter(*, revision: str | None = None) -> GepardAdapter: + engine = SimpleNamespace(model_config=SimpleNamespace(model=_MODEL, revision=revision)) + return GepardAdapter(SpeechServingContext(server=SimpleNamespace(engine_client=engine), engine_client=engine)) + + +def test_gepard_detection_and_registration() -> None: + assert resolve_adapter("gepard") is GepardAdapter + assert GepardAdapter.stage_keys == {"gepard"} + assert GepardAdapter.model_archs == {"GepardTalkerForConditionalGeneration"} + assert detect_tts_model_type("gepard", None) == "gepard" + assert detect_tts_model_type("gepard", "GepardTalkerForConditionalGeneration") == "gepard" + assert detect_tts_model_type(None, "GepardTalkerForConditionalGeneration") == "gepard" + + +def test_gepard_accepts_omitted_and_default_voice() -> None: + adapter = _adapter() + assert adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.")) is None + assert adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.", voice="default")) is None + assert adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.", voice="Default")) is None + assert ( + adapter.validate( + OpenAICreateSpeechRequest.model_validate({"input": "Hello from Gepard.", "speaker": "default"}) + ) + is None + ) + + +@pytest.mark.parametrize( + ("kwargs", "field"), + [ + ({"voice": "vivian"}, "voice"), + ({"speed": 1.2}, "speed"), + ({"extra_params": {"temperature": 0.3}}, "extra_params"), + ({"extra_params": {"top_p": 0.9}}, "extra_params"), + ({"extra_params": {"top_k": 20}}, "extra_params"), + ({"extra_params": {"foo": 1}}, "extra_params"), + ({"ref_audio": "https://example.com/ref.wav"}, "ref_audio"), + ({"ref_text": "transcript"}, "ref_text"), + ({"ref_audio_2": "https://example.com/ref2.wav"}, "ref_audio_2"), + ({"speaker_embedding": [0.1, 0.2]}, "speaker_embedding"), + ({"x_vector_only_mode": True}, "x_vector_only_mode"), + ({"task_type": "Base"}, "task_type"), + ({"non_streaming_mode": True}, "non_streaming_mode"), + ({"instructions": "speak happily"}, "instructions"), + ({"language": "English"}, "language"), + ({"ambient_sound": "rain"}, "ambient_sound"), + ({"duration_seconds": 2.0}, "duration_seconds"), + ({"initial_codec_chunk_frames": 4}, "initial_codec_chunk_frames"), + ({"word_timestamps": True}, "word_timestamps"), + ({"response_format": "opus"}, "opus"), + ], +) +def test_gepard_rejects_unsupported_fields(kwargs: dict, field: str) -> None: + adapter = _adapter() + err = adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.", **kwargs)) + assert err is not None + assert field in err + + +@pytest.mark.parametrize("text", ["", " ", "\n\t"]) +def test_gepard_rejects_empty_input(text: str) -> None: + adapter = _adapter() + err = adapter.validate(OpenAICreateSpeechRequest(input=text)) + assert err is not None + assert "empty" in err.lower() + + +@pytest.mark.parametrize("max_new_tokens", [0, -1]) +def test_gepard_schema_rejects_non_positive_max_new_tokens(max_new_tokens: int) -> None: + with pytest.raises(ValidationError, match="greater than or equal to 1"): + OpenAICreateSpeechRequest(input="Hello from Gepard.", max_new_tokens=max_new_tokens) + + +def test_gepard_rejects_max_new_tokens_above_frame_budget() -> None: + adapter = _adapter() + err = adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.", max_new_tokens=4097)) + assert err is not None + assert "max_new_tokens" in err + assert "frames" in err + + +def test_gepard_accepts_max_new_tokens_bounds() -> None: + adapter = _adapter() + assert adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.", max_new_tokens=1)) is None + assert adapter.validate(OpenAICreateSpeechRequest(input="Hello from Gepard.", max_new_tokens=4096)) is None + + +def test_gepard_apply_sampling_overrides_uses_frames() -> None: + adapter = _adapter() + stage_defaults = [SimpleNamespace(max_tokens=1000)] + overridden = adapter.apply_sampling_overrides( + stage_defaults, + OpenAICreateSpeechRequest(input="Hello from Gepard.", max_new_tokens=32), + ) + assert overridden[0].max_tokens == 32 + assert stage_defaults[0].max_tokens == 1000 + + +def test_gepard_build_matches_offline_prompt_builder(monkeypatch) -> None: + adapter = _adapter() + cfg = GepardConfig() + monkeypatch.setattr( + "vllm_omni.model_executor.models.gepard.configuration_gepard.GepardConfig.from_checkpoint", + staticmethod(lambda *_a, **_k: cfg), + ) + + class _Tok: + def __call__(self, text, add_special_tokens=False): + return {"input_ids": [11, 12, 13]} + + adapter._tokenizer = _Tok() + adapter._gepard_config = cfg + text = "Hello from Gepard." + prepared = asyncio.run(adapter.build(OpenAICreateSpeechRequest(input=text, seed=7), [], False)) + + expected = build_gepard_prompt_ids([11, 12, 13], config=cfg) + assert prepared.prompt["prompt_token_ids"] == expected + assert prepared.prompt["additional_information"] == {"text": [text]} + assert prepared.model_type == "gepard" + assert prepared.output_policy.accumulate_nonstreaming is True + + +def test_gepard_loads_tokenizer_and_config_at_served_revision(monkeypatch) -> None: + adapter = _adapter(revision="abc123") + seen: dict[str, object] = {} + + def fake_from_checkpoint(model, backbone_config=None, revision=None): + seen["config"] = (model, revision) + return GepardConfig() + + class _Tok: + def __init__(self, *args, **kwargs): + seen["tok"] = (args[0] if args else None, kwargs.get("revision")) + + def __call__(self, text, add_special_tokens=False): + return {"input_ids": [11, 12, 13]} + + monkeypatch.setattr( + "vllm_omni.model_executor.models.gepard.configuration_gepard.GepardConfig.from_checkpoint", + staticmethod(fake_from_checkpoint), + ) + monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", _Tok) + prepared = asyncio.run(adapter.build(OpenAICreateSpeechRequest(input="Hello from Gepard."), [], False)) + assert seen["config"] == (_MODEL, "abc123") + assert seen["tok"] == (_MODEL, "abc123") + assert prepared.prompt["prompt_token_ids"] == build_gepard_prompt_ids([11, 12, 13], config=GepardConfig()) + + +def _gepard_mock_output(sample_rate: int = 22050, num_samples: int = 2048) -> OmniRequestOutput: + class MockCompletionOutput: + def __init__(self) -> None: + self.index = 0 + self.text = "" + self.token_ids: list[int] = [] + self.finish_reason = "stop" + self.stop_reason = None + self.logprobs = None + + class MockRequestOutput: + def __init__(self) -> None: + self.request_id = "speech-gepard-mock" + self.outputs = [MockCompletionOutput()] + self.multimodal_output = { + "audio": torch.zeros(num_samples), + "sr": torch.tensor(sample_rate), + } + self.finished = True + self.prompt_token_ids = None + self.encoder_prompt_token_ids = None + self.num_cached_tokens = None + self.prompt_logprobs = None + self.kv_transfer_params = None + + output = OmniRequestOutput.from_stage_output(MockRequestOutput(), stage_id=0, final_output_type="audio") + output.metrics = {"stage_metrics": {"0": {"num_tokens_in": 0, "num_tokens_out": 2}}} + return output + + +@pytest.fixture +def gepard_server(mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value={"default"}, + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=22050.0, + ) + + async def _gen(*_args, **_kwargs): + yield _gepard_mock_output() + + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.model_config = mocker.MagicMock(model=_MODEL, async_chunk=False, revision=None) + mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=1000, seed=42, extra_args=None)] + mock_engine_client.tts_batch_max_items = 32 + mock_engine_client.generate = mocker.MagicMock(side_effect=lambda **_k: _gen()) + mock_engine_client.stage_configs = [ + SimpleNamespace( + engine_args=SimpleNamespace(model_stage="gepard", model_arch="GepardTalkerForConditionalGeneration"), + tts_args={}, + ) + ] + + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + cfg = GepardConfig() + server._adapter._gepard_config = cfg + server._adapter._tokenizer = lambda text, add_special_tokens=False: {"input_ids": [11, 12, 13]} + yield server + server.shutdown() + + +def test_gepard_prepare_uses_prompt_builder(gepard_server) -> None: + request_id, _generator, _ = asyncio.run( + gepard_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello from Gepard.")) + ) + assert request_id.startswith("speech-") + prompt = gepard_server.engine_client.generate.call_args.kwargs["prompt"] + assert prompt["prompt_token_ids"] == build_gepard_prompt_ids([11, 12, 13], config=GepardConfig()) + + +@pytest.mark.parametrize("response_format", ["wav", "pcm", "flac", "mp3"]) +def test_gepard_mocked_engine_encodes_response_formats(gepard_server, response_format: str) -> None: + audio, media_type = asyncio.run( + gepard_server._generate_audio_bytes( + OpenAICreateSpeechRequest(input="Hello from Gepard.", response_format=response_format) + ) + ) + assert isinstance(audio, (bytes, str)) + assert len(audio) > 0 + expected_media_type = {"wav": "audio/wav", "pcm": "audio/pcm", "flac": "audio/flac", "mp3": "audio/mpeg"} + assert media_type == expected_media_type[response_format] + + +def test_gepard_streaming_rejects_non_pcm_wav() -> None: + with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"): + OpenAICreateSpeechRequest(input="Hello from Gepard.", stream=True, response_format="mp3") diff --git a/tests/entrypoints/openai_api/test_serving_speech_metrics.py b/tests/entrypoints/openai_api/test_serving_speech_metrics.py index a5e284a01..a1fb74a46 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_speech_metrics.py @@ -46,6 +46,7 @@ def _serving(metrics: _MetricsStub, *, adapter=None) -> OmniOpenAIServingSpeech: serving._tts_model_type = "qwen3_tts" serving.engine_client = SimpleNamespace(mod_metrics=metrics, request_states={}) serving._get_tts_adapter = lambda: adapter + serving._speech_output_policies = {} serving.create_audio = lambda audio_obj: SimpleNamespace( audio_data=b"\0\0" * int(audio_obj.audio_tensor.size), media_type="audio/pcm", diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index edd2b4cd4..e0a165a55 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -69,6 +69,7 @@ EXPECTED_MODEL_TYPES = { "step_audio2", "indextts2", "indextts2_5", + "gepard", "dots_tts", } @@ -143,6 +144,17 @@ def _build_moss_tts_request(adapter_cls, mocker, *, request_seed): ) +@pytest.mark.parametrize( + ("adapter_cls", "expect_accumulate"), + [(MossTTSAdapter, False), (MossTTSNanoAdapter, True)], +) +def test_moss_tts_accumulate_nonstreaming_follows_adapter_flag(adapter_cls, expect_accumulate, mocker): + prepared = _build_moss_tts_request(adapter_cls, mocker, request_seed=7) + + assert adapter_cls.accumulate_nonstreaming is expect_accumulate + assert prepared.output_policy.accumulate_nonstreaming is expect_accumulate + + # Full-family coverage pins the adapter contract; only Nano consumes this seed end to end today. @pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter]) @pytest.mark.parametrize("request_seed", [0, 1234]) diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 6c7671659..90e7e0fee 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -148,6 +148,7 @@ _PIPELINE_STAGES = [ "dit", "fish_speech_slow_ar", "fused_thinker_talker", + "gepard", "glm_tts", "glm_tts_dit", "higgs_audio_v2", diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index be282da81..33abbb618 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -231,6 +231,7 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc serving._tts_model_type = "higgs_audio_v3" serving._get_tts_adapter = lambda: Adapter() serving._track_ref_audio_artifact_warmup = lambda *args, **kwargs: None + serving._speech_output_policies = {} _, stage_sampling_params, _ = asyncio.run(serving._prepare_speech_generation(request, request_id="speech-test")) stage0_params = stage_sampling_params[0] diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index d4aa72818..f54ca3e28 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -53,6 +53,7 @@ from vllm_omni.entrypoints.openai.speech_usage import ( qwen3_tts_input_token_details, ) from vllm_omni.entrypoints.openai.tts_adapters import ( + OutputPolicy, SpeechServingContext, TTSGenerationError, all_tts_stage_keys, @@ -396,6 +397,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._max_instructions_length = self._compute_max_instructions_length() self._tts_tokenizer = None + # Per-request output policy from the adapter's PreparedRequest. Keyed + # by request_id so concurrent speech requests cannot overwrite each + # other. Stored only for non-streaming requests and consumed (popped) + # by the non-streaming accumulator in `_generate_audio_bytes`. + self._speech_output_policies: dict[str, OutputPolicy] = {} # Batch configuration self._batch_max_items: int = getattr(self.engine_client, "tts_batch_max_items", 32) @@ -1872,6 +1878,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params = prepared.tts_params model_type = prepared.model_type qwen3_ref_audio_warmup_artifact_key = prepared.warmup_artifact_key + output_policy = prepared.output_policy else: # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker" # stage whose preprocess requires chat-templated tokens. The @@ -1893,6 +1900,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) tts_params = {} prompt = {"prompt": request.input} + output_policy = OutputPolicy() if model_type is None: if self._is_tts: @@ -1962,6 +1970,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): qwen3_ref_audio_warmup_artifact_key, x_vector_only=self._tts_x_vector_only(tts_params), ) + # Only the non-streaming path consumes output policies (in + # `_generate_audio_bytes`); streaming never reads them. Store late — + # after every fallible step above — and skip streaming requests + # entirely, so no caller has to discard entries (request ids are + # unique per request, so a stranded entry would never be reclaimed). + if not request.is_streaming(): + self._speech_output_policies[request_id] = output_policy return request_id, generator, tts_params async def _generate_pcm_chunks( @@ -2041,13 +2056,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): artifact_ready = False try: - # MOSS-TTS-Nano emits delta chunks per yield (single-stage, - # async_chunk=false). The engine surfaces each yield as its own - # RequestOutput, so we need to accumulate across the async-for loop — - # final_output alone only carries the last (often empty) sentinel. - is_moss = self._tts_model_type == "moss_tts_nano" - moss_chunks: list[Any] = [] - moss_sample_rate: int | None = None + # Sparse-audio models (MOSS-TTS-Nano, Gepard) emit delta chunks per + # yield with async_chunk=false. The engine surfaces each yield as its + # own RequestOutput, so we accumulate across the async-for loop when + # the adapter asks for it — final_output alone may only carry the + # last (often empty) sentinel. Prefer the engine's own concatenated + # waveform when FINAL_ONLY already produced one. + accumulate_nonstreaming = self._speech_output_policies.pop( + request_id, OutputPolicy() + ).accumulate_nonstreaming + delta_chunks: list[Any] = [] + delta_sample_rate: int | None = None final_output: OmniRequestOutput | None = None # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full @@ -2067,7 +2086,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _, audio_key = self._extract_audio_output(res) if audio_key is not None: audio_res = res - if not is_moss: + if not accumulate_nonstreaming: continue try: step_audio, step_key = self._extract_audio_output(res) @@ -2080,11 +2099,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): candidates = chunk if isinstance(chunk, list) else [chunk] for cand in candidates: if hasattr(cand, "numel") and cand.numel() > 0: - moss_chunks.append(cand) + delta_chunks.append(cand) sr_step = step_audio.get("sr") if sr_step is not None: sr_val_step = sr_step[-1] if isinstance(sr_step, list) and sr_step else sr_step - moss_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step) + delta_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step) if final_output is None: raise ValueError("No output generated from the model.") @@ -2117,7 +2136,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) - if is_moss: + if accumulate_nonstreaming: # Prefer the engine's own consolidated audio when present. After the # vllm 0.20 rebase non-stream requests resolve to FINAL_ONLY, so # final_output already carries the full concatenated waveform; the @@ -2133,12 +2152,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if final_audio is not None: audio_tensor = final_audio - elif moss_chunks: - audio_tensor = torch.cat(moss_chunks, dim=-1) + elif delta_chunks: + audio_tensor = torch.cat(delta_chunks, dim=-1) else: audio_tensor = np.zeros((0,), dtype=np.float32) - if moss_sample_rate is not None: - sample_rate = moss_sample_rate + if delta_sample_rate is not None: + sample_rate = delta_sample_rate elif isinstance(audio_tensor, list): async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False)) if async_chunk: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index 8ed54afd8..9f5387a7f 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -110,6 +110,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 covo_audio, dots_tts, fish_speech, + gepard, glm_tts, higgs_audio_v2, higgs_audio_v3, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/gepard.py b/vllm_omni/entrypoints/openai/tts_adapters/gepard.py new file mode 100644 index 000000000..255525a0d --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/gepard.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Gepard-1.0 serving adapter for ``/v1/audio/speech``. + +Zero-shot only: the trained ``null_prefix`` is the default voice. Prompt IDs +are assembled with the same layout helper the offline path uses. Frame budget +rides ``SamplingParams.max_tokens`` (one token = one 1024-sample frame). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from vllm.utils.async_utils import make_async + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + OutputPolicy, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig + +_DEFAULT_VOICE = "default" + +# Declared schema fields Gepard does not consume. Rejecting them here is +# load-bearing: the pydantic model silently drops undeclared keys, so a +# declared-but-ignored field would otherwise succeed and mislead. +_UNSUPPORTED_OPTIONAL_FIELDS: tuple[tuple[str, str], ...] = ( + ("ref_audio", "ref_audio"), + ("ref_text", "ref_text"), + ("ref_audio_2", "ref_audio_2"), + ("speaker_embedding", "speaker_embedding"), + ("x_vector_only_mode", "x_vector_only_mode"), + ("task_type", "task_type"), + ("non_streaming_mode", "non_streaming_mode"), + ("instructions", "instructions"), + ("language", "language"), + ("ambient_sound", "ambient_sound"), + ("duration_seconds", "duration_seconds"), + ("initial_codec_chunk_frames", "initial_codec_chunk_frames"), +) + + +@register_tts_adapter +class GepardAdapter(ARTTSAdapter): + name = "gepard" + stage_keys = frozenset({"gepard"}) + model_archs = frozenset({"GepardTalkerForConditionalGeneration"}) + + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._tokenizer: Any = None + self._gepard_config: GepardConfig | None = None + self._build_prompt_async = make_async(self._build_prompt, executor=getattr(ctx.server, "_tts_executor", None)) + + def _load_supported_speakers(self) -> set[str]: + return {_DEFAULT_VOICE} + + def validate(self, request: OpenAICreateSpeechRequest) -> str | None: + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.voice is not None: + voice = request.voice.lower() + if voice != _DEFAULT_VOICE: + return ( + f"Invalid voice '{request.voice}'. Supported: {_DEFAULT_VOICE}. " + "Voice cloning from reference audio is not available yet." + ) + + if request.speed is not None and request.speed != 1.0: + return "Gepard does not support 'speed' adjustments; audio is rendered at the native 22.05 kHz rate" + + if request.extra_params: + keys = ", ".join(sorted(request.extra_params)) + return ( + f"Gepard does not accept extra_params (got: {keys}). " + "Sampling is fixed in-model; temperature/top_p/top_k cannot be overridden." + ) + + if request.word_timestamps: + return "Gepard does not support 'word_timestamps'; there is no forced-aligner stage in this pipeline" + + fmt = (request.response_format or "").lower() + if fmt == "opus": + return ( + "Gepard does not support 'response_format'='opus'; " + "native 22.05 kHz is not an Opus sample rate (8000, 12000, 16000, 24000, 48000)" + ) + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min or request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens must be between {self.max_new_tokens_min} and {self.max_new_tokens_max} frames" + + for attr, field in _UNSUPPORTED_OPTIONAL_FIELDS: + if getattr(request, attr, None) is not None: + return f"Gepard does not support '{field}'" + + return None + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: OpenAICreateSpeechRequest, + prompt: dict | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) + + async def build( + self, + request: OpenAICreateSpeechRequest, + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + tts_params: dict = {} + prompt = await self._build_prompt_async(request.input) + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) + return PreparedRequest( + prompt=prompt, + tts_params=tts_params, + model_type=self.name, + output_policy=OutputPolicy(accumulate_nonstreaming=True), + ) + + def _build_prompt(self, text: str) -> dict[str, Any]: + from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids + + prompt_token_ids = build_gepard_prompt_ids(self._tokenize(text), config=self._config()) + return { + "prompt_token_ids": prompt_token_ids, + "additional_information": {"text": [text]}, + } + + def _model_config(self) -> Any: + engine = self.ctx.engine_client + assert engine is not None + return engine.model_config + + def _config(self) -> GepardConfig: + if self._gepard_config is None: + from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig + + model_config = self._model_config() + self._gepard_config = GepardConfig.from_checkpoint( + model_config.model, revision=getattr(model_config, "revision", None) + ) + return self._gepard_config + + def _tokenize(self, text: str) -> list[int]: + tokenizer = self._tokenizer + if tokenizer is None: + from transformers import AutoTokenizer + + model_config = self._model_config() + tokenizer = AutoTokenizer.from_pretrained( + model_config.model, + revision=getattr(model_config, "revision", None), + trust_remote_code=True, + ) + self._tokenizer = tokenizer + return tokenizer(text, add_special_tokens=False)["input_ids"] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 07e4bb197..bed20efc2 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -13,6 +13,7 @@ from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, + OutputPolicy, PreparedRequest, apply_max_new_tokens, conditioning_cache_salt, @@ -23,6 +24,8 @@ if TYPE_CHECKING: class _MossTTSAdapterBase(ARTTSAdapter): + accumulate_nonstreaming: bool = False + def __init__(self, ctx) -> None: super().__init__(ctx) self._moss_variant = None if self.name == "moss_tts_nano" else self._detect_moss_variant() @@ -369,7 +372,12 @@ class _MossTTSAdapterBase(ARTTSAdapter): prompt = tokens_input(prompt_token_ids=[1]) prompt["additional_information"] = tts_params prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) - return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + return PreparedRequest( + prompt=prompt, + tts_params=tts_params, + model_type=self.name, + output_policy=OutputPolicy(accumulate_nonstreaming=self.accumulate_nonstreaming), + ) def apply_sampling_overrides( self, @@ -385,6 +393,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): class MossTTSNanoAdapter(_MossTTSAdapterBase): stage_keys = frozenset({"moss_tts_nano"}) name = "moss_tts_nano" + accumulate_nonstreaming = True @register_tts_adapter From 6040bb97e2036927632151a87b0aaed94ccbce7d Mon Sep 17 00:00:00 2001 From: yzhautouskay Date: Tue, 15 Sep 2026 23:57:15 +0200 Subject: [PATCH 087/209] [Perf][Cosmos3] Add SeaCache support for Cosmos3 (#6922) Signed-off-by: Yuliya Zhautouskaya --- recipes/cosmos3/Cosmos3-Nano.md | 61 +- recipes/cosmos3/Cosmos3-Super.md | 27 +- tests/diffusion/cache/test_seacache.py | 635 ++++++++++++++++++ .../models/cosmos3/test_cosmos3_pipeline.py | 105 ++- .../cosmos3/test_cosmos3_transformer.py | 104 ++- .../diffusion/cache/seacache/__init__.py | 18 + vllm_omni/diffusion/cache/seacache/backend.py | 100 +++ vllm_omni/diffusion/cache/seacache/config.py | 36 + vllm_omni/diffusion/cache/seacache/hook.py | 397 +++++++++++ .../diffusion/cache/seacache/sea_filter.py | 93 +++ vllm_omni/diffusion/cache/seacache/state.py | 26 + vllm_omni/diffusion/cache/selector.py | 11 +- .../diffusion/cache/teacache/extractors.py | 100 ++- vllm_omni/diffusion/data.py | 8 + .../models/cosmos3/pipeline_cosmos3.py | 77 ++- .../models/cosmos3/transformer_cosmos3.py | 274 +++++--- .../worker/diffusion_model_runner.py | 49 +- vllm_omni/entrypoints/cli/serve.py | 5 +- 18 files changed, 1966 insertions(+), 160 deletions(-) create mode 100644 tests/diffusion/cache/test_seacache.py create mode 100644 vllm_omni/diffusion/cache/seacache/__init__.py create mode 100644 vllm_omni/diffusion/cache/seacache/backend.py create mode 100644 vllm_omni/diffusion/cache/seacache/config.py create mode 100644 vllm_omni/diffusion/cache/seacache/hook.py create mode 100644 vllm_omni/diffusion/cache/seacache/sea_filter.py create mode 100644 vllm_omni/diffusion/cache/seacache/state.py diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 6ef24bcf4..29088e94c 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -32,13 +32,13 @@ mode is selected per request: T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples. - **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks: - - `forward_dynamics` — given a first frame or video **and** an action trajectory, + - `forward_dynamics` — given a first frame or video **and** an action trajectory, roll out the resulting video. Synchronous: `POST /v1/videos/sync`. - - `policy` — given a first frame or video and a language instruction, + - `policy` — given a first frame or video and a language instruction, **predict** the action trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and read the predicted action from the top-level `action` field. - - `inverse_dynamics` — given a video, **recover** the action trajectory. Use + - `inverse_dynamics` — given a video, **recover** the action trajectory. Use the async `POST /v1/videos` endpoint and read the recovered action from the top-level `action` field (`{data, shape, dtype, raw_action_dim, domain_id}`). @@ -106,6 +106,25 @@ vllm serve nvidia/Cosmos3-Nano \ --init-timeout 1800 ``` +For Cosmos3, SeaCache is the recommended default choice when opting into +diffusion caching. Caching remains opt-in; add `--cache-backend sea_cache` to +the command above. Its default threshold and maximum cached-step streak are +tuned for Cosmos3, so no `--cache-config` is required. + +Override individual defaults with a JSON cache configuration; for example: + +```bash +vllm serve nvidia/Cosmos3-Nano \ + --omni \ + --cache-backend sea_cache \ + --cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}' +``` + +Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are +more conservative. Higher values allow more cached steps and may improve +speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to +`0` removes the streak cap. + To run **without** guardrails (you are responsible for license compliance), add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`; @@ -382,12 +401,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's #### Notes -- **Measured latency (1x B300, bf16, guardrails off):** - - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s** - - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s** - - I2V 1280×720, 189 frames @ 35 steps → ~**99 s** - - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**. - - Guardrails-on overhead: ~8% on T2I, negligible on video. +- **Measured latency (1x B300, bf16, guardrails off, diffusion cache off):** + - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s** + - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s** + - I2V 1280×720, 189 frames @ 35 steps → ~**99 s** + - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**. + - Guardrails-on overhead: ~8% on T2I, negligible on video. - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU; full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer) ~33 GB on disk. @@ -443,12 +462,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's `conditioning_fps`, `action_chunk_size`, `raw_action_dim`, `deterministic_seed`, and `session_id`. - **Known limitations:** - - Guardrails-on requires `cosmos-guardrail` **and** access to the gated + - Guardrails-on requires `cosmos-guardrail` **and** access to the gated `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise the server fails at pipeline build with a gated-repo / safety-checker error. - - A guardrail-blocked prompt currently returns HTTP 500 + - A guardrail-blocked prompt currently returns HTTP 500 (`"Guardrail blocked prompt"`). - - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported + - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported online. Use async `POST /v1/videos` when you need the predicted/recovered action payload under the top-level `action` field; sync `/v1/videos/sync` returns raw MP4 bytes and does not expose action metadata in the response body. @@ -579,14 +598,14 @@ so an md5 comparison across two runs also works as a smoke check. a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is **540 s cold vs 161 s warm, with byte-identical output**. -- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps +- **Measured on 1x MI350X (bf16, guardrails off, diffusion cache off, warm):** T2I 1024² @ 50 steps **~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the denoise loop. The optional flags act on different terms of the memory bill and are therefore complementary rather than redundant: | Flag | Latency | Peak reserved | Peak allocated | Acts on | - |---|---|---|---|---| + | --- | --- | --- | --- | --- | | *(none)* | 161 s | 120 GiB | 95 GiB | — | | `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations | | `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights | @@ -768,11 +787,11 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ #### Notes -- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):** - - T2I 1024² — 10 steps → ~8 s - - T2V 1280×720 @ 20 steps — 49 frames → ~55 s - - I2V 1280×720 @ 10 steps — 25 frames → ~25 s - - V2V 480×320 @ 10 steps — 17 frames → ~12 s +- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off, diffusion cache off):** + - T2I 1024² — 10 steps → ~8 s + - T2V 1280×720 @ 20 steps — 49 frames → ~55 s + - I2V 1280×720 @ 10 steps — 25 frames → ~25 s + - V2V 480×320 @ 10 steps — 17 frames → ~12 s - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU; full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer) ~33 GB on disk. @@ -785,9 +804,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ (for model loading), `--tensor-parallel-size 8` for multi-NPU, and `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class. - **Known limitations:** - - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a + - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a resolution-parsing bug; basic V2V without transfer hints works. - - FP8 online quantization and layerwise offload are not supported on NPU. + - FP8 online quantization and layerwise offload are not supported on NPU. ### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index c8bfb6bbd..82ae1f20a 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -62,6 +62,21 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs; generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V composition can shift at the same seed). +SeaCache is the recommended default choice when opting into diffusion caching, +which remains opt-in. Add `--cache-backend sea_cache` to either command above; +its defaults are tuned for Cosmos3. + +Override individual defaults with `--cache-config`; for example: + +```bash +--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}' +``` + +Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are +more conservative. Higher values allow more cached steps and may improve +speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to +`0` removes the streak cap. + #### Verification Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full @@ -109,12 +124,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ #### Notes -- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):** - - T2I 1024², 50 steps → **~6 s** - - T2V 1280×720, 189 frames, 35 steps → **~197 s** - - I2V 1280×720, 189 frames, 35 steps → **~200 s** - - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo** - - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.) +- **Measured (2x B300, bf16, guardrails off, diffusion cache off, official 2-GPU config above):** + - T2I 1024², 50 steps → **~6 s** + - T2V 1280×720, 189 frames, 35 steps → **~197 s** + - I2V 1280×720, 189 frames, 35 steps → **~200 s** + - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo** + - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.) - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk. - Same generation defaults, supported sizes, V2V reference-video controls (`condition_frame_indexes_vision`, `condition_video_keep`), and diff --git a/tests/diffusion/cache/test_seacache.py b/tests/diffusion/cache/test_seacache.py new file mode 100644 index 000000000..ea11537d1 --- /dev/null +++ b/tests/diffusion/cache/test_seacache.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import datetime +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers.runtime import get_distributed_init_method +from vllm_omni.diffusion.cache.seacache import ( + SeaCacheBackend, + SeaCacheConfig, + SeaCacheRootHook, + apply_sea_cache_hook, +) +from vllm_omni.diffusion.cache.seacache.sea_filter import ( + apply_sea_filter, + extrapolate_residual, + indicator_distance, +) +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext, get_extractor +from vllm_omni.diffusion.data import DiffusionCacheConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TinyCosmos3Transformer(torch.nn.Module): + """Small model that exposes a cache-neutral execution boundary.""" + + def _run_gen_layers(self, hidden_gen: torch.Tensor) -> torch.Tensor: + return hidden_gen * 2 + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor | None = None, + text_mask: torch.Tensor | None = None, + video_shape: tuple[int, int, int] | None = None, + noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + del timestep, text_ids, text_mask, video_shape, noisy_frame_mask + if kwargs: + raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}") + controls = ( + [] + if control_latents is None + else [control_latents] + if isinstance(control_latents, torch.Tensor) + else list(control_latents) + ) + inputs = [*controls, hidden_states] + gen_input = torch.cat( + [value.movedim(1, -1).flatten(1, 3) for value in inputs], + dim=1, + ) + return self._run_gen_layers(gen_input) + + +def _extract_tiny_cosmos3_context( + module: TinyCosmos3Transformer, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor | None = None, + text_mask: torch.Tensor | None = None, + video_shape: tuple[int, int, int] | None = None, + noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | torch.Tensor | None = None, + **kwargs, +) -> CacheContext: + del timestep, text_ids, text_mask, video_shape + if kwargs: + raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}") + controls = ( + [] + if control_latents is None + else [control_latents] + if isinstance(control_latents, torch.Tensor) + else list(control_latents) + ) + vision_items = [*controls, hidden_states] + gen_input = torch.cat( + [value.movedim(1, -1).flatten(1, 3) for value in vision_items], + dim=1, + ) + return CacheContext( + modulated_input=gen_input, + hidden_states=gen_input, + encoder_hidden_states=None, + temb=torch.zeros_like(gen_input[:, 0]), + run_transformer_blocks=lambda: (module._run_gen_layers(gen_input),), + postprocess=lambda output: output, + extra_states={ + "sea_cache_latents": vision_items, + "sea_cache_noisy_frame_mask": noisy_frame_mask, + }, + ) + + +class Cosmos3OmniDiffusersPipeline: + def __init__(self) -> None: + self.transformer = TinyCosmos3Transformer() + self._current_step_index: int | None = None + self._current_sigma: float | None = None + self._num_timesteps: int | None = None + + @property + def current_step_index(self) -> int | None: + return self._current_step_index + + @property + def current_sigma(self) -> float | None: + return self._current_sigma + + @property + def num_timesteps(self) -> int | None: + return self._num_timesteps + + +def _latent(value: float) -> torch.Tensor: + return torch.full((1, 2, 2, 2, 2), value) + + +def _run_step( + transformer: TinyCosmos3Transformer, + timestep: int, + value: float, + *, + hook: SeaCacheRootHook | None = None, + context: str = "cond", + control: float | None = None, + noisy_frame_mask: torch.Tensor | None = None, +) -> torch.Tensor: + controls = None if control is None else [_latent(control)] + cache_context = hook.cache_context(context) if hook is not None else nullcontext() + with torch.inference_mode(), cache_context: + return transformer( + hidden_states=_latent(value), + timestep=torch.tensor([timestep]), + noisy_frame_mask=noisy_frame_mask, + control_latents=controls, + ) + + +def _apply_test_hook( + transformer: TinyCosmos3Transformer, + metadata: SimpleNamespace, + config: SeaCacheConfig | None = None, +) -> SeaCacheRootHook: + return apply_sea_cache_hook( + transformer, + config or SeaCacheConfig(threshold=100.0), + current_step_callback=lambda: metadata.step, + current_sigma_callback=lambda: metadata.sigma, + num_inference_steps_callback=lambda: metadata.num_steps, + extractor_fn=_extract_tiny_cosmos3_context, + ) + + +def test_config_validation() -> None: + assert SeaCacheConfig().threshold == 0.25 + with pytest.raises(ValueError, match="residual_order"): + SeaCacheConfig(residual_order=-1) + with pytest.raises(ValueError, match="max_consecutive_cached"): + SeaCacheConfig(max_consecutive_cached=-1) + + +def test_sea_filter_matches_reference_equation() -> None: + hidden = torch.randn(3, 4, 5, 2, dtype=torch.float32) + sigma = 0.4 + power_exp = 3.0 + + spectrum = torch.fft.fftn(hidden, dim=(0, 1, 2)) + gain = None + for axis in (0, 1, 2): + frequencies = torch.fft.fftfreq(hidden.shape[axis], dtype=torch.float32) + clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16) + axis_gain = (1.0 - sigma) * clean_power / ((1.0 - sigma) ** 2 * clean_power + sigma**2 + 1e-16) + shape = [1] * hidden.ndim + shape[axis] = hidden.shape[axis] + gain = axis_gain.reshape(shape) if gain is None else gain * axis_gain.reshape(shape) + assert gain is not None + gain = gain / gain.mean() + expected = torch.fft.ifftn(spectrum * gain, dim=(0, 1, 2)).real + + torch.testing.assert_close( + apply_sea_filter(hidden, sigma=sigma, power_exp=power_exp), + expected, + ) + + +def test_indicator_distance_and_linear_extrapolation() -> None: + previous = [torch.ones(2, 2)] + current = [torch.full((2, 2), 1.5)] + assert indicator_distance(current, previous) == pytest.approx(0.5) + assert indicator_distance([torch.ones(3)], previous) == float("inf") + + history = [ + (0, torch.full((2, 2), 2.0)), + (2, torch.full((2, 2), 6.0)), + ] + torch.testing.assert_close( + extrapolate_residual(history, step=3, order=1), + torch.full((2, 2), 8.0), + ) + torch.testing.assert_close( + extrapolate_residual(history, step=3, order=0), + torch.full((2, 2), 6.0), + ) + quadratic_history = [ + (0, torch.zeros(2, 2)), + (1, torch.ones(2, 2)), + (2, torch.full((2, 2), 4.0)), + ] + torch.testing.assert_close( + extrapolate_residual(quadratic_history, step=3, order=2), + torch.full((2, 2), 9.0), + ) + + +def test_hook_skips_middle_steps_and_forces_endpoints() -> None: + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=4) + hook = _apply_test_hook( + transformer, + metadata, + SeaCacheConfig(threshold=100.0, max_consecutive_cached=2), + ) + hook.refresh(transformer) + + for step, timestep in enumerate((1000, 750, 500, 250)): + metadata.step = step + metadata.sigma = timestep / 1000 + _run_step(transformer, timestep, 1.0 - step * 0.01, hook=hook) + + assert hook.full_count == 2 + assert hook.skip_count == 2 + assert [step for step, _ in hook.state_manager._states["cond"].history] == [0, 3] + + +def test_parameter_sharded_hook_skips_when_all_shard_ranks_agree(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.distributed import parallel_state + + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3) + hook = _apply_test_hook(transformer, metadata) + hook._parameter_sharded = True + fs_group = object() + reduced_groups: list[object] = [] + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr( + parallel_state, + "get_fs_group", + lambda: SimpleNamespace(world_size=2, device_group=fs_group), + ) + monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1) + + def all_reduce(decision, *, op, group): + assert op == torch.distributed.ReduceOp.MAX + reduced_groups.append(group) + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + _run_step(transformer, 1000, 1.0, hook=hook) + metadata.step = 1 + metadata.sigma = 0.5 + _run_step(transformer, 500, 0.99, hook=hook) + + assert hook.full_count == 1 + assert hook.skip_count == 1 + assert reduced_groups == [fs_group, fs_group] + + +def test_parameter_sharded_peer_forces_full_compute(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.distributed import parallel_state + + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3) + hook = _apply_test_hook(transformer, metadata) + hook._parameter_sharded = True + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr( + parallel_state, + "get_fs_group", + lambda: SimpleNamespace(world_size=2, device_group=object()), + ) + monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1) + reduce_count = 0 + + def all_reduce(decision, *, op, group): + nonlocal reduce_count + del op, group + reduce_count += 1 + if reduce_count == 2: + decision.fill_(1) + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + _run_step(transformer, 1000, 1.0, hook=hook) + metadata.step = 1 + metadata.sigma = 0.5 + _run_step(transformer, 500, 0.99, hook=hook) + + assert hook.full_count == 2 + assert hook.skip_count == 0 + assert hook.state_manager._states["cond"].accumulated_distance == 0.0 + + +def _uneven_cfg_sharded_worker( + rank: int, + init_method: str, + result_queue: torch.multiprocessing.Queue, +) -> None: + from vllm_omni.diffusion.distributed import parallel_state + + torch.distributed.init_process_group( + "gloo", + init_method=init_method, + rank=rank, + world_size=4, + timeout=datetime.timedelta(seconds=10), + ) + try: + fs_groups = [ + (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10))) + for ranks in ([0, 1], [2, 3]) + ] + sp_groups = [ + (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10))) + for ranks in ([0, 1], [2, 3]) + ] + cfg_groups = [ + (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10))) + for ranks in ([0, 2], [1, 3]) + ] + fs_group = next(group for ranks, group in fs_groups if rank in ranks) + sp_group = next(group for ranks, group in sp_groups if rank in ranks) + cfg_group = next(group for ranks, group in cfg_groups if rank in ranks) + parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group) + parallel_state._SP = SimpleNamespace(world_size=2, device_group=sp_group) + + hook = SeaCacheRootHook(SeaCacheConfig()) + hook._parameter_sharded = True + hook._synchronize_compute(False, torch.device("cpu")) + if rank in (0, 1): + hook._synchronize_compute(False, torch.device("cpu")) + + gathered = [torch.zeros(1) for _ in range(2)] + torch.distributed.all_gather( + gathered, + torch.tensor([rank], dtype=torch.float32), + group=cfg_group, + ) + result_queue.put((rank, [int(value.item()) for value in gathered])) + finally: + parallel_state._FS = None + parallel_state._SP = None + torch.distributed.destroy_process_group() + + +def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None: + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + result_queue = manager.Queue() + torch.multiprocessing.spawn( + _uneven_cfg_sharded_worker, + args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue), + nprocs=4, + ) + + results = sorted(result_queue.get(timeout=2) for _ in range(4)) + assert results == [ + (0, [0, 2]), + (1, [1, 3]), + (2, [0, 2]), + (3, [1, 3]), + ] + + +def _hybrid_sp_worker( + rank: int, + init_method: str, + result_queue: torch.multiprocessing.Queue, +) -> None: + from vllm_omni.diffusion.distributed import parallel_state + + torch.distributed.init_process_group( + "gloo", + init_method=init_method, + rank=rank, + world_size=4, + timeout=datetime.timedelta(seconds=10), + ) + try: + fs_groups = [ + (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10))) + for ranks in ([0, 1], [2, 3]) + ] + ulysses_groups = [ + (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10))) + for ranks in ([0, 1], [2, 3]) + ] + ring_groups = [ + (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10))) + for ranks in ([0, 2], [1, 3]) + ] + sp_group = torch.distributed.new_group( + [0, 1, 2, 3], + timeout=datetime.timedelta(seconds=10), + ) + fs_group = next(group for ranks, group in fs_groups if rank in ranks) + ulysses_group = next(group for ranks, group in ulysses_groups if rank in ranks) + ring_group = next(group for ranks, group in ring_groups if rank in ranks) + parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group) + parallel_state._SP = SimpleNamespace( + world_size=4, + device_group=sp_group, + ulysses_world_size=2, + ulysses_group=ulysses_group, + ring_world_size=2, + ring_group=ring_group, + ) + + compute = rank == 3 + non_sharded_hook = SeaCacheRootHook(SeaCacheConfig()) + non_sharded_result = non_sharded_hook._synchronize_compute(compute, torch.device("cpu")) + + sharded_hook = SeaCacheRootHook(SeaCacheConfig()) + sharded_hook._parameter_sharded = True + sharded_result = sharded_hook._synchronize_compute(compute, torch.device("cpu")) + result_queue.put((rank, non_sharded_result, sharded_result)) + finally: + parallel_state._FS = None + parallel_state._SP = None + torch.distributed.destroy_process_group() + + +def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None: + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + result_queue = manager.Queue() + torch.multiprocessing.spawn( + _hybrid_sp_worker, + args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue), + nprocs=4, + ) + + results = sorted(result_queue.get(timeout=2) for _ in range(4)) + assert results == [ + (0, True, True), + (1, True, True), + (2, True, True), + (3, True, True), + ] + + +def test_parameter_sharded_hook_fails_open_without_distributed_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hook = SeaCacheRootHook(SeaCacheConfig()) + hook._parameter_sharded = True + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + + assert hook._synchronize_compute(False, torch.device("cpu")) is True + + +def test_hook_keeps_three_transfer_branches_separate() -> None: + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3) + hook = _apply_test_hook(transformer, metadata) + hook.refresh(transformer) + + contexts = ("cond", "cond_no_control", "uncond") + for step, (timestep, value) in enumerate(((1000, 1.0), (500, 0.9))): + metadata.step = step + metadata.sigma = timestep / 1000 + _run_step(transformer, timestep, value, hook=hook, context=contexts[0], control=0.5) + _run_step(transformer, timestep, value, hook=hook, context=contexts[1]) + _run_step(transformer, timestep, value, hook=hook, context=contexts[2], control=0.5) + + assert set(hook.state_manager._states) == set(contexts) + assert hook.full_count == 3 + assert hook.skip_count == 3 + assert len(hook.state_manager._states["cond"].previous_indicator) == 2 + assert len(hook.state_manager._states["cond_no_control"].previous_indicator) == 1 + + +def test_hook_fails_open_without_noisy_vision() -> None: + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3) + hook = _apply_test_hook(transformer, metadata) + hook.refresh(transformer) + all_clean = torch.zeros(1, 1, 2, 1, 1) + + _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean) + metadata.step = 1 + metadata.sigma = 0.5 + _run_step(transformer, 500, 0.9, hook=hook, noisy_frame_mask=all_clean) + + assert hook.full_count == 0 + assert hook.skip_count == 0 + + +def test_conditioning_only_model_failure_runs_forward_once(monkeypatch: pytest.MonkeyPatch) -> None: + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3) + hook = _apply_test_hook(transformer, metadata) + hook.refresh(transformer) + forward_calls = 0 + + def fail_forward(hidden_gen: torch.Tensor) -> torch.Tensor: + nonlocal forward_calls + forward_calls += 1 + raise RuntimeError("model forward failed") + + monkeypatch.setattr(transformer, "_run_gen_layers", fail_forward) + all_clean = torch.zeros(1, 1, 2, 1, 1) + + with pytest.raises(RuntimeError, match="model forward failed"): + _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean) + + assert forward_calls == 1 + + +def test_hook_fails_open_without_explicit_context() -> None: + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=2) + hook = _apply_test_hook(transformer, metadata) + hook.refresh(transformer) + + _run_step(transformer, 1000, 1.0) + + assert hook.full_count == 0 + assert hook.skip_count == 0 + assert hook.state_manager._states == {} + + +def test_hook_uses_exact_sigma_callback(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.cache.seacache import hook as hook_module + + transformer = TinyCosmos3Transformer() + metadata = SimpleNamespace(step=0, sigma=0.37, num_steps=2) + observed_sigmas: list[float] = [] + original_filter = hook_module.apply_sea_filter + + def recording_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float) -> torch.Tensor: + observed_sigmas.append(sigma) + return original_filter(hidden_states, sigma, power_exp) + + monkeypatch.setattr(hook_module, "apply_sea_filter", recording_filter) + hook = _apply_test_hook(transformer, metadata) + hook.refresh(transformer) + _run_step(transformer, timestep=999, value=1.0, hook=hook) + + assert observed_sigmas + assert all(sigma == pytest.approx(0.37) for sigma in observed_sigmas) + + +def test_backend_selector_and_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context) + + class FSDPTinyCosmos3Transformer(TinyCosmos3Transformer): + pass + + assert get_extractor(FSDPTinyCosmos3Transformer) is _extract_tiny_cosmos3_context + backend = get_cache_backend( + "sea_cache", + { + "sea_threshold": 0.4, + "sea_residual_order": 0, + }, + ) + assert isinstance(backend, SeaCacheBackend) + assert backend.config.sea_threshold == 0.4 + + pipeline = Cosmos3OmniDiffusersPipeline() + backend.enable(pipeline) + hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME) + assert isinstance(hook, SeaCacheRootHook) + assert callable(getattr(pipeline, "_cache_context_factory", None)) + for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"): + assert not hasattr(pipeline.transformer, name) + pipeline._current_step_index = 0 + pipeline._current_sigma = 1.0 + pipeline._num_timesteps = 7 + _run_step(pipeline.transformer, 1000, 1.0, hook=hook) + assert hook.full_count == 1 + + backend.refresh(pipeline, num_inference_steps=7) + assert hook.full_count == 0 + assert hook.skip_count == 0 + assert hook.state_manager._states == {} + + +def test_backend_uses_resolved_pipeline_metadata_not_refresh_argument(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context) + pipeline = Cosmos3OmniDiffusersPipeline() + backend = SeaCacheBackend(DiffusionCacheConfig()) + backend.enable(pipeline) + hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME) + assert isinstance(hook, SeaCacheRootHook) + backend.refresh(pipeline, num_inference_steps=35) + pipeline._num_timesteps = 50 + pipeline._current_step_index = 0 + pipeline._current_sigma = 1.0 + _run_step(pipeline.transformer, 1000, 1.0, hook=hook) + assert hook.full_count == 1 + assert hook.num_inference_steps_callback is not None + assert hook.num_inference_steps_callback() == 50 + + pipeline._current_step_index = 1 + pipeline._current_sigma = 0.98 + _run_step(pipeline.transformer, 980, 0.99, hook=hook) + assert hook.full_count == 1 + assert hook.skip_count == 1 + + +def test_shared_config_defaults() -> None: + config = DiffusionCacheConfig() + assert config.sea_threshold == 0.25 + assert config.sea_residual_order == 1 + assert config.sea_max_consecutive_cached == 2 + assert config.sea_power_exp == 3.0 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 9892a72d7..e73dd7ff6 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import json import sys import types +from contextlib import contextmanager from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -85,6 +86,7 @@ class StubScheduler: flow_shift: float = 1.0, ) -> None: self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64) + self.sigmas = self.timesteps.float() / 10 self.config = SimpleNamespace( num_train_timesteps=1000, flow_shift=flow_shift, @@ -113,6 +115,7 @@ class StubScheduler: else: assert num_inference_steps is not None self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device) + self.sigmas = self.timesteps.float() def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs): del kwargs @@ -250,7 +253,7 @@ def passthrough_progress_bar(iterable): @pytest.fixture(autouse=True) def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch): - module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails") + module: Any = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails") module.is_guardrails_enabled = lambda od_config, sampling_params=None: False module.ensure_initialized = lambda od_config: None module.check_text_safety = lambda text: None @@ -283,6 +286,8 @@ def make_cosmos3_pipeline(): pipeline.is_edge_model = False pipeline._guidance_scale = None pipeline._num_timesteps = None + pipeline._current_step_index = None + pipeline._current_sigma = None pipeline._cosmos3_branch_caches = None pipeline._cache_dit_requires_paired_cfg = False pipeline._sound_tokenizer = None @@ -300,7 +305,7 @@ def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None: def make_sampling_params(**overrides: Any) -> SimpleNamespace: - values = { + values: dict[str, Any] = { "height": None, "width": None, "num_frames": None, @@ -814,7 +819,7 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file( t_list = [1.0, 0.75, 0.5, 0.25] scheduler_dir = tmp_path / "scheduler" scheduler_dir.mkdir() - scheduler_config = {"_class_name": scheduler_class_name} + scheduler_config: dict[str, Any] = {"_class_name": scheduler_class_name} if expected_distilled: scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list} (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config)) @@ -1531,7 +1536,7 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke fused_add_rms_norm: list[str] custom_op: list[str] - fake_kernel = types.ModuleType("vllm.config.kernel") + fake_kernel: Any = types.ModuleType("vllm.config.kernel") fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel) @@ -2154,6 +2159,47 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result)) +def test_diffuse_publishes_exact_seacache_metadata_and_cfg_contexts(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.scheduler.sigmas = torch.tensor([0.91, 0.17]) + observations: list[tuple[str, int | None, float | None, int | None]] = [] + + class RecordingHook: + @contextmanager + def cache_context(self, name: str): + sigma = pipeline.current_sigma + observations.append( + ( + name, + pipeline.current_step_index, + None if sigma is None else float(sigma), + pipeline.num_timesteps, + ) + ) + yield + + pipeline._cache_context_factory = RecordingHook().cache_context + pipeline.diffuse( + latents=torch.zeros(1, 2, 1, 1, 1), + timesteps=torch.tensor([900, 100]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0}, + ) + + assert observations == [ + ("cond", 0, pytest.approx(0.91), 2), + ("uncond", 0, pytest.approx(0.91), 2), + ("cond", 1, pytest.approx(0.17), 2), + ("uncond", 1, pytest.approx(0.17), 2), + ] + assert pipeline.current_step_index is None + assert pipeline.current_sigma is None + + def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None: pipeline = make_cosmos3_pipeline() pipeline._use_session_state = True @@ -2213,6 +2259,53 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_ torch.testing.assert_close(result, torch.full_like(latents, 254.0)) +def test_diffuse_transfer_uses_named_seacache_contexts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.scheduler.sigmas = torch.tensor([0.42]) + contexts: list[tuple[str, int | None, float | None, int | None]] = [] + + class RecordingHook: + @contextmanager + def cache_context(self, name: str): + sigma = pipeline.current_sigma + contexts.append( + ( + name, + pipeline.current_step_index, + None if sigma is None else float(sigma), + pipeline.num_timesteps, + ) + ) + yield + + pipeline._cache_context_factory = RecordingHook().cache_context + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + pipeline.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.5, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + + assert contexts == [ + ("cond", 0, pytest.approx(0.42), 1), + ("cond_no_control", 0, pytest.approx(0.42), 1), + ("uncond", 0, pytest.approx(0.42), 1), + ] + assert pipeline.current_step_index is None + assert pipeline.current_sigma is None + + def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None: pipeline = make_cosmos3_pipeline() pipeline._use_session_state = True @@ -2583,7 +2676,7 @@ def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) - class TestForwardRouting: def _install_forward_stubs(self, pipeline): - captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []} + captured: dict[str, Any] = {"diffuse_calls": [], "prepare_calls": []} def fake_format( prompt, diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 44bd8e9a7..de14ee233 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -513,6 +513,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No assert tuple(output.shape) == (1, 2, 1, 2, 2) +def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + class TrackingNorm(nn.Module): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.calls += 1 + return hidden_states + 5.0 + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + norm = TrackingNorm() + model.norm_moe_gen = norm + captured: dict[str, torch.Tensor] = {} + + def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor: + del kwargs + captured["input"] = hidden_gen.detach().clone() + return hidden_gen + 2.0 + + monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers) + forward_kwargs = { + "hidden_states": torch.zeros(1, 2, 1, 2, 2), + "timestep": torch.tensor([1.0]), + "text_ids": torch.tensor([[1, 2]], dtype=torch.long), + "text_mask": torch.ones(1, 2, dtype=torch.long), + "video_shape": (1, 2, 2), + "fps": 24.0, + } + + full_output = model(**forward_kwargs) + + assert norm.calls == 1 + norm.calls = 0 + + ctx = extract_cosmos3_context(model, **forward_kwargs) + execution_input = ctx.hidden_states.detach().clone() + execution_output = ctx.run_transformer_blocks()[0] + residual = execution_output - execution_input + extracted_output = ctx.postprocess(execution_output) + + assert norm.calls == 1 + expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"] + assert torch.equal(residual, expected_residual) + assert torch.equal(extracted_output, full_output) + + norm.calls = 0 + + def fail_if_gen_layers_run(*args, **kwargs): + del args, kwargs + raise AssertionError("SeaCache hit unexpectedly executed GEN layers") + + monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run) + + cached_ctx = extract_cosmos3_context(model, **forward_kwargs) + cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual) + + assert norm.calls == 0 + assert torch.equal(cached_output, full_output) + for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"): + assert not hasattr(model, name) + + +def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + class CountingNorm(nn.Module): + def __init__(self, original: nn.Module) -> None: + super().__init__() + self.original = original + self.calls = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.calls += 1 + return self.original(hidden_states) + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + norm = CountingNorm(model.norm_moe_gen) + model.norm_moe_gen = norm + + output = model( + hidden_states=torch.zeros(1, 2, 1, 2, 2), + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + ) + + assert norm.calls == 1 + assert tuple(output.shape) == (1, 2, 1, 2, 2) + + def test_qwen_gen_mlp_remains_gated() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer diff --git a/vllm_omni/diffusion/cache/seacache/__init__.py b/vllm_omni/diffusion/cache/seacache/__init__.py new file mode 100644 index 000000000..2821c4de0 --- /dev/null +++ b/vllm_omni/diffusion/cache/seacache/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from vllm_omni.diffusion.cache.seacache.backend import SeaCacheBackend +from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig +from vllm_omni.diffusion.cache.seacache.hook import ( + SeaCacheRootHook, + apply_sea_cache_hook, +) +from vllm_omni.diffusion.cache.seacache.state import SeaCacheState + +__all__ = [ + "SeaCacheBackend", + "SeaCacheConfig", + "SeaCacheRootHook", + "SeaCacheState", + "apply_sea_cache_hook", +] diff --git a/vllm_omni/diffusion/cache/seacache/backend.py b/vllm_omni/diffusion/cache/seacache/backend.py new file mode 100644 index 000000000..e5e513a73 --- /dev/null +++ b/vllm_omni/diffusion/cache/seacache/backend.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig +from vllm_omni.diffusion.cache.seacache.hook import ( + SeaCacheRootHook, + apply_sea_cache_hook, +) +from vllm_omni.diffusion.data import DiffusionCacheConfig + +logger = init_logger(__name__) + + +def enable_cosmos3_seacache( + pipeline: Any, + config: DiffusionCacheConfig, +) -> SeaCacheRootHook: + transformer = getattr(pipeline, "transformer", None) + if transformer is None: + raise ValueError("SeaCache requires a pipeline with a transformer") + + sea_config = SeaCacheConfig( + threshold=config.sea_threshold, + residual_order=config.sea_residual_order, + max_consecutive_cached=config.sea_max_consecutive_cached, + power_exp=config.sea_power_exp, + ) + hook = apply_sea_cache_hook( + transformer, + sea_config, + current_step_callback=lambda: getattr(pipeline, "current_step_index", None), + current_sigma_callback=lambda: getattr(pipeline, "current_sigma", None), + num_inference_steps_callback=lambda: getattr( + pipeline, + "num_timesteps", + None, + ), + ) + logger.info( + "SeaCache enabled for %s (threshold=%s, residual_order=%d, max_consecutive_cached=%d, power_exp=%s)", + pipeline.__class__.__name__, + sea_config.threshold, + sea_config.residual_order, + sea_config.max_consecutive_cached, + sea_config.power_exp, + ) + return hook + + +CUSTOM_SEACACHE_ENABLERS = { + "Cosmos3OmniDiffusersPipeline": enable_cosmos3_seacache, + "Cosmos3OmniPipeline": enable_cosmos3_seacache, +} + + +class SeaCacheBackend(CacheBackend): + """Backend for spectral-evolution-aware diffusion caching.""" + + def __init__(self, config: DiffusionCacheConfig): + super().__init__(config) + self._transformer_id: int | None = None + + def enable(self, pipeline: Any) -> None: + pipeline_type = pipeline.__class__.__name__ + enabler = CUSTOM_SEACACHE_ENABLERS.get(pipeline_type) + if enabler is None: + raise ValueError(f"SeaCache does not support pipeline type {pipeline_type}") + hook = enabler(pipeline, self.config) + self._transformer_id = id(pipeline.transformer) + self.enabled = True + pipeline._cache_context_factory = hook.cache_context + + def refresh( + self, + pipeline: Any, + num_inference_steps: int, + verbose: bool = True, + ) -> None: + del num_inference_steps + transformer = getattr(pipeline, "transformer", None) + if transformer is None: + raise ValueError("SeaCache requires a pipeline with a transformer") + if not self.enabled or self._transformer_id != id(transformer): + self.enable(pipeline) + + registry = getattr(transformer, "_hook_registry", None) + hook = registry.get_hook(SeaCacheRootHook._HOOK_NAME) if registry is not None else None + if not isinstance(hook, SeaCacheRootHook): + raise RuntimeError("SeaCache hook is not installed on the pipeline transformer") + hook.refresh(transformer) + pipeline._cache_context_factory = hook.cache_context + if verbose: + logger.debug("SeaCache state refreshed") diff --git a/vllm_omni/diffusion/cache/seacache/config.py b/vllm_omni/diffusion/cache/seacache/config.py new file mode 100644 index 000000000..501d4653e --- /dev/null +++ b/vllm_omni/diffusion/cache/seacache/config.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SeaCacheConfig: + """Configuration for SeaCache. + + Defaults are tuned for Cosmos3 and may require adjustment for other models. + """ + + threshold: float = 0.25 + residual_order: int = 1 + max_consecutive_cached: int = 2 + power_exp: float = 3.0 + + def __post_init__(self) -> None: + if not math.isfinite(self.threshold) or self.threshold < 0: + raise ValueError(f"threshold must be finite and non-negative, got {self.threshold}") + if isinstance(self.residual_order, bool) or not isinstance(self.residual_order, int) or self.residual_order < 0: + raise ValueError(f"residual_order must be a non-negative integer, got {self.residual_order!r}") + if ( + isinstance(self.max_consecutive_cached, bool) + or not isinstance(self.max_consecutive_cached, int) + or self.max_consecutive_cached < 0 + ): + raise ValueError( + f"max_consecutive_cached must be a non-negative integer, got {self.max_consecutive_cached!r}" + ) + if not math.isfinite(self.power_exp) or self.power_exp <= 0: + raise ValueError(f"power_exp must be finite and positive, got {self.power_exp}") diff --git a/vllm_omni/diffusion/cache/seacache/hook.py b/vllm_omni/diffusion/cache/seacache/hook.py new file mode 100644 index 000000000..3fac58c09 --- /dev/null +++ b/vllm_omni/diffusion/cache/seacache/hook.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig +from vllm_omni.diffusion.cache.seacache.sea_filter import ( + apply_sea_filter, + extrapolate_residual, + indicator_distance, +) +from vllm_omni.diffusion.cache.seacache.state import SeaCacheState +from vllm_omni.diffusion.cache.teacache.extractors import CacheContext, get_extractor +from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager + +logger = init_logger(__name__) + + +def _is_parameter_sharded(module: torch.nn.Module) -> bool: + """Detect parameter-sharding runtimes whose collectives cannot be skipped.""" + for submodule in module.modules(): + module_type = type(submodule) + if callable(getattr(submodule, "_get_fsdp_state", None)): + return True + if module_type.__name__ == "FullyShardedDataParallel" and module_type.__module__.startswith( + "torch.distributed.fsdp" + ): + return True + for parameter in submodule.parameters(recurse=False): + parameter_type = type(parameter) + if ( + parameter_type.__name__ == "FlatParameter" + and parameter_type.__module__.startswith("torch.distributed.fsdp") + ) or ( + parameter_type.__name__ == "DTensor" + and parameter_type.__module__.startswith("torch.distributed.tensor") + ): + return True + return False + + +class SeaCacheRootHook(ModelHook): + """Drive SeaCache gating and transformer forward control.""" + + _HOOK_NAME = "sea_cache" + + def __init__( + self, + config: SeaCacheConfig, + *, + current_step_callback: Callable[[], int | torch.Tensor | None] | None = None, + current_sigma_callback: Callable[[], float | torch.Tensor | None] | None = None, + num_inference_steps_callback: Callable[[], int | torch.Tensor | None] | None = None, + extractor_fn: Callable[..., CacheContext] | None = None, + ) -> None: + super().__init__() + self.config = config + self.current_step_callback = current_step_callback + self.current_sigma_callback = current_sigma_callback + self.num_inference_steps_callback = num_inference_steps_callback + self.state_manager = StateManager(SeaCacheState) + self._warned_messages: set[str] = set() + self.full_count = 0 + self.skip_count = 0 + self.extractor_fn = extractor_fn + self._parameter_sharded = False + self._collective_skip_groups: list[torch.distributed.ProcessGroup] = [] + + def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + if self.extractor_fn is None: + self.extractor_fn = get_extractor(type(module)) + self._parameter_sharded = _is_parameter_sharded(module) + seen_groups: set[int] = set() + for block in getattr(module, "gen_layers", ()): + registry = getattr(block, "_hook_registry", None) + dlo_hook = registry.get_hook("distributed_layerwise_offload") if registry is not None else None + group = getattr(dlo_hook, "dp_group", None) + if group is not None and int(getattr(dlo_hook, "dp_size", 1)) > 1 and id(group) not in seen_groups: + seen_groups.add(id(group)) + self._collective_skip_groups.append(group) + return module + + def _warn_once(self, message: str) -> None: + if message not in self._warned_messages: + logger.warning(message) + self._warned_messages.add(message) + + @contextmanager + def cache_context(self, name: str) -> Iterator[None]: + previous_context = self.state_manager._context + self.state_manager.set_context(name) + try: + yield + finally: + self.state_manager.set_context(previous_context) + + def _build_indicator( + self, + vision_items: list[torch.Tensor] | None, + sigma: float, + ) -> list[torch.Tensor] | None: + if not vision_items: + return None + hidden_states = vision_items[-1] + if not isinstance(hidden_states, torch.Tensor) or hidden_states.ndim != 5: + return None + + # Controls precede the denoised target in packed vision-token order, + # and each vision item is filtered independently. + if any( + not isinstance(item, torch.Tensor) + or item.ndim != 5 + or item.shape[0] != hidden_states.shape[0] + or item.shape[1:] != hidden_states.shape[1:] + for item in vision_items + ): + return None + + indicator = [] + for batch_index in range(hidden_states.shape[0]): + for latent in vision_items: + thwc = latent[batch_index].movedim(0, -1) + indicator.append( + apply_sea_filter( + thwc, + sigma=sigma, + power_exp=self.config.power_exp, + ).detach() + ) + return indicator or None + + def _resolve_gate( + self, + state: SeaCacheState, + indicator: list[torch.Tensor] | None, + step: int, + num_inference_steps: int, + ) -> bool: + if state.last_step is not None and step != state.last_step + 1: + state.reset() + state.last_step = step + max_consecutive = bool( + self.config.max_consecutive_cached and state.consecutive_cached >= self.config.max_consecutive_cached + ) + forced_compute = ( + step < 1 + or step >= num_inference_steps - 1 + or max_consecutive + or not state.history + or indicator is None + or state.previous_indicator is None + ) + if forced_compute: + state.accumulated_distance = 0.0 + state.previous_indicator = None if indicator is None else [value.detach() for value in indicator] + return True + + assert indicator is not None + assert state.previous_indicator is not None + distance = indicator_distance(indicator, state.previous_indicator) + state.previous_indicator = [value.detach() for value in indicator] + if not math.isfinite(distance): + state.accumulated_distance = 0.0 + self._warn_once("SeaCache indicator history changed shape, device, or dtype; running full.") + return True + + state.accumulated_distance += distance + if state.accumulated_distance < self.config.threshold: + return False + state.accumulated_distance = 0.0 + return True + + def _synchronize_compute(self, compute: bool, device: torch.device) -> bool: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return True if self._parameter_sharded else compute + decision = torch.tensor(int(compute), dtype=torch.int32, device=device) + if self._parameter_sharded: + from vllm_omni.diffusion.distributed.parallel_state import ( + get_fs_group, + get_sequence_parallel_world_size, + get_sp_group, + ) + + fs_group = get_fs_group() + if fs_group.world_size > 1: + torch.distributed.all_reduce( + decision, + op=torch.distributed.ReduceOp.MAX, + group=fs_group.device_group, + ) + if get_sequence_parallel_world_size() > 1: + torch.distributed.all_reduce( + decision, + op=torch.distributed.ReduceOp.MAX, + group=get_sp_group().device_group, + ) + return bool(decision.item()) + + for group in self._collective_skip_groups: + torch.distributed.all_reduce( + decision, + op=torch.distributed.ReduceOp.MAX, + group=group, + ) + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_world_size, + get_sp_group, + ) + + if get_sequence_parallel_world_size() > 1: + torch.distributed.all_reduce( + decision, + op=torch.distributed.ReduceOp.MAX, + group=get_sp_group().device_group, + ) + return bool(decision.item()) + + @torch.compiler.disable + def new_forward( + self, + module: torch.nn.Module, + *args: Any, + **kwargs: Any, + ) -> Any: + if self.extractor_fn is None: + raise RuntimeError("SeaCache extractor was not initialized") + ctx = self.extractor_fn(module, *args, **kwargs) + + if torch.is_grad_enabled(): + self._warn_once("SeaCache is inference-only; autograd-enabled calls run in full.") + return self._run_uncached(ctx) + if self.state_manager._current_context is None: + self._warn_once("SeaCache requires an explicit cache context; running full.") + return self._run_uncached(ctx) + callbacks = ( + self.current_step_callback, + self.current_sigma_callback, + self.num_inference_steps_callback, + ) + if any(callback is None for callback in callbacks): + self._warn_once("SeaCache requires scheduler step, sigma, and step-count callbacks; running full.") + return self._run_uncached(ctx) + assert self.current_step_callback is not None + assert self.current_sigma_callback is not None + assert self.num_inference_steps_callback is not None + + try: + extra_states = ctx.extra_states or {} + vision_items = extra_states.get("sea_cache_latents") + if not isinstance(vision_items, list): + raise ValueError("extractor did not provide SeaCache vision inputs") + noisy_frame_mask = extra_states.get("sea_cache_noisy_frame_mask") + conditioning_only = isinstance(noisy_frame_mask, torch.Tensor) and not bool( + torch.any(noisy_frame_mask != 0).item() + ) + + step = self.current_step_callback() + sigma = self.current_sigma_callback() + num_inference_steps = self.num_inference_steps_callback() + if isinstance(step, torch.Tensor): + step = step.item() + if isinstance(sigma, torch.Tensor): + sigma = sigma.item() + if isinstance(num_inference_steps, torch.Tensor): + num_inference_steps = num_inference_steps.item() + if step is None or sigma is None or num_inference_steps is None: + raise ValueError("scheduler metadata is unavailable") + step = int(step) + sigma = float(sigma) + num_inference_steps = int(num_inference_steps) + if ( + step < 0 + or num_inference_steps <= 0 + or step >= num_inference_steps + or not math.isfinite(sigma) + or not 0.0 <= sigma <= 1.0 + ): + raise ValueError("expected a valid step index and exact sigma in [0, 1]") + except (IndexError, TypeError, ValueError, RuntimeError) as error: + self._warn_once(f"SeaCache metadata is invalid; running full: {error}") + return self._run_uncached(ctx) + + if conditioning_only: + self._warn_once("SeaCache requires noisy vision; conditioning-only calls run in full.") + return self._run_uncached(ctx) + + state: SeaCacheState = self.state_manager.get_state() + try: + indicator = self._build_indicator(vision_items, sigma) + except (TypeError, ValueError, RuntimeError) as error: + self._warn_once(f"SeaCache could not construct its vision indicator; running full: {error}") + indicator = None + + local_compute = self._resolve_gate(state, indicator, step, num_inference_steps) + should_compute = self._synchronize_compute(local_compute, ctx.hidden_states.device) + if should_compute and not local_compute: + state.accumulated_distance = 0.0 + + if should_compute: + self.full_count += 1 + output = self._run_full_stack(ctx) + result = ctx.postprocess(output) + self._record_execution(state, step, ctx.hidden_states, output) + return result + + residual = extrapolate_residual( + state.history, + step, + self.config.residual_order, + ) + if residual.device != ctx.hidden_states.device: + residual = residual.to(ctx.hidden_states.device) + state.consecutive_cached += 1 + self.skip_count += 1 + + can_reuse = ( + residual.shape == ctx.hidden_states.shape + and residual.device == ctx.hidden_states.device + and residual.dtype == ctx.hidden_states.dtype + ) + if can_reuse: + return ctx.postprocess(ctx.hidden_states + residual) + + output = self._run_full_stack(ctx) + result = ctx.postprocess(output) + self._record_execution(state, step, ctx.hidden_states, output) + return result + + @staticmethod + def _run_full_stack(ctx: CacheContext) -> torch.Tensor: + outputs = ctx.run_transformer_blocks() + if not outputs: + raise RuntimeError("Cache extractor returned no transformer outputs") + return outputs[0] + + @staticmethod + def _run_uncached(ctx: CacheContext) -> Any: + return ctx.postprocess(SeaCacheRootHook._run_full_stack(ctx)) + + def _record_execution( + self, + state: SeaCacheState, + step: int, + execution_input: torch.Tensor, + output: torch.Tensor, + ) -> None: + if ( + output.shape == execution_input.shape + and output.device == execution_input.device + and output.dtype == execution_input.dtype + ): + state.history.append((step, (output - execution_input).detach().clone())) + state.history = state.history[-(self.config.residual_order + 1) :] + state.consecutive_cached = 0 + return + + state.history.clear() + state.accumulated_distance = 0.0 + self._warn_once("SeaCache execution boundary returned an incompatible tensor; clearing cache history.") + + def reset_state(self, module: torch.nn.Module) -> torch.nn.Module: + self.state_manager.reset() + self.full_count = 0 + self.skip_count = 0 + return module + + def refresh(self, module: torch.nn.Module) -> None: + self.reset_state(module) + + +def apply_sea_cache_hook( + module: torch.nn.Module, + config: SeaCacheConfig, + *, + current_step_callback: Callable[[], int | torch.Tensor | None] | None = None, + current_sigma_callback: Callable[[], float | torch.Tensor | None] | None = None, + num_inference_steps_callback: Callable[[], int | torch.Tensor | None] | None = None, + extractor_fn: Callable[..., CacheContext] | None = None, +) -> SeaCacheRootHook: + registry = HookRegistry.get_or_create(module) + hook = SeaCacheRootHook( + config, + current_step_callback=current_step_callback, + current_sigma_callback=current_sigma_callback, + num_inference_steps_callback=num_inference_steps_callback, + extractor_fn=extractor_fn, + ) + registry.register_hook(SeaCacheRootHook._HOOK_NAME, hook) + return hook diff --git a/vllm_omni/diffusion/cache/seacache/sea_filter.py b/vllm_omni/diffusion/cache/seacache/sea_filter.py new file mode 100644 index 000000000..f68b5c288 --- /dev/null +++ b/vllm_omni/diffusion/cache/seacache/sea_filter.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math + +import torch + + +# SeaCache paper: "Spectral-Evolution-Aware Cache for Accelerating Diffusion Models" +# https://arxiv.org/abs/2602.18993 +def apply_sea_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float = 3.0) -> torch.Tensor: + """Apply the separable SEA Wiener filter over ``(T, H, W)``.""" + original_dtype = hidden_states.dtype + hidden_states = hidden_states.contiguous().float() + dimensions = (0, 1, 2) + spectrum = torch.fft.fftn(hidden_states, dim=dimensions) + + sigma = max(1e-6, min(1.0 - 1e-6, float(sigma))) + signal_scale = 1.0 - sigma + noise_scale = sigma + gain: torch.Tensor | None = None + for axis in dimensions: + frequencies = torch.fft.fftfreq( + hidden_states.shape[axis], + device=hidden_states.device, + dtype=torch.float32, + ) + clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16) + axis_gain = signal_scale * clean_power / (signal_scale**2 * clean_power + noise_scale**2 + 1e-16) + axis_shape = [1] * hidden_states.ndim + axis_shape[axis] = axis_gain.shape[0] + reshaped_gain = axis_gain.reshape(axis_shape) + gain = reshaped_gain if gain is None else gain * reshaped_gain + + assert gain is not None + gain = gain / gain.mean() + return torch.fft.ifftn(spectrum * gain, dim=dimensions).real.to(original_dtype) + + +def relative_l1_distance(current: torch.Tensor, previous: torch.Tensor, eps: float = 1e-16) -> float: + numerator = (current.float() - previous.float()).abs().mean() + denominator = previous.float().abs().mean() + eps + return float((numerator / denominator).detach().cpu()) + + +def indicator_distance( + current: list[torch.Tensor] | None, + previous: list[torch.Tensor] | None, +) -> float: + """Mean relative-L1 distance, or infinity for incompatible indicators.""" + if current is None or previous is None or not current or len(current) != len(previous): + return float("inf") + + total = 0.0 + for current_value, previous_value in zip(current, previous, strict=True): + if ( + current_value.shape != previous_value.shape + or current_value.device != previous_value.device + or current_value.dtype != previous_value.dtype + ): + return float("inf") + total += relative_l1_distance(current_value, previous_value) + distance = total / len(current) + return distance if math.isfinite(distance) else float("inf") + + +def extrapolate_residual( + history: list[tuple[int, torch.Tensor]], + step: int, + order: int, +) -> torch.Tensor: + """Evaluate a Newton residual polynomial at ``step``.""" + if not history: + raise ValueError("SeaCache residual extrapolation requires non-empty history") + + polynomial_order = max(0, min(order, len(history) - 1)) + window = history[-(polynomial_order + 1) :] + steps = [history_step for history_step, _ in window] + coefficients = [residual.clone() for _, residual in window] + + for level in range(1, polynomial_order + 1): + for index in range(polynomial_order, level - 1, -1): + denominator = steps[index] - steps[index - level] + if denominator == 0: + return window[-1][1].clone() + coefficients[index] = (coefficients[index] - coefficients[index - 1]) / float(denominator) + + result = coefficients[polynomial_order] + for index in range(polynomial_order - 1, -1, -1): + result = result * float(step - steps[index]) + coefficients[index] + return result diff --git a/vllm_omni/diffusion/cache/seacache/state.py b/vllm_omni/diffusion/cache/seacache/state.py new file mode 100644 index 000000000..bdf21c1f4 --- /dev/null +++ b/vllm_omni/diffusion/cache/seacache/state.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + + +@dataclass(slots=True) +class SeaCacheState: + """Per-cache-context trajectory state.""" + + last_step: int | None = None + accumulated_distance: float = 0.0 + previous_indicator: list[torch.Tensor] | None = None + history: list[tuple[int, torch.Tensor]] = field(default_factory=list) + consecutive_cached: int = 0 + + def reset(self) -> None: + self.last_step = None + self.accumulated_distance = 0.0 + self.previous_indicator = None + self.history.clear() + self.consecutive_cached = 0 diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py index d58bada6c..79bb86a4c 100644 --- a/vllm_omni/diffusion/cache/selector.py +++ b/vllm_omni/diffusion/cache/selector.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from typing import Any from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend from vllm_omni.diffusion.cache.magcache import MagCacheBackend +from vllm_omni.diffusion.cache.seacache import SeaCacheBackend from vllm_omni.diffusion.cache.stepcache import StepCacheBackend from vllm_omni.diffusion.cache.teacache import TeaCacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig @@ -15,11 +19,12 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface - mag_cache: Uses MagCacheBackend with enable()/refresh() interface + - sea_cache: Uses SeaCacheBackend for spectral residual caching - step_cache: Uses StepCacheBackend for DreamZero velocity step skip Args: cache_backend: Cache backend name ("cache_dit", "tea_cache", - "mag_cache", "step_cache", or None). + "mag_cache", "sea_cache", "step_cache", or None). cache_config: Cache configuration (dict or DiffusionCacheConfig instance). Returns: @@ -40,10 +45,12 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack return TeaCacheBackend(cache_config) elif cache_backend == "mag_cache": return MagCacheBackend(cache_config) + elif cache_backend == "sea_cache": + return SeaCacheBackend(cache_config) elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"): return StepCacheBackend(cache_config) else: raise ValueError( f"Unsupported cache backend: {cache_backend}. " - "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'" + "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'sea_cache', 'step_cache'" ) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 7f0242d57..3aca638ec 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -199,7 +199,7 @@ def extract_qwen_context( hidden_states, vid_freqs, txt_freqs = module.image_rope_prepare(hidden_states, img_shapes, txt_seq_lens) image_rotary_emb = (vid_freqs, txt_freqs) - timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) + timestep = torch.as_tensor(timestep, device=hidden_states.device, dtype=hidden_states.dtype) # Call modulate_index_prepare instead of handling timestep directly. # For zero_cond_t=False: timestep unchanged, modulate_index=None. @@ -618,6 +618,8 @@ def extract_flux2_klein_context( # ============================================================================ dtype = hidden_states.dtype + if encoder_hidden_states is None: + raise ValueError("Flux2Klein requires encoder_hidden_states") num_txt_tokens = encoder_hidden_states.shape[1] timestep = timestep.to(dtype=dtype) * 1000 @@ -1466,6 +1468,77 @@ def extract_minimax_h3_context( ) +def extract_cosmos3_context( + module: nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_start_frame_offset: int = 1, + action_fps: float | None = None, + sound_latents: torch.Tensor | None = None, + noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None, + transfer_share_vision_temporal_positions: bool = True, + **kwargs: Any, +) -> CacheContext: + """Build the shared cache execution context for Cosmos3's GEN pathway.""" + if kwargs: + raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") + + prep = module._gen_preprocess( + hidden_states, + timestep, + text_ids, + text_mask, + video_shape, + fps=fps, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_noisy_mask, + action_start_frame_offset=action_start_frame_offset, + action_fps=action_fps, + sound_latents=sound_latents, + noisy_frame_mask=noisy_frame_mask, + control_latents=control_latents, + control_weights=control_weights, + transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions, + ) + + def run_transformer_blocks() -> tuple[torch.Tensor, ...]: + return (module._run_gen_stack(prep),) + + def postprocess(hidden: torch.Tensor) -> Any: + return module._gen_postprocess(hidden, prep) + + if control_latents is None: + controls: list[torch.Tensor] = [] + elif isinstance(control_latents, torch.Tensor): + controls = [control_latents] + else: + controls = list(control_latents) + + return CacheContext( + # SeaCache uses the separate inputs in extra_states for its decision. + modulated_input=prep.hidden_gen, + hidden_states=prep.hidden_gen, + encoder_hidden_states=None, + temb=prep.time_embed, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + extra_states={ + "sea_cache_latents": [*controls, hidden_states], + "sea_cache_noisy_frame_mask": noisy_frame_mask, + }, + ) + + # Registry for model-specific extractors # Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -1475,6 +1548,8 @@ def extract_minimax_h3_context( EXTRACTOR_REGISTRY: dict[str, Callable] = { "QwenImageTransformer2DModel": extract_qwen_context, "Bagel": extract_bagel_context, + "Cosmos3EdgeVFMTransformer": extract_cosmos3_context, + "Cosmos3VFMTransformer": extract_cosmos3_context, "ZImageTransformer2DModel": extract_zimage_context, "Flux2Klein": extract_flux2_klein_context, "StableAudioDiTModel": extract_stable_audio_context, @@ -1521,16 +1596,15 @@ def register_extractor(transformer_cls_name: str, extractor_fn: Callable) -> Non EXTRACTOR_REGISTRY[transformer_cls_name] = extractor_fn -def get_extractor(transformer_cls_name: str) -> Callable: +def get_extractor(transformer_type: str | type) -> Callable: """ Get extractor function for given transformer class. - This function looks up the extractor based on the exact transformer_cls_name string, - which should match the transformer type in the pipeline (i.e., pipeline.transformer.__class__.__name__). + String lookups match an exact registered name. Type lookups also check base + classes so runtime wrappers such as FSDP retain the underlying model extractor. Args: - transformer_cls_name: Transformer class name (e.g., "QwenImageTransformer2DModel") - Must exactly match a key in EXTRACTOR_REGISTRY. + transformer_type: Transformer class name or runtime type. Returns: Extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -1543,14 +1617,18 @@ def get_extractor(transformer_cls_name: str) -> Callable: >>> extractor = get_extractor("QwenImageTransformer2DModel") >>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...) """ - # Direct lookup - no substring matching - if transformer_cls_name in EXTRACTOR_REGISTRY: - return EXTRACTOR_REGISTRY[transformer_cls_name] + candidate_names: tuple[str, ...] + if isinstance(transformer_type, str): + candidate_names = (transformer_type,) + else: + candidate_names = tuple(candidate.__name__ for candidate in transformer_type.__mro__) + for candidate_name in candidate_names: + if candidate_name in EXTRACTOR_REGISTRY: + return EXTRACTOR_REGISTRY[candidate_name] - # No match found available_types = list(EXTRACTOR_REGISTRY.keys()) raise ValueError( - f"Unknown model type: '{transformer_cls_name}'. " + f"Unknown model type: '{transformer_type}'. " f"Available types: {available_types}\n" f"To add support for a new model, use register_extractor() or add to EXTRACTOR_REGISTRY." ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 7a7de3bb1..c333fb11a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -541,6 +541,8 @@ class DiffusionCacheConfig: scm_steps_mask_policy, scm_steps_policy - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio, mag_ratios, mag_calibrate + - SeaCache: sea_threshold, sea_residual_order, + sea_max_consecutive_cached, sea_power_exp - step_cache: step_cache_dit_enabled, velocity_sim_thresholds, velocity_skip_countdowns, step_cache_dit_min_history @@ -560,6 +562,12 @@ class DiffusionCacheConfig: rel_l1_thresh: float | None = None coefficients: list[float] | None = None # Uses model-specific defaults if None + # SeaCache parameters [sea_cache only] + sea_threshold: float = 0.25 + sea_residual_order: int = 1 + sea_max_consecutive_cached: int = 2 + sea_power_exp: float = 3.0 + # MagCache parameters [mag_cache only] # Default: 0.24 threshold for accumulated magnitude error mag_threshold: float = 0.24 diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 5977e0eea..04f25ac58 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 text/image/video/sound/action pipeline for vllm-omni. One pipeline class serves the Cosmos3 family modes. Output modality is selected @@ -34,6 +34,7 @@ import math import os import time from collections.abc import Iterable, Mapping +from contextlib import nullcontext from dataclasses import fields from typing import Any, ClassVar @@ -1075,6 +1076,8 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = None self._num_timesteps = None + self._current_step_index = None + self._current_sigma = None self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None self._robolab_transforms: dict[bool, Any] = {} @@ -1291,17 +1294,24 @@ class Cosmos3OmniDiffusersPipeline( or a tuple in video, action, sound order for multimodal generation. """ cache_key = kwargs.pop("_cosmos3_cache_key", None) - if cache_key is None: - return self.transformer(**kwargs) + context_name = str(kwargs.pop("_cache_context", "cond")) - branch_caches = self._cosmos3_branch_caches - if branch_caches is None: - return self.transformer(**kwargs) - - cache_key = str(cache_key) - self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(cache_key, (None, None)) - prediction = self.transformer(**kwargs) - branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + context_factory = getattr(self, "_cache_context_factory", None) + context = context_factory(context_name) if callable(context_factory) else nullcontext() + with context: + if cache_key is None: + prediction = self.transformer(**kwargs) + else: + branch_caches = self._cosmos3_branch_caches + if branch_caches is None: + prediction = self.transformer(**kwargs) + else: + cache_key = str(cache_key) + self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get( + cache_key, (None, None) + ) + prediction = self.transformer(**kwargs) + branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) return prediction def combine_multi_branch_cfg_noise( @@ -1840,6 +1850,29 @@ class Cosmos3OmniDiffusersPipeline( def num_timesteps(self): return self._num_timesteps + @property + def current_step_index(self): + return self._current_step_index + + @property + def current_sigma(self): + return self._current_sigma + + def _set_denoise_step_metadata( + self, + step_index: int, + timesteps: torch.Tensor, + scheduler: Any, + ) -> None: + self._current_step_index = step_index + self._num_timesteps = len(timesteps) + sigmas = getattr(scheduler, "sigmas", None) + self._current_sigma = sigmas[step_index] if sigmas is not None and step_index < len(sigmas) else None + + def _clear_denoise_step_metadata(self) -> None: + self._current_step_index = None + self._current_sigma = None + def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: setter = getattr(self.transformer, "set_mixed_precision_step", None) if setter is not None: @@ -2848,6 +2881,7 @@ class Cosmos3OmniDiffusersPipeline( # else uncond), so session keying loads/stores only this rank's branch. cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0 for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, step_scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) # Outside the interval, scale=1 makes the combined output equal @@ -2859,6 +2893,7 @@ class Cosmos3OmniDiffusersPipeline( do_true_cfg=True, true_cfg_scale=step_scale, positive_kwargs=dict( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2868,6 +2903,7 @@ class Cosmos3OmniDiffusersPipeline( **shared_kwargs, ), negative_kwargs=dict( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -2888,13 +2924,15 @@ class Cosmos3OmniDiffusersPipeline( keep_uncond_for_cache = self._cache_requires_paired_cfg() for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, step_scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) cfg_active = _cfg_active_at(t) if not self._kv_load_und(kv_state, is_negative=False): self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache - noise_cond = self.transformer( + noise_cond = self.predict_noise( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2911,7 +2949,8 @@ class Cosmos3OmniDiffusersPipeline( if cfg_active or keep_uncond_for_cache: if not self._kv_load_und(kv_state, is_negative=True): self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache - noise_uncond = self.transformer( + noise_uncond = self.predict_noise( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -2943,10 +2982,12 @@ class Cosmos3OmniDiffusersPipeline( # No CFG: a single cond branch per step. Bespoke (state None) keeps # using the transformer-instance cache exactly as before. for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, step_scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) self._kv_load_und(kv_state, is_negative=False) - noise_pred = self.transformer( + noise_pred = self.predict_noise( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2959,6 +3000,7 @@ class Cosmos3OmniDiffusersPipeline( self._kv_capture_und(kv_state, is_negative=False) _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) finally: + self._clear_denoise_step_metadata() self._reset_mixed_precision() # Cosmos3 currently receives a unique request_id rather than a # reusable rollout session id. Retaining its state would only pin @@ -3145,6 +3187,7 @@ class Cosmos3OmniDiffusersPipeline( self._cosmos3_branch_caches = {} try: for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, self.scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 @@ -3153,6 +3196,7 @@ class Cosmos3OmniDiffusersPipeline( needs_control_cfg = step_control != 1.0 cond_full_kwargs = dict( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -3166,6 +3210,7 @@ class Cosmos3OmniDiffusersPipeline( branches_kwargs = [ cond_full_kwargs, dict( + _cache_context="cond_no_control", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -3175,6 +3220,7 @@ class Cosmos3OmniDiffusersPipeline( **shared_kwargs, ), dict( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -3200,6 +3246,7 @@ class Cosmos3OmniDiffusersPipeline( branches_kwargs = [ cond_full_kwargs, dict( + _cache_context="cond_no_control", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -3223,6 +3270,7 @@ class Cosmos3OmniDiffusersPipeline( branches_kwargs = [ cond_full_kwargs, dict( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -3257,6 +3305,7 @@ class Cosmos3OmniDiffusersPipeline( )[0] latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents finally: + self._clear_denoise_step_metadata() self._reset_mixed_precision() self._cosmos3_branch_caches = None self.transformer.reset_cache() diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 1ded3db2f..fc5a6f7cc 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 VFM Transformer for vllm-omni. Implements the Mixture-of-Transformers architecture with two pathways: @@ -14,7 +14,7 @@ from __future__ import annotations import math from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple import torch import torch.distributed as dist @@ -427,7 +427,7 @@ class TimestepEmbedder(nn.Module): max_period: int = 10000, ) -> None: super().__init__() - # Following diffusers naming pattern here for checkpoint compatibility. + # Preserve checkpoint-compatible parameter names. self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True) self.act = nn.SiLU() self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) @@ -1105,6 +1105,28 @@ class Cosmos3GenSPPrepare(nn.Module): return hidden_gen, freqs_cos, freqs_sin +class _GenPrepared(NamedTuple): + """GEN-pathway state shared by normal and cached execution.""" + + hidden_gen: torch.Tensor + time_embed: torch.Tensor + t: int + h: int + w: int + s_video: int + s_control: int + s_action: int + s_sound: int + has_control: bool + has_action: bool + has_sound: bool + action_domain_ids: torch.Tensor | None + ulysses_size: int + use_multi_control_attention: bool + multi_control_token_sizes: tuple[int, ...] | None + multi_control_weights: tuple[float, ...] | None + + class Cosmos3VFMTransformer(nn.Module): """Cosmos3 VFM Transformer: UND language model + GEN denoising layers. @@ -1675,6 +1697,61 @@ class Cosmos3VFMTransformer(nn.Module): pad = (-(base + sound_frames)) % ulysses_size return sound_frames + pad + def _run_gen_layers( + self, + hidden_gen: torch.Tensor, + *, + s_video: int, + s_control: int, + s_action: int, + s_sound: int, + has_action: bool, + has_sound: bool, + has_control: bool, + ulysses_size: int, + use_multi_control_attention: bool, + multi_control_token_sizes: tuple[int, ...] | None, + multi_control_weights: tuple[float, ...] | None, + ) -> torch.Tensor: + """Run the complete GEN decoder between full-layout boundaries.""" + if self.cached_kv is None or self.cached_freqs_gen is None: + raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") + freqs_cos, freqs_sin = self.cached_freqs_gen + if not use_multi_control_attention: + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + freqs_gen = (freqs_cos, freqs_sin) + + if len(self.gen_layers) == len(self.cached_kv): + for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): + hidden_gen = layer( + hidden_gen, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, + ) + # Cache-dit's block wrapper may return a tuple; unwrap it. + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + else: + # Cache-dit patches gen_layers to a grouped wrapper. + for layer in self.gen_layers: + hidden_gen = layer( + hidden_gen, + cached_kv=self.cached_kv, + freqs_gen=freqs_gen, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, + ) + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + + if not use_multi_control_attention: + hidden_gen = self.gen_sp_gather(hidden_gen) + return hidden_gen + # -- Forward ------------------------------------------------------------- def forward( @@ -1697,7 +1774,51 @@ class Cosmos3VFMTransformer(nn.Module): transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Run the shared Cosmos3 GEN preprocess, stack, and postprocess path.""" + if kwargs: + raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") + prep = self._gen_preprocess( + hidden_states, + timestep, + text_ids, + text_mask, + video_shape, + fps=fps, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_noisy_mask, + action_start_frame_offset=action_start_frame_offset, + action_fps=action_fps, + sound_latents=sound_latents, + noisy_frame_mask=noisy_frame_mask, + control_latents=control_latents, + control_weights=control_weights, + transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions, + ) + return self._gen_postprocess(self._run_gen_stack(prep), prep) + + def _gen_preprocess( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_start_frame_offset: int = 1, + action_fps: float | None = None, + sound_latents: torch.Tensor | None = None, + noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None, + transfer_share_vision_temporal_positions: bool = True, + ) -> _GenPrepared: """ + Prepare the packed GEN sequence before its cacheable execution region. + Args: hidden_states: [B, C, t, h, w] noisy latents timestep: [B] diffusion timestep @@ -1720,13 +1841,9 @@ class Cosmos3VFMTransformer(nn.Module): transfer controls. Values are normalized to sum to one. Returns: - [B, C, t, h, w] velocity prediction, or - tuple outputs in video, action, sound order when action/sound streams - are provided. Transfer-control streams condition the video prediction - and are not returned. + Packed GEN inputs and metadata consumed by the shared execution and + postprocessing helpers. """ - if kwargs: - raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") t, h, w = video_shape hp, wp, _, _ = self._pad_to_patch_size(h, w) text_lengths = text_mask.sum(dim=1) @@ -1926,78 +2043,81 @@ class Cosmos3VFMTransformer(nn.Module): hidden_parts.append(hidden_sound) hidden_gen = torch.cat(hidden_parts, dim=1) - # Run GEN layers. UND K/V (replicated) is passed to each layer; - # the Cosmos3CrossAttention forwards them as joint_key/value so the - # framework Attention handles the Ulysses head-slicing internally. - if self.cached_kv is None or self.cached_freqs_gen is None: - raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") - freqs_cos, freqs_sin = self.cached_freqs_gen - if not use_multi_control_attention: - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) - freqs_gen = (freqs_cos, freqs_sin) + return _GenPrepared( + hidden_gen=hidden_gen, + time_embed=time_embed, + t=t, + h=h, + w=w, + s_video=s_video, + s_control=s_control, + s_action=s_action, + s_sound=s_sound, + has_control=has_control, + has_action=has_action, + has_sound=has_sound, + action_domain_ids=action_domain_ids, + ulysses_size=ulysses_size, + use_multi_control_attention=use_multi_control_attention, + multi_control_token_sizes=multi_control_token_sizes, + multi_control_weights=multi_control_weights, + ) - if len(self.gen_layers) == len(self.cached_kv): - for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): - hidden_gen = layer( - hidden_gen, - k_und=k_und, - v_und=v_und, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - control_token_sizes=multi_control_token_sizes, - control_weights=multi_control_weights, - ) - # Cache-dit's block wrapper may return a tuple; unwrap it. - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] - else: - # Cache-dit patches gen_layers to a grouped wrapper. - for layer in self.gen_layers: - hidden_gen = layer( - hidden_gen, - cached_kv=self.cached_kv, - freqs_gen=freqs_gen, - control_token_sizes=multi_control_token_sizes, - control_weights=multi_control_weights, - ) - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] + def _run_gen_stack(self, prep: _GenPrepared) -> torch.Tensor: + """Execute the cacheable full-layout GEN stack, including final norm.""" + hidden_gen = self._run_gen_layers( + prep.hidden_gen, + s_video=prep.s_video, + s_control=prep.s_control, + s_action=prep.s_action, + s_sound=prep.s_sound, + has_action=prep.has_action, + has_sound=prep.has_sound, + has_control=prep.has_control, + ulysses_size=prep.ulysses_size, + use_multi_control_attention=prep.use_multi_control_attention, + multi_control_token_sizes=prep.multi_control_token_sizes, + multi_control_weights=prep.multi_control_weights, + ) + return self.norm_moe_gen(hidden_gen) - if not use_multi_control_attention: - hidden_gen = self.gen_sp_gather(hidden_gen) + def _gen_postprocess( + self, + hidden_gen: torch.Tensor, + prep: _GenPrepared, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Project an already-normalized packed GEN state to model outputs.""" + if not prep.has_action and not prep.has_sound and not prep.has_control: + return self.unpatchify(self.proj_out(hidden_gen), prep.t, prep.h, prep.w) - # Final norm and project back to latent space - hidden_gen = self.norm_moe_gen(hidden_gen) - if not has_action and not has_sound and not has_control: - return self.unpatchify(self.proj_out(hidden_gen), t, h, w) - - split_sizes = [] - if has_control: - split_sizes.append(s_control) - split_sizes.append(s_video) - if has_action: - split_sizes.append(s_action) - if has_sound: - split_sizes.append(s_sound) - split_hidden = hidden_gen.split(split_sizes, dim=1) - split_idx = 0 - if has_control: - split_idx += 1 - hidden_video = split_hidden[split_idx] + split_sizes = [] + if prep.has_control: + split_sizes.append(prep.s_control) + split_sizes.append(prep.s_video) + if prep.has_action: + split_sizes.append(prep.s_action) + if prep.has_sound: + split_sizes.append(prep.s_sound) + split_hidden = hidden_gen.split(split_sizes, dim=1) + split_idx = 0 + if prep.has_control: split_idx += 1 - video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) - if has_control: - return video_pred - outputs: list[torch.Tensor] = [video_pred] - if has_action: - hidden_action = split_hidden[split_idx] - split_idx += 1 - assert action_domain_ids is not None - outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids))) - if has_sound: - hidden_sound = split_hidden[split_idx] - outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) - return tuple(outputs) + hidden_video = split_hidden[split_idx] + split_idx += 1 + video_pred = self.unpatchify(self.proj_out(hidden_video), prep.t, prep.h, prep.w) + if prep.has_control: + return video_pred + + outputs: list[torch.Tensor] = [video_pred] + if prep.has_action: + hidden_action = split_hidden[split_idx] + split_idx += 1 + assert prep.action_domain_ids is not None + outputs.append(self.unpack_action(self.action_proj_out(hidden_action, prep.action_domain_ids))) + if prep.has_sound: + hidden_sound = split_hidden[split_idx] + outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) + return tuple(outputs) def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 9dfcd6f3f..5fb124a18 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -25,7 +25,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.cache.cachedit import cache_summary +from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, cache_summary from vllm_omni.diffusion.cache.prompt_embed_cache import ( install_prompt_embed_cache, resolve_prompt_embed_cache_config, @@ -409,7 +409,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # may explicitly adopt the enabled Cache-DiT backend and then # own all later request-boundary enable/disable transitions. self.cache_backend.enable(self.pipeline) - if str(self.od_config.cache_backend).lower() == "cache_dit" and adopt_request_scoped_cache_dit( + if isinstance(self.cache_backend, CacheDiTBackend) and adopt_request_scoped_cache_dit( self.pipeline, self.cache_backend, ): @@ -663,12 +663,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): num_inference_steps = getattr(self.pipeline, "default_num_inference_steps", None) if num_inference_steps is None and od_config.cache_backend in ( "tea_cache", + "sea_cache", "step_cache", ): # When num_inference_steps is None, some pipelines defer to their - # own defaults. TeaCache refresh ignores this value; step_cache - # refresh is a no-op because per-chunk state resets in the denoise - # loop. Use the pipeline default when available to keep refresh + # own defaults. These backends use refresh to reset request state; + # runtime step metadata is either unused or resolved in the + # pipeline. Use the pipeline default when available to keep refresh # behavior aligned with single-request execution. num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0 @@ -1046,6 +1047,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): new_request_ids: list[str], ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]: # process new reqs + pipeline = self.pipeline + assert pipeline is not None, "Model not loaded. Call load_model() first." prepared_states: list[StepRequestState] = [] error_outputs: list[RunnerOutput] = [] for state in states: @@ -1058,12 +1061,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): per_req_exc: BaseException | None = None try: self._initialize_generator(state.sampling) - clear_pipeline_stage_durations(self.pipeline) + clear_pipeline_stage_durations(pipeline) # encode - self.pipeline.prepare_encode(state) + pipeline.prepare_encode(state) merge_stage_durations( state, - consume_pipeline_stage_durations(self.pipeline), + consume_pipeline_stage_durations(pipeline), ) except Exception as exc: per_req_exc = exc @@ -1247,6 +1250,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: + pipeline = self.pipeline + assert pipeline is not None, "Model not loaded. Call load_model() first." had_active_states = bool(self.state_cache) states, new_request_ids = self._update_states(scheduler_output) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 @@ -1261,7 +1266,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids) if input_batch is None: return BatchRunnerOutput.from_list(runner_output_list) - attn_metadata = {} + attn_metadata: dict[str, Any] = {} kv_backend = getattr(self, "diffusion_kv_backend", None) paged_kv_runtime = kv_backend if getattr(kv_backend, "paged_attention_adapter", None) is not None else None @@ -1272,16 +1277,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): paged_kv_runtime=paged_kv_runtime, in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ): - clear_pipeline_stage_durations(self.pipeline) - noise_pred = self.pipeline.denoise_step(input_batch, states=states) - denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline) + clear_pipeline_stage_durations(pipeline) + noise_pred = pipeline.denoise_step(input_batch, states=states) + denoise_stage_durations = consume_pipeline_stage_durations(pipeline) for state in states: merge_stage_durations( state, denoise_stage_durations, ) - pipeline_interrupted = getattr(self.pipeline, "interrupt", False) + pipeline_interrupted = getattr(pipeline, "interrupt", False) if noise_pred is None and pipeline_interrupted: for state in states: runner_output_list.append( @@ -1296,9 +1301,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): else: offset = 0 for req in states: + if req.latents is None: + raise RuntimeError(f"Stepwise request {req.request_id} has no latent state.") row_num = req.latents.shape[0] try: - self.pipeline.step_scheduler( + pipeline.step_scheduler( req, noise_pred[offset : offset + row_num] if noise_pred is not None else None ) if self.od_config.streaming_output: @@ -1307,8 +1314,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): should_decode = req.denoise_completed if should_decode: - clear_pipeline_stage_durations(self.pipeline) - result = self.pipeline.post_decode(req) + clear_pipeline_stage_durations(pipeline) + result = pipeline.post_decode(req) if result is not None: result = self._prepare_output_for_transport(result, req.sampling) self._attach_stepwise_metadata( @@ -1317,8 +1324,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) # After consuming this chunk's interaction metadata, apply pending interactions and # prepare the next chunk (prepare_next_chunk may be a no-op---depending on pipeline) - if supports_interaction_apply(self.pipeline) and not req.request_denoise_completed: - pipe = cast(SupportsInteractionApply, self.pipeline) + if supports_interaction_apply(pipeline) and not req.request_denoise_completed: + pipe = cast(SupportsInteractionApply, pipeline) pipe.apply_interaction_at_chunk_boundary(req) pipe.prepare_next_chunk(req) else: @@ -1370,12 +1377,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): for runner_output in runner_output_list: if runner_output.result is None: continue - state = states_by_id.get(runner_output.request_id) - if state is None: + matched_state = states_by_id.get(runner_output.request_id) + if matched_state is None: continue runner_output.result.peak_memory_mb = max( runner_output.result.peak_memory_mb, - state.peak_memory_mb, + matched_state.peak_memory_mb, ) terminal_request_ids = [ diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 0166cb4c1..bc7f81132 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -712,7 +712,10 @@ class OmniServeCommand(CLISubcommand): "--cache-backend", type=str, default="none", - help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'", + help=( + "Cache backend for diffusion models, options: 'tea_cache', " + "'cache_dit', 'mag_cache', 'sea_cache', 'step_cache'" + ), ) omni_config_group.add_argument( "--cache-config", From bbf77cce69bff7dfd43234e938f1b8149e2a4359 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Tue, 15 Sep 2026 16:25:56 -0700 Subject: [PATCH 088/209] [CI][ROCm] Add non-blocking entrypoint GPU coverage (#7398) Signed-off-by: andyluo7 --- .buildkite/amd/test-amd-ready.yml | 65 +++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 181a451bb..505ed1b87 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -205,6 +205,71 @@ steps: pytest -s -v tests/engine/test_async_omni_engine_abort.py ' +- label: "ROCm · Entrypoints GPU Coverage (R2-02)" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 30 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + artifact_paths: + - "artifacts/rocm-r2-02/**/*" + commands: + # Native AMD tests run from /app/vllm-omni, while Buildkite uploads + # artifacts relative to its checkout directory. + - export R2_02_ARTIFACT_DIR="$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-02" + - mkdir -p "$$R2_02_ARTIFACT_DIR" + - >- + ps -eo pid=,ppid=,lstart=,stat=,comm= > + "$$R2_02_ARTIFACT_DIR/processes-before.txt" + # Verbose pytest output and JUnit retain selected node IDs without a + # second collection-only pass over the same GPU-backed test tree. + # The inner timeout leaves five minutes for teardown verification and + # artifact generation before the Buildkite step timeout. + - >- + timeout --signal=TERM --kill-after=1m 25m + pytest -s -v -ra + tests/entrypoints/ + -m 'core_model and cuda and not + (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' + --run-level "core_model" + --durations=0 + --junitxml="$$R2_02_ARTIFACT_DIR/pytest.xml" + 2>&1 | tee "$$R2_02_ARTIFACT_DIR/pytest.log" + - >- + tail -n 100 "$$R2_02_ARTIFACT_DIR/pytest.log" > + "$$R2_02_ARTIFACT_DIR/pytest-summary.txt" + - | + sleep 5 + ps -eo pid=,ppid=,lstart=,stat=,comm= > "$$R2_02_ARTIFACT_DIR/processes-after.txt" + # Compare PID plus process start time so PID reuse cannot hide a leak. + # Double dollars preserve awk field references through Buildkite interpolation. + awk ' + NR == FNR { + if ($$9 != "ps") { + baseline[$$1 "|" $$3 "|" $$4 "|" $$5 "|" $$6 "|" $$7] = 1 + } + next + } + $$9 != "ps" { + identity = $$1 "|" $$3 "|" $$4 "|" $$5 "|" $$6 "|" $$7 + if (!(identity in baseline)) { + print + leaked = 1 + } + } + END { + if (leaked) { + print "status=FAIL leaked_processes_detected=1" + exit 1 + } + print "status=PASS leaked_processes=0" + } + ' "$$R2_02_ARTIFACT_DIR/processes-before.txt" \ + "$$R2_02_ARTIFACT_DIR/processes-after.txt" \ + 2>&1 | tee "$$R2_02_ARTIFACT_DIR/process-cleanup.txt" + # - label: "Omni Model Test Qwen3-Omni" # agent_pool: mi300_2 # depends_on: amd-build From 65693b30fb52da47f1ce9129508e9cd524c86513 Mon Sep 17 00:00:00 2001 From: LOGO127 Date: Wed, 16 Sep 2026 07:50:32 +0800 Subject: [PATCH 089/209] [Bugfix][Core] Keep sample-rate snapshots bounded in delta output (#7448) Signed-off-by: luozijian --- .../engine/test_output_metadata_snapshots.py | 143 ++++++++++++++++++ vllm_omni/outputs/mm_outputs.py | 17 ++- 2 files changed, 157 insertions(+), 3 deletions(-) create mode 100644 tests/engine/test_output_metadata_snapshots.py diff --git a/tests/engine/test_output_metadata_snapshots.py b/tests/engine/test_output_metadata_snapshots.py new file mode 100644 index 000000000..688f1252d --- /dev/null +++ b/tests/engine/test_output_metadata_snapshots.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Snapshot metadata must not accumulate as if it were generated content.""" + +import pytest +import torch +from vllm.outputs import PoolingRequestOutput +from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine import FinishReason + +from vllm_omni.outputs.mm_outputs import MultimodalPayload +from vllm_omni.outputs.output_modality import OutputModality, TensorAccumulationStrategy +from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +RATE_KEYS = ("sr", "sample_rate", "audio_sample_rate") + + +def _state(kind: RequestOutputKind) -> OmniRequestState: + return OmniRequestState( + request_id="audio", + external_req_id="audio", + parent_req=None, + request_index=0, + lora_request=None, + prompt=None, + prompt_token_ids=[0], + prompt_embeds=None, + logprobs_processor=None, + detokenizer=None, + max_tokens_param=None, + arrival_time=0.0, + queue=None, + log_stats=False, + stream_interval=1, + output_kind=kind, + ) + + +@pytest.mark.parametrize("key", RATE_KEYS) +@pytest.mark.parametrize("representation", ["integer", "scalar_tensor", "vector_tensor", "metadata_tensor"]) +def test_rate_is_latest_snapshot_before_consolidation(key, representation): + payload = MultimodalPayload() + for index, rate in enumerate((16000, 24000, 48000)): + value = ( + rate if representation == "integer" else torch.tensor([rate] if representation == "vector_tensor" else rate) + ) + incoming = ( + MultimodalPayload(tensors={"audio": torch.full((2,), index)}, metadata={key: value}) + if representation == "metadata_tensor" + else MultimodalPayload.from_dict({"audio": torch.full((2,), index), key: value}) + ) + assert incoming is not None + payload = payload.merged_with(incoming) + snapshot = payload[key] + assert not isinstance(snapshot, list) + assert int(snapshot) == rate + assert payload.get(key) is payload.to_dict()[key] + payload = payload.merged_with(MultimodalPayload(tensors={"audio": torch.full((2,), 3)})) + assert int(payload[key]) == 48000 # Missing metadata means retain, not clear. + payload.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + payload.consolidate_metadata() + torch.testing.assert_close(payload["audio"], torch.arange(4).repeat_interleave(2)) + + +@pytest.mark.parametrize("key", RATE_KEYS) +@pytest.mark.parametrize("tensor_first", [False, True]) +def test_rate_representation_change_removes_stale_partition(key, tensor_first): + first = torch.tensor(16000) if tensor_first else 16000 + last = 24000 if tensor_first else torch.tensor(24000) + payload = MultimodalPayload.from_dict({key: first}) + incoming = MultimodalPayload.from_dict({key: last}) + assert payload is not None and incoming is not None + incoming_keys = (set(incoming.tensors), set(incoming.metadata)) + payload = payload.merged_with(incoming) + assert int(payload[key]) == 24000 + assert payload[key] is payload.to_dict()[key] + assert (key in payload.tensors) != (key in payload.metadata) + assert incoming[key] is last + assert incoming_keys == (set(incoming.tensors), set(incoming.metadata)) + + +@pytest.mark.parametrize("key", RATE_KEYS) +def test_self_merge_retains_snapshot_and_content(key): + payload = MultimodalPayload.from_dict({key: torch.tensor(24000), "audio": torch.tensor([1.0, 2.0])}) + assert payload is not None + payload = payload.merged_with(payload) + assert int(payload[key]) == 24000 + payload.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + torch.testing.assert_close(payload["audio"], torch.tensor([1.0, 2.0, 1.0, 2.0])) + + +@pytest.mark.parametrize("key", RATE_KEYS) +@pytest.mark.parametrize("kind", [RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE]) +def test_long_stream_keeps_rate_bounded_without_losing_audio(key, kind): + state = _state(kind) + for index in range(1000): + state.add_multimodal_tensor( + {"model_outputs": torch.full((3,), index / 1000), key: torch.tensor(24000)}, "audio" + ) + output = state.make_request_output([], None, None, None) + assert output is not None and not isinstance(output, PoolingRequestOutput) + payload = output.outputs[0].multimodal_output + assert isinstance(payload[key], torch.Tensor) and payload[key].numel() == 1 + audio = payload["audio"] + assert audio.numel() == (3 if kind == RequestOutputKind.DELTA else 3 * (index + 1)) + torch.testing.assert_close(audio[-3:], torch.full((3,), index / 1000)) + assert ("audio" in state.mm_accumulated) is (kind == RequestOutputKind.CUMULATIVE) + final = state.make_request_output([], None, FinishReason.STOP, None) + assert final is not None and final.finished + if kind == RequestOutputKind.DELTA: + assert "audio" not in final.outputs[0].multimodal_output + + +@pytest.mark.parametrize("pending", [0, 3]) +def test_abort_flushes_only_pending_content_with_latest_rate(pending): + state = _state(RequestOutputKind.DELTA) + processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False, output_modality=OutputModality.AUDIO) + processor.request_states[state.request_id] = state + processor.external_req_ids[state.external_req_id] = [state.request_id] + for rate in (16000, 24000): + state.add_multimodal_tensor({"model_outputs": torch.ones(3), "sr": torch.tensor(rate)}, "audio") + output = state.make_request_output([], None, None, None) + assert output is not None + if pending: + state.add_multimodal_tensor({"model_outputs": torch.ones(pending)}, "audio") + aborted, outputs = processor.abort_requests_collecting_outputs([state.request_id], internal=True) + assert aborted == [state.request_id] and len(outputs) == 1 and outputs[0].finished + payload = outputs[0].outputs[0].multimodal_output + assert int(payload["sr"]) == 24000 + assert (payload["audio"].numel() if "audio" in payload else 0) == pending + assert state.request_id not in processor.request_states + + +def test_non_snapshot_tensors_still_accumulate_and_empty_merge_keeps_identity(): + first = MultimodalPayload.from_dict({"latent": torch.tensor([1.0]), "audio": torch.tensor([2.0])}) + second = MultimodalPayload.from_dict({"latent": torch.tensor([3.0]), "audio": torch.tensor([4.0])}) + assert first is not None and second is not None + assert MultimodalPayload().merged_with(first) is first + merged = first.merged_with(second) + merged.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + torch.testing.assert_close(merged["latent"], torch.tensor([1.0, 3.0])) + torch.testing.assert_close(merged["audio"], torch.tensor([2.0, 4.0])) diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py index 8c5fe6442..b1393ba19 100644 --- a/vllm_omni/outputs/mm_outputs.py +++ b/vllm_omni/outputs/mm_outputs.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Multimodal output data structures for vLLM-Omni. This module defines structured types for multimodal outputs. @@ -113,7 +115,7 @@ class MultimodalPayload(Mapping): return self.tensors[key] return self.metadata.get(key, default) - def __contains__(self, key: str) -> bool: + def __contains__(self, key: object) -> bool: return key in self.tensors or key in self.metadata def __getitem__(self, key: str) -> Any: @@ -145,15 +147,24 @@ class MultimodalPayload(Mapping): def merged_with(self, incoming: MultimodalPayload) -> MultimodalPayload: """Merge *incoming* onto this payload and return the result. - Tensor values accumulate into lists for deferred concatenation; - non-tensor values are replaced with the latest. When this payload + Content tensors accumulate into lists for deferred concatenation; + known sample-rate keys are snapshots replaced immediately, including + in DELTA streams that do not consolidate each emission. Missing keys + retain their previous value. Other values keep the existing merge + behavior. When this payload is empty, *incoming* is returned as-is, so callers should use the return value: ``accumulated = accumulated.merged_with(incoming)``. """ if self.is_empty: return incoming + # Capture before merging: incoming may be this payload. Store these + # snapshots in only one partition even if their representation changes. + snapshots = {key: incoming[key] for key in _METADATA_TENSOR_KEYS if key in incoming} _append_entries(self.tensors, incoming.tensors) _append_entries(self.metadata, incoming.metadata) + for key, value in snapshots.items(): + self.tensors.pop(key, None) + self.metadata[key] = value return self def consolidate_tensors(self, strategy: TensorAccumulationStrategy) -> None: From f420cfcfbf5d7dff2ed471305e0ef145dc32f956 Mon Sep 17 00:00:00 2001 From: ChangXiong Deng <101780062+peterDengcx@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:06:41 +0800 Subject: [PATCH 090/209] [NPU][CosyVoice3] Add CPU STFT fallback (#7520) Signed-off-by: dengchangxiong Co-authored-by: dengchangxiong --- .../cosyvoice3/test_cosyvoice3_components.py | 68 +++++++++++++++++++ .../cosyvoice3/code2wav_core/hifigan.py | 42 ++++++++++-- 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 1b1df6367..7dfbe6e74 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -76,6 +76,74 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift): assert causal_hift.stft_window.device == waveform.device +def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch): + class NpuTensor: + device = SimpleNamespace(type="npu") + + tensor = NpuTensor() + stft_calls = [] + istft_calls = [] + monkeypatch.setattr( + causal_hift, + "_stft_on_cpu", + lambda value: stft_calls.append(value) or ("real", "imag"), + ) + monkeypatch.setattr( + causal_hift, + "_istft_on_cpu", + lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform", + ) + + assert causal_hift._stft(tensor) == ("real", "imag") + assert causal_hift._istft(tensor, tensor) == "waveform" + assert stft_calls == [tensor] + assert istft_calls == [(tensor, tensor)] + + +def test_causal_hift_stft_cpu_fallback_uses_float32(causal_hift, monkeypatch): + original_stft = torch.stft + + def assert_cpu_float32_stft(input, *args, window, **kwargs): + assert input.device.type == "cpu" + assert input.dtype == torch.float32 + assert window.device.type == "cpu" + assert window.dtype == torch.float32 + return original_stft(input, *args, window=window, **kwargs) + + monkeypatch.setattr(torch, "stft", assert_cpu_float32_stft) + + waveform = torch.randn(1, 64, dtype=torch.float64) + real, imag = causal_hift._stft_on_cpu(waveform) + + assert real.device == waveform.device + assert imag.device == waveform.device + assert real.dtype == waveform.dtype + assert imag.dtype == waveform.dtype + + +def test_causal_hift_istft_cpu_fallback_uses_float32(causal_hift, monkeypatch): + original_istft = torch.istft + + def assert_cpu_complex64_istft(input, *args, window, **kwargs): + assert input.device.type == "cpu" + assert input.dtype == torch.complex64 + assert window.device.type == "cpu" + assert window.dtype == torch.float32 + return original_istft(input, *args, window=window, **kwargs) + + monkeypatch.setattr(torch, "istft", assert_cpu_complex64_istft) + + waveform = torch.randn(1, 64, dtype=torch.float32) + spec = torch.stft(waveform, 16, 4, 16, window=causal_hift.stft_window, return_complex=True) + magnitude = torch.abs(spec).to(torch.float64) + phase = torch.angle(spec).to(torch.float64) + reconstructed = causal_hift._istft_on_cpu(magnitude, phase) + + assert reconstructed.device == magnitude.device + assert reconstructed.dtype == magnitude.dtype + torch.testing.assert_close(reconstructed, waveform.to(torch.float64), rtol=1e-5, atol=1e-5) + + class TestPreLookaheadLayer: """Tests for PreLookaheadLayer.""" diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 2a5d472a0..7cbd31e9c 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -489,6 +489,9 @@ class HiFTGenerator(nn.Module): block.remove_weight_norm() def _stft(self, x): + if x.device.type == "npu": + return self._stft_on_cpu(x) + spec = torch.stft( x, self.istft_params["n_fft"], @@ -500,18 +503,49 @@ class HiFTGenerator(nn.Module): spec = torch.view_as_real(spec) # [B, F, TT, 2] return spec[..., 0], spec[..., 1] + def _stft_on_cpu(self, x): + target_device = x.device + target_dtype = x.dtype + spec = torch.stft( + x.float().cpu(), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.float().cpu(), + return_complex=True, + ) + spec = torch.view_as_real(spec).to(device=target_device, dtype=target_dtype) + return spec[..., 0], spec[..., 1] + def _istft(self, magnitude, phase): + if magnitude.device.type == "npu": + return self._istft_on_cpu(magnitude, phase) + magnitude = torch.clip(magnitude, max=1e2) real = magnitude * torch.cos(phase) - img = magnitude * torch.sin(phase) - inverse_transform = torch.istft( - torch.complex(real, img), + imag = magnitude * torch.sin(phase) + return torch.istft( + torch.complex(real, imag), self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self._get_stft_window(magnitude), ) - return inverse_transform + + def _istft_on_cpu(self, magnitude, phase): + target_device = magnitude.device + target_dtype = magnitude.dtype + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + imag = magnitude * torch.sin(phase) + waveform = torch.istft( + torch.complex(real.float().cpu(), imag.float().cpu()), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.float().cpu(), + ) + return waveform.to(device=target_device, dtype=target_dtype) def _get_stft_window(self, tensor: torch.Tensor) -> torch.Tensor: if self.stft_window.device != tensor.device: From 5527f3f407e224dc3908098a1513cca974c62b53 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Wed, 16 Sep 2026 09:07:44 +0800 Subject: [PATCH 091/209] [CI][Bugfix] Update Qwen3-TTS CI Validation Message (#7276) Signed-off-by: Weiming Liao --- .../qwen3_tts/batch_speech_client.py | 5 ++- .../text_to_speech/qwen3_tts/gradio_demo.py | 39 ++++++++++++------- .../qwen3_tts/openai_speech_client.py | 7 +++- .../qwen3_tts/run_gradio_demo.sh | 5 ++- .../test_invalid_audio_speech.py | 10 +---- 5 files changed, 40 insertions(+), 26 deletions(-) diff --git a/examples/online_serving/text_to_speech/qwen3_tts/batch_speech_client.py b/examples/online_serving/text_to_speech/qwen3_tts/batch_speech_client.py index 47fdc3691..1b91905b7 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/batch_speech_client.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/batch_speech_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Batch speech client for Qwen3-TTS via /v1/audio/speech/batch endpoint. This script demonstrates how to synthesize multiple texts in a single request. @@ -7,7 +10,7 @@ the reference for each item. Start the server (with batch-optimized stage settings for best throughput): - vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --omni \ --trust-remote-code \ --stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2}, diff --git a/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py b/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py index 3c99a82d3..105571f38 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Gradio demo for Qwen3-TTS with gapless streaming audio playback. Uses a custom AudioWorklet-based player for gap-free streaming, @@ -351,7 +354,7 @@ def generate_speech( raise gr.Error(f"Failed to decode audio: {e}") -def create_app(api_base: str): +def create_app(api_base: str, task_type: str = "CustomVoice"): """Create the FastAPI app with streaming proxy + Gradio UI.""" fastapi_app = FastAPI() @@ -463,10 +466,11 @@ def create_app(api_base: str): lines=4, ) with gr.Row(): - task_type = gr.Radio( + task_type_input = gr.Radio( choices=TASK_TYPES, - value="CustomVoice", + value=task_type, label="Task Type", + interactive=False, scale=2, ) language = gr.Dropdown( @@ -479,17 +483,21 @@ def create_app(api_base: str): choices=voices, value=voices[0] if voices else None, label="speaker", - visible=True, + visible=task_type == "CustomVoice", allow_custom_value=True, ) instructions = gr.Textbox( label="Instructions", placeholder="e.g., Speak with excitement / A warm, friendly female voice", lines=2, - visible=True, - info="Optional style/emotion instructions", + visible=task_type != "Base", + info=( + "Required: describe the voice style" + if task_type == "VoiceDesign" + else "Optional style/emotion instructions" + ), ) - with gr.Column(visible=False) as ref_group: + with gr.Column(visible=task_type == "Base") as ref_group: ref_audio = gr.Audio( label="Reference Audio (upload for voice cloning)", type="numpy", @@ -563,7 +571,7 @@ def create_app(api_base: str): autoplay=True, visible=False, ) - with gr.Column(visible=True) as examples_cv: + with gr.Column(visible=task_type == "CustomVoice") as examples_cv: gr.Examples( examples=[ [ @@ -588,7 +596,7 @@ def create_app(api_base: str): inputs=[text_input, voice, language, instructions], label="examples", ) - with gr.Column(visible=False) as examples_vd: + with gr.Column(visible=task_type == "VoiceDesign") as examples_vd: gr.Examples( examples=[ [ @@ -605,7 +613,7 @@ def create_app(api_base: str): inputs=[text_input, language, instructions], label="examples", ) - with gr.Column(visible=False) as examples_base: + with gr.Column(visible=task_type == "Base") as examples_base: gr.Examples( examples=[ [ @@ -653,9 +661,9 @@ def create_app(api_base: str): gr.update(visible=is_base), # examples_base ) - task_type.change( + task_type_input.change( fn=on_task_change, - inputs=[task_type], + inputs=[task_type_input], outputs=[ voice, instructions, @@ -703,7 +711,7 @@ def create_app(api_base: str): all_inputs = [ text_input, - task_type, + task_type_input, voice, language, instructions, @@ -788,6 +796,9 @@ def parse_args(): description="Gradio demo for Qwen3-TTS with gapless AudioWorklet streaming.", ) add_common_args(parser) + parser.add_argument( + "--task-type", choices=TASK_TYPES, default="CustomVoice", help="Task matching the served checkpoint" + ) return parser.parse_args() @@ -796,7 +807,7 @@ def main(): logging.basicConfig(level=logging.INFO) print(f"Connecting to vLLM server at: {args.api_base}") - app = create_app(args.api_base) + app = create_app(args.api_base, task_type=args.task_type) import uvicorn diff --git a/examples/online_serving/text_to_speech/qwen3_tts/openai_speech_client.py b/examples/online_serving/text_to_speech/qwen3_tts/openai_speech_client.py index 7b2fecef5..a3f45bce8 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/openai_speech_client.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/openai_speech_client.py @@ -8,25 +8,28 @@ to generate audio from text using Qwen3-TTS models. Examples: # CustomVoice task (predefined speaker) - python openai_speech_client.py --text "Hello, how are you?" --voice vivian + python openai_speech_client.py --text "Hello, how are you?" --speaker vivian # CustomVoice with emotion instruction - python openai_speech_client.py --text "I'm so happy!" --voice vivian \ + python openai_speech_client.py --text "I'm so happy!" --speaker vivian \ --instructions "Speak with excitement" # VoiceDesign task (voice from description) python openai_speech_client.py --text "Hello world" \ + --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \ --task-type VoiceDesign \ --instructions "A warm, friendly female voice" # Base task (voice cloning) python openai_speech_client.py --text "Hello world" \ + --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --task-type Base \ --ref-audio "https://example.com/reference.wav" \ --ref-text "This is the reference transcript" # Base task with pre-computed speaker embedding python openai_speech_client.py --text "Hello world" \ + --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --task-type Base \ --speaker-embedding embedding.json """ diff --git a/examples/online_serving/text_to_speech/qwen3_tts/run_gradio_demo.sh b/examples/online_serving/text_to_speech/qwen3_tts/run_gradio_demo.sh index d79be3c2a..650797beb 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/run_gradio_demo.sh +++ b/examples/online_serving/text_to_speech/qwen3_tts/run_gradio_demo.sh @@ -1,4 +1,7 @@ #!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Launch both vLLM server and Gradio demo for Qwen3-TTS # # Usage: @@ -179,7 +182,7 @@ fi echo "" echo "Starting Gradio demo..." cd "$SCRIPT_DIR" -GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--host" "$GRADIO_IP" "--port" "$GRADIO_PORT") +GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--host" "$GRADIO_IP" "--port" "$GRADIO_PORT" "--task-type" "$TASK_TYPE") if [ "$GRADIO_SHARE" = true ]; then GRADIO_CMD+=("--share") fi diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py index 9861b609e..98b4184d9 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py @@ -157,14 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client: "ref_audio": None, "ref_text": None, }, - # NOTE: This is delegating down the TTS adapter, so the message is both specific - # to the Qwen3TTS adapter and dependent on whether or not the model has any speakers - # (as opposed to having valida speakers, but getting an invalid speaker). it would - # be more ideal to validate that this is landing on `.validate()` on the adapter - # and ensuring that the returned message matches what we would get from calling the - # adapter validate directly. - ("no speakers configured",), - id="customvoice_invalid_voice", + ("Base checkpoint does not support task_type='CustomVoice'",), + id="base_checkpoint_customvoice_task_mismatch", ), pytest.param( {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid" From 45a3e77351b9ef2b36612159bb0d5a3df812466b Mon Sep 17 00:00:00 2001 From: leo <92263211+liuyao0322@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:08:06 +0800 Subject: [PATCH 092/209] [Bugfix] Route Music3 float32 attention to SDPA (#7354) Signed-off-by: liuyao0322 --- .../test_minimax_music3_dit_attention.py | 117 ++++++++++++++++++ .../models/minimax_music3/dit.py | 47 ++++--- 2 files changed, 150 insertions(+), 14 deletions(-) create mode 100644 tests/model_executor/models/test_minimax_music3_dit_attention.py diff --git a/tests/model_executor/models/test_minimax_music3_dit_attention.py b/tests/model_executor/models/test_minimax_music3_dit_attention.py new file mode 100644 index 000000000..cdbb57796 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_dit_attention.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Attention dtype routing for MiniMax Music 3's float32 acoustic stage.""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.minimax_music3 import dit + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_FLOAT32_INCOMPATIBLE_BACKENDS = [ + "FLASH_ATTN", + "FLASH_ATTN_HUB", + "FLASH_ATTN_3_HUB", + "CUDNN_ATTN", + "FLASHINFER_ATTN", +] + + +class _BackendType: + def __init__(self, name: str) -> None: + self.name = name + + def get_name(self) -> str: + return self.name + + +class _Backend(nn.Module): + def __init__(self, *, name: str, explicit: bool) -> None: + super().__init__() + self.attn_backend = _BackendType(name) + self.backend_explicit = explicit + self.calls = 0 + + def forward(self, q, k, v): + del k, v + self.calls += 1 + return torch.full_like(q, 7) + + +def _attention(monkeypatch, *, backend_name: str, explicit: bool): + backend = _Backend(name=backend_name, explicit=explicit) + monkeypatch.setattr(dit, "_build_native_attention", lambda **_kwargs: backend) + return dit.Attention(8, head_dim=4), backend + + +@pytest.mark.parametrize( + "backend_name", + _FLOAT32_INCOMPATIBLE_BACKENDS, +) +def test_float32_uses_sdpa_for_automatic_incompatible_backend(monkeypatch, backend_name): + attention, backend = _attention(monkeypatch, backend_name=backend_name, explicit=False) + q = torch.randn(2, 5, 2, 4) + + actual = attention._attend(q, q, q) + + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + q.transpose(1, 2), + q.transpose(1, 2), + is_causal=False, + scale=attention.softmax_scale, + ).transpose(1, 2) + torch.testing.assert_close(actual, expected) + assert backend.calls == 0 + assert attention.backend_name == f"TORCH_SDPA(float32)/{backend_name}(lower precision)" + + +@pytest.mark.parametrize("backend_name", _FLOAT32_INCOMPATIBLE_BACKENDS) +@pytest.mark.parametrize( + ("dtype", "explicit"), + [ + (torch.float16, False), + (torch.bfloat16, False), + (torch.float32, True), + ], +) +def test_lower_precision_or_explicit_backend_remains_authoritative(monkeypatch, backend_name, dtype, explicit): + attention, backend = _attention(monkeypatch, backend_name=backend_name, explicit=explicit) + q = torch.zeros(1, 3, 2, 4, dtype=dtype) + + actual = attention._attend(q, q, q) + + assert backend.calls == 1 + torch.testing.assert_close(actual, torch.full_like(q, 7)) + + +def test_float32_uses_sdpa_when_native_attention_is_unavailable(monkeypatch): + monkeypatch.setattr(dit, "_build_native_attention", lambda **_kwargs: None) + attention = dit.Attention(8, head_dim=4) + q = torch.randn(1, 3, 2, 4) + + actual = attention._attend(q, q, q) + + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + q.transpose(1, 2), + q.transpose(1, 2), + is_causal=False, + scale=attention.softmax_scale, + ).transpose(1, 2) + torch.testing.assert_close(actual, expected) + assert attention.backend_name == "TORCH_SDPA" + + +def test_float32_keeps_compatible_automatic_backend(monkeypatch): + attention, backend = _attention(monkeypatch, backend_name="SDPA", explicit=False) + q = torch.zeros(1, 3, 2, 4) + + actual = attention._attend(q, q, q) + + assert backend.calls == 1 + torch.testing.assert_close(actual, torch.full_like(q, 7)) diff --git a/vllm_omni/model_executor/models/minimax_music3/dit.py b/vllm_omni/model_executor/models/minimax_music3/dit.py index cf180ac52..1e7af63fc 100644 --- a/vllm_omni/model_executor/models/minimax_music3/dit.py +++ b/vllm_omni/model_executor/models/minimax_music3/dit.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Condition encoder and flow-matching transformer for MiniMax Music 3. Stage 1 turns the AR stage's conditioning frames into a DAC latent in two @@ -16,13 +16,13 @@ a single ``to_qkv`` so each block runs one GEMM instead of three. :func:`remap_transformer_state` performs that fusion, and is the only place the checkpoint's layout is reinterpreted. -Attention runs through the project's own diffusion attention layer, so this -model picks up the same backend selection, sequence-parallel plumbing and -KV-quant policy as every other transformer in the tree. That layer resolves -its backend from the diffusion config, which an ``LLM_GENERATION`` stage such -as this one does not set; it tolerates that and falls through to the platform -default. If it cannot be built at all the block runs plain SDPA instead, a -choice made once when the blocks are built and never inside ``forward``. +Attention uses the project's diffusion attention layer when its selected +backend supports the runtime dtype. This ``LLM_GENERATION`` stage does not set +a diffusion backend and therefore gets the platform default. CUDA commonly +resolves that default to FlashAttention, cuDNN, or FlashInfer, so the model +routes its native float32 decode through SDPA while retaining the selected +backend for supported lower-precision inputs. An explicitly selected backend +remains authoritative. """ from __future__ import annotations @@ -75,6 +75,19 @@ _TRANSFORMER_IN_DIM = DIT_LATENT_CHANNELS * 2 + _CONDITION_DIM # reaches exactly zero, so the noise term never fully vanishes. _SIGMA_FLOOR = 1e-6 +# These backends reject float32 Q/K/V. MiniMax Music 3 decodes in float32, so +# automatically selected incompatible backends must not receive its attention +# tensors. This includes cuDNN and FlashInfer selected on Blackwell. Explicit +# backend choices retain the shared layer's fail-fast contract instead of +# being silently replaced here. +_FLOAT32_UNSUPPORTED_BACKENDS = { + "FLASH_ATTN", + "FLASH_ATTN_HUB", + "FLASH_ATTN_3_HUB", + "CUDNN_ATTN", + "FLASHINFER_ATTN", +} + class FourierFeatures(nn.Module): """Random-Fourier timestep features: ``[B, 1] -> [B, out_features]``.""" @@ -149,10 +162,7 @@ def _build_native_attention(*, num_heads: int, head_size: int, softmax_scale: fl except Exception as exc: # noqa: BLE001 - any failure means: use SDPA _NATIVE_ATTENTION_UNAVAILABLE = repr(exc) logger.warning( - "MiniMax Music 3 DiT could not build the diffusion attention layer " - "(%s); running torch SDPA instead. The stage decodes in float32, " - "for which the diffusion layer dispatches to SDPA anyway, so this " - "changes performance bookkeeping and not the audio.", + "MiniMax Music 3 DiT could not build the diffusion attention layer (%s); running torch SDPA instead.", exc, ) return None @@ -180,10 +190,19 @@ class Attention(nn.Module): softmax_scale=self.softmax_scale, prefix=prefix, ) - self.backend_name = self.backend.attn_backend.get_name() if self.backend is not None else "TORCH_SDPA" + native_backend_name = self.backend.attn_backend.get_name() if self.backend is not None else None + self._auto_float32_sdpa = ( + self.backend is not None + and not getattr(self.backend, "backend_explicit", False) + and native_backend_name in _FLOAT32_UNSUPPORTED_BACKENDS + ) + if self._auto_float32_sdpa: + self.backend_name = f"TORCH_SDPA(float32)/{native_backend_name}(lower precision)" + else: + self.backend_name = native_backend_name or "TORCH_SDPA" def _attend(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: - if self.backend is not None: + if self.backend is not None and not (q.dtype == torch.float32 and self._auto_float32_sdpa): return self.backend(q, k, v) q, k, v = (t.transpose(1, 2) for t in (q, k, v)) out = F.scaled_dot_product_attention(q, k, v, is_causal=False, scale=self.softmax_scale) From a1a234a0b439c2c00469e8c841a0286604156a8d Mon Sep 17 00:00:00 2001 From: Clodagh Walsh <42538837+clodaghwalsh17@users.noreply.github.com> Date: Wed, 16 Sep 2026 02:26:08 +0100 Subject: [PATCH 093/209] [Voxtral] Raise BadRequest for ref_audio (#7012) Signed-off-by: Clodagh Walsh --- recipes/mistralai/Voxtral-TTS.md | 2 +- .../test_voxtral_tts_expansion.py | 21 +++++++++++++++++ .../openai_api/test_serving_speech.py | 2 ++ .../openai/tts_adapters/voxtral.py | 23 +++++++++++++++++-- .../models/voxtral_tts/voxtral_tts.py | 5 ++++ .../voxtral_tts_audio_tokenizer.py | 6 ++++- vllm_omni/worker/base.py | 13 +++++++++++ 7 files changed, 68 insertions(+), 4 deletions(-) diff --git a/recipes/mistralai/Voxtral-TTS.md b/recipes/mistralai/Voxtral-TTS.md index 8868ed3c3..c4eb2d12f 100644 --- a/recipes/mistralai/Voxtral-TTS.md +++ b/recipes/mistralai/Voxtral-TTS.md @@ -178,7 +178,7 @@ mistralai/Voxtral-4B-TTS-2603 checkpoint. running the voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream) ) or adding a `ref_audio` request parameter will fail because the public checkpoint does not provide the encoder weights needed to -convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.` +convert reference audio into conditioning features. The expected failure is a 400 BadRequest `Voice cloning with 'ref_audio' requires the full Voxtral checkpoint with encoder weights. The open-source variant only supports preset voices via the 'voice' parameter.`. ### 1x NVIDIA RTX A6000 48GB diff --git a/tests/e2e/online_serving/test_voxtral_tts_expansion.py b/tests/e2e/online_serving/test_voxtral_tts_expansion.py index 278e563e8..40d419a91 100644 --- a/tests/e2e/online_serving/test_voxtral_tts_expansion.py +++ b/tests/e2e/online_serving/test_voxtral_tts_expansion.py @@ -12,11 +12,14 @@ import os os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") import pytest +from openai import BadRequestError as OpenAIBadRequestError from tests.helpers.mark import hardware_test +from tests.helpers.media import generate_synthetic_audio from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +# Note this model is the public checkpoint without the encoder weights required for voice cloning MODEL = "mistralai/Voxtral-4B-TTS-2603" STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml") EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"] @@ -173,3 +176,21 @@ def test_speech_batches(omni_server, online_client) -> None: "timeout": 120.0, } ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_ref_audio(omni_server, online_client) -> None: + """Test ref_audio failure is handled as a BadRequest instead of a RuntimeError""" + audio_data = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" + with pytest.raises(OpenAIBadRequestError) as exc_info: + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "how are you", + "ref_audio": audio_data, + "language": "English", + "response_format": "wav", + } + ) + + assert exc_info.value.status_code == 400 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index e7666ee33..f910c68af 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -5133,6 +5133,7 @@ class TestTTSAsyncOffloading: def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mocker: MockerFixture): """Voxtral path in _prepare_speech_generation should call the async wrapper.""" + voxtral_server._adapter._encoder_loaded = mocker.AsyncMock(return_value=False) voxtral_server._adapter._build_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], @@ -5288,6 +5289,7 @@ class TestTTSAsyncOffloading: """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only.""" voxtral_server.engine_client.model_config.async_chunk = False mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None) + voxtral_server._adapter._encoder_loaded = mocker.AsyncMock(return_value=False) voxtral_server._adapter._build_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py index f9d7639ca..5ea82babd 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py @@ -22,6 +22,7 @@ class VoxtralTTSAdapter(ARTTSAdapter): def __init__(self, ctx) -> None: super().__init__(ctx) self._build_prompt_async = make_async(self._build_prompt, executor=getattr(ctx.server, "_tts_executor", None)) + self._encoder_loaded = None def _build_prompt(self, request: "OpenAICreateSpeechRequest") -> dict[str, Any]: """Build a Voxtral prompt for either a preset voice or inline reference audio.""" @@ -31,13 +32,24 @@ class VoxtralTTSAdapter(ARTTSAdapter): ref_audio = request.ref_audio if not request.voice and not ref_audio: raise ValueError("Voxtral requires either a voice name or ref_audio.") - if isinstance(ref_audio, str) and ref_audio.startswith("data:"): - _, _, ref_audio = ref_audio.partition(",") + + if ref_audio is not None and not self._encoder_loaded: + validation_error = ( + "Voice cloning with 'ref_audio' requires the full Voxtral checkpoint " + "with encoder weights. The open-source variant only supports " + "preset voices via the 'voice' parameter." + ) + raise ValueError(validation_error) + server = self.ctx.server if server._tts_tokenizer is None: from vllm.tokenizers import cached_tokenizer_from_config server._tts_tokenizer = cached_tokenizer_from_config(self.ctx.engine_client.model_config).instruct + + if isinstance(ref_audio, str) and ref_audio.startswith("data:"): + _, _, ref_audio = ref_audio.partition(",") + if request.voice is not None: tokenized = server._tts_tokenizer.encode_speech_request( SpeechRequest(input=request.input, voice=request.voice) @@ -45,6 +57,7 @@ class VoxtralTTSAdapter(ARTTSAdapter): prompt = tokens_input(prompt_token_ids=tokenized.tokens) prompt["additional_information"] = {"voice": [request.voice]} return prompt + tokenized = server._tts_tokenizer.encode_speech_request(SpeechRequest(input=request.input, ref_audio=ref_audio)) audio = tokenized.audios[0] return { @@ -84,6 +97,12 @@ class VoxtralTTSAdapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: + # Cache encoder_loaded property to avoid RPC call on every request with ref_audio + if self._encoder_loaded is None: + loaded = await self.ctx.server.engine_client.collective_rpc("encoder_loaded") + # Flatten results before checking if encoder loaded property is set + self._encoder_loaded = any(sum(loaded, [])) + prompt = await self._build_prompt_async(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts") diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index c61aeaaad..de79c4cba 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -174,6 +174,11 @@ class VoxtralTTSForConditionalGeneration( return self.model.get_language_model() return self.model + def encoder_loaded(self): + if self.audio_tokenizer is not None: + return self.audio_tokenizer.encoder_loaded + return False + def _enable_acoustic_transformer_cudagraph(self): """Initialize and capture CUDA graphs for compute_mm_logits.""" if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"): diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py index 9be16c0f1..ac8e09625 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py @@ -900,6 +900,10 @@ class VoxtralTTSAudioTokenizer(nn.Module): """List of size of each codebook""" return self.quantizer.codebook_sizes + @property + def encoder_loaded(self) -> bool: + return self._encoder_loaded + def load_weight(self, weight: tuple[str, torch.Tensor]) -> str: params_dict = dict(self.named_parameters()) name, loaded_weight = weight @@ -985,7 +989,7 @@ class VoxtralTTSAudioTokenizer(nn.Module): def encode_waveforms(self, x: list[torch.Tensor]) -> list[torch.Tensor]: if not self._encoder_loaded: - raise RuntimeError( + raise ValueError( "encode_waveforms requires encoder weights which are not available in the open-source checkpoint." ) audio_codes = [] diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 2a7e9c876..44ec3dad5 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -308,3 +308,16 @@ class OmniGPUWorkerBase(GPUWorker): pass tid = task.task_id if hasattr(task, "task_id") else "unknown" return OmniACK(task_id=tid, status="ERROR", error_msg=str(e)) + + def encoder_loaded(self) -> bool: + """Check if encoder weights are loaded in the model. + + This method is exposed via collective_rpc to check encoder availability + for models that support voice cloning with reference audio. + + Returns: + bool: True if encoder weights are available, False otherwise. + """ + if hasattr(self.model_runner, "model") and hasattr(self.model_runner.model, "encoder_loaded"): + return self.model_runner.model.encoder_loaded() + return False From c197cd242207eaa1d3dc13622ceac8a4fc273ef5 Mon Sep 17 00:00:00 2001 From: Deep Shah Date: Tue, 15 Sep 2026 18:26:43 -0700 Subject: [PATCH 094/209] [Bugfix] Fix Qwen3-Omni streaming crops in mixed-length batches (#7340) Signed-off-by: Deep Shah --- .../test_qwen3_omni_code2wav_streaming.py | 64 +++++++++++++------ .../models/qwen3_omni/qwen3_omni_code2wav.py | 21 +++--- 2 files changed, 52 insertions(+), 33 deletions(-) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py index 7be4f6afa..69dd4e9ce 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for ``Qwen3OmniMoeCode2Wav.chunked_decode_streaming``. Background @@ -210,26 +210,50 @@ def test_no_shift_when_decoder_returns_nominal_length(): torch.testing.assert_close(streamed, gt) -def test_batched_streaming_only_shifts_the_boundary_row(): - """Per-row tail: in a shared window the shorter row had right context from - the longer row's frames, so it lost nothing (tail==0, no shift); only the - row at the decoded boundary (code_seq_len == window length) is shifted.""" - model = _ShortTrimDecoder() - # One window of W=30 frames, two rows. codes[:, 0, :] are absolute indices. - idx = torch.arange(30, dtype=torch.long) - window = idx.view(1, 1, 30).expand(2, _Q, 30).contiguous() +@pytest.mark.parametrize("cudagraph", [False, True]) +@pytest.mark.parametrize("trim", [_TRIM, 0]) +@pytest.mark.parametrize("short_frames,left_context", [(1, 0), (4, 0), (29, 4), (26, 25), (37, 25)]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_streaming_batch_matches_independent_requests(cudagraph, trim, short_frames, left_context, reverse): + """Padding a request beside a longer one cannot supply valid future codec frames. - long_frames, short_frames = 30, 20 # row0 spans the full window, row1 is shorter - lc = 5 - out = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( + Compare values as well as lengths: an incorrect start and end can shift + the whole chunk while preserving its duration. The independent request + is the oracle, rather than another copy of the slicing arithmetic. + """ + model = _ShortTrimDecoder(trim=trim, cudagraph=cudagraph) + short = _make_codes(short_frames) + long = _make_codes(50) + 100 + windows = [(short, left_context), (long, 25)] + if reverse: + windows.reverse() + batch = torch.zeros(2, _Q, 50, dtype=torch.long) + for i, (window, _) in enumerate(windows): + batch[i, :, : window.shape[-1]] = window[0] + actual = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( model, - window, - left_context_size=[lc, lc], - seq_token_counts=[long_frames * _Q, short_frames * _Q], + batch, + left_context_size=[lc for _, lc in windows], + seq_token_counts=[window.numel() for window, _ in windows], ) + assert len(actual) == len(windows) + for output, (window, lc) in zip(actual, windows): + expected = _decode_chunk(model, window, lc) + torch.testing.assert_close(output.reshape(-1), expected, rtol=0, atol=0) - batch_len = 30 * _UP - _TRIM # actual decoder output width - # Row 0 (boundary row): tail == TRIM -> start = lc*up - TRIM - assert out[0].shape[-1] == batch_len - (lc * _UP - _TRIM) - # Row 1 (shorter): code_seq_len*up < batch_len -> tail == 0 -> start = lc*up - assert out[1].shape[-1] == short_frames * _UP - lc * _UP + +@pytest.mark.parametrize("cudagraph", [False, True]) +def test_streaming_preserves_samples_when_batch_composition_changes(cudagraph): + """A singleton -> mixed batch -> singleton stream cannot skip then repeat samples.""" + model = _ShortTrimDecoder(cudagraph=cudagraph) + codes = _make_codes(54) + first = _decode_chunk(model, codes[:, :, :4], 0) + middle_batch = torch.zeros(2, _Q, 50, dtype=torch.long) + middle_batch[0, :, :29] = codes[0, :, :29] + middle_batch[1] = _make_codes(50)[0] + 100 + middle = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( + model, middle_batch, left_context_size=[4, 25], seq_token_counts=[29 * _Q, 50 * _Q] + )[0].reshape(-1) + last = _decode_chunk(model, codes[:, :, 4:54], 25) + streamed = torch.cat((first, middle, last)) + torch.testing.assert_close(streamed, _single_shot(model, codes), rtol=0, atol=0) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 67352eecb..bc8f6b3b4 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -316,21 +316,16 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): else: # Fallback: assume all batch elements share the same sequence length. code_seq_lens = [codes.shape[-1]] * codes.shape[0] + # The decoder trims the same right-edge tail from every batch row. + # Infer it from the decoded window, not a shorter request's length: + # padding that request with codec zeros does not provide valid context. + tail = max(0, int(codes.shape[-1] * self.total_upsample) - batch_wav.shape[-1]) for idx, code_seq_len in enumerate(code_seq_lens): - # The eager decoder (self(codes)) emits a fixed `tail` fewer samples than - # code_seq_len*total_upsample: the causal conv stack trims the sequence end - # because it lacks right context. The original slice assumed full length, so - # under async_chunk every streaming chunk silently dropped its last `tail` - # samples -> a ~23ms gap at every chunk boundary + cumulative ~1.2% time - # compression. Shift the start back by `tail` so this chunk refills the gap the - # previous chunk left (those frames are re-decoded here with proper context). - # `tail` is measured per call, so this is a no-op only when the decoder returns - # the nominal full length (tail==0, e.g. a C==0 decoder like Qwen3-TTS); for - # Qwen3-Omni Code2Wav both the eager forward AND the CUDA-graph wrapper - # (_trim_replay_output -> actual*upsample - 555) are short, so it applies to both. - tail = max(0, int(code_seq_len * self.total_upsample) - batch_wav.shape[-1]) + # Refill the previous chunk's tail using the re-decoded left context, + # and exclude the current row's padded tail just as singleton decode does. start = max(0, left_context_size[idx] * self.total_upsample - tail) - wav_chunk = batch_wav[idx, :, start : code_seq_len * self.total_upsample] + end = max(0, code_seq_len * self.total_upsample - tail) + wav_chunk = batch_wav[idx, :, start:end] wavs.append(wav_chunk) return wavs From 2185c4e26f9aded75553760bf2d29bf175a2d137 Mon Sep 17 00:00:00 2001 From: tlysanhuo <166924864+tlysanhuo@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:51:49 +0800 Subject: [PATCH 095/209] [Bugfix] Tolerate a bounded TTS tail in the audio-text similarity gate (#7358) Signed-off-by: tly <2200895168@qq.com> --- tests/helpers/assertions.py | 43 ++++++++++++++++++++++- tests/helpers/tests/test_assertions.py | 47 ++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 1 deletion(-) diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index dd25aca71..7dc7a6ab3 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -171,6 +171,38 @@ def _short_transcript_contains_expected(transcript: str, expected: str) -> bool: return short_text and small_word_delta and expected_clean in transcript_clean +_MAX_TAIL_WORDS_FLOOR = 2 +_TAIL_WORD_RATIO = 0.2 + + +def _transcript_has_bounded_tail(transcript: str, expected: str) -> bool: + """Pass when the expected text is spoken verbatim, only trailed by noise. + + A short, unrelated tail after the complete expected sentence is the exact + case #6828's rationale anticipated: a different valid TTS sample can append + an audible tail even when the spoken text is identical, and Whisper can + hallucinate brief words on a non-speech tail. The length-penalized n-gram + cosine stays below the gate for such tails, so accept the match when the + expected text is a word-aligned prefix of the transcript and the trailing + words stay within the floor/ratio bounds. + """ + transcript_clean = preprocess_text(transcript) + expected_clean = preprocess_text(expected) + if not transcript_clean or not expected_clean: + return False + + transcript_words = transcript_clean.split() + expected_words = expected_clean.split() + if not transcript_words or not expected_words: + return False + + if transcript_words[: len(expected_words)] != expected_words: + return False + tail_words = len(transcript_words) - len(expected_words) + max_tail_words = max(_MAX_TAIL_WORDS_FLOOR, math.ceil(_TAIL_WORD_RATIO * len(expected_words))) + return tail_words <= max_tail_words + + def assert_image_diffusion_response( response: "DiffusionResponse", request_config: dict[str, Any], @@ -834,7 +866,16 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve text_output.lower(), ) print(f"similarity is: {similarity}") - assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE + if similarity <= similarity_threshold and _transcript_has_bounded_tail(transcript, text_output): + # The full answer is spoken verbatim and only a short + # noise tail follows it; see #6828's rationale. + print( + "bounded-tail containment check passed: " + f"text={text_output!r} is a word-aligned prefix of " + f"transcript={transcript!r}" + ) + else: + assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE if audio_ref_text: assert transcript is not None, "No audio transcript for reference-text validation" audio_similarity = cosine_similarity_text( diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index 98deaa511..e9b1a771c 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -126,3 +126,50 @@ def test_escalated_transcript_keeps_declared_language(monkeypatch): assert captured["model_size"] == "large-v3" assert captured["language"] == "en" + + +def test_bounded_tail_after_complete_answer_passes(): + # The exact shape from the #6815 nightly failure: the full answer is spoken + # verbatim, then a two-word unrelated tail trips the cosine gate. + assert assertions._transcript_has_bounded_tail( + "The squares in this image are black. nack shit.", + "The squares in this image are black.", + ) + + +def test_bounded_tail_exact_match_passes(): + assert assertions._transcript_has_bounded_tail( + "The squares in this image are black.", + "The squares in this image are black.", + ) + + +def test_bounded_tail_rejects_different_content(): + assert not assertions._transcript_has_bounded_tail( + "The circles in this picture are black and white.", + "The squares in this image are black.", + ) + + +def test_bounded_tail_rejects_long_tail(): + # Eight expected words allow at most max(2, ceil(0.2 * 8)) = 2 tail words. + assert not assertions._transcript_has_bounded_tail( + "The squares in this image are black and some extra words follow here", + "The squares in this image are black.", + ) + + +def test_bounded_tail_rejects_word_boundary_split(): + assert not assertions._transcript_has_bounded_tail( + "The squares in this image are blacks.", + "The squares in this image are black.", + ) + + +def test_bounded_tail_allows_ratio_bound_for_long_answers(): + expected = " ".join(["word"] * 20) + four_word_tail = expected + " one two three four" + five_word_tail = expected + " one two three four five" + # Twenty expected words allow up to max(2, ceil(0.2 * 20)) = 4 tail words. + assert assertions._transcript_has_bounded_tail(four_word_tail, expected) + assert not assertions._transcript_has_bounded_tail(five_word_tail, expected) From 3eb5c3b4e6c9a26c5d37858bbf3f75bdb6d17bf0 Mon Sep 17 00:00:00 2001 From: HAAZZZEEEE <49724853+HAAZZZEEEE@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:06:41 +0800 Subject: [PATCH 096/209] [Bugfix][Wan2.2] Apply UniPC flow shift once (#7576) Signed-off-by: HAAZZZEEEE --- tests/diffusion/models/wan2_2/conftest.py | 5 +- .../wan2_2/fixtures/native_unipc_shift12.json | 97 +++++++++++++++++++ .../models/wan2_2/test_wan22_i2v_pipeline.py | 34 +++++-- .../wan2_2/test_wan22_pipeline_diffuse.py | 33 ++++++- .../models/wan2_2/test_wan22_s2v_pipeline.py | 44 ++++++++- .../models/wan2_2/test_wan22_scheduler.py | 54 +++++++++++ .../models/wan2_2/test_wan22_vace_pipeline.py | 40 ++++++-- .../models/wan2_2/pipeline_wan2_2.py | 9 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 5 +- .../models/wan2_2/pipeline_wan2_2_s2v.py | 6 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 5 +- 11 files changed, 298 insertions(+), 34 deletions(-) create mode 100644 tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json create mode 100644 tests/diffusion/models/wan2_2/test_wan22_scheduler.py diff --git a/tests/diffusion/models/wan2_2/conftest.py b/tests/diffusion/models/wan2_2/conftest.py index f836fa545..cc95914c0 100644 --- a/tests/diffusion/models/wan2_2/conftest.py +++ b/tests/diffusion/models/wan2_2/conftest.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from contextlib import contextmanager @@ -33,7 +36,7 @@ class StubScheduler: self.config = SimpleNamespace(num_train_timesteps=1000) self.set_timesteps_calls: list[tuple[int, torch.device]] = [] - def set_timesteps(self, num_steps: int, device: torch.device) -> None: + def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None: self.set_timesteps_calls.append((num_steps, device)) diff --git a/tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json b/tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json new file mode 100644 index 000000000..9b6fd2433 --- /dev/null +++ b/tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json @@ -0,0 +1,97 @@ +{ + "source_url": "https://github.com/Wan-Video/Wan2.2/blob/42bf4cfaa384bc21833865abc2f9e6c0e67233dc/wan/utils/fm_solvers_unipc.py", + "source_sha256": "0dec8c7ed17f6f2049275c6848113314da6ccec1c8db5bdc89df43c05c6038d9", + "git_head": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc", + "num_train_timesteps": 1000, + "use_dynamic_shifting": false, + "reproduction": "Load FlowUniPCMultistepScheduler from source_url; construct with num_train_timesteps=1000, shift=1.0, use_dynamic_shifting=False; call set_timesteps(40, device='cpu', shift=12.0); record timesteps.tolist() and sigmas.tolist().", + "constructor_shift": 1.0, + "runtime_shift": 12.0, + "steps": 40, + "timesteps": [ + 999, + 997, + 995, + 993, + 990, + 988, + 985, + 982, + 979, + 976, + 972, + 969, + 965, + 961, + 956, + 952, + 947, + 941, + 936, + 929, + 922, + 915, + 907, + 898, + 888, + 877, + 865, + 852, + 837, + 819, + 799, + 776, + 749, + 717, + 678, + 631, + 571, + 492, + 386, + 235 + ], + "sigmas": [ + 0.9999166131019592, + 0.9977826476097107, + 0.9955461621284485, + 0.9931996464729309, + 0.9907346963882446, + 0.988142192363739, + 0.9854119420051575, + 0.9825326800346375, + 0.9794917702674866, + 0.9762753844261169, + 0.9728676676750183, + 0.9692512154579163, + 0.9654061794281006, + 0.961310088634491, + 0.9569376707077026, + 0.95225989818573, + 0.9472436308860779, + 0.9418508410453796, + 0.9360373020172119, + 0.9297519326210022, + 0.9229347705841064, + 0.9155154228210449, + 0.9074103832244873, + 0.8985199332237244, + 0.888724148273468, + 0.8778772950172424, + 0.865800678730011, + 0.8522725701332092, + 0.8370144367218018, + 0.8196719288825989, + 0.7997865080833435, + 0.7767546772956848, + 0.7497654557228088, + 0.717703104019165, + 0.6789888143539429, + 0.6313128471374512, + 0.5711563229560852, + 0.49288034439086914, + 0.38684695959091187, + 0.23510953783988953, + 0.0 + ], + "implementation": "Wan-Video/Wan2.2 wan/utils/fm_solvers_unipc.py" +} diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py index 26e78ddab..70bb146ab 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +import numpy as np import pytest import torch from PIL import Image from torch import nn from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar -from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( Wan22I2VPipeline, get_wan22_i2v_post_process_func, @@ -50,7 +51,7 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: def _make_i2v_sampling(**overrides): - values = { + values: dict[str, object] = { "height": 16, "width": 16, "num_frames": 5, @@ -275,16 +276,18 @@ def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None: assert first_frame_mask.shape == (2, 1, 2, 2, 2) -def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None: +@pytest.mark.parametrize("solver", ["unipc", "euler"]) +@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0]) +def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch, solver: str, shift: float) -> None: monkeypatch.setattr( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform", SimpleNamespace(is_available=lambda: False), ) pipeline = _make_i2v_pipeline(expand_timesteps=True) - pipeline.scheduler = StubScheduler([9]) - pipeline.od_config = SimpleNamespace(flow_shift=5.0) - pipeline._sample_solver = "unipc" - pipeline._flow_shift = 5.0 + pipeline.scheduler = build_wan_scheduler(solver, shift) + pipeline.od_config = SimpleNamespace(flow_shift=shift) + pipeline._sample_solver = solver + pipeline._flow_shift = shift pipeline.boundary_ratio = 0.875 pipeline.has_image_encoder = False pipeline._guidance_scale = None @@ -327,6 +330,8 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) - generator=gen_a, latents=latents_a, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), SimpleNamespace( @@ -336,12 +341,25 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) - generator=gen_b, latents=latents_b, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), ] ) outputs = pipeline.forward(batch) + if solver == "unipc": + sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1] + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + torch.testing.assert_close( + pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0 + ) + assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist() + else: + reference = build_wan_scheduler("euler", shift) + reference.set_timesteps(50, device="cpu") + torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0) assert prepare_call["batch_size"] == 4 assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 6f7833ff8..526eeeb15 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -10,7 +10,7 @@ import torch from torch import nn from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange -from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -19,6 +19,12 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.fixture(autouse=True) +def _cpu_platform(monkeypatch): + module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2") + monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False) + + class _StubTransformer(nn.Module): @property def dtype(self) -> torch.dtype: @@ -51,10 +57,10 @@ class _StubScheduler: def __init__(self, timesteps: list[int]) -> None: self.timesteps = torch.tensor(timesteps, dtype=torch.int64) self.config = SimpleNamespace(num_train_timesteps=1000) - self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = [] - def set_timesteps(self, num_steps: int, device: torch.device) -> None: - self.set_timesteps_calls.append((num_steps, device)) + def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None: + self.set_timesteps_calls.append((num_steps, device, shift)) @contextmanager @@ -184,7 +190,24 @@ def test_forward_delegates_denoising_to_diffuse( assert captured["boundary_timestep"] == pytest.approx(875.0) assert captured["latent_condition"] is None assert captured["first_frame_mask"] is None - assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))] + assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)] + + +@pytest.mark.parametrize("solver", ["unipc", "euler"]) +def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None: + pipeline = _make_pipeline() + pipeline.diffuse = lambda **kwargs: kwargs["latents"] + for shift in (3.0, 12.0, 5.0): + sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift}) + request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling) + pipeline.forward(DiffusionRequestBatch(requests=[request])) + reference = build_wan_scheduler(solver, shift) + if solver == "unipc": + reference.set_timesteps(5, device="cpu", shift=shift) + assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 + else: + reference.set_timesteps(5, device="cpu") + torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0) def test_forward_batches_text_generators_latents_and_splits_outputs() -> None: diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index 3b49d7edb..94f8b86fd 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -7,6 +10,7 @@ import pytest import torch import torch.nn as nn +from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( Wan22S2VPipeline, _make_clip_generators, @@ -14,11 +18,32 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.mark.parametrize("configured_shift", [None, 12.0]) +def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None: + module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v" + config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False) + + def init_components(pipeline, *args): + pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)) + + with ( + patch(f"{module}.get_local_device", return_value=torch.device("cpu")), + patch(f"{module}._resolve_model_path", return_value="unused"), + patch(f"{module}._is_diffusers_format", return_value=True), + patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components), + patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"), + ): + pipeline = Wan22S2VPipeline(od_config=config) + assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift) + assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 + assert pipeline.scheduler.sigma_max == float(np.float32(0.999)) def _make_s2v_sampling(**overrides): - values = { + values: dict[str, object] = { "height": 16, "width": 16, "num_frames": 8, @@ -211,7 +236,8 @@ def test_s2v_forward_rejects_different_raw_audio_shapes() -> None: pipeline.forward(batch) -def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: +@pytest.mark.parametrize("shift", [3.0, 12.0]) +def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None: pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") @@ -227,7 +253,8 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: enable_cpu_offload=False, parallel_config=SimpleNamespace(use_hsdp=False), ) - pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0])) + pipeline._flow_shift = shift + pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.0) pipeline.vae_scale_factor_spatial = 8 pipeline.resolution_divisor = 16 pipeline.motion_frames = 7 @@ -262,6 +289,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: "multi_modal_data": {"image": image, "audio": audio_a}, }, sampling_params=_make_s2v_sampling( + num_inference_steps=5, num_outputs_per_prompt=2, generator=generators_a, latents=latents_a, @@ -275,6 +303,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: "multi_modal_data": {"image": image, "audio": audio_b}, }, sampling_params=_make_s2v_sampling( + num_inference_steps=5, num_outputs_per_prompt=2, generator=generators_b, latents=latents_b, @@ -287,6 +316,12 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: platform.is_available.return_value = False outputs = pipeline.forward(batch) + expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1] + expected = shift * expected / (1.0 + (shift - 1.0) * expected) + torch.testing.assert_close( + pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0 + ) + assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 assert len(outputs) == 2 assert outputs[0].output[0].shape[0] == 2 assert outputs[1].output[0].shape[0] == 2 @@ -519,6 +554,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): mock_scheduler.timesteps = torch.linspace(999, 0, 5) mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),)) pipeline.scheduler = mock_scheduler + pipeline._flow_shift = 3.0 # -- Bind methods from the real class -- pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline) diff --git a/tests/diffusion/models/wan2_2/test_wan22_scheduler.py b/tests/diffusion/models/wan2_2/test_wan22_scheduler.py new file mode 100644 index 000000000..87cc24670 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_scheduler.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic.""" + +import json +from pathlib import Path + +import numpy as np +import numpy.typing as npt +import pytest +import torch + +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + FASTWAN_DMD_SCHEDULER_SHIFT, + build_wan_scheduler, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_wan_unipc_matches_captured_native_full_schedule() -> None: + # The fixture pins the official Wan source and records how to regenerate it. + golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text()) + scheduler = build_wan_scheduler("unipc", golden["runtime_shift"]) + for steps in (golden["steps"], 20, golden["steps"]): + scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"]) + if steps == golden["steps"]: + assert scheduler.timesteps.tolist() == golden["timesteps"] + torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0) + + +@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0]) +@pytest.mark.parametrize("steps", [20, 40]) +def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None: + # Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0; + # interpolate in float64, apply the runtime shift, then cast sigma to FP32. + unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1] + shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted) + expected_timesteps = (shifted * 1000).astype(np.int64) + expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32) + scheduler = build_wan_scheduler("unipc", shift) + scheduler.set_timesteps(steps, device="cpu", shift=shift) + assert scheduler.config.shift == scheduler.config["shift"] == 1.0 + torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0) + torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0) + + +def test_wan_euler_keeps_dmd_shift_semantics() -> None: + scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT) + scheduler.set_timesteps(3, device="cpu") + # The unchanged Euler branch starts at sigma=1, with shift=8 for DMD. + expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0]) + torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7) + torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4) diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 5f712bc8a..75cf868ce 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -1,15 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass, field from types import SimpleNamespace +import numpy as np import pytest import torch from PIL import Image from torch import nn from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( Wan22VACEPipeline, create_vace_transformer_from_config, @@ -33,6 +35,8 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp ) enable_layerwise_offload: bool = True post_init_calls: int = 0 + _resolved_diffusion_offload: object = None + _diffusion_offload_flags_materialized: bool = False def __post_init__(self) -> None: self.post_init_calls += 1 @@ -41,7 +45,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp resolved_offload = object() config._resolved_diffusion_offload = resolved_offload config._diffusion_offload_flags_materialized = True - captured = {} + captured: dict[str, object] = {} def capture_base_init(self, *, od_config, prefix="") -> None: captured.update(od_config=od_config, prefix=prefix) @@ -51,6 +55,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp Wan22VACEPipeline(od_config=config, prefix="stage") forwarded = captured["od_config"] + assert isinstance(forwarded, RuntimeConfig) assert forwarded is not config assert config.flow_shift is None assert forwarded.flow_shift == 3.0 @@ -75,7 +80,7 @@ def _make_vace_pipeline() -> Wan22VACEPipeline: def _make_vace_sampling(**overrides): - values = { + values: dict[str, object] = { "height": 16, "width": 16, "num_frames": 5, @@ -123,7 +128,7 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None: def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None: - captured = {} + captured: dict[str, object] = {} class FakeVACETransformer: def __init__(self, **kwargs) -> None: @@ -203,7 +208,9 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None: torch.testing.assert_close(result, torch.ones_like(latents)) -def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None: +@pytest.mark.parametrize("solver", ["unipc", "euler"]) +@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0]) +def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> None: monkeypatch.setattr( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", SimpleNamespace(is_available=lambda: False), @@ -211,10 +218,10 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N pipeline = _make_vace_pipeline() pipeline.transformer.vace_patch_embedding = None pipeline.transformer_2 = None - pipeline.scheduler = StubScheduler([9]) - pipeline.od_config = SimpleNamespace(flow_shift=3.0) - pipeline._sample_solver = "unipc" - pipeline._flow_shift = 3.0 + pipeline.scheduler = build_wan_scheduler(solver, shift) + pipeline.od_config = SimpleNamespace(flow_shift=shift) + pipeline._sample_solver = solver + pipeline._flow_shift = shift pipeline.boundary_ratio = None pipeline._guidance_scale = None pipeline._num_timesteps = None @@ -247,6 +254,8 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N generator=gen_a, latents=latents_a, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), SimpleNamespace( @@ -256,12 +265,25 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N generator=gen_b, latents=latents_b, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), ] ) outputs = pipeline.forward(batch) + if solver == "unipc": + sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1] + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + torch.testing.assert_close( + pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0 + ) + assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist() + else: + reference = build_wan_scheduler("euler", shift) + reference.set_timesteps(50, device="cpu") + torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0) assert prepare_call["batch_size"] == 4 assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index a285e0009..15b036864 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -59,9 +59,11 @@ FASTWAN_DMD_SCHEDULER_SHIFT = 8.0 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any: if sample_solver == "unipc": + # Keep native Wan's unshifted training endpoints; set_timesteps applies + # the requested shift to the interpolated inference sigmas. return FlowUniPCMultistepScheduler( num_train_timesteps=1000, - shift=flow_shift, + shift=1.0, prediction_type="flow_prediction", ) if sample_solver == "euler": @@ -774,7 +776,10 @@ class Wan22Pipeline( self._sample_solver = sample_solver self._flow_shift = flow_shift - self.scheduler.set_timesteps(num_steps, device=device) + if sample_solver == "unipc": + self.scheduler.set_timesteps(num_steps, device=device, shift=flow_shift) + else: + self.scheduler.set_timesteps(num_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) boundary_timestep = None diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 908a35020..01bd53bb8 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -629,7 +629,10 @@ class Wan22I2VPipeline( self._flow_shift = flow_shift # Timesteps - self.scheduler.set_timesteps(num_steps, device=device) + if sample_solver == "unipc": + self.scheduler.set_timesteps(num_steps, device=device, shift=flow_shift) + else: + self.scheduler.set_timesteps(num_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index ea99e065f..1d28b425d 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -545,10 +545,10 @@ class Wan22S2VPipeline( self._init_original_format(model_path, dtype) # -- Scheduler -- - flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0 # S2V default + self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0 # S2V default self.scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, - shift=flow_shift, + shift=1.0, prediction_type="flow_prediction", ) @@ -1342,7 +1342,7 @@ class Wan22S2VPipeline( ) # -- Scheduler -- - self.scheduler.set_timesteps(num_steps, device=device) + self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 0d48dffb5..aeecff608 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -690,7 +690,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): ) # Set up scheduler - self.scheduler.set_timesteps(num_inference_steps, device=device) + if self._sample_solver == "unipc": + self.scheduler.set_timesteps(num_inference_steps, device=device, shift=self._flow_shift) + else: + self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) From 3319f6ca7891c53b267ae6a6beb65d128362e8f5 Mon Sep 17 00:00:00 2001 From: Genghong Hu <111997394+GodHu777777@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:37:50 +0800 Subject: [PATCH 097/209] [CI/Build][Hardware][Ascend] Add torchmetrics to the NPU CI image (#7570) Signed-off-by: Genghong Hu --- docker/Dockerfile.npu.ci | 1 + 1 file changed, 1 insertion(+) diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index b65b803d6..a25155111 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -44,6 +44,7 @@ RUN python3 -m pip install -U pip && \ "onnx>=1.16.0" \ "librosa>=0.10.0" \ "datasets>=3.0.0" \ + "torchmetrics==1.9.0" \ --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ --default-timeout=60 From 7bceb1a76b99652b7a7abb315c7cc6f1f952b967 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:49:07 +0800 Subject: [PATCH 098/209] [Bugfix] Align video frame consumption reporting with prompt sampling (#7329) --- docs/serving/video_stream_api.md | 8 + .../openai_api/test_serving_video_stream.py | 194 ++++++++++++++++++ .../openai/serving_video_stream.py | 20 +- .../entrypoints/openai/video_stream_base.py | 32 ++- 4 files changed, 232 insertions(+), 22 deletions(-) diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index 101289fc0..23aa4cb08 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -50,6 +50,14 @@ WebSocket /v1/video/chat/stream | Server -> Client | `session.done` | none | Session closed. | | Server -> Client | `error` | `message` | Recoverable protocol or generation error. | +### Frame Consumption Reporting + +When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt. + +Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path. + +**Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`. + ### `session.config` Fields | Field | Type | Default | Description | diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index f7267a061..c9c4f6a4b 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -375,6 +375,200 @@ async def test_video_frames_consumed_is_emitted_after_engine_uses_frame_prompt() ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "num_frames,bad_index,expected_indices", + [ + pytest.param(2, 4, [0], id="bad-last-no-refill"), + pytest.param(1, 4, [], id="all-selected-frames-bad"), + pytest.param(8, 4, [0, 1, 2, 3], id="below-sampling-limit"), + pytest.param(2, 1, [0, 4], id="bad-frame-not-selected"), + pytest.param(2, 0, [4], id="bad-first-frame"), + ], +) +async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot( + monkeypatch, num_frames, bad_index, expected_indices +): + """A query queued before bad-frame cleanup must report only its prompt images.""" + query_queued = asyncio.Event() + decode_failed = asyncio.Event() + release_decode = asyncio.Event() + pong_blocked = asyncio.Event() + release_pong = asyncio.Event() + turn_done = asyncio.Event() + original_to_thread = asyncio.to_thread + bad_image = b"not-a-jpeg" + frames = [_make_jpeg(r=index * 40) for index in range(5)] + frames[bad_index] = bad_image + engine_prompts = [] + + class ObservedQueue(asyncio.Queue): + def put_nowait(self, item): + super().put_nowait(item) + if isinstance(item, dict): + if item.get("type") == "video.query": + query_queued.set() + elif item.get("type") == "_internal.frame_decode_failed": + decode_failed.set() + + async def gated_to_thread(function, *args, **kwargs): + if function is video_stream_base._decode_frame_bytes and args[0] == bad_image: + await release_decode.wait() + return await original_to_thread(function, *args, **kwargs) + + class BackpressuredWebSocket(TimedWebSocket): + async def send_json(self, data): + await super().send_json(data) + if data["type"] == "pong": + pong_blocked.set() + await release_pong.wait() + elif data["type"] == "response.text.done": + turn_done.set() + + class OneOutputEngine: + async def generate(self, **kwargs): + engine_prompts.append(kwargs["prompt"]) + yield _text_result("visible") + + class CapturingHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"messages": request.messages} + + # Construct the client queue before observing the server's message queue. + ws = BackpressuredWebSocket() + monkeypatch.setattr(asyncio, "Queue", ObservedQueue) + monkeypatch.setattr(asyncio, "to_thread", gated_to_thread) + handler = CapturingHandler(chat_service=object(), engine_client=OneOutputEngine(), idle_timeout=5.0) + ws.put( + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "num_frames": num_frames, + "enable_frame_filter": False, + } + ) + for index, frame in enumerate(frames): + ws.put({"type": "video.frame", "data": _b64(frame), "frame_id": f"frame-{index}", "pts_ms": index * 100}) + ws.put({"type": "ping"}) + ws.put({"type": "video.query", "text": "describe"}) + task = asyncio.create_task(handler.handle_session(ws)) + try: + # Hold the processor at a WebSocket send. Queue the query first, then + # finish the real failing decode so its cleanup lands behind the query. + await asyncio.wait_for(pong_blocked.wait(), timeout=5.0) + await asyncio.wait_for(query_queued.wait(), timeout=5.0) + release_decode.set() + await asyncio.wait_for(decode_failed.wait(), timeout=5.0) + release_pong.set() + await asyncio.wait_for(turn_done.wait(), timeout=5.0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=5.0) + finally: + release_decode.set() + release_pong.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert len(engine_prompts) == 1 + images = [part for part in engine_prompts[0]["messages"][-1]["content"] if part["type"].startswith("image_")] + # Expected positions preserve sample-then-filter: do not refill a dropped + # position or shift the remaining metadata when a frame fails to decode. + assert len(images) == len(expected_indices) + for image, index in zip(images, expected_indices): + if image["type"] == "image_pil": + assert image["image_pil"].tobytes() == Image.open(io.BytesIO(frames[index])).convert("RGB").tobytes() + else: + assert image["image_url"]["url"] == f"data:image/jpeg;base64,{_b64(frames[index])}" + consumed_events = [message for message in ws.sent if message.get("type") == "video.frames.consumed"] + assert len(consumed_events) == 1 + consumed = consumed_events[0] + expected_ids = [f"frame-{index}" for index in expected_indices] + assert consumed["frame_ids"] == expected_ids + assert [frame["frame_id"] for frame in consumed["frames"]] == expected_ids + assert [frame["pts_ms"] for frame in consumed["frames"]] == [index * 100 for index in expected_indices] + assert consumed["latest_pts_ms"] == (expected_indices[-1] * 100 if expected_indices else None) + assert {"type": "error", "message": "Frame decode failed"} in ws.sent + assert ws.sent[-1]["type"] == "session.done" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "num_frames,frame_ids,repeat_image,expected_indices", + [ + pytest.param(2, ["f0", "f1", "f2", "f3", "f4"], False, [0, 4], id="stride-and-last"), + pytest.param(1, ["f0", "f1", "f2"], False, [2], id="last-only"), + pytest.param(8, ["f0", "f1", "f2"], False, [0, 1, 2], id="below-sampling-limit"), + pytest.param(2, [None, "f1", None, "f3", None], False, [0, 4], id="unidentified-selected-frames"), + pytest.param(2, [None, None, None], False, [0, 2], id="no-frame-ids"), + pytest.param(2, ["same", "same", "same"], False, [0, 2], id="repeated-frame-id"), + pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"), + ], +) +async def test_video_frame_selection_preserves_uncached_frames( + monkeypatch, num_frames, frame_ids, repeat_image, expected_indices +): + """Prewarm misses still use image URLs, with metadata paired by buffer position.""" + engine_prompts = [] + pending_decode = asyncio.Event() + + async def blocked_to_thread(*args, **kwargs): + await pending_decode.wait() + + monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread) + + class OneOutputEngine: + async def generate(self, **kwargs): + engine_prompts.append(kwargs["prompt"]) + yield _text_result("visible") + + class CapturingHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"messages": request.messages} + + frames = [_b64(_make_jpeg(r=0 if repeat_image else index * 40)) for index in range(len(frame_ids))] + messages = [ + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "num_frames": num_frames, + "enable_frame_filter": False, + }, + *[ + {"type": "video.frame", "data": frame, "frame_id": frame_ids[index], "pts_ms": index * 100} + for index, frame in enumerate(frames) + ], + {"type": "video.query", "text": "describe"}, + {"type": "video.done"}, + ] + ws = MockWebSocket([json.dumps(message) for message in messages]) + handler = CapturingHandler(chat_service=object(), engine_client=OneOutputEngine()) + await asyncio.wait_for(handler.handle_session(ws), timeout=5.0) + + assert not [message for message in ws.sent if message["type"] == "error"] + assert len(engine_prompts) == 1 + content = engine_prompts[0]["messages"][-1]["content"] + assert content == [ + *[ + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}} + for index in expected_indices + ], + {"type": "text", "text": "describe"}, + ] + consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"] + if not any(frame_ids): + assert consumed_events == [] + else: + assert len(consumed_events) == 1 + consumed = consumed_events[0] + assert consumed["frame_ids"] == [frame_ids[index] for index in expected_indices if frame_ids[index] is not None] + assert [frame["frame_id"] for frame in consumed["frames"]] == [frame_ids[index] for index in expected_indices] + assert [frame["pts_ms"] for frame in consumed["frames"]] == [index * 100 for index in expected_indices] + assert consumed["latest_pts_ms"] == expected_indices[-1] * 100 + + @pytest.mark.asyncio async def test_audio_in_video_sets_mm_processor_kwargs(): captured_requests = [] diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index 5b30ff0fa..5f7334bd2 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -30,7 +30,6 @@ from collections.abc import Mapping from typing import Any from vllm_omni.entrypoints.openai.video_stream_base import ( - _BAD_FRAME, _DEFAULT_CONFIG_TIMEOUT, _DEFAULT_IDLE_TIMEOUT, PrewarmedFrame, @@ -62,22 +61,19 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): message_history: list[dict[str, Any]], query_text: str, prewarmed_frames: Mapping[str, PrewarmedFrame], + *, + frame_indices: list[int] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - n_buf = len(frame_buffer) - if n_buf <= config.num_frames: - frames = list(frame_buffer) - else: - stride = max(1, n_buf // config.num_frames) - idx = [i * stride for i in range(config.num_frames - 1)] + [n_buf - 1] - frames = [frame_buffer[i] for i in idx] - prewarmed = prewarmed_frames or {} + if frame_indices is None: + frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed) user_content: list[dict] = [] - for frame_b64 in frames: + for index in frame_indices: + frame_b64 = frame_buffer[index] cached = prewarmed.get(frame_b64) - if cached is _BAD_FRAME: - continue if cached is not None: + # The shared selector excludes failed decodes before this point. + assert isinstance(cached, tuple) pil, pil_uuid = cached user_content.append( { diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index 8885e9a33..52ae9d894 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -93,8 +93,10 @@ class VideoStreamPipelineHooks(Protocol): message_history: list[dict[str, Any]], query_text: str, prewarmed_frames: Mapping[str, PrewarmedFrame], + *, + frame_indices: list[int] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Build OpenAI-style messages and the current user message.""" + """Build messages using the supplied frame selection, or sample if omitted.""" ... def on_turn_complete( @@ -179,6 +181,8 @@ class OmniStreamingVideoHandler: message_history: list[dict[str, Any]], query_text: str, prewarmed_frames: Mapping[str, PrewarmedFrame], + *, + frame_indices: list[int] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: raise NotImplementedError @@ -639,6 +643,9 @@ class OmniStreamingVideoHandler: ChatCompletionRequest, ) + # Select once from this turn's snapshot so prompt images and the + # consumed event refer to exactly the same frames, including decode failures. + frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames) messages, user_message = self.build_engine_prompt( config, frame_buffer, @@ -646,6 +653,7 @@ class OmniStreamingVideoHandler: message_history, query_text, prewarmed_frames, + frame_indices=frame_indices, ) request_kwargs: dict[str, Any] = { @@ -678,7 +686,7 @@ class OmniStreamingVideoHandler: await self._send_error(websocket, f"Preprocess failed: {e}") return decoded_ready_ts_ms = _time.monotonic() * 1000 - selected_metadata = self._sample_frame_metadata(frame_metadata or [], config.num_frames) + selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else [] model_selected_ts_ms = _time.monotonic() * 1000 await websocket.send_json({"type": "response.start"}) @@ -1071,15 +1079,19 @@ class OmniStreamingVideoHandler: pass @staticmethod - def _sample_frame_metadata( - frame_metadata: list[dict[str, Any]], + def _sample_frame_indices( + frame_buffer: list[str], num_frames: int, - ) -> list[dict[str, Any]]: - if len(frame_metadata) <= num_frames: - return list(frame_metadata) - stride = max(1, len(frame_metadata) // num_frames) - indices = [index * stride for index in range(num_frames - 1)] + [len(frame_metadata) - 1] - return [frame_metadata[index] for index in indices] + prewarmed_frames: Mapping[str, PrewarmedFrame], + ) -> list[int]: + """Stride-sample with the last frame, then drop known bad frames without refilling.""" + n_buf = len(frame_buffer) + if n_buf <= num_frames: + indices = list(range(n_buf)) + else: + stride = max(1, n_buf // num_frames) + indices = [index * stride for index in range(num_frames - 1)] + [n_buf - 1] + return [index for index in indices if prewarmed_frames.get(frame_buffer[index]) is not _BAD_FRAME] @staticmethod async def _send_frame_ack( From 5dbae2b182c1680586a25928ea5089971fd5b515 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:08:49 -0700 Subject: [PATCH 099/209] [Bugfix] Per-key accumulation strategy for Qwen3-TTS codec frame outputs (#7608) Signed-off-by: Yueqian Lin --- tests/engine/test_output_modality.py | 88 +++++++++++++++++++ .../models/qwen3_tts/pipeline.py | 14 +++ vllm_omni/outputs/mm_outputs.py | 49 +++++++---- vllm_omni/outputs/output_modality.py | 39 +++++++- vllm_omni/outputs/output_processor.py | 9 +- 5 files changed, 174 insertions(+), 25 deletions(-) diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py index ff15e5824..e5c5d36c1 100644 --- a/tests/engine/test_output_modality.py +++ b/tests/engine/test_output_modality.py @@ -41,6 +41,7 @@ _mm_mod = _load_module( OutputModality = _om_mod.OutputModality TensorAccumulationStrategy = _om_mod.TensorAccumulationStrategy get_accumulation_strategy = _om_mod.get_accumulation_strategy +register_key_accumulation_strategy = _om_mod.register_key_accumulation_strategy MultimodalPayload = _mm_mod.MultimodalPayload MultimodalCompletionOutput = _mm_mod.MultimodalCompletionOutput @@ -156,3 +157,90 @@ def test_output_modality_printed_examples(capsys): captured = capsys.readouterr() assert "OutputModality Parsing" in captured.out assert "MultimodalPayload" in captured.out + + +def test_get_accumulation_strategy_key_override_precedence(): + """A registered per-key override wins over the modality-wide default; + any other key under the same modality keeps using that default.""" + assert get_accumulation_strategy(OutputModality.AUDIO, "unregistered.key") == TensorAccumulationStrategy.CONCAT_LAST + + register_key_accumulation_strategy("__test_override__.frames", TensorAccumulationStrategy.CONCAT_DIM0) + assert ( + get_accumulation_strategy(OutputModality.AUDIO, "__test_override__.frames") + == TensorAccumulationStrategy.CONCAT_DIM0 + ) + # A plain waveform key under the same AUDIO modality is unaffected. + assert get_accumulation_strategy(OutputModality.AUDIO, "waveform") == TensorAccumulationStrategy.CONCAT_LAST + assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST + + +def _talker_codec_frame_payload(audio_key: str, ref_key: str) -> MultimodalPayload: + """Synthetic stand-in for a TTS talker's accumulated per-step output. + + ``audio_key`` grows one ``[1, num_codebooks]`` row per decode step on top + of a ``[5, num_codebooks]`` prefill block (Qwen3-TTS's ``codes.audio``); + ``ref_key`` re-emits the same constant ``[100, num_codebooks]`` + reference-context matrix unchanged at every step (Qwen3-TTS's + ``codes.ref``). Neither is a waveform chunk. + """ + payload = MultimodalPayload() + payload.tensors[audio_key] = [ + torch.zeros(5, 16), + torch.full((1, 16), 1.0), + torch.full((1, 16), 2.0), + torch.full((1, 16), 3.0), + ] + ref_block = torch.arange(100 * 16, dtype=torch.float32).reshape(100, 16) + payload.tensors[ref_key] = [ref_block.clone(), ref_block.clone(), ref_block.clone()] + return payload + + +def test_consolidate_tensors_default_audio_strategy_corrupts_codec_frame_keys(): + """Reproduces both failure modes of forcing codec-frame keys through the + AUDIO modality's waveform-tuned default (CONCAT_LAST) with no per-key + override registered: the audio-frames key's dim-0 growth makes + CONCAT_LAST's dim=-1 concat raise (mismatched dim 0), which the + flatten-chunks fallback turns into a 1-D blob instead of the intended + ``[8, 16]``; the ref-frames key's matching last dim lets CONCAT_LAST + "succeed" by silently concatenating three identical ``[100, 16]`` blocks + into ``[100, 48]`` instead of keeping a single copy. + """ + payload = _talker_codec_frame_payload("unregistered.codes.audio", "unregistered.codes.ref") + + payload.consolidate_tensors(OutputModality.AUDIO) + + assert tuple(payload.tensors["unregistered.codes.audio"].shape) != (8, 16) + assert tuple(payload.tensors["unregistered.codes.ref"].shape) == (100, 48) + + +def test_consolidate_tensors_with_qwen3_tts_key_overrides_is_correct(): + """With the per-key overrides Qwen3-TTS's pipeline module registers + (``codes.audio`` -> CONCAT_DIM0, ``codes.ref`` -> REPLACE), both keys + consolidate correctly under the AUDIO modality despite not being + waveform chunks.""" + register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0) + register_key_accumulation_strategy("codes.ref", TensorAccumulationStrategy.REPLACE) + + payload = _talker_codec_frame_payload("codes.audio", "codes.ref") + expected_ref = payload.tensors["codes.ref"][0] + + payload.consolidate_tensors(OutputModality.AUDIO) + + assert tuple(payload.tensors["codes.audio"].shape) == (8, 16) + assert tuple(payload.tensors["codes.ref"].shape) == (100, 16) + assert torch.equal(payload.tensors["codes.ref"], expected_ref) + + +def test_consolidate_tensors_raises_with_key_name_instead_of_silently_keeping_last(): + """A concat failure under a non-CONCAT_LAST strategy must be surfaced + with the offending key name rather than resolved by silently keeping + only the last chunk. Previously the except-branch only special-cased + the literal key ``"audio"``, so any other mismatched key (e.g. a + codec-frame key routed to CONCAT_DIM0) was corrupted instead of + failing loudly. + """ + payload = MultimodalPayload() + payload.tensors["latent.frames"] = [torch.zeros(2, 8), torch.zeros(3, 9)] + + with pytest.raises(RuntimeError, match="latent.frames"): + payload.consolidate_tensors(OutputModality.LATENT) diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 6dfe8736d..dd33aeca8 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -10,9 +10,23 @@ from vllm_omni.config.stage_config import ( StageExecutionType, StagePipelineConfig, ) +from vllm_omni.outputs.output_modality import ( + TensorAccumulationStrategy, + register_key_accumulation_strategy, +) _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" +# The talker's codec-frame keys are not waveform chunks, so they must not go +# through the AUDIO modality's waveform-tuned default (CONCAT_LAST) whenever +# a stage carrying them consolidates its output under that modality (e.g. an +# AUDIO-typed final stage that forwards them directly instead of routing +# through code2wav): ``codes.audio`` grows one row per decode step along +# dim 0, and ``codes.ref`` re-emits the constant reference-context matrix +# unchanged at every step. +register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0) +register_key_accumulation_strategy("codes.ref", TensorAccumulationStrategy.REPLACE) + QWEN3_TTS_PIPELINE = PipelineConfig( model_type="qwen3_tts", default_deploy_config_name="qwen3_tts.yaml", diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py index b1393ba19..45d1fc6fc 100644 --- a/vllm_omni/outputs/mm_outputs.py +++ b/vllm_omni/outputs/mm_outputs.py @@ -16,7 +16,11 @@ import torch from vllm.logger import init_logger from vllm.outputs import CompletionOutput -from vllm_omni.outputs.output_modality import TensorAccumulationStrategy +from vllm_omni.outputs.output_modality import ( + OutputModality, + TensorAccumulationStrategy, + get_accumulation_strategy, +) from vllm_omni.outputs.utils import _is_tensor_list, _to_cpu logger = init_logger(__name__) @@ -46,21 +50,24 @@ def _consolidate_tensor_list( tensor_list: list[torch.Tensor], strategy: TensorAccumulationStrategy, ) -> torch.Tensor: - """Concatenate a deferred tensor list, with fallbacks on shape mismatch.""" + """Concatenate a deferred tensor list, with a fallback for CONCAT_LAST only. + + A shape mismatch under CONCAT_LAST is expected for genuine audio waveform + chunks (their non-last dims can vary between async-chunk steps), so that + case falls back to flattening each chunk before concatenating. Any other + strategy failing to concatenate means the tensors are not shaped the way + that strategy assumes (e.g. a key was assigned the wrong strategy), so it + is surfaced instead of silently keeping only the last chunk. + """ try: return _cat_tensors(tensor_list, strategy) - except RuntimeError: - # [TODO: this part is for async chunk, not for audio only. can we come up with a design - # that identify this by whether using async chunk instead of modality?] - if key != "audio": - logger.warning("Error concatenating tensor for key %s; keeping last tensor", key) - return tensor_list[-1] - # Audio chunks may have mismatched shapes; retry along the last dim, - # then fall back to flattening each chunk. - try: - return torch.cat(tensor_list, dim=-1) - except RuntimeError: - return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) + except RuntimeError as exc: + if strategy is not TensorAccumulationStrategy.CONCAT_LAST: + raise RuntimeError( + f"Error concatenating tensor list for key {key!r} under strategy {strategy}: {exc}" + ) from exc + logger.warning("CONCAT_LAST failed for key %s; flattening each chunk before concatenating", key) + return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None: @@ -167,12 +174,17 @@ class MultimodalPayload(Mapping): self.metadata[key] = value return self - def consolidate_tensors(self, strategy: TensorAccumulationStrategy) -> None: + def consolidate_tensors(self, modality: OutputModality) -> None: """Concatenate deferred tensor lists into single tensors. - Tensors are generated content accumulated as chunks, so lists are - concatenated according to *strategy* (e.g. audio chunks along the - time dimension, latent frames along the batch dimension). + Tensors are generated content accumulated as chunks, so each key's + list is concatenated according to the strategy + ``get_accumulation_strategy(modality, key)`` resolves for it (e.g. + audio waveform chunks along the time dimension, latent frames along + the batch dimension). Most keys share *modality*'s default, but a + key can be registered for a different strategy -- e.g. a codec-frame + matrix that grows along dim 0 rather than the waveform-tuned default + for its modality -- via ``register_key_accumulation_strategy``. """ # Relocate scalar metadata tensors (e.g. sample rate) that from_dict # routed into .tensors, so they take the REPLACE path via .metadata @@ -185,6 +197,7 @@ class MultimodalPayload(Mapping): for key, value in list(self.tensors.items()): if _is_tensor_list(value): + strategy = get_accumulation_strategy(modality, key) self.tensors[key] = _consolidate_tensor_list(key, value, strategy) def consolidate_metadata(self) -> None: diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 7cdbc1f19..412e61619 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/vllm_omni/outputs/output_modality.py @@ -115,8 +115,43 @@ class TensorAccumulationStrategy(Enum): """Replace previous tensor entirely with the latest one.""" -def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy: - """Determine tensor merge strategy from the multimodal flags.""" +# Per-tensor-key overrides for the modality-wide default below. A modality +# is a coarse hint (e.g. AUDIO tensors are usually waveform chunks that grow +# along the last dimension), but some models emit tensors under that same +# modality whose accumulation semantics differ -- e.g. discrete codec-frame +# matrices that grow along dim 0, or a constant reference/context tensor +# re-emitted unchanged at every step. Pipelines that produce such keys +# register the correct strategy here via ``register_key_accumulation_strategy`` +# instead of forcing every tensor of the modality through one default. +_KEY_ACCUMULATION_OVERRIDES: dict[str, TensorAccumulationStrategy] = {} + + +def register_key_accumulation_strategy(key: str, strategy: TensorAccumulationStrategy) -> None: + """Register a per-key override for ``get_accumulation_strategy``. + + Args: + key: The flattened multimodal output key this override applies to, + e.g. ``"codes.audio"`` as produced by + ``vllm_omni.data_entry_keys.flatten_payload`` from a nested + ``{"codes": {"audio": tensor}}`` payload. + strategy: The accumulation strategy to use for this key, regardless + of the modality its request is otherwise associated with. + """ + _KEY_ACCUMULATION_OVERRIDES[key] = strategy + + +def get_accumulation_strategy( + modality: OutputModality, + key: str | None = None, +) -> TensorAccumulationStrategy: + """Determine the tensor merge strategy for one output key. + + A registered per-key override (see ``register_key_accumulation_strategy``) + always wins; otherwise the strategy falls back to the modality-wide + default. + """ + if key is not None and key in _KEY_ACCUMULATION_OVERRIDES: + return _KEY_ACCUMULATION_OVERRIDES[key] if OutputModality.AUDIO in modality: return TensorAccumulationStrategy.CONCAT_LAST if OutputModality.IMAGE in modality or OutputModality.LATENT in modality: diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index bd60e4793..5d2ede43d 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -28,7 +28,7 @@ from vllm_omni.outputs.multimodal_accumulation import ( is_non_final_delta_audio_chunk, replace_snapshot_keys, ) -from vllm_omni.outputs.output_modality import OutputModality, get_accumulation_strategy +from vllm_omni.outputs.output_modality import OutputModality logger = init_logger(__name__) @@ -140,8 +140,8 @@ class OmniRequestState(RequestState): """Consolidate accumulated tensor lists into single tensors. Uses TensorAccumulationStrategy derived from the output modality - to determine concatenation behavior. Metadata values always use - REPLACE (keep latest). + (with any per-key override applied) to determine concatenation + behavior. Metadata values always use REPLACE (keep latest). """ if self.mm_accumulated.is_empty: return @@ -150,10 +150,9 @@ class OmniRequestState(RequestState): modality = OutputModality.from_string(self.mm_type) except (ValueError, KeyError): modality = OutputModality.TEXT - strategy = get_accumulation_strategy(modality) try: - self.mm_accumulated.consolidate_tensors(strategy) + self.mm_accumulated.consolidate_tensors(modality) self.mm_accumulated.consolidate_metadata() except (RuntimeError, TypeError, KeyError): logger.exception("Error consolidating multimodal tensors") From 275720b3782a321427b0a58a7efc2827bf583a62 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Wed, 16 Sep 2026 11:43:30 +0800 Subject: [PATCH 100/209] [Bugfix] Keep async video jobs queued until inference and abort on delete (#6759) Signed-off-by: ZhengWG Signed-off-by: wuli666 <421774554@qq.com> Co-authored-by: wuli666 <421774554@qq.com> --- docs/configuration/environment_variables.md | 1 + docs/serving/videos_api.md | 10 + tests/config/test_environment_variables.py | 2 +- tests/diffusion/test_diffusion_engine.py | 29 ++- .../test_diffusion_output_formatter.py | 11 + .../test_inline_stage_diffusion_client.py | 46 +++- tests/diffusion/test_stage_diffusion_proc.py | 44 ++++ .../openai_api/test_video_server.py | 202 +++++++++++++++++- tests/entrypoints/test_async_omni.py | 55 ++++- tests/entrypoints/test_omni_entrypoints.py | 3 +- .../config/environment_variable_inventory.py | 1 + vllm_omni/diffusion/data.py | 15 ++ vllm_omni/diffusion/diffusion_engine.py | 24 ++- .../inline_stage_diffusion_client.py | 13 +- vllm_omni/diffusion/output_formatter.py | 4 +- vllm_omni/diffusion/stage_diffusion_proc.py | 21 +- vllm_omni/engine/orchestrator.py | 7 +- vllm_omni/entrypoints/async_omni.py | 35 ++- vllm_omni/entrypoints/openai/api_server.py | 32 ++- vllm_omni/entrypoints/openai/serving_video.py | 27 ++- .../openai/video/generation/helpers.py | 6 +- vllm_omni/inputs/data.py | 5 + 22 files changed, 549 insertions(+), 44 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 3a37b52d2..331487fac 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -82,6 +82,7 @@ depends on the installed kernels and model path. | --- | --- | --- | --- | --- | | `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable | | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | +| `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental | | `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental | | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index c6e91df60..1152ebce3 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -109,6 +109,16 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 The final content is available from `/v1/videos/{video_id}/content` after the job status becomes `completed`. +`queued` means the request is still waiting for diffusion scheduler admission. +The status changes to `in_progress` when the scheduler first selects the +request for execution. + +`DELETE /v1/videos/{video_id}` issues a bounded engine abort +(`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend +task. Cancellation cleanup is also bounded and best-effort: it confirms +the abort was queued, and the current request batch may still drain. +The job is then re-read so a completed save is not orphaned. + ### Synchronous Response `POST /v1/videos/sync` blocks until generation finishes and returns raw video diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index ce608dca3..b73062689 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 27, + EnvironmentVariableCategory.PUBLIC_OMNI: 28, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 64, diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 1ee7424e9..77a6da5db 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -15,7 +15,12 @@ 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.data import ( + DIFFUSION_REQUEST_LIFECYCLE_KEY, + DIFFUSION_REQUEST_STARTED, + DiffusionOutput, + OmniDiffusionConfig, +) from vllm_omni.diffusion.diffusion_engine import ( DiffusionEngine, DiffusionExecutionMode, @@ -183,6 +188,28 @@ def _make_request_mode_sched_output(*request_ids: str) -> RealDiffusionScheduler ) +@pytest.mark.cpu +def test_request_started_output_is_emitted_only_for_opted_in_requests() -> None: + sched_output = _make_request_mode_sched_output("tracked", "untracked") + sched_output.scheduled_new_reqs[0].req.sampling_params.emit_request_lifecycle = True + engine = object.__new__(DiffusionEngine) + emitted = [] + engine._put_output = lambda request_id, output: emitted.append((request_id, output)) + + engine._emit_request_started_outputs(sched_output) + + assert len(emitted) == 1 + request_id, output = emitted[0] + assert request_id == "tracked" + assert output.request_started is True + assert output.finished is False + [formatted] = engine.postprocess_output(sched_output.scheduled_new_reqs[0].req, output) + assert formatted.custom_output == { + DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED, + } + assert formatted.finished is False + + class TestRequestBatchCapability: pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index 4578d8368..a1dfec637 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -443,3 +443,14 @@ def test_format_empty_diffusion_outputs_preserves_empty_response_shape() -> None assert [result.prompt for result in results] == ["prompt-0"] assert [result.images for result in results] == [[]] assert [result.metrics for result in results] == [{}] + + +def test_format_empty_diffusion_outputs_preserves_custom_output() -> None: + [result] = format_empty_diffusion_outputs( + _request("prompt-0"), + finished=False, + custom_output={"lifecycle": "started"}, + ) + + assert result.custom_output == {"lifecycle": "started"} + assert result.finished is False diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index e39c18ff4..3eadc53c1 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -7,7 +10,11 @@ from unittest.mock import MagicMock, patch import pytest from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + DIFFUSION_REQUEST_LIFECYCLE_KEY, + DIFFUSION_REQUEST_STARTED, + OmniDiffusionConfig, +) from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -80,6 +87,43 @@ async def test_inline_dispatch_request_success(client, mock_engine): assert output.request_id == "req-1" +@pytest.mark.asyncio +async def test_inline_non_streaming_dispatches_lifecycle_before_final(client, mock_engine): + lifecycle = OmniRequestOutput.from_diffusion( + request_id="req-lifecycle", + images=[], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + finished=False, + ) + intermediate = OmniRequestOutput.from_diffusion( + request_id="req-lifecycle", + images=[], + custom_output={"chunk": 0}, + finished=False, + ) + final = OmniRequestOutput.from_diffusion(request_id="req-lifecycle", images=[MagicMock()]) + + async def _step_streaming(_request): + yield [lifecycle] + yield [intermediate] + yield [final] + + mock_engine.step_streaming = _step_streaming + + await client.add_request_async("req-lifecycle", "A test prompt", OmniDiffusionSamplingParams()) + + outputs = [] + for _ in range(20): + output = client.get_diffusion_output_nowait() + if output is not None: + outputs.append(output) + if output.finished: + break + await asyncio.sleep(0.01) + + assert outputs == [lifecycle, final] + + @pytest.mark.asyncio async def test_inline_requests_clone_shared_sampling_params(client, mock_engine): requests = [] diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index 90e753c9b..c6b5761d9 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -9,6 +9,7 @@ import pytest import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc import vllm_omni.plugins as omni_plugins +from vllm_omni.diffusion.data import DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -103,6 +104,49 @@ async def test_proc_streaming_request_yields_each_engine_chunk(): assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}} +@pytest.mark.asyncio +async def test_proc_non_streaming_forwards_lifecycle_before_final_output(): + lifecycle = OmniRequestOutput.from_diffusion( + request_id="", + images=[], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + finished=False, + ) + intermediate = OmniRequestOutput.from_diffusion( + request_id="", + images=[], + custom_output={"chunk": 0}, + finished=False, + ) + final = OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True) + + class _LifecycleEngine: + async def step_streaming(self, request): + del request + yield [lifecycle] + yield [intermediate] + yield [final] + + stage_proc = object.__new__(StageDiffusionProc) + stage_proc._engine = _LifecycleEngine() + intermediate_outputs = [] + + async def _capture(output): + intermediate_outputs.append(output) + + result = await stage_proc._process_request( + request_id="req-lifecycle", + prompt="prompt", + sampling_params_dict=asdict(OmniDiffusionSamplingParams()), + on_request_started=_capture, + ) + + assert intermediate_outputs == [lifecycle] + assert lifecycle.request_id == "req-lifecycle" + assert result is final + assert result.request_id == "req-lifecycle" + + @pytest.mark.asyncio async def test_proc_process_request_with_batching_async_output(): stage_proc = object.__new__(StageDiffusionProc) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 57b2daaf8..2de5c4056 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -25,6 +25,7 @@ from PIL import Image from pytest_mock import MockerFixture from vllm import envs +from vllm_omni.diffusion.data import DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes from vllm_omni.entrypoints.openai import api_server, video_api_utils from vllm_omni.entrypoints.openai.api_server import router @@ -58,6 +59,7 @@ class MockVideoResult: multimodal_output=None, stage_durations=None, peak_memory_mb=0.0, + custom_output=None, ): self.multimodal_output = dict(multimodal_output or {"video": videos}) if audios is not None: @@ -66,6 +68,7 @@ class MockVideoResult: self.multimodal_output["audio_sample_rate"] = sample_rate self.stage_durations = stage_durations or {} self.peak_memory_mb = peak_memory_mb + self.custom_output = custom_output or {} class FakeAsyncOmni: @@ -102,9 +105,17 @@ class FakeAsyncOmni: ): self.captured_reference_video_bytes = [Path(item).read_bytes() for item in reference_videos] num_outputs = sampling_params_list[0].num_outputs_per_prompt + if sampling_params_list[0].emit_request_lifecycle: + yield MockVideoResult( + [], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + ) videos = [object() for _ in range(num_outputs)] yield MockVideoResult(videos) + async def abort(self, request_id, *, timeout=None): + del request_id, timeout + def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture): engine = FakeAsyncOmni() @@ -224,9 +235,18 @@ class BlockingVideoHandler: self.stage_configs = stage_configs async def generate_video_bytes( - self, request, reference_id, *, reference_image=None, reference_video=None, reference_audio=None + self, + request, + reference_id, + *, + reference_image=None, + reference_video=None, + reference_audio=None, + on_started=None, ): del request, reference_id, reference_image, reference_video, reference_audio + if on_started is not None: + await on_started() self.started.set() try: await asyncio.Future() @@ -234,6 +254,104 @@ class BlockingVideoHandler: self.cancelled.set() raise + async def abort_request(self, request_id): + del request_id + + +class HangingAbortHandler(BlockingVideoHandler): + """Engine abort never returns; DELETE must time out and still cancel.""" + + async def abort_request(self, request_id): + del request_id + await asyncio.sleep(30) + + +class CompletingDuringAbortHandler(BlockingVideoHandler): + """Finishes and persists output while DELETE is still awaiting abort.""" + + def __init__(self): + super().__init__() + self._finish = asyncio.Event() + + async def generate_video_bytes( + self, + request, + reference_id, + *, + reference_image=None, + reference_video=None, + reference_audio=None, + on_started=None, + ): + del request, reference_image, reference_video, reference_audio + if on_started is not None: + await on_started() + self.started.set() + await self._finish.wait() + return b"late-complete", {}, 0.0, None + + async def abort_request(self, request_id): + self._finish.set() + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline: + job = await api_server.VIDEO_STORE.get(request_id) + if job is not None and job.status is VideoGenerationStatus.COMPLETED and job.file_name is not None: + return + await asyncio.sleep(0.01) + raise RuntimeError(f"video job {request_id} did not complete during abort") + + +class SchedulerQueuedVideoHandler(BlockingVideoHandler): + """Parks after engine submission but before scheduler admission.""" + + def __init__(self): + super().__init__() + self.admit = threading.Event() + self.in_progress = threading.Event() + + async def generate_video_bytes( + self, + request, + reference_id, + *, + reference_image=None, + reference_video=None, + reference_audio=None, + on_started=None, + ): + del request, reference_id, reference_image, reference_video, reference_audio + self.started.set() + try: + while not self.admit.is_set(): + await asyncio.sleep(0.01) + if on_started is not None: + await on_started() + self.in_progress.set() + await asyncio.Future() + except asyncio.CancelledError: + self.cancelled.set() + raise + + +class AbortTrackingOmni(FakeAsyncOmni): + def __init__(self): + super().__init__() + self.aborted: list[str] = [] + self.entered = threading.Event() + + async def generate(self, prompt, request_id, sampling_params_list): + del prompt, request_id, sampling_params_list + self.entered.set() + yield MockVideoResult( + [], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + ) + await asyncio.Future() + + async def abort(self, request_id, *, timeout=None): + assert timeout is not None + self.aborted.append(request_id) + class FakeServerSocket: def __init__(self): @@ -592,12 +710,12 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): engine = FakeAsyncOmni() - engine.get_diffusion_od_config = lambda: SimpleNamespace( + engine.get_diffusion_od_config = lambda: SimpleNamespace( # type: ignore[method-assign] model="org/model", model_class_name="ExamplePipeline", revision="pinned-revision", ) - captured = {} + captured: dict[str, str | None] = {} def fake_policy(model_class_name, *, model, revision=None): captured.update( @@ -625,6 +743,7 @@ def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): ) ) + assert engine.captured_prompt is not None input_image = engine.captured_prompt["multi_modal_data"]["image"] assert isinstance(input_image, Image.Image) assert input_image.size == (48, 32) @@ -718,6 +837,7 @@ def test_video_generation_bridges_request_fields(generation_request, expected_nu asyncio.run(handler._run_and_extract(generation_request, "field-bridge")) + assert engine.captured_sampling_params_list is not None sampling = engine.captured_sampling_params_list[0] # Top-level ``seconds`` bridges into extra_args["duration"]; num_frames is # passed through (or derived as seconds x fps when omitted). No private @@ -1503,7 +1623,7 @@ def test_video_generation_response_exposes_action_payload(mocker: MockerFixture) }, ) - engine.generate = _generate + engine.generate = _generate # type: ignore[method-assign] mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", return_value="encoded-video", @@ -2096,6 +2216,80 @@ def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client): assert retrieve_resp.status_code == 404 +def test_async_video_stays_queued_until_scheduler_admission(test_client): + handler = SchedulerQueuedVideoHandler() + test_client.app.state.openai_serving_video = handler + + create_resp = test_client.post("/v1/videos", data={"prompt": "Queue this video"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert handler.started.wait(timeout=2.0) + + queued = test_client.get(f"/v1/videos/{video_id}") + assert queued.status_code == 200 + assert queued.json()["status"] == VideoGenerationStatus.QUEUED.value + + handler.admit.set() + assert handler.in_progress.wait(timeout=2.0) + in_progress = test_client.get(f"/v1/videos/{video_id}") + assert in_progress.status_code == 200 + assert in_progress.json()["status"] == VideoGenerationStatus.IN_PROGRESS.value + + assert test_client.delete(f"/v1/videos/{video_id}").status_code == 200 + + +def test_delete_times_out_engine_abort_and_still_cancels(test_client, monkeypatch): + monkeypatch.setattr(api_server, "VIDEO_ABORT_TIMEOUT_S", 0.05) + handler = HangingAbortHandler() + test_client.app.state.openai_serving_video = handler + + create_resp = test_client.post("/v1/videos", data={"prompt": "Hang abort"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert handler.started.wait(timeout=2.0) + + started = time.monotonic() + delete_resp = test_client.delete(f"/v1/videos/{video_id}") + assert time.monotonic() - started < 2.0 + assert delete_resp.status_code == 200 + assert handler.cancelled.wait(timeout=2.0) + + +def test_delete_removes_artifact_if_job_completes_during_abort(test_client): + handler = CompletingDuringAbortHandler() + test_client.app.state.openai_serving_video = handler + + create_resp = test_client.post("/v1/videos", data={"prompt": "Complete during delete"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert handler.started.wait(timeout=2.0) + + assert test_client.delete(f"/v1/videos/{video_id}").status_code == 200 + file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id) + assert not os.path.exists(file_path) + assert asyncio.run(api_server.VIDEO_STORE.get(video_id)) is None + + +def test_delete_aborts_engine_request_before_cancelling_task(test_client): + engine = AbortTrackingOmni() + test_client.app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion( + engine, + model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + ) + + create_resp = test_client.post("/v1/videos", data={"prompt": "Abort this video"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert engine.entered.wait(timeout=2.0) + _wait_for_status(test_client, video_id, VideoGenerationStatus.IN_PROGRESS.value) + + delete_resp = test_client.delete(f"/v1/videos/{video_id}") + assert delete_resp.status_code == 200 + assert delete_resp.json()["deleted"] is True + _wait_until(lambda: engine.aborted == [video_id]) + assert asyncio.run(api_server.VIDEO_STORE.get(video_id)) is None + + def test_video_response_file_extension_is_robust(): response = VideoResponse(model="test-model", prompt="Make something beautiful") assert response.file_extension == "mp4" diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index a69bcce09..2379db78e 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -12,6 +12,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.entrypoints import async_omni as async_omni_mod from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.outputs import OmniRequestOutput @@ -38,7 +39,8 @@ def get_fake_add_request(submitted_request_ids, submitted_lora_requests=None): def get_fake_abort(aborted_request_batches): - async def fake_abort_async(request_ids): + async def fake_abort_async(request_ids, timeout=None): + del timeout aborted_request_batches.append(list(request_ids)) return fake_abort_async @@ -196,8 +198,8 @@ def test_abort_keeps_request_states_until_generate_cleanup(): release = asyncio.Event() seen_during_wait: list[int] = [] - async def slow_abort_async(request_ids): - del request_ids + async def slow_abort_async(request_ids, timeout=None): + del request_ids, timeout seen_during_wait.append(len(omni.request_states)) await release.wait() @@ -224,8 +226,8 @@ def test_abort_keeps_request_states_until_generate_cleanup(): @pytest.mark.cpu def test_abort_propagates_engine_errors_without_popping_state(): async def run(): - async def failing_abort_async(request_ids): - del request_ids + async def failing_abort_async(request_ids, timeout=None): + del request_ids, timeout raise RuntimeError("orchestrator abort failed") omni = get_async_omni_instance(fake_abort_request=failing_abort_async) @@ -253,7 +255,8 @@ def test_abort_enqueues_prefix_tokens_from_engine(): queue: asyncio.Queue = asyncio.Queue() - async def abort_with_prefix(request_ids): + async def abort_with_prefix(request_ids, timeout=None): + del timeout rid = request_ids[0] engine_output = OmniRequestOutput( request_id=rid, @@ -308,8 +311,8 @@ def test_abort_enqueues_synthetic_finished_when_engine_returns_empty(): async def run(): queue: asyncio.Queue = asyncio.Queue() - async def empty_abort_async(request_ids): - del request_ids + async def empty_abort_async(request_ids, timeout=None): + del request_ids, timeout return [] omni = get_async_omni_instance(fake_abort_request=empty_abort_async) @@ -467,6 +470,42 @@ def test_generate_accepts_request_after_repeated_cancellations(): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_cancel_bounds_cleanup_abort(monkeypatch): + """Cancelled generate() must not wait forever on a wedged orchestrator abort.""" + monkeypatch.setattr(async_omni_mod, "ABORT_TIMEOUT_S", 0.05) + + async def run(): + seen_timeouts: list[float | None] = [] + + async def hanging_abort_async(request_ids, timeout=None): + del request_ids + seen_timeouts.append(timeout) + await asyncio.Event().wait() + + omni = get_async_omni_instance(fake_abort_request=hanging_abort_async) + + async def collect(): + async for _ in omni.generate( + prompt={"prompt": "prompt"}, + request_id="cancel-hang-abort", + sampling_params_list=[SimpleNamespace()], + output_modalities=["image"], + ): + pass + + task = asyncio.create_task(collect()) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1.0) + assert seen_timeouts + assert seen_timeouts[-1] == 0.05 + assert omni.request_states == {} + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker): async def run_test(): diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 95059a06e..01264eaea 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -167,7 +167,8 @@ class FakeAsyncOmniEngine: def abort(self, request_ids: list[str]) -> None: self.aborted.append(list(request_ids)) - async def abort_async(self, request_ids: list[str]) -> None: + async def abort_async(self, request_ids: list[str], timeout=None) -> None: + del timeout self.abort(request_ids) async def collective_rpc_async(self, **_: Any) -> list[Any]: diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index b5214720a..97782328d 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -63,6 +63,7 @@ _PUBLIC_OMNI = ( "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS", "SPEAKER_MAX_UPLOADED", "SPEAKER_SAMPLES_DIR", + "VLLM_OMNI_ABORT_TIMEOUT", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", "VLLM_OMNI_EVENT_DRIVEN_ORCH", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c333fb11a..10f8642c6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1637,6 +1637,18 @@ class OmniDiffusionConfig: return cls(**cls.normalize_init_kwargs(kwargs)) +DIFFUSION_REQUEST_LIFECYCLE_KEY = "_diffusion_request_lifecycle" +DIFFUSION_REQUEST_STARTED = "started" + + +def is_diffusion_request_started_output(output: Any) -> bool: + custom_output = getattr(output, "custom_output", None) + return ( + isinstance(custom_output, dict) + and custom_output.get(DIFFUSION_REQUEST_LIFECYCLE_KEY) == DIFFUSION_REQUEST_STARTED + ) + + @dataclass class DiffusionOutput: """ @@ -1685,6 +1697,9 @@ class DiffusionOutput: # mode) and the receiving side must not initialise a stray CUDA context. to_cpu: bool = False + # Internal control-plane event emitted on first scheduler admission. + request_started: bool = False + # Typed video-media contract. Declared last so the pre-existing positional # constructor order (output, trajectory_timesteps, ...) that out-of-tree # pipelines rely on is preserved. Mutually exclusive with ``output``. diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 417090503..6e53b5909 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -26,6 +26,8 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.kv_cache_interface import KVCacheConfig from vllm_omni.diffusion.data import ( + DIFFUSION_REQUEST_LIFECYCLE_KEY, + DIFFUSION_REQUEST_STARTED, DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig, @@ -55,7 +57,7 @@ from vllm_omni.diffusion.registry import ( ) from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest from vllm_omni.diffusion.sched import BaseScheduler, RequestScheduler, StepScheduler -from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus +from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus, DiffusionSchedulerOutput from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput from vllm_omni.errors import client_error_from_metadata, is_client_error_status from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt @@ -533,7 +535,12 @@ class DiffusionEngine: error_type=output.error_type, ) raise RuntimeError(output.error) - logger.debug("Generation completed successfully.") + if output.request_started: + return format_empty_diffusion_outputs( + request, + finished=False, + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + ) if output.media is not None: if output.output is not None: @@ -599,6 +606,8 @@ class DiffusionEngine: sched_output = self.scheduler.schedule() self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0) + self._emit_request_started_outputs(sched_output) + if sched_output.is_empty: self._emit_finished_outputs(sched_output.finished_req_ids, None) continue @@ -629,6 +638,17 @@ class DiffusionEngine: # Engine is stopping: fail any RPCs still queued so callers don't hang. self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down.")) + def _emit_request_started_outputs(self, sched_output: DiffusionSchedulerOutput) -> None: + """Notify opted-in callers when requests leave the scheduler queue.""" + for new_req in sched_output.scheduled_new_reqs: + request = getattr(new_req, "req", None) + sampling_params = getattr(request, "sampling_params", None) + if getattr(sampling_params, "emit_request_lifecycle", False): + self._put_output( + new_req.request_id, + DiffusionOutput(finished=False, request_started=True), + ) + def _wait_for_admission_if_needed_locked(self) -> None: """Apply scheduler admission policy while holding the engine condition. diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 6c8cc236c..f583aba14 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -19,7 +19,10 @@ from typing import TYPE_CHECKING, Any from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.diffusion.data import DiffusionRequestAbortedError +from vllm_omni.diffusion.data import ( + DiffusionRequestAbortedError, + is_diffusion_request_started_output, +) from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.engine.stage_client import StageClientBase @@ -158,7 +161,13 @@ class InlineStageDiffusionClient(StageClientBase): # only publish the final output. result = None async for results in self._engine.step_streaming(request): - result = results[0] + output = results[0] + if is_diffusion_request_started_output(output): + if not output.request_id: + output.request_id = request_id + self._output_queue.put_nowait(output) + continue + result = output if result is None: raise RuntimeError("Diffusion execution finished without output.") if not result.request_id: diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 62c2330d5..4074806e1 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -5,7 +5,7 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass, field -from typing import TypeGuard, cast +from typing import Any, TypeGuard, cast from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.io_support import get_diffusion_output_type, supports_audio_output @@ -106,6 +106,7 @@ def format_empty_diffusion_outputs( request: OmniDiffusionRequest, *, finished: bool = True, + custom_output: dict[str, Any] | None = None, ) -> list[OmniRequestOutput]: return [ OmniRequestOutput.from_diffusion( @@ -114,6 +115,7 @@ def format_empty_diffusion_outputs( prompt=request.prompt, metrics={}, latents=None, + custom_output=custom_output, finished=finished, ) ] diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index aedb730e3..220158f1f 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -13,7 +13,7 @@ import asyncio import contextlib import multiprocessing.connection import signal -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from typing import TYPE_CHECKING, Any @@ -33,7 +33,10 @@ from vllm.v1.engine.utils import ( ) from vllm.v1.utils import shutdown -from vllm_omni.diffusion.data import DiffusionRequestAbortedError +from vllm_omni.diffusion.data import ( + DiffusionRequestAbortedError, + is_diffusion_request_started_output, +) from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.utils.serialization import ( @@ -163,6 +166,7 @@ class StageDiffusionProc: prompt: Any, sampling_params_dict: dict, kv_sender_info: dict[str, Any] | None = None, + on_request_started: Callable[[OmniRequestOutput], Awaitable[None]] | None = None, kv_transfer_params: dict[str, Any] | None = None, ) -> OmniRequestOutput: """Build a diffusion request and consume DiffusionEngine.step_streaming() to completion.""" @@ -180,7 +184,13 @@ class StageDiffusionProc: # return the final output. result = None async for results in self._engine.step_streaming(request): - result = results[0] + output = results[0] + if is_diffusion_request_started_output(output) and on_request_started is not None: + if not output.request_id: + output.request_id = request_id + await on_request_started(output) + continue + result = output if result is None: raise RuntimeError("Diffusion execution finished without output.") if not result.request_id: @@ -354,11 +364,16 @@ class StageDiffusionProc: """Process a single diffusion request and send the response.""" try: if not self._od_config.streaming_output: + + async def _send_request_started(output: OmniRequestOutput) -> None: + await response_socket.send(encoder.encode({"type": "result", "output": output})) + result = await self._process_request( request_id, prompt, sampling_params_dict, kv_sender_info=kv_sender_info, + on_request_started=_send_request_started, kv_transfer_params=kv_transfer_params, ) await response_socket.send(encoder.encode({"type": "result", "output": result})) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 498a9dc89..75620bec9 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -35,6 +35,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker @@ -1199,7 +1200,8 @@ class Orchestrator: idle = False continue - pool.record_output_timestamps([diffusion_output]) + if not is_diffusion_request_started_output(diffusion_output): + pool.record_output_timestamps([diffusion_output]) processed = [diffusion_output] else: raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001) @@ -1441,7 +1443,8 @@ class Orchestrator: if self._absorb_diffusion_metrics(stage_id, replica_id, payload): self._orch_monitor.note_loop(idle=False) continue - pool.record_output_timestamps([payload]) + if not is_diffusion_request_started_output(payload): + pool.record_output_timestamps([payload]) processed = [payload] else: processed = await self._process_llm_stage_outputs( diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 0012039e0..993f0ca1e 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -11,6 +11,7 @@ StageEngineCoreClient instances) instead of OmniStage with worker processes. from __future__ import annotations import asyncio +import os import time import uuid from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence @@ -64,6 +65,8 @@ _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 # the janus queue's condition variable; this timeout only bounds how often the # orchestrator liveness check runs while the pipeline is idle. _FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 +# Shared DELETE / generate() cleanup abort bound. Env is the documented knob. +ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) class AsyncEventResolver: @@ -661,12 +664,12 @@ class AsyncOmni(EngineClient, OmniBase): except (asyncio.CancelledError, GeneratorExit): self._record_request_failure_once(request_id, reason="client_disconnect") - await self._abort_internal_requests(request_id) + await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: self._record_request_failure_once(request_id, reason="stage_error") - await self._abort_internal_requests(request_id) + await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise finally: @@ -1120,14 +1123,14 @@ class AsyncOmni(EngineClient, OmniBase): return all(bool(item) for item in result) return bool(result) - async def abort(self, request_id: str | Iterable[str]) -> None: + async def abort(self, request_id: str | Iterable[str], *, timeout: float | None = None) -> None: """Abort request(s) via the Orchestrator.""" request_ids = [request_id] if isinstance(request_id, str) else list(request_id) # Map the external user request IDs to internal IDs used by the Orchestrator. # NOTE: If the user request_id matches multiple requests, all of them will be # aborted. This is also what happens in this case in vLLM's output processor. internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id in request_ids] - await self._abort(internal_ids) + await self._abort(internal_ids, timeout=timeout) async def submit_interaction_async( self, @@ -1169,16 +1172,32 @@ class AsyncOmni(EngineClient, OmniBase): if self.log_stats: logger.info("[AsyncOmni] Queued interaction for request %s", request_id) - async def _abort_internal_requests(self, request_id: str | Iterable[str]): + async def _abort_internal_requests( + self, + request_id: str | Iterable[str], + *, + timeout: float = ABORT_TIMEOUT_S, + ): """Abort request(s) via the Orchestrator given internal request IDs, which take the format -. """ request_ids = [request_id] if isinstance(request_id, str) else list(request_id) # Request IDs are already internal, so we just need to get the matching states. internal_req_ids = [rid for rid in request_ids if rid in self.request_states] - await self._abort(internal_req_ids) + try: + # Unbind generate() if abort_async blocks in the executor. + await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) + except TimeoutError: + logger.warning( + "[AsyncOmni] Timed out aborting %s after %.1fs; " + "engine abort is best-effort until the current batch drains", + ",".join(internal_req_ids), + timeout, + ) + except Exception: + logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) - async def _abort(self, request_ids: list[str]) -> None: + async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: """Abort request IDs via the engine and enqueue terminal abort outputs. Waits for orchestrator abort acknowledgment, enqueues any AR terminal @@ -1191,7 +1210,7 @@ class AsyncOmni(EngineClient, OmniBase): registered yet, unbound replica, or orchestrator id drop), enqueue a synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. """ - abort_outputs = await self.engine.abort_async(request_ids) or [] + abort_outputs = await self.engine.abort_async(request_ids, timeout=timeout) or [] delivered: set[str] = set() for output_msg in abort_outputs: req_id = getattr(output_msg, "request_id", None) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 8acd1d6fa..36b7544ee 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -86,7 +86,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.endpoint_policy import ( shutdown_unsupported_routes, ) -from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni from vllm_omni.entrypoints.duplex.capability import should_enable_duplex_endpoint from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime @@ -201,6 +201,8 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName logger = init_logger(__name__) router = APIRouter() +VIDEO_ABORT_TIMEOUT_S = ABORT_TIMEOUT_S + profiler_router = APIRouter() @@ -2518,11 +2520,11 @@ async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse: @router.delete("/v1/videos/{video_id}") -async def delete_video(video_id: str) -> VideoDeleteResponse: +async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteResponse: """Delete a stored video job and any generated output. - If the job is still queued or running, this endpoint first attempts to - cancel the in-flight generation task before removing the stored metadata. + In-flight jobs get a bounded engine abort, then frontend cancel. The job + is re-read afterwards so a completed save is not orphaned. Args: video_id: Identifier of the video job to delete. @@ -2539,19 +2541,37 @@ async def delete_video(video_id: str) -> VideoDeleteResponse: raise HTTPException(status_code=404, detail="Video not found") if job.status in (VideoGenerationStatus.QUEUED, VideoGenerationStatus.IN_PROGRESS): + handler = raw_request.app.state.openai_serving_video + try: + await asyncio.wait_for(handler.abort_request(video_id), timeout=VIDEO_ABORT_TIMEOUT_S) + except asyncio.TimeoutError: + logger.warning( + "Timed out aborting video request %s after %.1fs; " + "engine abort is best-effort until the current batch drains", + video_id, + VIDEO_ABORT_TIMEOUT_S, + ) + except Exception: + logger.exception("Failed to abort in-flight video request %s", video_id) task = await VIDEO_TASKS.get(video_id) if task is not None: task.cancel() try: - await asyncio.wait_for(task, timeout=2.0) + # Cancel cleanup may spend a full abort budget; +2s covers scheduling slack. + await asyncio.wait_for(task, timeout=VIDEO_ABORT_TIMEOUT_S + 2.0) except asyncio.TimeoutError: raise HTTPException(status_code=409, detail="Cancellation in progress. Please try again later.") except asyncio.CancelledError: pass + job = await VIDEO_STORE.get(video_id) + if job is None: + return VideoDeleteResponse(id=video_id, deleted=True) + if job.status in (VideoGenerationStatus.QUEUED, VideoGenerationStatus.IN_PROGRESS): await VIDEO_STORE.pop(video_id) return VideoDeleteResponse(id=job.id, deleted=True) - elif job.status is VideoGenerationStatus.FAILED: + + if job.status is VideoGenerationStatus.FAILED: if job.file_name is not None: try: await STORAGE_MANAGER.delete(video_id) diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 5ff055fa2..9a24a1ef9 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -6,7 +6,7 @@ from __future__ import annotations import copy import math import time -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from functools import cached_property from http import HTTPStatus @@ -19,9 +19,10 @@ from PIL import Image from vllm.engine.protocol import EngineClient from vllm.logger import init_logger +from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.media_utils import count_mp4_frames, normalize_preencode_batch_frames -from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni from vllm_omni.entrypoints.openai.protocol.videos import ( VideoAction, VideoData, @@ -256,6 +257,9 @@ class OmniOpenAIServingVideo: def shutdown(self) -> None: self._video_frame_converter.shutdown() + async def abort_request(self, request_id: str) -> None: + await self._engine_client.abort(request_id, timeout=ABORT_TIMEOUT_S) + async def _run_and_extract( self, request: VideoGenerationRequest, @@ -264,6 +268,7 @@ class OmniOpenAIServingVideo: reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, + on_started: Callable[[], Awaitable[None]] | None = None, ) -> VideoGenerationArtifacts: """Run the generation pipeline and extract video/audio/profiler outputs.""" prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"]) @@ -444,7 +449,12 @@ class OmniOpenAIServingVideo: gen_params.seed, ) - result = await self._run_generation(prompt, gen_params, reference_id) + result = await self._run_generation( + prompt, + gen_params, + reference_id, + on_started=on_started, + ) multimodal_output = self._extract_multimodal_output(result) metadata = multimodal_output.get("metadata") if isinstance(multimodal_output, dict) else {} common_metadata = metadata.get("common") if isinstance(metadata, dict) else {} @@ -528,6 +538,7 @@ class OmniOpenAIServingVideo: reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, + on_started: Callable[[], Awaitable[None]] | None = None, ) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: """Generate a video and return raw MP4 bytes, bypassing base64 encoding.""" artifacts = await self._run_and_extract( @@ -536,6 +547,7 @@ class OmniOpenAIServingVideo: reference_image=reference_image, reference_video=reference_video, reference_audio=reference_audio, + on_started=on_started, ) if len(artifacts.videos) > 1: logger.warning( @@ -641,6 +653,8 @@ class OmniOpenAIServingVideo: prompt: OmniTextPrompt, gen_params: OmniDiffusionSamplingParams, request_id: str, + *, + on_started: Callable[[], Awaitable[None]] | None = None, ) -> object: stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None) @@ -660,6 +674,7 @@ class OmniOpenAIServingVideo: # Common generation logic for both paths engine_client = cast(AsyncOmni, self._engine_client) + gen_params.emit_request_lifecycle = on_started is not None sampling_params_list = build_stage_sampling_params_list( list(stage_configs), get_default_sampling_params_list(engine_client), @@ -668,11 +683,17 @@ class OmniOpenAIServingVideo: ) result = None + started_notified = False async for output in engine_client.generate( prompt=prompt, request_id=request_id, sampling_params_list=sampling_params_list, ): + if is_diffusion_request_started_output(output): + if on_started is not None and not started_notified: + await on_started() + started_notified = True + continue result = output if result is None: diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py index e23d876e5..4473fce07 100644 --- a/vllm_omni/entrypoints/openai/video/generation/helpers.py +++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py @@ -330,9 +330,12 @@ async def _run_video_generation_job( _cleanup_video_references(reference_video, reference_audio, control_path) return - await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) started_at = time.perf_counter() try: + + async def _mark_started() -> None: + await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) + video_bytes, stage_durations, peak_memory_mb, action, video_metadata = _unpack_video_generation_result( await handler.generate_video_bytes( request, @@ -340,6 +343,7 @@ async def _run_video_generation_job( reference_image=reference_image, reference_video=reference_video, reference_audio=reference_audio, + on_started=_mark_started, ) ) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index f9f0a9a9a..287ddb3fb 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -349,6 +349,11 @@ class OmniDiffusionSamplingParams: # results output: torch.Tensor | None = None + # Internal control-plane event used by asynchronous APIs. When enabled, + # the diffusion engine emits a non-terminal output when the scheduler + # first admits this request for execution. + emit_request_lifecycle: bool = False + def __post_init__(self) -> None: if self.quality is not None and self.quality not in DIFFUSION_QUALITY_LEVELS: raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {self.quality!r}") From a2138b830fe01230f39f0149cce3d14eec0beec6 Mon Sep 17 00:00:00 2001 From: GXIN <37653830+gxxx-hum@users.noreply.github.com> Date: Wed, 16 Sep 2026 12:50:00 +0800 Subject: [PATCH 101/209] [Feature][Qwen3-TTS] Make Code2Wav honor the stage dtype (#6059) Signed-off-by: XIN GAO <1037396230@qq.com> --- tests/config/test_config_factory.py | 20 +++++++ .../qwen3_tts/test_qwen3_tts_code2wav.py | 22 ++++++-- tests/platforms/npu/test_310p_patches.py | 52 ++++++++++++++++--- vllm_omni/deploy/qwen3_tts.yaml | 2 + .../deploy/qwen3_tts_high_concurrency.yaml | 2 + .../models/qwen3_tts/qwen3_tts_code2wav.py | 4 +- vllm_omni/platforms/npu/models/qwen3_tts.py | 19 ++++--- 7 files changed, 100 insertions(+), 21 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 82c1fb32e..d1dc0fd6d 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2587,6 +2587,26 @@ class TestPlatformOverrides: assert rocm.stages[0].enforce_eager is None assert rocm.stages[1].enforce_eager is True + @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"]) + @pytest.mark.parametrize("platform", ["cuda", "npu", "rocm"]) + def test_qwen3_tts_default_code2wav_dtype_is_bf16(self, deploy_name, platform): + deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) + deploy = _apply_platform_overrides(deploy, platform=platform) + stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy) + + assert stages[1].yaml_engine_args["dtype"] == "bfloat16" + + @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"]) + @pytest.mark.parametrize("platform", ["cuda", "npu"]) + @pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"]) + def test_qwen3_tts_propagates_explicit_code2wav_dtype(self, deploy_name, platform, dtype): + deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) + deploy.stages[1].engine_extras["dtype"] = dtype + deploy = _apply_platform_overrides(deploy, platform=platform) + stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy) + + assert stages[1].yaml_engine_args["dtype"] == dtype + def test_higgs_audio_v3_rocm_uses_triton_attention(self): deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml")) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 64f9ac8de..8ad2b519f 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -23,7 +23,7 @@ class _FakeDecoder(nn.Module): def __init__(self, total_upsample: int = _TOTAL_UPSAMPLE): super().__init__() self.total_upsample = total_upsample - self.decode_calls: list[dict[str, int]] = [] + self.decode_calls: list[dict[str, int | tuple[int, ...]]] = [] self.batched_decode_calls: list[dict[str, object]] = [] self.decode_codes: list[torch.Tensor] = [] self.cudagraph_calls: list[dict[str, int | torch.device]] = [] @@ -63,7 +63,7 @@ class _FakeDecoder(nn.Module): chunk_size: int = 300, left_context_size: int = 25, max_batch_size: int = 0, - ) -> torch.Tensor: + ) -> list[torch.Tensor]: self.batched_decode_calls.append( { "chunk_size": chunk_size, @@ -91,8 +91,8 @@ class _FakeDecoder(nn.Module): wav_len = frames * self.total_upsample + 6 wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1) offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000 - outputs = wav.expand(batch, 1, wav_len) + offsets - return [outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] + batched_wav = wav.expand(batch, 1, wav_len) + offsets + return [batched_wav[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] def enable_cudagraph(self, **kwargs): self.cudagraph_calls.append(kwargs) @@ -131,6 +131,7 @@ def _make_model( load_config=SimpleNamespace(), model_config=SimpleNamespace( model="unused", + dtype=torch.bfloat16, revision=None, stage_connector_config=stage_connector_config, async_chunk=async_chunk, @@ -708,3 +709,16 @@ def test_invalid_decode_chunking_is_rejected(): with pytest.raises(ValueError, match="decode_chunk_frames=0"): _load_weights_noop(model) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_load_weights_uses_model_dtype_before_precomputing_caches(dtype): + model = _make_model() + model.vllm_config.model_config.dtype = dtype + model.decoder = nn.Linear(2, 2) + cache_dtypes = [] + model.decoder.precompute_snake_caches = lambda: cache_dtypes.append(model.decoder.weight.dtype) + + assert _load_weights_noop(model) == {"decoder.fake_weight"} + assert model.decoder.weight.dtype is dtype + assert cache_dtypes == [dtype] diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index c19dcd997..0db4b9d41 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/tests/platforms/npu/test_310p_patches.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for 310P patch wiring. The tests load patch modules from source with fake Qwen3-TTS dependencies, so @@ -644,7 +644,11 @@ def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest. torch.testing.assert_close(norm_out, torch.tensor([[2.0, 3.0]])) -def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None: +def _load_qwen3_tts_code2wav_npu_patch( + monkeypatch: pytest.MonkeyPatch, + *, + platform_dtype_override: bool = True, +): linear_weights = [] conv_weights = [] @@ -674,15 +678,22 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes self.prefix = prefix self.decoder = FakeDecoder() - def _npu_decoder_runtime_dtype(self, _device): - return torch.float16 - def load_weights(self, weights): assert list(weights) == [] + self.decoder.to(dtype=self.vllm_config.model_config.dtype) + self.decoder.precompute_snake_caches() return {"loaded"} - logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None) + if platform_dtype_override: + + def _npu_decoder_runtime_dtype(self, _device: torch.device) -> torch.dtype: + return torch.float16 + + FakeCode2Wav._npu_decoder_runtime_dtype = _npu_decoder_runtime_dtype # type: ignore[attr-defined] + + logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None, warning=lambda *_: None) current_platform = SimpleNamespace(is_npu=lambda: False) + target = _install_fake_module( monkeypatch, "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav", @@ -702,11 +713,23 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts.py" module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_npu_patch", path) module.apply_qwen3_tts_code2wav_patch() + return module, target, linear_weights, conv_weights - model = target.Qwen3TTSCode2Wav( - vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))), + +def _make_patched_code2wav(target, *, dtype=torch.bfloat16): + return target.Qwen3TTSCode2Wav( + vllm_config=SimpleNamespace( + device_config=SimpleNamespace(device=torch.device("cpu")), + model_config=SimpleNamespace(dtype=dtype), + ), prefix="stage1", ) + + +def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None: + module, target, linear_weights, conv_weights = _load_qwen3_tts_code2wav_npu_patch(monkeypatch) + model = _make_patched_code2wav(target, dtype=torch.bfloat16) + assert model.load_weights(iter(())) == {"loaded"} assert model.prefix == "stage1" @@ -717,4 +740,17 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes model.decoder.deconv.weight.data_ptr(), } assert all(fmt == module._ACL_FORMAT_FRACTAL_Z for _, fmt in conv_weights) + assert model.decoder.cache_precompute_calls == 2 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_qwen3_tts_code2wav_npu_patch_preserves_stage_dtype( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, +) -> None: + _, target, _, _ = _load_qwen3_tts_code2wav_npu_patch(monkeypatch, platform_dtype_override=False) + model = _make_patched_code2wav(target, dtype=dtype) + + assert model.load_weights(iter(())) == {"loaded"} + assert all(parameter.dtype is dtype for parameter in model.decoder.parameters()) assert model.decoder.cache_precompute_calls == 1 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 747b68852..4019f10aa 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -95,6 +95,8 @@ stages: top_p: 1.0 - stage_id: 1 + # Decoder compute dtype on GPU/NPU; 310P requires FP16 and overrides this. + dtype: bfloat16 # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1. max_num_seqs: 64 gpu_memory_utilization: 0.3 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 7a73dd50a..7914ecd52 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -85,6 +85,8 @@ stages: top_p: 1.0 - stage_id: 1 + # Decoder compute dtype on GPU/NPU; 310P requires FP16 and overrides this. + dtype: bfloat16 max_num_seqs: 10 gpu_memory_utilization: 0.3 enforce_eager: true diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index f1692f73d..8ab5c5e34 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -499,7 +499,7 @@ class Qwen3TTSCode2Wav(nn.Module): if wav.numel() == 0: continue if wav.shape[0] > 0: - # Decoder already runs in fp32, so the .to(float32) is a redundant dispatch. + # Return FP32 audio regardless of the decoder compute dtype. audios[idx] = (wav if wav.dtype == torch.float32 else wav.to(torch.float32)).reshape(-1) for req_id, finished, segment_finished in zip( @@ -560,7 +560,7 @@ class Qwen3TTSCode2Wav(nn.Module): ) device = self.vllm_config.device_config.device - self.decoder.to(device=device, dtype=torch.float32) + self.decoder.to(device=device, dtype=self.vllm_config.model_config.dtype) # Precompute SnakeBeta exp caches (benefits both Triton and eager paths) if hasattr(self.decoder, "precompute_snake_caches"): diff --git a/vllm_omni/platforms/npu/models/qwen3_tts.py b/vllm_omni/platforms/npu/models/qwen3_tts.py index a91d2cf80..124e60624 100644 --- a/vllm_omni/platforms/npu/models/qwen3_tts.py +++ b/vllm_omni/platforms/npu/models/qwen3_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Patch Qwen3-TTS NPU runtime setup and weights. @@ -68,11 +68,16 @@ def _patched_load_weights(self, weights): assert _original_load_weights is not None loaded = _original_load_weights(self, weights) device = self.vllm_config.device_config.device - runtime_dtype = getattr(self, "_npu_decoder_runtime_dtype", lambda _: torch.float32)(device) - self.decoder.to(device=device, dtype=runtime_dtype) + platform_dtype_override = getattr(self, "_npu_decoder_runtime_dtype", None) + if platform_dtype_override is not None: + # 310P requires FP16 and deliberately overrides the stage dtype. + runtime_dtype = platform_dtype_override(device) + self.decoder.to(device=device, dtype=runtime_dtype) + if hasattr(self.decoder, "precompute_snake_caches"): + self.decoder.precompute_snake_caches() _prepare_npu_decoder_weights(self.decoder) - if runtime_dtype != torch.float32 and hasattr(self.decoder, "precompute_snake_caches"): - self.decoder.precompute_snake_caches() + actual_dtype = next(self.decoder.parameters()).dtype + logger.info("Prepared NPU Code2Wav decoder with runtime dtype %s", actual_dtype) return loaded @@ -108,8 +113,8 @@ def _apply_a5_prompt_embeds_builder_patch() -> None: # The talker imports the builder by name at module scope, so its reference # must be swapped as well (same as the 310P patch). - prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilderA5 - qwen3_tts_talker.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilderA5 + setattr(prompt_embeds_builder, "Qwen3TTSPromptEmbedsBuilder", _Qwen3TTSPromptEmbedsBuilderA5) + setattr(qwen3_tts_talker, "Qwen3TTSPromptEmbedsBuilder", _Qwen3TTSPromptEmbedsBuilderA5) _A5_PATCHED = True logger.debug("Applied A5 prompt-embeds builder patch for Qwen3-TTS") From 21d86ec9239d5aba08ea5fcac7e20a8f0679b0c9 Mon Sep 17 00:00:00 2001 From: vOv <124694036+cr-gao@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:01:26 +0800 Subject: [PATCH 102/209] [Diffusion] MAGI-2: regional compile coverage (#7174) Signed-off-by: cr-gao --- .../models/magi2/test_native_compile.py | 137 ++++++++++++ .../models/magi2/test_native_compile_cuda.py | 77 +++++++ .../magi2/test_native_compile_distributed.py | 179 ++++++++++++++++ .../models/magi2/test_native_packing.py | 118 +++++++++++ .../diffusion/models/magi2/modeling_magi2.py | 198 +++++++++++++----- .../diffusion/models/magi2/pipeline_magi2.py | 18 ++ .../diffusion/models/magi2/sampler_magi2.py | 28 +-- 7 files changed, 674 insertions(+), 81 deletions(-) create mode 100644 tests/diffusion/models/magi2/test_native_compile.py create mode 100644 tests/diffusion/models/magi2/test_native_compile_cuda.py create mode 100644 tests/diffusion/models/magi2/test_native_compile_distributed.py create mode 100644 tests/diffusion/models/magi2/test_native_packing.py diff --git a/tests/diffusion/models/magi2/test_native_compile.py b/tests/diffusion/models/magi2/test_native_compile.py new file mode 100644 index 000000000..0abac75f8 --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_compile.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch + +from tests.diffusion.models.magi2.test_native_packing import ( + _longer_text_tensors, + _sampler_tensors, + _tiny_model, + _tiny_sampler, +) +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def test_layer_regions_bracket_the_eager_kernels() -> None: + moe_layer, dense_layer = _tiny_model().block.layers + + assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output") + assert dense_layer.region_methods == ("_attention_input", "_dense_output") + + +def test_compile_regions_fullgraph_matches_eager_bitwise() -> None: + model = _tiny_model() + sampler = _tiny_sampler(model) + steps = [ + sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig()) + for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0]))) + ] + expected = [sampler.forward(step) for step in steps] + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + first = sampler.forward(steps[0]) + graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before + with torch._dynamo.config.patch(error_on_recompile=True): + second = sampler.forward(steps[1]) + + for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True): + assert torch.equal(video, video_ref) + assert torch.equal(audio, audio_ref) + # pre adapter, post adapter, three MoE-layer regions, two dense-layer regions + assert graphs == 7 + + +def test_pipeline_setup_compile_forwards_config_dynamic() -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline + + calls: list[dict[str, object]] = [] + + class _Transformer(torch.nn.Module): + def compile_regions(self, **compile_kwargs: object) -> None: + calls.append(compile_kwargs) + + pipeline = object.__new__(Magi2Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = OmniDiffusionConfig( + model="sand-ai/MAGI-2-preview", + model_class_name="Magi2Pipeline", + diffusion_compile_dynamic=False, + ) + pipeline.transformer = _Transformer() + + pipeline.setup_compile() + + assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}] + + +def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline + + calls: list[dict[str, object]] = [] + + class _Transformer(torch.nn.Module): + def compile_regions(self, **compile_kwargs: object) -> None: + calls.append(compile_kwargs) + + pipeline = object.__new__(Magi2Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = OmniDiffusionConfig( + model="sand-ai/MAGI-2-preview", + model_class_name="Magi2Pipeline", + diffusion_compile_granularity="full", + ) + pipeline.transformer = _Transformer() + + pipeline.setup_compile() + + assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}] + assert "diffusion_compile_granularity='full'" in caplog.text + + +def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None: + model = _tiny_model(moe_layers=2) + sampler = _tiny_sampler(model) + step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig()) + expected = sampler.forward(step) + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + actual = sampler.forward(step) + + assert torch.equal(actual[0], expected[0]) + # Two MoE layers and two dense layers still compile one graph per region + # kind: pre adapter, post adapter, three MoE regions, two dense regions. + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7 + + +def test_compiled_regions_survive_new_requests_without_recompiling() -> None: + model = _tiny_model() + sampler = _tiny_sampler(model) + t = torch.tensor([900.0]) + first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig()) + second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig()) + longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig()) + expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)] + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + actual = [sampler.forward(first_request)] + with torch._dynamo.config.patch(error_on_recompile=True): + actual.append(sampler.forward(second_request)) + actual.append(sampler.forward(longer_request)) + + for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True): + assert torch.equal(video, video_ref) + assert torch.equal(audio, audio_ref) + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7 diff --git a/tests/diffusion/models/magi2/test_native_compile_cuda.py b/tests/diffusion/models/magi2/test_native_compile_cuda.py new file mode 100644 index 000000000..1db394750 --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_compile_cuda.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch + +from tests.diffusion.models.magi2.test_native_packing import ( + _longer_text_tensors, + _sampler_tensors, + _tiny_model, + _tiny_sampler, +) +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig + +pytestmark = [ + pytest.mark.diffusion, + pytest.mark.core_model, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"), +] + + +def _cuda_steps(sampler, tensors=_sampler_tensors): + timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0])) + return [ + sampler.prepare_model_input( + **{name: value.cuda() for name, value in tensors(seed).items()}, + t=t, + cfg_config=CFGConfig(), + ) + for seed, t in enumerate(timesteps) + ] + + +def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MAGI2_DETERMINISTIC", "1") + model = _tiny_model(params_dtype=torch.bfloat16).cuda() + sampler = _tiny_sampler(model) + steps = _cuda_steps(sampler) + expected = [sampler.forward(step) for step in steps] + new_request_step = _cuda_steps(sampler)[0] + longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0] + expected_new_request = sampler.forward(new_request_step) + expected_longer_text = sampler.forward(longer_text_step) + + torch._dynamo.reset() + graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"] + # The pipeline compiles with inductor emulating eager's intermediate + # low-precision casts, so fused bf16 chains keep eager's rounding. + model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True}) + first = sampler.forward(steps[0]) + with torch._dynamo.config.patch(error_on_recompile=True): + compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])] + repeated = [sampler.forward(step) for step in steps] + # A second request repacks the same shapes; a third changes the + # text length and therefore every token count. + compiled_new_request = sampler.forward(new_request_step) + compiled_longer_text = sampler.forward(longer_text_step) + + # Emulated precision casts restore eager's rounding boundaries; what is + # left is one-ulp transcendental noise that some Triton backends carry. + for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True): + torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0) + torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0) + for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True): + assert torch.equal(video, video_again) + assert torch.equal(audio, audio_again) + for (video, audio), (video_ref, audio_ref) in ( + (compiled_new_request, expected_new_request), + (compiled_longer_text, expected_longer_text), + ): + torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0) + torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0) + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7 diff --git a/tests/diffusion/models/magi2/test_native_compile_distributed.py b/tests/diffusion/models/magi2/test_native_compile_distributed.py new file mode 100644 index 000000000..03f6f601c --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_compile_distributed.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import os +import tempfile +from contextlib import ExitStack, contextmanager +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn +from torch.distributed import init_device_mesh +from torch.distributed.fsdp import MixedPrecisionPolicy + +import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dlo_module +from tests.diffusion.models.magi2.test_native_distributed_parity import ( + _current_group, + _new_groups, + _patched_groups, +) +from tests.diffusion.models.magi2.test_native_packing import _tiny_config, _tiny_model, _tiny_sampler +from tests.diffusion.offloader.helpers import patch_offload_runtime +from vllm_omni.diffusion.distributed.hsdp import shard_model +from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer +from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup +from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig +from vllm_omni.diffusion.offloader.base import OffloadConfig +from vllm_omni.diffusion.offloader.config import OffloadStrategy + +_WORLD_SIZE = 4 +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +@contextmanager +def _cpu_offload_runtime(): + with pytest.MonkeyPatch.context() as monkeypatch: + patch_offload_runtime(monkeypatch, dlo_module.current_omni_platform, synchronize=True) + yield + + +def _enable_dlo( + model: Magi2PreviewTransformer, + *, + dp_group: dist.ProcessGroup | None, + dp_size: int, + dp_rank: int, +) -> None: + pipeline = Magi2Pipeline.__new__(Magi2Pipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = model + backend = dlo_module.DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=dp_size, + dlo_use_allgather=dp_size > 1, + ), + torch.device("cpu"), + ) + backend.dp_group = dp_group + backend.rank = dp_rank + backend.enable(pipeline) + assert backend.enabled + + +def _enable_hsdp(model: Magi2PreviewTransformer) -> None: + mesh = init_device_mesh("cpu", mesh_shape=(1, _WORLD_SIZE), mesh_dim_names=("replicate", "shard")) + ignored = { + parameter for name in model._hsdp_ignored_modules for parameter in model.get_submodule(name).parameters() + } + shard_model( + model, + mesh=mesh, + mp_policy=MixedPrecisionPolicy(cast_forward_inputs=False), + hsdp_shard_conditions=model._hsdp_shard_conditions, + ignored_params=ignored, + ) + model.requires_grad_(False) + + +def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]: + # 14 tokens per CFG branch, so the packed sequence splits evenly across + # two Ulysses ranks; gloo cannot gather uneven shards. + generator = torch.Generator(device="cpu").manual_seed(seed) + return { + "latent": torch.randn(1, 4, 2, 1, 3, generator=generator), + "audio_latent": torch.randn(1, 5, 4, generator=generator), + "txt_feat": torch.randn(1, 3, 4, generator=generator), + "null_txt_feat": torch.randn(1, 3, 4, generator=generator), + } + + +def _steps(sampler): + return [ + sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig()) + for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0]))) + ] + + +def _run(model: Magi2PreviewTransformer, *, compile_regions: bool) -> list[tuple[torch.Tensor, torch.Tensor]]: + sampler = _tiny_sampler(model) + steps = _steps(sampler) + if not compile_regions: + return [sampler.forward(step) for step in steps] + torch._dynamo.reset() + model.compile_regions(fullgraph=True, backend="eager", dynamic=True) + outputs = [sampler.forward(steps[0])] + with torch._dynamo.config.patch(error_on_recompile=True): + outputs.append(sampler.forward(steps[1])) + return outputs + + +def _assert_same(actual, expected, *, exact: bool, label: str) -> None: + for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True): + if exact: + assert torch.equal(video, video_ref), label + assert torch.equal(audio, audio_ref), label + else: + torch.testing.assert_close(video, video_ref, atol=2e-5, rtol=2e-4, msg=label) + torch.testing.assert_close(audio, audio_ref, atol=2e-5, rtol=2e-4, msg=label) + + +def _distributed_worker(rank: int, rendezvous: str) -> None: + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=_WORLD_SIZE) + try: + singleton = Magi2ParallelGroup(None, world_size=1, rank=0) + with _patched_groups(singleton, singleton): + oracle = _tiny_model() + checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()] + expected = _run(oracle, compile_regions=False) + + sp2_groups = _new_groups(((0, 1), (2, 3))) + dp2_groups = _new_groups(((0, 2), (1, 3))) + sp_group = _current_group(rank, sp2_groups) + dp_ranks, dp_process_group = next((ranks, group) for ranks, group in dp2_groups if rank in ranks) + + def dlo_rank_local(model: Magi2PreviewTransformer) -> None: + _enable_dlo(model, dp_group=None, dp_size=1, dp_rank=0) + + def dlo_allgather(model: Magi2PreviewTransformer) -> None: + _enable_dlo(model, dp_group=dp_process_group, dp_size=len(dp_ranks), dp_rank=dp_ranks.index(rank)) + + variants = ( + ("dlo_rank_local_sp2", sp_group, dlo_rank_local), + ("dlo_allgather_dp2sp2", sp_group, dlo_allgather), + ("hsdp4", singleton, _enable_hsdp), + ) + for name, ulysses_group, enable in variants: + with ExitStack() as stack: + stack.enter_context(_patched_groups(singleton, ulysses_group)) + stack.enter_context(_cpu_offload_runtime()) + outputs = {} + for compile_regions in (False, True): + model = Magi2PreviewTransformer(_tiny_config()) + assert model.load_weights(checkpoint) == set(model.state_dict()) + enable(model) + outputs[compile_regions] = _run(model, compile_regions=compile_regions) + _assert_same(outputs[False], expected, exact=False, label=f"{name} eager vs single-rank oracle") + _assert_same(outputs[True], outputs[False], exact=True, label=f"{name} compiled vs eager") + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="requires torch.distributed gloo", +) +def test_compiled_regions_match_eager_under_dlo_and_hsdp() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}" + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + mp.spawn(_distributed_worker, args=(rendezvous,), nprocs=_WORLD_SIZE, join=True) diff --git a/tests/diffusion/models/magi2/test_native_packing.py b/tests/diffusion/models/magi2/test_native_packing.py new file mode 100644 index 000000000..cedabebcb --- /dev/null +++ b/tests/diffusion/models/magi2/test_native_packing.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.diffusion.models.magi2.configuration_magi2 import ( + Magi2MHCConfig, + Magi2MoEConfig, + Magi2PreviewConfig, +) +from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm +from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE +from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer +from vllm_omni.diffusion.models.magi2.preview_data_proxy import ( + Magi2DataProxy, + Magi2PreviewDataProxyConfig, +) +from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig: + # The leading layers are multimodal with MoE, the rest single-modality + # dense, so every layer belongs to one of two kinds. + moe_indices = tuple(range(moe_layers)) + return Magi2PreviewConfig( + num_layers=2 * moe_layers, + hidden_size=16, + head_dim=8, + num_query_groups=2, + video_in_channels=4, + audio_in_channels=4, + text_in_channels=4, + intermediate_factor=2, + multimodal_layers=moe_indices, + params_dtype=params_dtype, + mhc=Magi2MHCConfig(num_streams=2), + moe=Magi2MoEConfig( + num_heads=2, + num_experts=4, + top_k=2, + expert_intermediate_size=8, + shared_expert_intermediate_size=8, + modality_shared_expert_intermediate_size=8, + layers=moe_indices, + ), + ) + + +def _tiny_model( + seed: int = 11, + params_dtype: torch.dtype = torch.float32, + *, + moe_layers: int = 1, +) -> Magi2PreviewTransformer: + model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers)) + generator = torch.Generator(device="cpu").manual_seed(seed) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) + for module in model.modules(): + if isinstance(module, MultiModalityRMSNorm): + module.weight.zero_() + elif isinstance(module, Magi2MultiHeadMoE): + module.router.expert_bias.zero_() + module.router.expert_bias_ema.zero_() + return model + + +def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler: + return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8))) + + +def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(seed) + return { + "latent": torch.randn(1, 4, 2, 1, 3, generator=generator), + "audio_latent": torch.randn(1, 5, 4, generator=generator), + "txt_feat": torch.randn(1, 3, 4, generator=generator), + "null_txt_feat": torch.randn(1, 2, 4, generator=generator), + } + + +def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]: + tensors = _sampler_tensors(seed) + generator = torch.Generator(device="cpu").manual_seed(seed + 100) + tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator) + return tensors + + +def test_prepare_model_input_keeps_lengths_on_host() -> None: + sampler = Magi2PreviewSampler(torch.nn.Identity()) + model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig()) + + assert model_input.audio_feat_len == [5, 5] + assert model_input.txt_feat_len == [3, 2] + assert model_input.ref_audio_feat_len == [0, 0] + assert model_input.ref_video_feat_len == [0, 0] + + positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input) + assert positive.txt_feat_len == [3] + assert negative.txt_feat_len == [2] + + +def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None: + model = _tiny_model(params_dtype=torch.bfloat16) + packed = torch.randn(6, 4) + indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5]) + + with torch.no_grad(): + hidden = model.pre_adapter(packed, *indices) + + assert hidden.dtype == torch.bfloat16 + assert hidden.shape == (6, model.config.virtual_width) diff --git a/vllm_omni/diffusion/models/magi2/modeling_magi2.py b/vllm_omni/diffusion/models/magi2/modeling_magi2.py index 87d3402d3..44e282aa0 100644 --- a/vllm_omni/diffusion/models/magi2/modeling_magi2.py +++ b/vllm_omni/diffusion/models/magi2/modeling_magi2.py @@ -16,6 +16,7 @@ import os from collections.abc import Iterable from enum import IntEnum from functools import partial +from typing import Any import torch import torch.nn as nn @@ -28,6 +29,7 @@ from .configuration_magi2 import Magi2PreviewConfig from .layers import ( ElementWiseFourierEmbed, MHCHandler, + MHCTensorTuple, ModalityDispatcher, MultiModalityRMSNorm, make_grouped_linear, @@ -130,14 +132,12 @@ class Magi2Attention(nn.Module): start = self.tp_group.rank * self.num_heads_q return checkpoint_tensor[:, start : start + self.num_heads_q] - def forward( + def project( self, hidden_states: torch.Tensor, rope: torch.Tensor, - varlen_handler: VarlenHandler, modality_dispatcher: ModalityDispatcher, - cp_split_sizes: list[int] | torch.Tensor, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: normalized = self.pre_norm(hidden_states, modality_dispatcher) gates = self.linear_g(normalized, modality_dispatcher) qkv = self.linear_qkv(normalized, modality_dispatcher) @@ -159,8 +159,17 @@ class Magi2Attention(nn.Module): q = apply_rotary_emb(q, cos, sin).squeeze(0).to(self.config.params_dtype) k = apply_rotary_emb(k, cos, sin).squeeze(0).to(self.config.params_dtype) v = v.squeeze(0).to(self.config.params_dtype) + return q, k, v, gates - output = self.packed_attention( + def attend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + varlen_handler: VarlenHandler, + cp_split_sizes: list[int] | torch.Tensor, + ) -> torch.Tensor: + return self.packed_attention( q, k, v, @@ -172,7 +181,14 @@ class Magi2Attention(nn.Module): } ), ) - output = modality_dispatcher.permute(output) + + def output( + self, + attention: torch.Tensor, + gates: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + output = modality_dispatcher.permute(attention) output = output * torch.sigmoid(gates) output = output.reshape(-1, self.q_size).to(self.config.params_dtype) return self.linear_proj(output, modality_dispatcher) @@ -303,12 +319,21 @@ class Magi2MultiHeadMoELayer(nn.Module): modality.contiguous(), dispatcher ) - def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor: + def route_input( + self, + hidden_states: torch.Tensor, + dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, torch.Tensor]: normalized = self.pre_norm(hidden_states, dispatcher) - routed = self.split_linear(normalized) - routed = self.moe_mlp(routed) - routed = self.merge_linear(routed) - return routed + self._shared_experts(normalized, dispatcher) + return normalized, self.split_linear(normalized) + + def combine( + self, + normalized: torch.Tensor, + routed: torch.Tensor, + dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + return self.merge_linear(routed) + self._shared_experts(normalized, dispatcher) class Magi2PreAdapter(nn.Module): @@ -339,37 +364,28 @@ class Magi2PreAdapter(nn.Module): def forward( self, packed: torch.Tensor, - coords_mapping: torch.Tensor, video_indices: torch.Tensor, audio_indices: torch.Tensor, text_indices: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - rope = self.rope(coords_mapping) - output = torch.zeros( + ) -> torch.Tensor: + # Every token belongs to exactly one modality, so the three copies + # below cover the whole buffer. Embedding in fp32 and rounding each + # modality once matches the released cast at the block boundary. + output = torch.empty( packed.shape[0], self.adapter_dim, - dtype=torch.float32, + dtype=self.config.params_dtype, device=packed.device, ) - if text_indices.numel(): - output.index_copy_( - 0, - text_indices, - self.text_embedder(packed.index_select(0, text_indices)[:, : self.config.text_in_channels].float()), - ) - if audio_indices.numel(): - output.index_copy_( - 0, - audio_indices, - self.audio_embedder(packed.index_select(0, audio_indices)[:, : self.config.audio_in_channels].float()), - ) - if video_indices.numel(): - output.index_copy_( - 0, - video_indices, - self.video_embedder(packed.index_select(0, video_indices)[:, : self.config.video_in_channels].float()), - ) - return output, rope + for indices, embedder, in_channels in ( + (text_indices, self.text_embedder, self.config.text_in_channels), + (audio_indices, self.audio_embedder, self.config.audio_in_channels), + (video_indices, self.video_embedder, self.config.video_in_channels), + ): + if indices.numel(): + tokens = packed[:, :in_channels].index_select(0, indices).float() + output.index_copy_(0, indices, embedder(tokens).to(output.dtype)) + return output class Magi2PostAdapter(nn.Module): @@ -425,8 +441,10 @@ class Magi2TransformerLayer(nn.Module): self.mlp: nn.Module if layer_index in config.moe.layers: self.mlp = Magi2MultiHeadMoELayer(config) + self.region_methods = ("_attention_input", "_moe_input", "_moe_output") else: self.mlp = Magi2MLP(config, num_modality=num_modality) + self.region_methods = ("_attention_input", "_dense_output") self._init_mhc(num_modality) def _init_mhc(self, num_modality: int) -> None: @@ -525,6 +543,74 @@ class Magi2TransformerLayer(nn.Module): ) return self.mhc_handler.hyper_connect(streams, output, post, residual) + # The layer forward is split into compile regions around the two eager + # kernels: packed attention and the multi-head MoE. + + def _attention_input( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size) + attention_logits = self._branch_logits(streams, "attn", modality_dispatcher) + attention_input = self._branch_input(streams, "attn", attention_logits) + q, k, v, gates = self.attention.project(attention_input, rope, modality_dispatcher) + return streams, attention_logits, q, k, v, gates + + def _mlp_input( + self, + streams: torch.Tensor, + attention: torch.Tensor, + gates: torch.Tensor, + attention_logits: MHCTensorTuple, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor]: + attention_output = self.attention.output(attention, gates, modality_dispatcher) + streams = self._connect(streams, attention_output, "attn", attention_logits) + mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher) + return streams, mlp_logits, self._branch_input(streams, "mlp", mlp_logits) + + def _dense_output( + self, + streams: torch.Tensor, + attention: torch.Tensor, + gates: torch.Tensor, + attention_logits: MHCTensorTuple, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + streams, mlp_logits, mlp_input = self._mlp_input( + streams, attention, gates, attention_logits, modality_dispatcher + ) + streams = self._connect(streams, self.mlp(mlp_input, modality_dispatcher), "mlp", mlp_logits) + return streams.reshape(streams.shape[0], -1) + + def _moe_input( + self, + streams: torch.Tensor, + attention: torch.Tensor, + gates: torch.Tensor, + attention_logits: MHCTensorTuple, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor, torch.Tensor]: + streams, mlp_logits, mlp_input = self._mlp_input( + streams, attention, gates, attention_logits, modality_dispatcher + ) + normalized, routed = self.mlp.route_input(mlp_input, modality_dispatcher) + return streams, mlp_logits, normalized, routed + + def _moe_output( + self, + streams: torch.Tensor, + mlp_logits: MHCTensorTuple, + normalized: torch.Tensor, + routed: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + mlp_output = self.mlp.combine(normalized, routed, modality_dispatcher) + streams = self._connect(streams, mlp_output, "mlp", mlp_logits) + return streams.reshape(streams.shape[0], -1) + def forward( self, hidden_states: torch.Tensor, @@ -533,22 +619,15 @@ class Magi2TransformerLayer(nn.Module): modality_dispatcher: ModalityDispatcher, cp_split_sizes: list[int] | torch.Tensor, ) -> torch.Tensor: - streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size) - attention_logits = self._branch_logits(streams, "attn", modality_dispatcher) - attention_input = self._branch_input(streams, "attn", attention_logits) - attention_output = self.attention( - attention_input, - rope, - varlen_handler, - modality_dispatcher, - cp_split_sizes, + streams, attention_logits, q, k, v, gates = self._attention_input(hidden_states, rope, modality_dispatcher) + attention = self.attention.attend(q, k, v, varlen_handler, cp_split_sizes) + if not isinstance(self.mlp, Magi2MultiHeadMoELayer): + return self._dense_output(streams, attention, gates, attention_logits, modality_dispatcher) + streams, mlp_logits, normalized, routed = self._moe_input( + streams, attention, gates, attention_logits, modality_dispatcher ) - streams = self._connect(streams, attention_output, "attn", attention_logits) - mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher) - mlp_input = self._branch_input(streams, "mlp", mlp_logits) - mlp_output = self.mlp(mlp_input, modality_dispatcher) - streams = self._connect(streams, mlp_output, "mlp", mlp_logits) - return streams.reshape(streams.shape[0], -1) + routed = self.mlp.moe_mlp(routed) + return self._moe_output(streams, mlp_logits, normalized, routed, modality_dispatcher) class Magi2TransformerBlock(nn.Module): @@ -570,7 +649,6 @@ class Magi2TransformerBlock(nn.Module): modality_dispatcher: ModalityDispatcher, cp_split_sizes: list[int] | torch.Tensor, ) -> torch.Tensor: - hidden_states = hidden_states.to(self.config.params_dtype) hidden_states = modality_dispatcher.permute(hidden_states) for layer in self.layers: hidden_states = layer( @@ -630,6 +708,15 @@ class Magi2PreviewTransformer(nn.Module): return self.block.layers + def compile_regions(self, **compile_kwargs: Any) -> None: + """Compile the dense compute between the eager attention and MoE kernels.""" + + self.pre_adapter.forward = torch.compile(self.pre_adapter.forward, **compile_kwargs) + self.post_adapter.forward = torch.compile(self.post_adapter.forward, **compile_kwargs) + for layer in self.block.layers: + for name in layer.region_methods: + setattr(layer, name, torch.compile(getattr(layer, name), **compile_kwargs)) + def forward( self, x: torch.Tensor, @@ -647,6 +734,7 @@ class Magi2PreviewTransformer(nn.Module): assert dispatcher.split_sizes is not None cp_split_sizes = dispatcher.split_sizes + rope = self.pre_adapter.rope(coords_mapping) time_mask = modality_mapping == int(Modality.TIME) modality_mapping = torch.where(time_mask, int(Modality.TEXT), modality_mapping) modality_dispatcher = ModalityDispatcher(modality_mapping, 3) @@ -654,13 +742,7 @@ class Magi2PreviewTransformer(nn.Module): audio_indices = torch.nonzero(modality_mapping == int(Modality.AUDIO)).flatten() text_indices = torch.nonzero(modality_mapping == int(Modality.TEXT)).flatten() - hidden_states, rope = self.pre_adapter( - x, - coords_mapping, - video_indices, - audio_indices, - text_indices, - ) + hidden_states = self.pre_adapter(x, video_indices, audio_indices, text_indices) if time_token_sequence is not None and time_token_sequence.shape[-1] > 0: hidden_states[:, : time_token_sequence.shape[-1]] = time_token_sequence.to(hidden_states.dtype) hidden_states = self.block( diff --git a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py index c7c480876..910e0d074 100644 --- a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py +++ b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py @@ -695,6 +695,24 @@ class Magi2Pipeline( return self.video_decoder.module + def setup_compile(self) -> None: + """Compile the transformer regions; the attention and MoE kernels stay eager.""" + + granularity = self.od_config.diffusion_compile_granularity + if granularity != "regional": + logger.warning( + "MAGI-2 compiles the transformer regions itself; diffusion_compile_granularity=%r is ignored.", + granularity, + ) + + # The mHC connections chain bf16 ops that eager rounds after every op. + # Emulating those casts keeps the compiled rounding boundaries equal. + self.transformer.compile_regions( + fullgraph=True, + dynamic=self.od_config.diffusion_compile_dynamic, + options={"emulate_precision_casts": True}, + ) + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], diff --git a/vllm_omni/diffusion/models/magi2/sampler_magi2.py b/vllm_omni/diffusion/models/magi2/sampler_magi2.py index 994fdd3e4..35dd2dab9 100644 --- a/vllm_omni/diffusion/models/magi2/sampler_magi2.py +++ b/vllm_omni/diffusion/models/magi2/sampler_magi2.py @@ -308,29 +308,11 @@ class Magi2PreviewSampler(CFGParallelMixin): per_token_video_t = t_batch.view(-1, 1, 1, 1, 1).expand(cfg_batch, 1, video_t, video_h, video_w) per_token_audio_t = t_batch.view(-1, 1, 1).expand(cfg_batch, audio_t, 1) - audio_lengths = torch.full( - (cfg_batch,), - audio_latent_len, - device=latent.device, - dtype=torch.long, - ) - text_lengths = torch.tensor( - [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size, - device=latent.device, - dtype=torch.long, - ) - ref_audio_lengths = torch.full( - (cfg_batch,), - ref_audio_feat_len, - device=latent.device, - dtype=torch.long, - ) - ref_video_lengths = torch.full( - (cfg_batch,), - ref_video_feat_len, - device=latent.device, - dtype=torch.long, - ) + # Lengths stay host integers; the data proxy reads them as Python ints. + audio_lengths = [audio_latent_len] * cfg_batch + text_lengths = [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size + ref_audio_lengths = [ref_audio_feat_len] * cfg_batch + ref_video_lengths = [ref_video_feat_len] * cfg_batch return ModelInput( x_t=torch.cat((latent, latent), dim=0), From 5e9024171d4a8ceb9089b7f8306e266d0eb5526a Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Wed, 16 Sep 2026 14:43:18 +0800 Subject: [PATCH 103/209] [Bugfix] Skip generic text-only warmup for DreamZero (#7548) Signed-off-by: tzhouam --- .../test_diffusion_engine_dummy_run.py | 17 +++++++++++++++++ .../models/dreamzero/pipeline_dreamzero.py | 5 +++++ 2 files changed, 22 insertions(+) diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 8d524c572..1f1e8cbab 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -7,6 +7,7 @@ from unittest.mock import Mock import pytest from vllm_omni.diffusion import io_support +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest @@ -16,6 +17,22 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.mark.parametrize("model_class_name", ["DreamZeroPipeline", "LingBotWorldCausalDMDPipeline"]) +def test_observation_conditioned_startup_skips_generic_warmup(model_class_name: str) -> None: + """Observation-conditioned pipelines must not receive generic text requests.""" + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = OmniDiffusionConfig.__new__(OmniDiffusionConfig) + engine.od_config.model_class_name = model_class_name + engine.od_config.diffusion_load_format = "default" + engine.add_req_and_wait_for_response = Mock(side_effect=AssertionError("generic text warmup submitted")) + engine.close = Mock() + + engine.run_startup_warmup() + + engine.add_req_and_wait_for_response.assert_not_called() + engine.close.assert_not_called() + + @pytest.mark.parametrize("step_execution", [False, True]) def test_dummy_run_uses_enough_steps_for_execution_mode( monkeypatch: pytest.MonkeyPatch, diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 09bd536cd..939544a9f 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -17,6 +17,7 @@ import re as re_module from collections import OrderedDict from collections.abc import Iterable from contextlib import contextmanager +from typing import ClassVar import numpy as np import torch @@ -121,6 +122,10 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): methods below. The runner binds one session state only for ``forward()``. """ + # Generic warmup cannot synthesize robot observations. AR-Diffusion uses + # ar_diffusion_warmup_requests() for model-specific warmup instead. + dummy_run_num_frames: ClassVar[int] = 0 + _POSITIVE_BRANCH = "positive" _NEGATIVE_BRANCH = "negative" _ar_diffusion_kv_state = None From 9005d789033b8c3ec5876a7a68c4e2d9238f5c69 Mon Sep 17 00:00:00 2001 From: linzhenpl07 <60651557+linzhenpl07@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:01:11 +0800 Subject: [PATCH 104/209] [LingBot World] Reuse the text encode for a prompt already encoded (#6845) Signed-off-by: linzhenpl07 Co-authored-by: Claude Opus 5 --- .../test_pipeline_lingbot_world.py | 122 ++++++++++++++++++ .../models/lingbot_world/pipeline.py | 34 ++++- 2 files changed, 154 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index fb9e93a5d..9d5e958c5 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -2711,3 +2711,125 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No assert "from .pipeline import" in lingbot_init assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init + + +# --------------------------------------------------------------------------- +# a prompt is encoded once, whichever path asks for it +# --------------------------------------------------------------------------- +def _counting_text_encoder(pipeline) -> list[str]: + """Give ``pipeline`` its real encode_prompt over stubs that record each encode.""" + del pipeline.encode_prompt # _pipeline() replaces it; these tests need the real one + encoded: list[str] = [] + + def tokenizer(texts, **kwargs): + del kwargs + encoded.extend(texts) + return SimpleNamespace( + input_ids=torch.arange(512).view(1, 512), + attention_mask=torch.ones(1, 512, dtype=torch.long), + ) + + def text_encoder(input_ids, attention_mask): + del input_ids, attention_mask + return SimpleNamespace(last_hidden_state=torch.full((1, 512, 8), float(len(encoded)))) + + pipeline.tokenizer = tokenizer + pipeline.text_encoder = text_encoder + return encoded + + +def test_the_text_encoder_runs_once_per_distinct_prompt() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + encoded = _counting_text_encoder(pipeline) + + def encode(prompt, *, length=512, dtype=torch.float32): + return pipeline.encode_prompt(prompt, max_sequence_length=length, dtype=dtype) + + first = encode("move through the room") + assert encode("move through the room") is first + # The key is the text the tokenizer actually sees, so spacing alone is not new. + assert encode(" move through the room ") is first + assert encoded == ["move through the room"] + + encode("enter the snowy valley") + # dtype and sequence length change the tensor, so each is its own entry. + encode("move through the room", dtype=torch.float16) + encode("move through the room", length=256) + assert encoded == [ + "move through the room", + "enter the snowy valley", + "move through the room", + "move through the room", + ] + + +def test_the_prompt_encode_cache_is_bounded() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + encoded = _counting_text_encoder(pipeline) + size = module._PROMPT_EMBEDS_CACHE_SIZE + + for index in range(size + 1): + pipeline.encode_prompt(f"prompt {index}", max_sequence_length=512, dtype=torch.float32) + assert len(pipeline._prompt_embeds_cache) == size + + # The oldest entry was the one evicted, so it is the one encoded again. + pipeline.encode_prompt("prompt 0", max_sequence_length=512, dtype=torch.float32) + pipeline.encode_prompt(f"prompt {size}", max_sequence_length=512, dtype=torch.float32) + assert encoded == [f"prompt {index}" for index in range(size + 1)] + ["prompt 0"] + + +def test_realtime_ticks_reuse_the_encode_and_still_invalidate_on_a_prompt_change( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + pipeline._ar_diffusion_kv_state = object() + encoded = _counting_text_encoder(pipeline) + invalidations: list[bool] = [] + + def ar_text_caches(prompt_embeds, *, invalidate): + del prompt_embeds + invalidations.append(invalidate) + return [SimpleNamespace()] + + monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches) + monkeypatch.setattr( + pipeline, + "_generate_block", + lambda **kwargs: torch.randn((1, 16, 3, 2, 2), generator=kwargs["generator"]), + ) + + for chunk_index in range(3): + pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=chunk_index)))) + switched = _prompt() + switched["prompt"] = "enter the snowy valley" + pipeline( + _request( + sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=3, prompt=switched["prompt"])), + prompt=switched, + ) + ) + + assert encoded == ["move through the room", "enter the snowy valley"] + # A new prompt must still drop the cross-attention K/V built from the old one. + assert invalidations == [False, False, False, True] + + +def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module, transformer=_RecordingTransformer()) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + encoded = _counting_text_encoder(pipeline) + + states = [_stepwise_state(request_id=request_id, num_frames=21) for request_id in ("req-a", "req-b")] + for state in states: + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + + assert encoded == ["move through the room"] + assert states[0].prompt_embeds is states[1].prompt_embeds diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index ccaa677e0..8e8d5dea4 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -6,6 +6,7 @@ from __future__ import annotations import math import os +from collections import OrderedDict from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager, nullcontext from dataclasses import dataclass, replace @@ -99,6 +100,10 @@ _CAMERA_SPATIAL_FOLD = 8 _MAX_PIXEL_AREA = 480 * 832 _MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096 _MAX_SEQUENCE_LENGTH = 512 +# Distinct prompts whose text encodes are kept. One entry is a single padded +# sequence -- a few MiB at 512 tokens -- so a handful covers a session's prompt +# plus the switches it makes, at no meaningful cost next to its KV. +_PROMPT_EMBEDS_CACHE_SIZE = 4 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" _PREPROCESSED_CAMERA_ACTIONS_KEY = "_lingbot_camera_actions" @@ -569,6 +574,8 @@ class LingBotWorldCausalDMDPipeline( self._ar_width = int(model_config.get("ar_diffusion_width", 832)) self._ar_diffusion_kv_state: ARDiffusionKVState | None = None self._ar_sessions: dict[str, _LingBotARSessionState] = {} + # Text encodes keyed by what they depend on; see encode_prompt. + self._prompt_embeds_cache: OrderedDict[tuple[str, int, torch.dtype], torch.Tensor] = OrderedDict() # One temporal decoder cache per stepwise session, keyed the way the # runner keys AR sessions (session_id == request_id) so both are # released together rather than through two independent lifecycles. @@ -1221,8 +1228,27 @@ class LingBotWorldCausalDMDPipeline( max_sequence_length: int, dtype: torch.dtype, ) -> torch.Tensor: + """Encode ``prompt`` with UMT5, reusing the result for a prompt already encoded. + + The encode depends only on the whitespace-normalised text, the sequence + length and the dtype -- the tokenizer and text encoder are fixed once + loaded -- so those three are the key. That covers every caller the same + way: a realtime tick repeats its session's prompt on every block, a + stepwise request encodes once per request, and a server tends to open + many sessions with the same scene prompt. + + A hit returns the stored tensor itself. Callers only read it: the + cross-attention projection and the DMD transformer both work out of + place, so a shared encode is never changed under another session. + """ + text = " ".join(prompt.strip().split()) + key = (text, int(max_sequence_length), dtype) + cached = self._prompt_embeds_cache.get(key) + if cached is not None: + self._prompt_embeds_cache.move_to_end(key) + return cached text_inputs = self.tokenizer( - [" ".join(prompt.strip().split())], + [text], padding="max_length", max_length=max_sequence_length, truncation=True, @@ -1234,7 +1260,11 @@ class LingBotWorldCausalDMDPipeline( attention_mask = text_inputs.attention_mask.to(self.device) prompt_embeds = self.text_encoder(input_ids, attention_mask).last_hidden_state prompt_embeds = prompt_embeds.to(device=self.device, dtype=dtype) - return prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype) + prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype) + self._prompt_embeds_cache[key] = prompt_embeds + if len(self._prompt_embeds_cache) > _PROMPT_EMBEDS_CACHE_SIZE: + self._prompt_embeds_cache.popitem(last=False) + return prompt_embeds def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: inputs = self._parse_request(req) From 99ff4f307040e85ba99e849ef013fdee89525378 Mon Sep 17 00:00:00 2001 From: chickeyton Date: Wed, 16 Sep 2026 15:35:36 +0800 Subject: [PATCH 105/209] [Core][Frontend] Unified Full-duplex Framework (#7413) Signed-off-by: chickeyton Co-authored-by: Claude Opus 5 (1M context) --- .../common/ci_source_file_dependencies.yml | 10 +- .buildkite/cuda/test-merge.yml | 4 +- docs/api/README.md | 7 +- docs/configuration/stage_configs.md | 3 +- docs/design/fullduplex.md | 1415 +-- docs/design/module/engine_orchestration.md | 11 +- docs/design/module/entrypoints.md | 11 +- docs/mkdocs/hooks/generate_api_readme.py | 4 +- docs/models/supported_models.md | 2 +- docs/serving/README.md | 17 +- docs/serving/full_duplex_api.md | 81 +- docs/serving/realtime_duplex_api.md | 349 +- examples/online_serving/barge_in_client.py | 228 +- .../online_serving/barge_in_client_flow.md | 15 +- .../online_serving/minicpmo/barge_in_serve.sh | 10 +- .../minicpmo/realtime_duplex_demo.py | 2 - .../minicpmo/realtime_web/app/static/app.js | 2 - recipes/OpenBMB/MiniCPM-o-4_5.md | 75 +- recipes/README.md | 2 +- tests/benchmarks/patch/test_patch.py | 68 +- tests/benchmarks/test_omniinteract.py | 27 +- tests/benchmarks/test_realtime_tts_probe.py | 179 + tests/clients/test_duplex_client.py | 514 +- tests/clients/test_model_session_configs.py | 7 +- tests/dfx/perf/tests/test_minicpmo_4_5.json | 86 +- .../test_minicpmo_4_5_duplex_seed_tts.json | 38 +- .../perf/tests/test_minicpmo_4_5_ready.json | 40 +- .../test_diffusion_streaming_output.py | 3 +- .../test_prompt_interaction.py | 3 +- .../test_diffusion_engine_metrics.py | 26 +- .../fullduplex/test_core_contracts.py | 11 +- .../minicpmo_realtime_duplex_scenarios.py | 25 +- .../run_minicpmo_realtime_duplex_demo_pair.py | 3 - ..._minicpmo_realtime_duplex_multi_session.py | 34 +- ...minicpmo_realtime_duplex_soft_interrupt.py | 3 - .../online_serving/test_duplex_client_live.py | 2 - .../test_minicpmo_4_5_duplex.py | 386 +- .../test_minicpmo_4_5_duplex_expansion.py | 2 +- .../test_minicpmo_realtime_duplex_drivers.py | 26 +- .../duplex/test_duplex_audio.py | 2 +- tests/engine/duplex/test_duplex_commands.py | 453 + .../duplex/test_duplex_control_client.py | 110 - .../duplex/test_duplex_control_plane.py | 1270 --- tests/engine/duplex/test_duplex_events.py | 386 + tests/engine/duplex/test_duplex_lease.py | 206 +- tests/engine/duplex/test_duplex_plugin.py | 324 + tests/engine/duplex/test_duplex_runtime.py | 563 - tests/engine/duplex/test_engine_session.py | 612 ++ tests/engine/duplex/test_overlap_policy.py | 229 + tests/engine/duplex/test_session_manager.py | 1569 +++ tests/engine/duplex/test_session_runner.py | 1195 +++ tests/engine/duplex/test_vad_backend.py | 379 + .../test_async_omni_engine_abort_ack.py | 8 +- .../engine/test_async_omni_engine_outputs.py | 336 +- .../test_async_omni_engine_stage_init.py | 4 +- .../test_diffusion_streaming_output_kwarg.py | 2 +- tests/engine/test_duplex_import_boundary.py | 83 +- tests/engine/test_duplex_omni_engine.py | 210 + tests/engine/test_duplex_orchestrator.py | 517 + tests/engine/test_orchestrator.py | 887 -- .../test_orchestrator_error_handling.py | 38 +- .../test_orchestrator_segment_stage_keying.py | 44 +- .../test_orchestrator_stage_input_bridge.py | 107 +- .../engine/test_output_metadata_snapshots.py | 8 +- tests/engine/test_single_stage_mode.py | 15 +- .../entrypoints/duplex/test_duplex_serving.py | 609 ++ .../duplex/test_runtime_adapter_boundary.py | 145 - .../duplex/test_runtime_control_redaction.py | 115 - .../duplex/test_websocket_actor.py | 145 - .../openai/test_duplex_protocol.py | 601 -- .../openai/test_duplex_session_attachment.py | 129 +- .../openai_api/test_api_server_guards.py | 1 - .../openai_api/test_duplex_api_server.py | 414 + .../openai_api/test_duplex_handler.py | 9491 ----------------- .../test_qwen3_omni_realtime_websocket.py | 2 +- .../entrypoints/openai_api/test_server_vad.py | 429 - .../test_async_omni_diffusion_config.py | 2 +- tests/entrypoints/test_async_omni_duplex.py | 601 -- .../test_disable_log_stats_routing.py | 2 +- .../test_duplex_fence_propagation.py | 126 - tests/entrypoints/test_duplex_omni.py | 326 + tests/entrypoints/test_omni_base_profiler.py | 19 +- tests/entrypoints/test_omni_entrypoints.py | 5 +- .../test_minicpmo_realtime_web_static.py | 2 +- .../minicpmo_4_5/duplex/test_commit_policy.py | 14 +- .../duplex/test_data_plane_audio_cursor.py | 78 + .../models/minicpmo_4_5/test_pipeline.py | 12 +- .../nemotron_voicechat/duplex/conftest.py | 20 + .../duplex/test_data_plane.py | 2 +- .../models/personaplex/duplex/conftest.py | 21 + .../duplex/test_unified_runtime.py | 12 +- tests/worker/test_native_duplex_hooks.py | 38 +- tools/pre_commit/check_forbidden_imports.py | 11 +- vllm_omni/benchmarks/omniinteract.py | 27 +- vllm_omni/benchmarks/patch/patch.py | 316 +- vllm_omni/clients/__init__.py | 6 +- vllm_omni/clients/duplex.py | 365 +- vllm_omni/clients/inline_duplex.py | 136 + vllm_omni/clients/minicpmo_4_5.py | 2 +- vllm_omni/config/stage_config.py | 21 +- vllm_omni/deploy/minicpmo_4_5.yaml | 4 +- vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 4 +- vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 4 +- vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 5 +- vllm_omni/engine/async_omni_engine.py | 1167 +- vllm_omni/engine/duplex/__init__.py | 3 +- .../{entrypoints => engine}/duplex/audio.py | 8 +- vllm_omni/engine/duplex/commands.py | 324 + vllm_omni/engine/duplex/config.py | 815 ++ vllm_omni/engine/duplex/contracts.py | 146 +- vllm_omni/engine/duplex/control_client.py | 228 - vllm_omni/engine/duplex/control_plane.py | 1063 -- vllm_omni/engine/duplex/events.py | 756 ++ vllm_omni/engine/duplex/messages.py | 128 +- vllm_omni/engine/duplex/plugin.py | 329 + vllm_omni/engine/duplex/realtime_commands.py | 887 ++ vllm_omni/engine/duplex/realtime_events.py | 1526 +++ vllm_omni/engine/duplex/runtime.py | 108 - vllm_omni/engine/duplex/session.py | 534 - vllm_omni/engine/duplex/session/__init__.py | 38 + .../engine/duplex/session/append_task.py | 195 + .../duplex/session}/commit_policy.py | 4 +- vllm_omni/engine/duplex/session/context.py | 161 + vllm_omni/engine/duplex/session/control.py | 328 + vllm_omni/engine/duplex/session/emitter.py | 249 + .../duplex/session/engine_session.py} | 1207 +-- vllm_omni/engine/duplex/session/helpers.py | 195 + .../engine/duplex/{ => session}/lease.py | 13 - vllm_omni/engine/duplex/session/manager.py | 872 ++ .../engine/duplex/session/model_channel.py | 1100 ++ .../engine/duplex/session/overlap_policy.py | 388 + .../engine/duplex/session/playback_ledger.py | 137 + vllm_omni/engine/duplex/session/runner.py | 1772 +++ vllm_omni/engine/duplex/turn_detection.py | 301 + vllm_omni/engine/duplex/vad.py | 488 + vllm_omni/engine/duplex_omni_engine.py | 299 + vllm_omni/engine/duplex_orchestrator.py | 352 + vllm_omni/engine/omni_engine_base.py | 1229 +++ vllm_omni/engine/orchestrator.py | 947 +- vllm_omni/entrypoints/async_omni.py | 724 +- vllm_omni/entrypoints/async_omni_base.py | 557 + .../entrypoints/cli/benchmark/cli_args.py | 6 +- vllm_omni/entrypoints/duplex/__init__.py | 11 +- .../entrypoints/duplex/audio_encoding.py | 55 + vllm_omni/entrypoints/duplex/capability.py | 34 - vllm_omni/entrypoints/duplex/chat_fallback.py | 302 - .../entrypoints/duplex/realtime_input.py | 1723 +-- .../entrypoints/duplex/realtime_output.py | 1007 -- .../entrypoints/duplex/realtime_session.py | 147 - .../entrypoints/duplex/realtime_state.py | 245 - .../entrypoints/duplex/runtime_adapter.py | 331 - .../entrypoints/duplex/runtime_bridge.py | 1454 --- vllm_omni/entrypoints/duplex/server_vad.py | 528 - vllm_omni/entrypoints/duplex/serving.py | 2431 +---- .../entrypoints/duplex/session_attachment.py | 151 +- .../entrypoints/duplex/session_runner.py | 2520 ----- vllm_omni/entrypoints/duplex/warmup.py | 7 +- vllm_omni/entrypoints/duplex/websocket.py | 290 +- vllm_omni/entrypoints/duplex_omni.py | 476 + .../entrypoints/duplex_request_client.py | 383 - vllm_omni/entrypoints/omni.py | 6 +- vllm_omni/entrypoints/omni_base.py | 18 +- vllm_omni/entrypoints/openai/api_server.py | 219 +- .../entrypoints/openai/audio_utils_mixin.py | 153 +- vllm_omni/experimental/fullduplex/README.md | 29 +- vllm_omni/model_executor/duplex_sampling.py | 6 - .../models/minicpmo_4_5/duplex/__init__.py | 6 +- .../models/minicpmo_4_5/duplex/adapter.py | 335 - .../minicpmo_4_5/duplex/capabilities.py | 14 +- .../models/minicpmo_4_5/duplex/compat.py | 6 +- .../models/minicpmo_4_5/duplex/data_plane.py | 30 +- .../models/minicpmo_4_5/duplex/input.py | 12 +- .../models/minicpmo_4_5/duplex/plugin.py | 773 ++ .../models/minicpmo_4_5/duplex/policy.py | 7 +- .../models/minicpmo_4_5/duplex/runtime.py | 439 - .../minicpmo_4_5/duplex/serving_adapter.py | 102 - .../models/minicpmo_4_5/duplex/session.py | 11 +- .../models/minicpmo_4_5/duplex/stage0.py | 107 +- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 14 +- .../models/minicpmo_4_5/pipeline.py | 8 +- .../minicpmo_4_5_omni.py | 3 - vllm_omni/utils/audio_resample.py | 167 + 182 files changed, 26951 insertions(+), 34743 deletions(-) create mode 100644 tests/benchmarks/test_realtime_tts_probe.py rename tests/{entrypoints => engine}/duplex/test_duplex_audio.py (95%) create mode 100644 tests/engine/duplex/test_duplex_commands.py delete mode 100644 tests/engine/duplex/test_duplex_control_client.py delete mode 100644 tests/engine/duplex/test_duplex_control_plane.py create mode 100644 tests/engine/duplex/test_duplex_events.py create mode 100644 tests/engine/duplex/test_duplex_plugin.py delete mode 100644 tests/engine/duplex/test_duplex_runtime.py create mode 100644 tests/engine/duplex/test_engine_session.py create mode 100644 tests/engine/duplex/test_overlap_policy.py create mode 100644 tests/engine/duplex/test_session_manager.py create mode 100644 tests/engine/duplex/test_session_runner.py create mode 100644 tests/engine/duplex/test_vad_backend.py create mode 100644 tests/engine/test_duplex_omni_engine.py create mode 100644 tests/engine/test_duplex_orchestrator.py create mode 100644 tests/entrypoints/duplex/test_duplex_serving.py delete mode 100644 tests/entrypoints/duplex/test_runtime_adapter_boundary.py delete mode 100644 tests/entrypoints/duplex/test_runtime_control_redaction.py delete mode 100644 tests/entrypoints/duplex/test_websocket_actor.py delete mode 100644 tests/entrypoints/openai/test_duplex_protocol.py create mode 100644 tests/entrypoints/openai_api/test_duplex_api_server.py delete mode 100644 tests/entrypoints/openai_api/test_duplex_handler.py delete mode 100644 tests/entrypoints/openai_api/test_server_vad.py delete mode 100644 tests/entrypoints/test_async_omni_duplex.py delete mode 100644 tests/entrypoints/test_duplex_fence_propagation.py create mode 100644 tests/entrypoints/test_duplex_omni.py create mode 100644 tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py create mode 100644 tests/model_executor/models/nemotron_voicechat/duplex/conftest.py create mode 100644 tests/model_executor/models/personaplex/duplex/conftest.py create mode 100644 vllm_omni/clients/inline_duplex.py rename vllm_omni/{entrypoints => engine}/duplex/audio.py (96%) create mode 100644 vllm_omni/engine/duplex/commands.py create mode 100644 vllm_omni/engine/duplex/config.py delete mode 100644 vllm_omni/engine/duplex/control_client.py delete mode 100644 vllm_omni/engine/duplex/control_plane.py create mode 100644 vllm_omni/engine/duplex/events.py create mode 100644 vllm_omni/engine/duplex/plugin.py create mode 100644 vllm_omni/engine/duplex/realtime_commands.py create mode 100644 vllm_omni/engine/duplex/realtime_events.py delete mode 100644 vllm_omni/engine/duplex/runtime.py delete mode 100644 vllm_omni/engine/duplex/session.py create mode 100644 vllm_omni/engine/duplex/session/__init__.py create mode 100644 vllm_omni/engine/duplex/session/append_task.py rename vllm_omni/{entrypoints/duplex => engine/duplex/session}/commit_policy.py (91%) create mode 100644 vllm_omni/engine/duplex/session/context.py create mode 100644 vllm_omni/engine/duplex/session/control.py create mode 100644 vllm_omni/engine/duplex/session/emitter.py rename vllm_omni/{entrypoints/duplex/protocol.py => engine/duplex/session/engine_session.py} (56%) create mode 100644 vllm_omni/engine/duplex/session/helpers.py rename vllm_omni/engine/duplex/{ => session}/lease.py (92%) create mode 100644 vllm_omni/engine/duplex/session/manager.py create mode 100644 vllm_omni/engine/duplex/session/model_channel.py create mode 100644 vllm_omni/engine/duplex/session/overlap_policy.py create mode 100644 vllm_omni/engine/duplex/session/playback_ledger.py create mode 100644 vllm_omni/engine/duplex/session/runner.py create mode 100644 vllm_omni/engine/duplex/turn_detection.py create mode 100644 vllm_omni/engine/duplex/vad.py create mode 100644 vllm_omni/engine/duplex_omni_engine.py create mode 100644 vllm_omni/engine/duplex_orchestrator.py create mode 100644 vllm_omni/engine/omni_engine_base.py create mode 100644 vllm_omni/entrypoints/async_omni_base.py create mode 100644 vllm_omni/entrypoints/duplex/audio_encoding.py delete mode 100644 vllm_omni/entrypoints/duplex/capability.py delete mode 100644 vllm_omni/entrypoints/duplex/chat_fallback.py delete mode 100644 vllm_omni/entrypoints/duplex/realtime_output.py delete mode 100644 vllm_omni/entrypoints/duplex/realtime_session.py delete mode 100644 vllm_omni/entrypoints/duplex/realtime_state.py delete mode 100644 vllm_omni/entrypoints/duplex/runtime_adapter.py delete mode 100644 vllm_omni/entrypoints/duplex/runtime_bridge.py delete mode 100644 vllm_omni/entrypoints/duplex/server_vad.py delete mode 100644 vllm_omni/entrypoints/duplex/session_runner.py create mode 100644 vllm_omni/entrypoints/duplex_omni.py delete mode 100644 vllm_omni/entrypoints/duplex_request_client.py delete mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py delete mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py delete mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py create mode 100644 vllm_omni/utils/audio_resample.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 33def4de4..e29001d9a 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -38,9 +38,13 @@ x-minicpmo-4-5: &minicpmo_4_5 x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/deploy/minicpmo_4_5_duplex.yaml - vllm_omni/engine/duplex/ + - vllm_omni/engine/duplex_omni_engine.py + - vllm_omni/engine/duplex_orchestrator.py + - vllm_omni/engine/omni_engine_base.py - vllm_omni/entrypoints/duplex/ + - vllm_omni/entrypoints/duplex_omni.py + - vllm_omni/entrypoints/async_omni_base.py - vllm_omni/clients/ - - vllm_omni/entrypoints/duplex_request_client.py - vllm_omni/outputs/duplex.py - vllm_omni/model_executor/duplex_sampling.py - vllm_omni/experimental/fullduplex/ @@ -249,7 +253,11 @@ source_file_dependencies: - tests/assets/minicpmo_4_5/ - tests/clients/ - tests/engine/duplex/ + - tests/engine/test_duplex_omni_engine.py + - tests/engine/test_duplex_orchestrator.py - tests/entrypoints/duplex/ + - tests/entrypoints/test_duplex_omni.py + - tests/entrypoints/openai_api/test_duplex_api_server.py - tests/model_executor/models/minicpmo_4_5/duplex/ omni_minicpmo_4_5_duplex_perf: diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index d600e8fbe..87d4c5362 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -66,7 +66,9 @@ steps: # 0.26 startup regression is understood and fixed. timeout_in_minutes: 90 commands: - - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" + # Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5 + # (see the P0 thread on #7413); the test stays in the tree and runs again there. + - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" --deselect tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_server_vad_multi_turn_without_client_commit mirror_hardwares: h100_2 - label: "Entrypoint Test with L4" diff --git a/docs/api/README.md b/docs/api/README.md index b2d4ea1ba..db760506e 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -5,6 +5,9 @@ Main entry points for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.async_omni.AsyncOmni][] +- [vllm_omni.entrypoints.async_omni_base.AsyncOmniBase][] +- [vllm_omni.entrypoints.duplex_omni.DuplexOmni][] +- [vllm_omni.entrypoints.duplex_omni.DuplexSessionHandle][] - [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][] - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][] - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][] @@ -17,12 +20,13 @@ Main entry points for vLLM-Omni inference and serving. ## Clients -Client libraries that talk to a running vLLM-Omni server (the full-duplex Realtime client and its per-model session presets). +Client libraries for full-duplex sessions: the Realtime WebSocket client, the in-process client over `DuplexOmni`, and the per-model session presets. - [vllm_omni.clients.duplex.AudioDelta][] - [vllm_omni.clients.duplex.AudioFormat][] - [vllm_omni.clients.duplex.ConnectionResumed][] - [vllm_omni.clients.duplex.DuplexClient][] +- [vllm_omni.clients.duplex.DuplexClientBase][] - [vllm_omni.clients.duplex.DuplexClientError][] - [vllm_omni.clients.duplex.DuplexConnectionError][] - [vllm_omni.clients.duplex.DuplexEvent][] @@ -44,6 +48,7 @@ Client libraries that talk to a running vLLM-Omni server (the full-duplex Realti - [vllm_omni.clients.duplex.TextDelta][] - [vllm_omni.clients.duplex.TranscriptDelta][] - [vllm_omni.clients.duplex.WebSocketTransport][] +- [vllm_omni.clients.inline_duplex.InlineDuplexClient][] ## Inputs diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 07beef2de..aa81d42c6 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -17,6 +17,7 @@ Common `PipelineConfig` fields include: | ------- | ------------- | | `model_type` | Pipeline identifier used during model and config resolution. | | `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. | +| `duplex_plugin` | Dotted path of the model's `DuplexModelPlugin`. Set only for full-duplex models; it makes `vllm-omni serve` run the model through `DuplexOmni` (a server whose every surface runs on a duplex session) and the engine host a `DuplexOrchestrator` with the plugin loaded. | | `model_arch` | Default Hugging Face architecture for the pipeline. | | `hf_architectures` | Architecture names used to identify checkpoints whose `model_type` is shared. | | `hf_config_predicate` | Optional predicate used to select between pipelines with otherwise identical HF metadata. | @@ -52,7 +53,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze | ------- | ------ | ---------- | --------- | ------------- | | `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. | | `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. | -| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. | +| `session_mode` | str | optional | `"turn"` | Session behavior. Duplex models (pipelines with a `duplex_plugin`, e.g. MiniCPM-o 4.5) set `"duplex"`; `vllm-omni serve` then serves it over `/v1/realtime?duplex=1` and `/v1/chat/completions`, both on duplex sessions. | | `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. | | `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. | | `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). | diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 7d3d5ec95..4646b85cd 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -1,1070 +1,373 @@ -# MiniCPM-o 4.5 Full-Duplex Runtime Architecture +# Unified Full-Duplex Framework -> Relocation note: after the H20 validation recorded below, the full-duplex -> implementation moved out of `vllm_omni/experimental/fullduplex/` into the -> stable tree (`vllm_omni/engine/duplex/`, `vllm_omni/entrypoints/duplex/`, -> `vllm_omni/model_executor/models/minicpmo_4_5/duplex/`, and companion -> modules). Module paths in this document reflect the new layout; the runtime -> evidence was produced on the pre-move tree, which is file-identical apart -> from module paths. +This document is the design record of the full-duplex serving framework: +the engine-resident session model, the typed command/event contract, the +single model plugin, and the serving/API surfaces built on top of them. It is +the owner document for `vllm_omni/engine/duplex/**`, +`vllm_omni/engine/duplex_omni_engine.py`, +`vllm_omni/engine/duplex_orchestrator.py`, +`vllm_omni/entrypoints/duplex_omni.py`, `vllm_omni/entrypoints/duplex/**`, +`vllm_omni/clients/**` and the per-model `duplex/` packages. + +The public wire contract lives in +[`docs/serving/realtime_duplex_api.md`](../serving/realtime_duplex_api.md); +this page describes how it is produced. ## Purpose -This document records the runtime architecture implemented by PR #3907, the -cleanup applied after published commit `e011d936`, and the remaining work that -must not be advertised as complete. Implemented contracts and future -architecture work are called out separately below. +A full-duplex model listens and speaks at the same time: audio streams into +the model while it produces audio, the model decides when to speak, and the +user can interrupt. Serving such a model needs a long-lived session with its +own identity (epoch, turn), ledgers (input, response, playback, history), a +lease (idle TTL, disconnect grace, resume) and a binding to the resumable +stage requests that carry the session's audio through the pipeline. -The checkpoint preserves the active runtime path already exercised on H20. It does -not introduce another reducer, controller, worker provider, or shadow runtime. +The framework gives that session exactly one owner, keeps the generic +turn-based stack free of duplex code, and exposes the same session to three +consumers: the OpenAI Realtime WebSocket route, the in-process Python API and +the model plugins. -## Review Snapshot - -- PR: `vllm-project/vllm-omni#3907` -- Published head before this refactor: `e011d936` -- Published base snapshot: `62589203` -- Published diff: 93 files, approximately `+29.2k/-0.6k` -- Local runtime and pytest validation: intentionally not used -- Required validation environment: an NVIDIA H20 CUDA host -- PR head reviewed for the current cleanup: - `f4f78fa555af4d65e46b003022d352ca2ea9d703` -- Validated tree: the uncommitted refactor synchronized file-for-file to the - isolated H20 worktree on 2026-07-20 - -The current tree has received fresh H20 focused and E2E validation against vLLM -0.25.0. The final affected matrix includes the pre-response continuation and -Stage1 CUDA Graph padding regressions. The checked-in MiniCPM profile admits -two concurrent sessions through an engine-owned limit; a third -session was rejected and capacity returned after one session closed. The two -accepted sessions independently completed their audio, transcript, response, -playback, and close lifecycles. This is isolation and admission evidence for the -tested two-session deployment, not a production fairness, arbitrary-capacity, -or failure-recovery claim. The validated paths include the MiniCPM duplex -runtime extension, typed resumable policy, extracted control plane and clients, request -preregistration, typed direct-output decision, engine-managed two-session -capacity in the checked-in deploy profile, separate public/runtime configuration -channels, private Session ledgers, ordered `session.update`, the MiniCPM runner -fast path, continuous -browser input, and the final Realtime input-lifecycle fixes. Exact evidence and -scope limits are recorded below. - -## Scope - -The checkpoint keeps these verified contracts: - -- MiniCPM Stage0 conversation KV continuity; -- Stage1 TTS and Token2Wav continuity; -- model-owned listen/speak decisions on the normal auto-response path; -- continuous browser PCM upload during assistant playback, without browser VAD - or browser-generated input commits; -- segment EOS and turn EOS as different boundaries; -- transcript/audio cursors scoped to a response and turn; -- playback acknowledgement and history commit; -- Stage0 request metadata routed per batch row rather than inherited from the - first request in the batch; -- Stage0 token timing snapshots and client-observed audio chunk cadence; -- scheduler data-plane append over a resumable request; -- stale epoch/turn/response fencing; -- `/v1/duplex` and OpenAI Realtime projection; -- optional OpenAI `server_vad` turn detection backed by a process-shared Silero - model and per-session VAD state; -- existing JoyVL behavior. - -The checkpoint does not claim: - -- scheduler-native KV append; -- an acoustic-onset-to-cancel latency target across arbitrary client chunk sizes; -- production multi-session admission, fairness, capacity, or failure recovery; -- bounded long-session KV; -- audio/video synchronization beyond one camera frame per model unit (camera - frames ride the audio appends as `video_frames`; the wire contract is in the - [Realtime Duplex API](../serving/realtime_duplex_api.md) guide). - -## Active Runtime Path - -The actual path is: +## Architecture ```text -WebSocket - -> OmniDuplexSessionHandler - -> DuplexSession + MiniCPMO45ServingSessionState - -> DuplexRequestClient - -> AsyncOmni thin open/append/signal/close proxies - -> DuplexControlClient over the engine-owned correlated RPC transport - -> DuplexControlPlane + DuplexSessionRuntimeManager - -> MiniCPM DuplexRuntimeExtension + engine session/stage bindings - -> StagePool - -> resumable scheduler request - -> MiniCPM Stage0 - -> Stage1 TTS / Token2Wav - -> output processor - -> Realtime protocol projection - -> WebSocket writer + Python users WebSocket clients + ┌──────────────────────┐ ┌────────────────────────┐ + │ InlineDuplexClient │ │ DuplexClient │ + │ (DuplexClientBase) │ │ (DuplexClientBase) │ + └──────────┬───────────┘ └───────────┬────────────┘ + │ DuplexSessionHandle │ /v1/realtime?duplex=1 (alias /v1/duplex) + │ ▼ + │ ┌──────────────────────────────────────┐ + │ │ OmniDuplexSessionHandler (thin) │ + │ │ websocket I/O, command_from_realtime│ + │ │ event.to_realtime(), attachment / │ + │ │ resume tokens / replay journal │ + │ └──────────────────┬───────────────────┘ + ▼ ▼ DuplexSessionHandle + ┌──────────────────────────────────────────────────────────────────────────────┐ + │ DuplexOmni(AsyncOmni) [entrypoints/duplex_omni.py, thin] │ + │ open_session / get_session / resume_session / detach_session / close_session│ + │ DuplexSessionHandle: submit(DuplexCommand) / events() -> DuplexEvent │ + └──────────────────────────────────────┬───────────────────────────────────────┘ + │ engine.open/close/resume/touch_session_async (RPC) + │ engine.submit_command_async (one-way) + ▼ + ┌──────────────────────────────────────────────────────────────────────────────┐ + │ DuplexOmniEngine(AsyncOmniEngine) [engine/duplex_omni_engine.py, thin] │ + │ loads the plugin, validates session_mode, builds queue messages │ + │ └─ DuplexOrchestrator(Orchestrator) [engine/duplex_orchestrator.py]│ + │ template seams + DuplexStagePort over the shared stage machinery │ + │ ├─ DuplexSessionManager admission, backpressure, lease reaper, │ + │ │ open/close/resume/touch, command dispatch │ + │ └─ DuplexSessionRunner ×N ONE session state (DuplexEngineSession), │ + │ ordered mailbox, append planning + stage │ + │ submit, output projection, response / │ + │ playback / history, VAD, typed events │ + │ plugin: DuplexModelPlugin (PipelineConfig.duplex_plugin, one/model) │ + └──────────────────────────────────────────────────────────────────────────────┘ ``` -The removed typed reducer/facade and worker-provider runtimes were not in this -path. This refactor continues to modify only the path above. +In each layer the duplex class extends the turn-based one, which extends +the shared base. The dependency runs duplex -> turn-based and never the +reverse, so the turn-based stack is unchanged and the duplex stack reuses +its request machinery instead of duplicating it: -## State Ownership +```text +OmniBase OmniEngineBase OrchestratorBase + ├─ Omni (sync, turn) └─ AsyncOmniEngine (turn requests) └─ Orchestrator (turn-based admission) + └─ AsyncOmniBase └─ DuplexOmniEngine (session msgs) └─ DuplexOrchestrator (hosts DuplexSessionManager, + └─ AsyncOmni (EngineClient) implements DuplexStagePort) + └─ DuplexOmni -The architecture allows several state objects only when their facts are -orthogonal. +AsyncOmni._create_engine -> AsyncOmniEngine AsyncOmniEngine._create_orchestrator -> Orchestrator +DuplexOmni._create_engine -> DuplexOmniEngine DuplexOmniEngine._create_orchestrator -> DuplexOrchestrator +``` -| Owner | Authoritative facts | +Each concrete class constructs its collaborator in an explicit factory +method the base calls at the right moment: `OmniBase.__init__` calls +`self._create_engine(...)`; `OmniEngineBase._bootstrap_orchestrator` calls +`self._create_orchestrator(...)` on the orchestrator thread once the stages +are initialized. `DuplexOmniEngine._create_orchestrator` passes the loaded +plugin and the session runtime config to `DuplexOrchestrator` directly. + +### Design rules + +1. **One session authority.** The whole duplex session (config, epoch/turn, + ledgers, lease, stage request resources, model state, projection state) is + one object, `DuplexEngineSession`, owned by one `DuplexSessionRunner` on + the orchestrator loop. Nothing above the engine keeps session state beyond + a handle (session id, event queue, capabilities, public session object). +2. **Typed contract.** Commands into a session are `DuplexCommand` + dataclasses (`engine/duplex/commands.py`); outputs are `DuplexEvent` + dataclasses (`engine/duplex/events.py`). `command_from_realtime()` and + `DuplexEvent.to_realtime()` derive the OpenAI Realtime JSON, so the wire + format is never hand-built, and the websocket handler and the inline client + share one conversion. +3. **Generic bases and the turn-based classes have zero duplex vocabulary.** + `OmniBase`, `AsyncOmniBase`, `AsyncOmni`, `OmniEngineBase`, + `AsyncOmniEngine`, `OrchestratorBase` and `Orchestrator` carry only + template seams; `tests/engine/test_duplex_import_boundary.py` checks that + importing the turn-based stack loads no duplex module. +4. **Every surface a duplex server exposes is backed by a session.** + `vllm-omni serve` constructs `DuplexOmni` when the pipeline declares + `duplex_plugin`; the server exposes `/v1/realtime?duplex=1` (alias + `/v1/duplex`), `POST /v1/chat/completions`, `/v1/models` and `/health`, and + every other turn-based route reports "not available". The chat route is + served by `DuplexChatCompletionsAdapter`, a Realtime client that runs one + short-lived session per request: speech becomes a committed turn, text is + seeded as the session's opening turn because a model-native model takes a + turn only when it hears speech. Whether a model can be reached with text is + `DuplexCapabilities.supports_chat_completions`; one that cannot is refused + at once rather than left to idle out. Turn-based use of the same model stays + available offline through `Omni` / `AsyncOmni`. +5. **Serving is transport only.** The websocket handler does socket I/O, + wire-envelope validation, command translation, event rendering and the + attachment/resume/replay bookkeeping; it holds no session state. +6. **No `typing.Protocol`** in the duplex surfaces: the plugin, data plane, + session state, PCM buffer, stage port and client transport seams are ABCs. + +## Package layout + +```text +vllm_omni/ +├── entrypoints/ +│ ├── omni_base.py OmniBase (+ abstract _create_engine) +│ ├── async_omni_base.py AsyncOmniBase (output pump, _route_engine_message seam) +│ ├── async_omni.py AsyncOmni (turn-based requests) +│ ├── duplex_omni.py DuplexOmni, DuplexSessionHandle (thin pipe) +│ ├── duplex/ SERVING (thin) +│ │ ├── serving.py OmniDuplexSessionHandler +│ │ ├── realtime_input.py RealtimeEnvelope (query rules, first message), parse_resume_request +│ │ ├── session_attachment.py DuplexSessionAttachmentRegistry (resume tokens, replay journal) +│ │ ├── audio_encoding.py encode_audio, injected into DuplexOmniEngine for the plugin's data plane +│ │ ├── chat_completions.py DuplexChatCompletionsAdapter (/v1/chat/completions on a session per request) +│ │ └── websocket.py websocket send/close/receive helpers +│ └── openai/api_server.py builds DuplexOmni for duplex models; session-backed app state +├── engine/ +│ ├── omni_engine_base.py OmniEngineBase (stage processes, orchestrator thread, queues, RPC) +│ ├── async_omni_engine.py AsyncOmniEngine (turn-based request building) +│ ├── duplex_omni_engine.py DuplexOmniEngine (session message surface, creates DuplexOrchestrator) +│ ├── orchestrator.py OrchestratorBase + Orchestrator +│ ├── duplex_orchestrator.py DuplexOrchestrator (+ DuplexOrchestratorRequestState; implements DuplexStagePort) +│ └── duplex/ +│ ├── commands.py DuplexCommand dataclasses, command_from_realtime +│ ├── realtime_commands.py Realtime client event -> DuplexCommand translation +│ ├── events.py DuplexEvent dataclasses (+ to_realtime), REALTIME_ERROR_TYPES_BY_CODE +│ ├── realtime_events.py RealtimeProjectionState: internal event -> typed events +│ ├── messages.py queue envelopes (Open/Close/Resume/Touch/Command/Result/Event), DuplexSessionError +│ ├── config.py DuplexSessionConfig, DuplexCapabilities, ResponseCreateOptions +│ ├── contracts.py DuplexFence (session_id, epoch, turn_id), stage request records, DuplexStagePort +│ ├── plugin.py DuplexModelPlugin, DuplexModelSessionState, DuplexDataPlane, PcmAppendBuffer ABCs +│ ├── turn_detection.py server-side VAD turn detector used by the session +│ ├── audio.py / vad.py / intermediate.py +│ └── session/ one engine-resident session and everything that runs it +│ ├── engine_session.py DuplexEngineSession: ledgers, lease, fence, stage resources, append sequencing +│ ├── runner.py DuplexSessionRunner (per-session mailbox on the orchestrator loop) +│ ├── manager.py DuplexSessionManager (admission, backpressure, reaper, dispatch) +│ ├── context.py DuplexSessionContext / DuplexRunState: what the runner shares with its components +│ ├── emitter.py SessionEmitter: projection, epoch filter, domain effects of a terminal event +│ ├── model_channel.py ModelChannel: submit an append, project stage output, continue a turn +│ ├── control.py SessionControl: server VAD and the events that reconfigure it +│ ├── append_task.py AppendAttempt: one append in flight and the rollback its failure owes +│ ├── helpers.py pure reads and payload builders over a session +│ ├── lease.py DuplexLeaseState (idle TTL, disconnect grace, resume generation) +│ └── overlap_policy.py / commit_policy.py / playback_ledger.py +├── config/stage_config.py PipelineConfig.duplex_plugin; DuplexSessionRuntimeConfig +├── model_executor/models/minicpmo_4_5/duplex/plugin.py MiniCPMO45DuplexPlugin (+ data_plane, input, policy, ...) +└── clients/ + ├── duplex.py DuplexClientBase (ABC), DuplexClient (websocket), client-side events + ├── inline_duplex.py InlineDuplexClient (in-process, over DuplexOmni) + ├── minicpmo_4_5.py MiniCPM-o 4.5 session preset + └── personaplex.py PersonaPlex session preset +``` + +## Session identity + +- **The server allocates every session id.** `DuplexOmni.open_session` + generates `duplex-`; a `session_id` / `id` key inside the + Realtime session object is ignored (Realtime clients echo the session object + back). The allocated id reaches the client in `session.created` and is the + only handle for `session.resume`, `close` and every command. +- **Ids are never reused** within an engine's lifetime, so the id alone + identifies a session. There is no incarnation counter: stage request ids are + `duplex-s..e..r.`, model-side per-session state + is keyed by `session_id` alone, and a command, close, touch or resume for an + id the manager no longer holds is answered with `unknown_session`. +- **Staleness inside a live session is epoch-based.** `DuplexFence + (session_id, epoch, turn_id)` is an engine-internal identity: every stage + request is bound to the fence it was submitted under, cancels advance the + epoch, and tracked tasks re-check the epoch after each `await`. +- **Resume identity** is `(session_id, resume_token, expected_lease_generation)`: + the token belongs to the attachment registry, the lease generation is the + engine-side compare-and-swap. + +## Session lifecycle + +```text +open DuplexOmni.open_session(config) + -> engine.open_session_async(session_id, DuplexSessionConfig) [RPC] + -> DuplexSessionManager.open: admission (max_sessions, closing sessions keep their slot), + plugin.validate_client_extra_body / prepare_runtime_config, DuplexEngineSession, + Stage0 request id reserved (ensure_stage_request), DuplexSessionRunner.start() + -> first event: SessionCreated (announces the allocated id) +command handle.submit(DuplexCommand) + -> engine.submit_command_async -> DuplexSessionCommandMessage on the request queue [one-way] + -> DuplexSessionManager.dispatch: unknown_session / input_backpressure checks, then runner mailbox +output DuplexOrchestrator._intercept_stage_output -> runner.on_stage_output -> mailbox -> typed events + -> output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message -> handle.events() +detach DuplexOmni.detach_session -> touch(DETACH): engine-owned disconnect grace; expiry -> SessionExpired +resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS; the existing handle is re-entered +close DuplexOmni.close_session -> close RPC; the manager tears the runner down, then the stage cleanup + (abort submitted requests, release reserved ids), then SessionClosed, then the RPC result: seeing + the event means the admission slot is free (it is held until the cleanup succeeded) +reap DuplexSessionManager.reaper_loop: idle TTL / disconnect grace expiry, cleanup retries +``` + +### Concurrency and ordering model of the runner + +Everything below runs on the orchestrator asyncio loop; there is no lock. + +- **Single writer per session.** `DuplexEngineSession` is mutated only by its + runner. Inputs reach the runner through one mailbox in this order: client + commands (from `_request_handler`, in request-queue order), stage outputs + (`on_stage_output` enqueues an internal item; it does not mutate the session + inline), timers (silence continuation, bounded auto-response) and internal + items (deferred overlap promotion, VAD-synthesized commits). The mailbox + worker processes items one at a time. +- **Slow work is not awaited inline.** An `AppendAudio` handler reserves + bytes, then schedules a tracked append task on the per-session append tail: + the task awaits `_offload(decode/resample/VAD)` on a dedicated + `ThreadPoolExecutor` owned by the manager, then `plugin.plan_append`, then + `stage_port.submit`. The worker returns to the mailbox immediately, so a + later `CancelResponse` is not blocked behind an in-flight append. +- **Re-validation after every `await`.** A tracked task captures `(epoch, + turn_id)` when it starts and re-checks them before the stage submit and + before committing ledgers; a task that finds the epoch advanced rolls back + its PCM reservation and exits. +- **Cancel is atomic at the session.** `CancelResponse` / `BargeIn` advance + the epoch, cancel tracked append tasks, abort the owned stage bindings + through the stage port and emit the cancelled terminal events before + yielding. An append task already past its last check is harmless: its stage + request is bound to the old epoch's request id, which the abort covers, and + the base never emits session-owned outputs to a client. +- **Ordered event stream.** Because the epoch bump and every emission happen + on one loop in program order, nothing can be emitted for an old epoch after + its terminal event; `runner.emit()` still drops a terminal carrying a stale + epoch and stamps `epoch` on every event for clients that filter + defensively. This is a statement about order, not about latency: the engine + output queue and each `DuplexSessionHandle` outbox are unbounded and FIFO, + so a cancellation is delivered behind whatever audio was already emitted for + the response it cancels. A client that must stop quickly cancels its own + playback on `response.done` / `audio.cancelled` rather than waiting for the + stream to drain. Bounding those buffers and letting an accepted invalidation + skip undelivered media is left to the follow-up RFC. +- **Backpressure before the mailbox.** `DuplexSessionManager.dispatch` + checks `max_pending_input_bytes_per_session` and reserves a pending turn for + `Commit` before the put; a rejected command is answered with + `ErrorEvent(code="input_backpressure")`. +- **Per-response stage metrics.** `_route_output` computes the per-segment + `StageMetrics` before calling `_intercept_stage_output`; the runner + accumulates the snapshot into the active response, so + `ResponseDone.stage_metrics` and `metadata.vllm_omni.stage_metrics` on the + wire keep their meaning. + +## Orchestrator seams + +`OrchestratorBase` keeps the generic machinery (request handler skeleton, +abort, collective RPC, output loops, stage error and dead-replica handling, +`_route_output`, stage forwarding, prewarm, cleanup, shutdown) and offers +these seams. `DuplexOrchestrator` extends `Orchestrator` rather than sitting +beside it: a seam it does not override falls through to the turn-based +implementation, and `_dispatch_message` hands anything that is not a session +message to `super()`: + +| Seam | `Orchestrator` | `DuplexOrchestrator` | +| --- | --- | --- | +| `_dispatch_message(msg) -> bool` | `add_request`, `streaming_update`, `add_companion_request`, `interaction` | open/command/close/resume/touch messages -> `DuplexSessionManager` | +| `_background_tasks()` | — | `[manager.reaper_loop()]` | +| `_shutdown_extensions()` | — | `await manager.shutdown()` | +| `_on_stage_submitted(stage_id, request_id, replica_id, req_state)` | — | bind the forwarded stage request to the owning session | +| `_intercept_stage_output(...) -> bool` | — | hand every session-owned output (Stage0 decision, Stage1 audio) plus its metric snapshot to `runner.on_stage_output`; `True` = do not forward | +| `_handle_forward_failure(...) -> bool` | — | fail the owning session's response and close the session, keep the loop alive | +| `_cleanup_request_ids(ids, abort=, release_owners=)` | ignores `release_owners` | brackets `super()` with `close_sessions_for_request_ids` / `defer_request_cleanups` / `finalize_closed_sessions` | +| `OrchestratorRequestState.session_owned` | never set | set by `ensure_request`: no synthetic terminal, no client emission, no terminal re-forward, no auto-cleanup on finish | + +`DuplexOrchestrator` also implements the `DuplexStagePort` ABC the manager and +runner use (`stage_count`, `sampling_defaults`, `ensure_request`, `submit`, +`cleanup`, `abort_requests`): the resumable Stage0 request is preregistered at +open with `DuplexOrchestratorRequestState(session_owned=True, session_id, +fence)`, submitted with `submit_initial` on the first unit and +`submit_update` afterwards, and the request's `streaming.bridge_states["duplex"]` +carries the session identity (`session_id`, `fence`, `epoch`, `turn_id`, +`model_turn_id`), the public session config and the server-owned runtime +config for the worker-side model hooks. + +## Model plugin + +`DuplexModelPlugin` (ABC, `engine/duplex/plugin.py`) is the one class a model +provides, selected by `PipelineConfig.duplex_plugin`; "is a duplex model" is +`duplex_plugin is not None`. It merges the engine policy and the session +policy that used to be two separately configured objects: + +| Half | Members | | --- | --- | -| `DuplexSession` | session state, epoch, turn, active/last response, pending input, overlap duration, playback cursor, history | -| `MiniCPMO45ServingSessionState` | MiniCPM PCM buffer, deferred overlap payload, data-plane task, continuation effect counter | -| `DuplexWebSocketActor` | inbound mailbox, outbound queue, writer, transport close state, effect task handles | -| `NativeRealtimeSessionProtocol` | wire-only response/item projection and input-buffer projection | -| `DuplexSessionRuntimeManager` | fence snapshot, append reservation, completed-operation cache, resource reservations, stage bindings | -| `DuplexControlPlane` | control-message validation, transactional open/append/update/close, extension invocation, typed segment/output decisions | - -Important distinctions: - -- Realtime response IDs are projection caches, not domain decision sources. -- Engine fences duplicate identity across a process boundary by design; they - are immutable validation snapshots. -- Stage bindings are engine resources, not serving response state. -- `continuation_owner_id` scopes an effect counter either to a visible response - or, before `response.created`, to an explicit model turn. A pre-response - continuation also revalidates request, epoch, turn state, and model-turn - identity immediately before append; it cannot survive turn completion. -- runtime open/close acknowledgements are handler-local effect bookkeeping, - not Session lifecycle. - -## Implemented Concurrency Contracts - -### Single inbound mailbox - -The Actor now uses one: - -```python -asyncio.Queue[event] -``` - -This replaces competing control/input/event `Queue.get()` tasks and the -deferred-event list while preserving WebSocket arrival order. A later close, -cancel, or clear must not overtake earlier append, playback ACK, or -session-update events. Slow runtime appends remain background tasks. Cancelling -their asyncio task does not stop a synchronous callable already running in the -default executor, so cancellation correctness is provided by the engine fence, -not by `Task.cancel()`. - -Properties covered by focused tests: - -- every enqueued event is delivered exactly once; -- terminal control preserves wire order while cancelling background work; -- clear and cancel never scan or delete later mailbox events; -- one writer owns all WebSocket sends. - -### Correlated RPC result router - -`AsyncOmniEngine` now has one consumer for `rpc_output_queue`. Results route by: - -```text -("duplex", control_id) -("collective", rpc_id) -``` - -The router replaces the global RPC lock. It supports out-of-order replies, -timeout unregister, late-result rejection, fatal-error broadcast, and close -unblocking. - -Additional lifecycle rules: - -- uncorrelated non-fatal errors are not broadcast to unrelated waiters; -- a latched fatal rejects new calls before they enter the request queue; -- `EngineDeadError` is broadcast directly to the RPC queue even though the - orchestrator consumes that exception during teardown; -- shutdown closes the router before joining the orchestrator; -- failure to enqueue shutdown closes the request queue and still attempts a - bounded orchestrator join; -- `StageRuntime` cleanup runs only after the orchestrator has stopped, avoiding - concurrent client teardown. - -The router removes wait-side head-of-line blocking. The orchestrator request -handler remains ordered, but executor submission order is not itself a session -ordering guarantee. - -### Extracted control ownership - -The generic engine mechanism and the duplex plugin are split by -responsibility: - -| Module | Responsibility | -| --- | --- | -| `engine.messages` | ordinary engine queue messages and collective RPC envelopes | -| `engine.rpc_result_router` | generic correlation routing through each result message's `rpc_correlation_key` | -| `engine.duplex.contracts` | immutable duplex DTOs, extension/stage protocols, typed decisions, request-ID codec | -| `engine.duplex.lease` | duplex lease activity, TTL/grace configuration, and expiry records | -| `engine.duplex.messages` | identity fence and duplex control message/result envelopes | -| `engine.duplex.runtime` | extension loading, validation, and compatibility exports | -| `engine.duplex.session` | session resources, fences, append reservations, idempotency cache, binding cleanup | -| `engine.duplex.control_plane` | control algorithms and the narrow `DuplexStagePort` used to submit and clean up stage requests | -| `engine.duplex.control_client` | typed control-message construction and calls through the engine-owned correlated RPC transport | -| `entrypoints.duplex_request_client` | deterministic request identity, client-state preregistration, acknowledgement validation, data-plane collection, rollback | - -`Orchestrator`, `AsyncOmniEngine`, and `AsyncOmni` retain explicit wiring and -compatibility proxies. They no longer own the corresponding duplex algorithms. -The ControlClient does not own waiters or consume the result queue; -`RpcResultRouter` remains the single correlation owner for both duplex and -collective RPC. The ControlPlane receives a narrow stage port rather than the -Orchestrator object. - -Control-plane enablement is independent from extension presence. -`PipelineConfig.duplex_control_enabled` determines whether an Orchestrator -constructs a ControlPlane. This preserves extension-free -`TURN_COMMIT_ONLY` deployments while allowing ordinary deployments to use a -single `None` fast-path check. The MiniCPM pipeline explicitly enables control -and separately selects its runtime extension. - -### Transactional open and append - -Opening a native session registers its engine state and reserves the Stage0 -request resource as one externally atomic operation. If capability, sampling, -extension, or request-state initialization fails, the ControlPlane removes the -session, reserved request resources, and any preregistered request state before -returning the error. Cleanup distinguishes a reserved resource from a submitted -request, so it does not decrement the running counter for unrelated work. -Closing a session that never appended also removes the preregistered resource. - -Expiry and request-triggered cleanup use a two-step close. The runtime manager -first marks the session closing but retains it in the admission set. Only a -successful stage cleanup finalizes and removes the session. A blocked or failed -cleanup therefore continues to consume its server-owned slot and is retried; -replacement admission cannot overlap request or KV resources that are still -owned by the closing session. - -Append uses prepare/submit/commit ordering. `prepare_append()` computes the -next sequence values without mutating the session. Prompt planning and stage -submission use that reservation, and only a successful submit commits -`input_seq`, `input_turn_seq`, and the accepted fence. Serving uses a matching -PCM reservation and consumes bytes only after the append acknowledgement; -failure or a closed predecessor rolls the bytes back. An `operation_id` caches -completed append results so a retry cannot submit the same physical append -twice. - -### Ordered append and configuration effects - -Native append effects form a per-session tail. Each append awaits its -predecessor before entering the engine, and `session.update` awaits that same -tail before changing either the serving aggregate or engine configuration. -Consequently, an append received before an update uses the old immutable -sampling snapshot, while the first accepted append after the update uses -the new generation. - -The engine increments `config_generation` when it accepts the replacement -configuration. At the next append boundary, the control plane rebuilds sampling -parameters from stage defaults through the model extension. It never mutates -the snapshot of a segment already in progress. - -### Irreversible cancellation fence - -Every cancel control carries the cancelled fence and the next session fence. -The `AsyncOmni` facade preserves both values across the serving, -entrypoint, engine-message, and orchestrator boundaries. The orchestrator then -atomically: - -1. advances the runtime state to the next fence; -2. releases and aborts bindings owned by the cancelled fence; -3. permanently rejects an append that arrives later with the old fence. - -Missing `next_fence` is rejected without releasing bindings; the engine never -advances independently from serving. A late cancel cannot move a session -backward when the runtime is already on a newer fence. Idle playback clear does -not send an engine cancel because no runtime resource is being cancelled. - -`model_native_duplex` also validates its engine-client contract before opening -the session. The open, append, signal, and close methods must exist and declare -the required `fence` arguments; signal must additionally accept `next_fence`. -An older or custom client that cannot preserve this boundary fails with -`runtime_contract_invalid` instead of silently running an unfenced native -session. The generic serving-session adapter retains its compatibility behavior. - -`DuplexFence` also carries a server-side session incarnation. Reusing a public -session ID after close increments that incarnation, and resource request IDs -include it after the first incarnation. A synchronous append from a closed -session therefore cannot be accepted by a newly opened session with the same -public ID. Stale close validates the full fence before removing the runtime -registry entry or releasing bindings. - -## Implemented Serving Ownership - -### Session-owned response and overlap state - -The Actor no longer stores: - -- a Session reference; -- lifecycle mirror state; -- active or last response ID; -- overlap duration; -- response-in-progress or playback business predicates; -- runtime open/close acknowledgements. - -`DuplexSession.begin_response()` owns both active and last response identity. -The Session owns overlap accumulation/reset and playback state. - -One `native_response_in_progress()` predicate now covers: - -- active response ID; -- active engine request ID; -- uncommitted ACK-only playback; -- active generic response task; -- response-bound append tasks; -- MiniCPM data-plane task; -- assistant-generating turn state. - -The write-only `MiniCPMO45ServingSessionState.response_emitted` field and a dead -data-plane task mirror were removed. - -### Response-scoped options - -`response.create` is represented by an immutable `ResponseCreateOptions` value. -The Session reserves it for one response, applies it to a response-local copy -when that response begins, and restores the base session configuration on -completion, listen, cancellation, or failure. Instructions, voice, modalities, -temperature, token limit, format, tools, and conversation metadata therefore do -not mutate the permanent session configuration or leak into a later response. - -### Explicit event effects and pure writer - -Business paths call `emit_event()`. It serializes domain effects and outbound -queue insertion with one session-local lock. Deferred append starts outside the -lock so recursive event emission cannot deadlock. - -The writer only: - -1. dequeues an event; -2. applies late stale filtering to streaming model output; -3. projects the event to the selected wire protocol; -4. sends JSON. - -The writer does not promote overlap, clear input, advance Session state, call -the engine, or commit history. - -### Single terminal acceptance point - -Terminal events are accepted or rejected before domain effects run: - -```text -response.done -response.listen -audio.cancelled -input.cancelled -session.closed -``` - -For a terminal carrying an epoch, that epoch must match the current Session. -Stale cancellation therefore cannot reset overlap belonging to a later turn. -Once a terminal is accepted and its effect is applied, the writer cannot revoke -it because close or epoch state changed while it waited in the output queue. - -Streaming audio/text deltas still receive late stale filtering at send time, -which prevents queued old audio from leaking after a cancel or epoch bump. - -## Runtime Type Boundary - -The stable kernel is deliberately limited to generic engine and scheduler -mechanisms: - -```text -vllm_omni.engine.messages # ordinary queue and collective RPC envelopes -vllm_omni.engine.rpc_result_router # generic correlation routing -vllm_omni.engine.async_engine_utils # ordinary engine lifecycle helpers -``` - -The full-duplex implementation lives in dedicated duplex modules: - -```text -vllm_omni.engine.duplex.runtime - # extension loading and validation -vllm_omni.engine.duplex.contracts - # immutable DTOs and narrow protocols -vllm_omni.engine.duplex.lease - # duplex lease primitives -vllm_omni.engine.duplex.messages - # duplex identity and queue envelopes -vllm_omni.engine.duplex.session - # session transaction implementation -vllm_omni.engine.duplex.control_plane - # engine-side control algorithms -vllm_omni.engine.duplex.control_client - # correlated control client -vllm_omni.entrypoints.duplex_request_client - # entrypoint request/output lifecycle -``` - -Stable engine and entrypoint modules import the duplex kernel eagerly at -module load (top-of-file imports). Duplex behavior remains opt-in: the -ControlPlane, ControlClient, and request client are *constructed* only on -duplex-enabled deployments, and the model runtime extension still loads only -through the dotted-string plugin path. The client-side package -(`vllm_omni/clients/`, home of the public `DuplexClient`) is never imported -by runtime code. - -`PipelineConfig.duplex_control_enabled` enables the generic mechanism, while -`PipelineConfig.duplex_runtime_extension` separately selects the model adapter. -The MiniCPM pipeline opts into both and installs -`MiniCPMO45DuplexRuntimeExtension`; ordinary pipelines construct neither a -ControlPlane nor a ControlClient. The model-neutral protocol exposes -sampling-parameter configuration, append planning, and a typed stage-output -decision. MiniCPM owns -stage-specific overrides, PCM-to-token budgeting, force-listen payload policy, -special-token interpretation, and `listen` response metadata. The generic -Orchestrator does not parse MiniCPM session keys, `listen_token_id`, or construct -MiniCPM metadata, and it does not import the adapter. - -Extension loading is fail-fast. Startup validates every required callable and -the number and type of returned stage sampling parameters. Invalid custom -extensions fail before a session can open. - -Direct listen/speak outcomes use `DuplexOutputDecision` on -`OmniRequestOutput`. That typed envelope is authoritative over inner completion -metadata for native projection, so a processed output cannot hide a listen -decision merely because it retains unrelated multimodal metadata. - -Generic multimodal chunk accumulation lives in -`vllm_omni.outputs.multimodal_accumulation`. It owns snapshot replacement, -DELTA draining, and audio-finality interpretation; `OmniRequestState` only -invokes that policy while assembling request output. Model output processors no -longer need to add chunk-retention rules to the main output orchestration file. - -The strict chunk-transfer schema still contains `hidden_states.tts` as the -generic AR-stage-to-TTS handoff and several MiniCPM streaming token IDs as -compatibility wire fields. The latter are not a preferred generic model -contract. Moving them behind a model-namespaced metadata envelope requires a -versioned producer/consumer migration because unknown msgspec fields are -rejected; silently deleting or renaming them in this checkpoint would break -existing Stage0-to-Stage1 transfer. - -### Public and runtime configuration channels - -Control messages carry two independent snapshots: - -- `session_config` contains the public Realtime/session values used by serving - and model prompt policy; -- `runtime_config` contains server-derived MiniCPM sampling, scheduler, context - reserve, prefix-budget, and resolved reference-audio values. - -The MiniCPM serving adapter derives `runtime_config` from deploy/model defaults -and validated public options. Clients cannot provide private runtime keys in -`session.create` or `session.update`; response-scoped `extra_body` drops those -keys rather than turning them into engine policy. The ControlPlane constructs -sampling parameters only from `runtime_config`, -while append planning receives both snapshots explicitly. Raw Realtime -`extra_body` is not used as the engine's sampling-policy channel. - -`session.update` prepares both candidates, validates the runtime candidate in -the engine, and replaces the serving snapshots only after the engine ACK. A -rejected candidate therefore leaves both the public configuration and runtime -sampling generation unchanged. - -`DeployConfig.duplex_session` is an immutable, server-owned -`DuplexSessionRuntimeConfig`. These limits are not read from client -`extra_body` or session payloads: - -| Field | Default | Owner | Lifecycle or resource effect | -| --- | ---: | --- | --- | -| `idle_ttl_s` | `300.0` | Engine lease manager | Expires an attached or detached session after inactivity; `null` disables only idle expiry. | -| `disconnect_grace_s` | `30.0` | Engine lease manager and API attachment registry | Bounds how long a detached session may retain engine resources before cleanup. | -| `reaper_interval_s` | `5.0` | Orchestrator | Sets the cadence for lease expiry and pending cleanup retries. | -| `resume_replay_ttl_s` | `60.0` | API attachment registry | Expires replay events retained for resume or takeover. | -| `resume_replay_max_bytes_per_session` | `8388608` | API attachment registry | Bounds each session's replay buffer. | -| `max_pending_input_bytes_per_session` | `16777216` | API session input ledger | Applies per-session byte backpressure before input is admitted. | -| `max_pending_turns_per_session` | `4` | API session input ledger | Bounds queued turns that have not completed runtime processing. | -| `max_sessions` | `1` | Engine session manager | Atomically enforces admission; closing sessions retain capacity until cleanup finalizes. | -| `completed_append_cache_size` | `256` | Engine session manager | Bounds completed append idempotency records per session. | - -Deploy profiles may lower or raise these values to match scheduler capacity, -but clients cannot override them. The engine remains authoritative for -admission and cleanup ownership even when the API performs an early check. - -### Client-side companion package - -The public client surface ships under `vllm_omni/clients/` and stays on the -consumer side of the boundary above: `vllm_omni.clients.duplex` provides the -model-agnostic `DuplexClient` (typed events, per-response demux via -`ResponseHandle`, `session.resume` with acknowledgement-based replay dedup) -and the `EventCollector` helper used for timing and request-metric summaries. -Model-specific session defaults live in thin per-model preset modules — -`vllm_omni.clients.minicpmo_4_5.create_duplex_session_config()` and -`vllm_omni.clients.personaplex.create_duplex_session_config()` — so the core -client never imports model packages and the model duplex plugin packages -never import the client (both directions are enforced by tests: -`tests/engine/test_duplex_import_boundary.py`, -`tests/entrypoints/openai_api/test_duplex_capability.py`, and the subprocess -purity guards in `tests/clients/test_model_session_configs.py`). -`examples/online_serving/barge_in_client.py` (with its flow companion -`barge_in_client_flow.md`) is the runnable reference for the client library -against a live duplex deployment. The wire vocabulary the client speaks — -every `/v1/realtime?duplex=1` event with its OpenAI-compatibility tier and a -JSON example, plus the native-to-Realtime name map — is catalogued in the -[Realtime Duplex API](../serving/realtime_duplex_api.md) serving guide; that -page derives from the normative contract table above and must be updated -with it. - -## Generic-Path Cleanup - -Model-specific `MINICPMO45_PROFILE_LOGS` probes were removed from: - -- the generic AR scheduler; -- orchestrator; -- AsyncOmni; -- the generic GPU AR runner. - -The cleanup removes the temporary serving and generic-runtime probes. The -generic AR runner builds `DuplexSamplingRow` values only when the model exposes -the optional `prepare_duplex_sampling()` hook, then invokes that hook exactly -once before the normal sampler. Ordinary models do not scan request rows. -MiniCPM owns force-listen logits, turn-ended state, and its row-local sampling -policy behind that hook; the runner no longer reads `_minicpmo45_*` attributes -or retries sampling after matching `TypeError` text. - -The scheduler no longer parses serving `session_config`, `extra_body`, or -`duplex_stage_sampling_params`. Stage-specific overrides are materialized as -ordinary `SamplingParams` by the model runtime extension before scheduler -admission. Existing vLLM stop-token, EOS, and max-token handling defines each -segment boundary; `resumable=True` parks the request for the next streaming -update. The scheduler does not inspect a duplex dictionary, Realtime -configuration, or MiniCPM sampling fields. - -The generic runner's unused `_duplex_force_listen_applied_segments` cleanup and -the obsolete `streaming_accumulated_keys` accumulation hook were removed. The -latter had no model producer after its upstream Qwen removal; streaming input -still performs the ordinary replacement merge needed by active models. - -Local `turn.signal` transitions such as `user_started` and -`assistant_started` no longer perform an engine RPC. Only cancellation and -`session.update` cross the engine boundary because those operations change -runtime resource identity or runtime configuration. - -During a stage transition, `StreamingInputState.source_token_decoder` exposes a -short-lived, model-neutral decoder sourced from the upstream output processor. -The MiniCPM Stage1 input processor consumes that capability for token-faithful -delta transcripts. The generic orchestrator no longer installs or names a -MiniCPM-specific bridge-state key. - -The stable API server imports the duplex serving package eagerly and inspects -the deployment `session_mode` before constructing the duplex handler. Ordinary -deployments do not create its registries and background state. -For an enabled deployment, the client still selects the Realtime duplex route -with `?duplex=1` (or an equivalent explicit true value). Model-name matching is -not used for routing or native-runtime activation. MiniCPM clients explicitly -set `extra_body.native_duplex=true`; the canonical way to do so is -the `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config()` preset, -which the repository's barge-in example -(`examples/online_serving/barge_in_client.py`) uses in its session payload. - -`DuplexRequestClient` derives the data-plane request identity from the accepted -fence and preregisters a resumable `ClientRequestState` before the append crosses -the engine boundary. The acknowledgement must return the same identity or the -operation fails. Unregistered outputs, including names that merely use the old -`duplex-` prefix, are dropped. Timeout, cancel, and close remove the -preregistered state. - -## Why Scheduler Changes Remain - -Serving cannot preserve model KV after a segment ends. The scheduler must keep -one resumable request in a waiting state and accept a later update: - -```text -RUNNING - -> segment stop -WAITING_FOR_STREAMING_REQ (KV retained) - -> next append -RUNNING - -> session close -FINISHED (KV released) -``` - -Scheduler responsibilities are limited to resumable request state, runtime -context update, stop/boundary handling, and final release. Scheduler must not -own response IDs, playback, overlap, Realtime events, or model policy. - -## Serving and Compatibility Boundaries - -### Normative Realtime event contract - -The public contract in this section applies to `/v1/realtime?duplex=1`. Internal -duplex event names such as `response.output_audio.delta` are projection inputs, -not additional public aliases. Events for one attachment preserve mailbox -order; replay after resume preserves the original event order and `event_id`. - -Accepted compatibility aliases are normalized before mailbox admission: - -| Canonical input event | Accepted aliases | Contract | -| --- | --- | --- | -| `input.text.append` | `input_text.append`, `push_text` | Appends text to the open input item. | -| `input_audio_buffer.append` | `input.audio.append`, `push_chunk` | Appends audio; legacy aliases default to WAV when no format is supplied. | -| `turn.signal` | `signal_turn` | Sends an explicit turn signal; it does not imply response creation. | -| `playback.ack` | `audio.playback_ack` | Advances playback/history only for the identified response and attachment. | -| `session.close` | `close_session`, `close` | Closes the session after previously admitted mailbox events. | - -The Realtime projector emits the following response lifecycle. Cardinality is -scoped to one `response_id` unless stated otherwise. - -| Public event | Trigger and ordering | Cardinality | -| --- | --- | --- | -| `response.created` | First event for a visible assistant response; precedes its output-item, content-part, speak, text, and audio events. | Exactly once for every visible response. | -| `response.speak` | The model selected speak; emitted after `response.created` and no later than the first `response.output_audio.delta`. It carries decision metadata, not transcript text. | At most once. | -| `response.output_audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.output_audio_transcript.delta` for the same chunk. | Zero or more. | -| `response.output_audio_transcript.delta` | Append-only transcript contribution paired with an audio delta; it may be empty for a text-less audio unit and must not repeat or overlap earlier text. | One per audio delta. | -| `response.output_audio.done` | Closes the audio stream after its final delta and before response terminal events. It is omitted for a response with no audio or transcript. | At most once. | -| `response.output_audio_transcript.done` | Contains the exact concatenation of all non-overlapping transcript deltas for the response. | At most once, and only for a non-empty transcript. | -| `response.output_item.done` / `conversation.item.done` | Finalize the assistant item after content-part terminal events and before `response.done`. | At most once each. | -| `response.done` | Terminal event for a created response; follows audio, transcript, content-part, and item terminal events. | Exactly once for every created response that reaches a terminal state. | -| `rate_limits.updated` | Compatibility event emitted immediately after `response.done`; the current payload has an empty rate-limit list. | Once per emitted `response.done`. | -| `response.listen` | The model selected listen at a model-unit boundary. It may be emitted without `response.created`; clients must not infer a missing or pending response from it. | Zero or more per session, at most once per model-unit decision. | - -Session lifecycle events (`session.created`, `session.updated`, -`session.resumed`, `session.replaced`, `session.expired`, and -`session.resync_required`) are session-scoped rather than response-scoped. -`session.created` is emitted before held response output on a new attachment. -Cancellation and truncation use terminal status/details on the affected -response and never reuse that response ID for later model turns. - -Compatibility changes must update this table and its protocol contract tests in -the same change. The golden transcript test requires joined -`response.output_audio_transcript.delta` values to equal -`response.output_audio_transcript.done.transcript`; response tests separately enforce -at-most-once `response.speak` and terminal event cardinality. - -### Session internal ledgers - -`DuplexSession` remains the serving aggregate root and now composes -`InputBufferState`, `ResponseState`, `PlaybackLedger`, and -`ConversationHistory`. The ledgers are private. Read-only properties return -tuples, snapshots, or mapping proxies, and serving modules mutate them through -Session transition methods for request, response-turn, lifecycle, playback, -capability replacement, and history. The Realtime session-update parser still -normalizes individual fields before atomically retaining or rolling back the -aggregate configuration. Further work may narrow the read-only compatibility -surface, but it must not introduce another reducer or a second session state -machine. - -### Realtime input ownership - -The browser is a continuous PCM producer. While the session is open and the -microphone is unmuted, it sends `input_audio_buffer.append` every 200 ms, -including while assistant audio is being generated or played. It does not run -VAD and does not send `input_audio_buffer.commit`. Native Stage0 consumes the -stream in approximately one-second model units; MiniCPM listen, speak, and turn -EOS decisions advance the model conversation. - -Explicit Realtime clients may still use `input_audio_buffer.commit` to create a -conversation item. The translator validates that wire input before producing a -commit carrying `item_id`. Native auto-response may already have -streamed those PCM samples into the runtime, so the runner accepts a validated -commit even when no runtime-side chunk remains. A truly empty explicit buffer -continues to return `input_audio_buffer_empty`. - -The commit-to-response contract is: - -| Session state at `input_audio_buffer.commit` | Runtime action | Required event outcome | -| --- | --- | --- | -| Auto-response disabled, valid idle input | Commit and retain the input item. | Emit `input_audio_buffer.committed`; emit no response until an explicit `response.create`. | -| Auto-response enabled, valid idle speech | Submit the final native append for the current model turn. | Emit `input_audio_buffer.committed` before model output. Create a response only on the first visible text/audio output; a model listen decision may end with `response.listen` and no response. | -| Auto-response enabled, short overlap during an active response | Discard the short overlap as an acknowledgement. | Emit the committed/no-response acknowledgement; do not create or defer a response. | -| Auto-response enabled, meaningful overlap during an active response | Retain one deferred payload and promote it after the active response terminates. | Emit `input_audio_buffer.committed` with deferred metadata. A later response follows the same visible-output rule and never reuses the prior response ID. | -| Explicit silence/noise (`is_speech=false`) | Clear the physical input without a runtime append. | Emit the committed/no-response acknowledgement followed by `response.listen`. | -| Empty or invalid input | Reject without changing response ownership. | Emit one typed `error`; do not emit `response.created`. | -| Previous response cancelled or truncated and cleanup complete | Apply the corresponding idle rule above. | The prior response remains terminal; any later visible output receives a new response ID. | - -Append tasks, explicit commits, cancellation, and deferred promotion are -serialized by the session actor. Tests cover post-response commits, active -playback deferral, truncated-response ownership, and two-turn -response-required execution; the H20 validation record below is the golden -multi-turn runtime evidence. - -During auto-response overlap, `preserve_realtime_input` distinguishes "do not -append this silent chunk to the native buffer" from "clear the open Realtime -item". Silent overlap no longer discards earlier user PCM. - -The first chunk of one overlapping input item also reserves its target model -turn. A later Realtime commit uses that reserved identity even if response EOS -has already advanced the Session turn; clear, cancel, close, and successful -promotion release the reservation. - -### Physical input and model-turn identity - -A Realtime input item, an optional `input_audio_buffer.commit`, and a MiniCPM -model turn are intentionally different identities. The browser normally keeps -one input stream open without commits. Native Stage0 evaluates that stream in -approximately one-second model units. A sampled `<|turn_eos|>` closes the -current model turn and allows a later streamed unit to start another turn. -Explicitly committed input may also produce only model-listen decisions and no -spoken response. - -Therefore neither `model_turn_id` nor response cardinality is derived from the -number of physical inputs. `model_turn_id` advances at model EOS. The -`response-required` fixture may require exactly one response per requested -turn, but `model-policy` validation must accept additional model-owned responses -and cannot assert cross-turn transcript independence through physical-input -attribution. -Before closing, the model-policy demo drains every created response through -`response.done`, acknowledges every completed audio playback, and requires a -bounded quiet interval with no newly created response. This validates lifecycle -completion without suppressing or overriding the model's listen/speak policy. - -### Non-cancel signal semantics - -`turn.signal` remains a public compatibility surface for local Session and -Realtime transitions. Serving forwards only cancellation and `session.update` -to the engine. The engine rejects arbitrary non-cancel signal names instead of -returning a misleading supported acknowledgement. `session.update` remains a -real engine operation because it replaces the configuration consumed by later -append plans. - -## Remaining Architecture Work - -The following work is deliberately not hidden inside this checkpoint. - -### Serving composition - -The generic runner and runtime bridge now depend on the model-neutral -`ServingRuntimeAdapter` protocol and no longer import MiniCPM modules. The -`MiniCPMO45ServingRuntimeAdapter` owns MiniCPM serving state, PCM preparation, -data-plane projection, client/runtime configuration validation, prefix policy, -and capability projection. `PipelineConfig.duplex_serving_adapter` explicitly -selects its import path; `AsyncOmniEngine` and `AsyncOmni` carry that path to the -API server, which passes it into the generic handler. The handler has no model -default and fails startup when the duplex endpoint is enabled without an -adapter. Tests may inject an adapter instance directly. Import-boundary tests -verify that loading the generic runner, bridge, or handler does not load MiniCPM -modules. - -Behavior is still assembled through `DuplexSessionRunnerMixin`, -`NativeRuntimeBridgeMixin`, and `ChatFallbackProjectorMixin`, so explicit -adapter composition is only the model-policy boundary, not a completed removal -of host-method coupling. A later refactor may replace the remaining mixins with -explicit protocol and effect-runner dependencies while preserving one actor -mailbox and one `DuplexSession` transition owner. Splitting those components -must not create a second reducer or concurrent state machine. - -The four private Session ledgers are data partitions, not independent owners. -`DuplexSession` intentionally remains the aggregate root. Further encapsulation -may narrow compatibility properties, but serving code must continue to perform -state changes through Session transitions. - -### Plugin descriptor and typed payloads - -The pipeline currently selects the MiniCPM engine extension and Serving adapter -through two explicit fields, but there is not yet one versioned plugin descriptor -that binds those identities in an open handshake. Realtime client payloads do -not select either implementation. Adding a second native model should first -introduce that descriptor and reject serving/engine plugin mismatches explicitly. - -The stable boundary has typed fences, plans, policies, decisions, and control -messages, but `session_config`, `runtime_config`, and append payloads still use -generic mappings or objects at the extension boundary. Model-neutral typed -session and input-chunk contracts remain follow-up work, especially before -adding video. - -### Output modality policy - -The current Realtime duplex contract supports text and audio output. Unknown -output modalities fail closed with `unsupported_response_modality`; they are -never silently projected as text. This checkpoint deliberately does not add a -modality registry because there is only one native text/audio projection path -to register. - -A registry becomes justified when a second non-audio/non-text implementation -uses the same validation, encoding, lifecycle, and projection contract. At that -point the abstraction must be extracted from two working adapters and retain -the same fail-closed behavior. A speculative registry before that condition -would add selection state without removing any existing branch. - -### Performance measurement contract - -Stage0 token timing comes only from the engine's per-segment metric snapshot: -`vllm_ttft_ms`, `vllm_tpot_ms`, and `vllm_itls_ms`. A resumable Stage0 request -exports its segment-local `num_generation_tokens`; ordinary requests retain the -existing output-token fallback. Realtime response state clears its metric -accumulator at `response.created`, combines only model units owned by that -response, and ignores session-level events with no matching `response_id`. -Realtime projection carries the result under -`metadata.vllm_omni.stage_metrics`; clients do not estimate TPOT from transcript -deltas or subtract cumulative snapshots. - -Audio timing uses the client's monotonic receive clock and is partitioned by -`response_id`. `response_created_to_first_audio_ms` measures projection startup -after `response.created`; `commit_to_first_audio_ms` measures end-to-end output -from an explicit input commit. Neither field is called TTFP without naming its -origin. Reports also include inter-chunk and chunk-duration -count/mean/p50/p95/max plus maximum chunk gap. Engine and client measurements -have different clock origins and must not be merged into one timeline. - -### Stage1 CUDA Graph boundary - -Native Stage1 streaming advances Python-owned request generators and turn -state on every scheduler step. A full-model CUDA Graph would execute that -Python control flow only while capturing and then replay stale chunk and turn -metadata, even if its host-to-device tensor construction were made -capture-safe. The duplex deploy profile therefore keeps `enforce_eager=false` -but selects `cudagraph_mode=PIECEWISE` for Stage1. Python orchestration runs on -every step while fixed-topology compiled regions remain eligible for CUDA -Graph capture. This is a model execution constraint, not an E2E-only override. - -### Resource capabilities - -At the time of this snapshot the checked-in MiniCPM duplex deploy profile set -Stage0 and Stage1 `max_num_seqs` to two and configured the engine runtime -manager with `max_sessions=2`; the duplex settings have since been folded into -`vllm_omni/deploy/minicpmo_4_5.yaml`, which admits four sessions -(`duplex_session.max_sessions: 4`, matching its stage `max_num_seqs`). The -Realtime capability response advertises -multi-session and same-replica multi-session support with -`session_admission_mode="engine_managed"`. Client session fields cannot raise -this server-owned limit. Fresh E2E evidence covers two concurrent sessions in -both model-policy and response-required modes, rejection of a third session, -and admission of a replacement after one accepted session closes. It does not -establish behavior beyond that configured limit. - -Session leases, lease TTL, resumable attachment, orphan reaping, and bounded -Serving input/backpressure are implemented. A production admission controller, -KV-aware capacity budgeting, fairness, starvation metrics, and multi-session -worker-failure recovery remain follow-up work. Capability claims must stay -limited to the mechanisms and two-session execution shape validated here. - -## Validation Evidence - -All pytest and runtime evidence for this branch must run on the remote H20. - -### Current synchronized tree - -The current dirty tree was synchronized to the isolated H20 worktree at -`/home/admin/workspace/aop_lab/model_runner_v2/vllm-omni-worktrees/pr3907-boundary-cg-0721`. - -- Final affected matrix: 661 passed, one network-dependent config test - deselected, 20 warnings. It includes the complete MiniCPM native-duplex hook - file and all control/session/handler/Orchestrator/import-boundary suites. The - run is task `be221f67` (`/tmp/remote_gpu_logs/be221f67.log`). The deselected - `test_resolve_when_autodetect_resolves_none` attempted a Hugging Face model - lookup and hung in the isolated environment; the duplex deploy-config test - passed independently. -- Stage1 PIECEWISE CUDA Graph validation: the deploy-config assertion first - failed with no Stage1 compilation config, then passed in task `3ed6091d`. - Server task `95827027` started with `enforce_eager=False`, logged - `cudagraph_mode=PIECEWISE`, and captured batch sizes 1, 2, and 4 without the - CPU-to-CUDA capture error. The first concurrent E2E exposed graph padding as - `request_token_spans` covering three valid rows in a four-row capture. A - focused RED/GREEN now verifies both single-request and batched-request - padding, and the full MiniCPM hook file passed 47 tests in task `e89dba88`. - After the fix, two concurrent response-required sessions each completed two - spoken responses, playback acknowledgements, and history lifecycles with no - error or stale event. Artifacts are under - `/tmp/pr3907_piecewise_response_required_2x2_fixed_20260721`. -- An exploratory two-session, two-turn model-policy run still left one session - waiting after its second input commit while the other session completed both - turns. The same liveness symptom was reproduced with Stage1 eager mode before - the padding fix, and no Stage1 fatal occurred after the fix. This run is not - counted as PIECEWISE acceptance evidence; the remaining Stage0 resumable - reactivation path requires separate diagnosis. Artifacts are under - `/tmp/pr3907_piecewise_model_policy_2x2_fixed_20260721`. -- Latest broad affected matrix before the final pre-response continuation fix: - 608 passed, 19 warnings. This includes the - Stage0 row-routing regression, cleanup-held admission regressions, - segment-local token accounting, response-owned metric filtering, the generic - output-processor suite, direct E2E script execution, and continuous-response - overlap ownership. The run is task `e2146318` - (`/tmp/remote_gpu_logs/e2146318.log`). -- The pre-response continuation regression group passed 5 focused tests. The - response-required two-turn scenario then passed 5/5 repeated runs in task - two-session semantic isolation plus resume, takeover, and - admission passed in task `49225794`; model-unit overlap passed in task - `3888387d`. The final affected matrix above supersedes the earlier pending - rerun requirement. -- Serving plugin selection was first reproduced as three RED tests: the generic - handler silently selected MiniCPM and the pipeline had no explicit Serving - adapter field. The final pipeline/import/fail-fast group passed 9 tests in task - `4db6d959` (`/tmp/remote_gpu_logs/4db6d959.log`). The token-metric and client - audio-cadence producer-to-consumer group passed 8 tests in task `cf80f6b2` - (`/tmp/remote_gpu_logs/cf80f6b2.log`). -- Stable-import isolation was first reproduced as a RED that loaded twelve - `experimental.fullduplex` modules. After extracting the kernel and making - implementation loading conditional, the isolated import regression passed. - The focused ControlPlane/ControlClient/session/runtime/Orchestrator/AsyncOmni - suite then passed 139 tests in task `fad613ea` - (`/tmp/remote_gpu_logs/fad613ea.log`). - After the final contract-import cleanup, the import-boundary, runtime, and - lease suites passed 45 tests in task `bc138887` - (`/tmp/remote_gpu_logs/bc138887.log`). -- Concurrency and failure-path hardening now has direct regressions for - single-flight cancel/close cleanup, cleanup retry without dropping request - resources, failed-replica affinity cleanup and same-address reattachment, - append ordering across resume, takeover ownership of pending-turn - reservations, and a closed-transport send immediately after resumable - credentials are delivered. The final focused resume/takeover group passed 5 - tests in task `dc2ac96e` (`/tmp/remote_gpu_logs/dc2ac96e.log`). -- Stage1 batching regression: the focused test first failed because only the - first request in `runtime_additional_information` reached the talker. After - routing every request by `request_token_spans`, preserving per-row terminal - flags, and emitting request-keyed sparse audio, the complete MiniCPM duplex - hook file passed 42 tests. -- Stage0 batching regression: two requests in one batch now retain different - `duplex_prompt_token_ids` and `special_token_ids`. The model emits row-aligned - payload elements, so Stage0 no longer copies request zero's duplex metadata to - later sessions. -- Admission finalization regressions cover `max_sessions=1` with blocked and - failed expiry cleanup plus request-triggered cleanup. Every replacement is - rejected with `resource_exhausted` until stage resources are cleaned and the - old session is explicitly finalized. -- Resumable token accounting exports each model unit's native generation-token - count instead of recounting the request's cumulative token IDs. Realtime - response state resets and accumulates those segment snapshots, and the client - ignores unowned session-level metrics while summarizing a response. -- Engine-owned admission and two-session model-policy E2E: two random session - IDs were admitted and remained isolated, a third was rejected with - `resource_exhausted`, closing one accepted session returned capacity, and a - random replacement session was admitted. The capability projection reported - `configured_limit=2`, `advertised_multi_session=true`, and - `session_admission_mode=engine_managed`. In the same run one session selected - speak and the other selected listen without identity leakage. Detach/resume - advanced attachment generation from one to two and rotated its token; active - takeover rejected all four writes from the replaced connection. The run is - task `1d9dd1ed` (`/tmp/remote_gpu_logs/1d9dd1ed.log`), with artifacts under - `/tmp/pr3907_resume_takeover_native_metrics`. -- Concurrent two-session, two-turn response-required E2E: both sessions produced - two complete spoken responses, committed two playback acknowledgements and - history records, and emitted no error, cancel, truncate, or stale events. The - semantically distinct inputs produced isolated transcripts containing - "当时" and "一加一" respectively. Per-response Stage0 output counts were 20, - 38, 32, and 33 tokens rather than cumulative request totals. The run is task - `a16f1467` (`/tmp/remote_gpu_logs/a16f1467.log`), with artifacts under - `/tmp/pr3907_semantic_2turn_native_metrics`. -- Disconnect-grace E2E: a detached session expired after the configured - 30-second engine grace, and a later resume was rejected with - `session_resume_expired`. The run is task `9872d50c`, with artifacts under - `/tmp/pr3907_0719_disconnect_grace_expiry`. -- A separate black-box takeover client exercised the current backend with two - semantically distinct concurrent inputs. Active resume replayed 17 events - including five audio deltas and committed playback history; takeover rejected - all four stale-attachment sends, emitted `session.replaced`, and preserved - attachment/session isolation. Both transcripts contained only their expected - semantic token. The client task is `6773bd34` - (`/tmp/remote_gpu_logs/6773bd34.log`), with artifacts under - `/tmp/pr3907_gaohan_final_takeover`. The backend task was `5abcd1ee`. -- Across the final response-required, overlap, and resume/takeover runs, engine - Stage0 TTFT ranged from 31.17 to 310.18 ms, TPOT from 15.18 to 21.18 ms, and - reported ITL p95 from 16.35 to 32.83 ms. Explicit - `commit_to_first_audio_ms` ranged from 386.80 to 1033.25 ms, while the - client-monotonic maximum audio chunk gap ranged from 388.30 to 1316.37 ms. - These are response-scoped observations from a small, warm-state-sensitive - sample, not an SLO claim. The slowest response still does not satisfy a 200 ms - end-to-end first-output target. -- Post-merge focused handler/demo/Web/Stage0/runner/video suite: 324 passed, - 19 warnings. The run is task `f78cb7a8` - (`/tmp/remote_gpu_logs/f78cb7a8.log`). The Web contract explicitly rejects - browser VAD markers and browser-generated input commits while the merged - video path retains camera-frame admission. -- No-response listen and residual accounting: the focused handler regressions - passed after first reproducing both failures. A model-listen unit without a - Realtime response no longer defers the next committed input, and three - consecutive 1.2-second `pcm_f32le` inputs each returned pending input bytes - to zero. -- Model-unit overlap E2E: the second physical input streamed while the first - response remained active. Seven overlap decisions had - `defer_runtime_append=false`; the one-second model-unit boundary occurred - before the first `response.done`, and no barge-in, cancel, truncate, stale, or - error event occurred. The model-policy contract accepts a later unit ending - the continuous response instead of requiring a second physical-input-owned - response. The two response-owned Stage0 snapshots independently reported 20 - tokens / 152.8 ms TTFT and 19 tokens / 43.3 ms TTFT. The run is task - `a4977622` (`/tmp/remote_gpu_logs/a4977622.log`), with artifacts under - `/tmp/pr3907_overlap_native_metrics_scoped`. -- The split multi-session E2E runner is directly executable from the repository - even when another installed package owns the top-level `tests` name. Its - script-import and continuous-terminal regressions passed in task `f11b106c`. -- Pinned response-required E2E: the first 1400 ms of the pinned fixture produced - one complete spoken response, five audio/transcript deltas, and one playback - history commit, with symmetric created/audio-done/done lifecycle and no - duplicate `response.speak`. The client task is `a76dec52` - (`/tmp/remote_gpu_logs/a76dec52.log`). -- Continuous-input model-policy E2E: 36 browser-like 200 ms chunks produced - seven model-listen decisions with zero input commits and no - error/cancel/truncate/stale events. The client task is `fb256595`; artifacts - are in `/tmp/minicpmo_pr3907_continuous_model_policy_20260718`. -- Three-response continuous-input E2E: three consecutive 9469 ms streams, - composed of the original fixture plus four seconds of silence, produced three - spoken responses in one session. All three completed `response.done`, audio - transcript, playback ACK, and history commit; the run had 22 audio/transcript - deltas, zero input commits, and no error/cancel/truncate/stale events. This is - lifecycle evidence, not semantic cross-input independence, because the same - fixture was repeated. The client task is `2a5f29f0`; artifacts are in - `/tmp/minicpmo_pr3907_continuous_three_response_20260718`. -- The fresh validation backend was stopped after the E2E runs and GPU0/1 returned - to baseline. Audio-quality and ASR evaluation were not rerun for these - lifecycle-only changes. Other GPU workloads were not modified. - -### Regression lineage - -Earlier synchronized checkpoints passed broader affected matrices before the -extension, policy, and ledger consolidation. Those runs were useful while -developing the refactor but are superseded by the current file-identical H20 -matrix and E2E evidence above; they are intentionally omitted to avoid mixing -different trees in the acceptance argument. - -### Focused suites - -- RPC router and AsyncOmni shutdown/error routing; -- orchestrator fatal and duplex control paths; -- WebSocket Actor mailbox and writer; -- duplex protocol and Session ownership; -- Realtime handler, terminal ordering, overlap, playback ACK; -- scheduler resumable boundary; -- MiniCPM runner and Stage0/Stage1 lifecycle. - -### E2E contract - -Run two deliberately separate fixture contracts: - -- model-policy uses three distinct inputs, no transcript hints, 200 ms - real-time pacing, and continuous input with no browser commit. It allows - natural listen or speak decisions and assumes no 1:1 mapping between - physical inputs, model turns, and responses; -- response-required uses `minicpmo_pr3907_jiayan_16k.wav` pinned to SHA-256 - `2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941`, - sends its first 1400 ms, and requires one complete spoken response. - It is an output-lifecycle fixture, not evidence that arbitrary no-hint, - real-time input must select speak; -- late and old-response playback ACK handling in the affected test matrix; -- final model-policy response drain before `session.close`; -- close/disconnect with no post-terminal audio; -- transcript delta/final consistency; -- generated WAV validity and external ASR sanity. - -Distinct-input fixtures must also be semantically distinct. Different files or -waveform hashes containing the same greeting can legitimately produce the same -reply and are not sufficient evidence for cross-turn transcript independence. - -For each run record the exact SHA/worktree, command, server log, client log, -WAV path, and cleanup result. - -## Acceptance Criteria - -This checkpoint is ready to publish only when: - -- one mailbox passes exactly-once and wire-order tests; -- cancel advances the engine fence even when the old append continues in an - executor, and the next-epoch append remains accepted; -- request completion removes MiniCPM helper and model-owned force-listen state; -- RPC fatal/close wakes every waiter and no new work enters after terminal; -- Session is the only serving response/overlap authority; -- stale terminal events cannot mutate a newer epoch; -- an accepted terminal cannot be dropped by the writer; -- no duplicate `response.speak` or transcript delta appears; -- Stage0 KV continuity and Stage1 turn reset pass; -- the real H20 multi-turn audio E2E passes; -- ordinary non-duplex paths in the affected matrix still pass; -- no local path, proxy token, temporary profiling probe, or test-only switch is present. - -Passing this checkpoint supports the statement: - -> Single-session, model-owned MiniCPM-o 4.5 native duplex is reviewable on the -> validated H20 configuration. - -It does not support claims for an acoustic-onset barge-in latency target across -arbitrary client chunk sizes, multi-session production concurrency, bounded -long-session KV, or scheduler-native append. Camera-frame input -(`video_frames`) landed after this snapshot and is covered by its own E2E -test rather than by the evidence above. +| engine policy | `configure_sampling_params(runtime_config, defaults)`, `plan_append(...) -> DuplexAppendPlan` (the resumable Stage0 prompt for one unit), `decide_output(...) -> DuplexOutputDecision \| None` (e.g. the listen decision on a finished Stage0 segment) | +| session policy | `capabilities(max_sessions)`, `validate_client_extra_body`, `prepare_runtime_config(config, model_config)` (server-owned runtime keys, reference audio resolution), `runtime_config_for_update`, `runtime_config_for_function_output`, `create_session_state() -> DuplexModelSessionState`, `data_plane: DuplexDataPlane` (projects raw stage outputs into internal events), `data_plane_context(...)` | + +`DuplexOmniEngine._validate_deployment` loads the plugin before any stage +starts; `DuplexSessionManager.__init__` validates it against the stage +sampling defaults once the stage pools exist. Plugin hooks that may block +(`prepare_runtime_config` fetching `ref_audio`) are awaited in `open()` and +offloaded from the loop. + +The MiniCPM-o 4.5 plugin (`model_executor/models/minicpmo_4_5/duplex/plugin.py`) +is the only integration in this framework version. PersonaPlex and Nemotron +VoiceChat still carry their pre-framework duplex code (runtime extension plus +serving adapter) and are therefore not served over this framework yet: their +pipelines declare no `duplex_plugin`, so they run turn-based until the +follow-up PRs port them (RFC vllm-omni#7181, PR 2/3). + +## Serving + +`OmniDuplexSessionHandler.handle_realtime_session(websocket)`: + +1. accept; build the `RealtimeEnvelope` (autostart / default session from + query params, `event_id` correlation, error rendering through the typed + `ErrorEvent`); +2. handshake: first message `session.update` -> `omni.open_session(session)` + (the id inside the payload is ignored; the allocated id is announced in + `session.created` together with `resume_token` / `attachment_generation` + when the model supports resume); `session.resume(session_id, resume_token, + last_received_server_event_seq)` -> `attachment.authenticate_resume` -> + `omni.resume_session`; +3. reader loop: JSON -> `RealtimeEnvelope.translate` (`translate_realtime_command` + with the session's declared audio defaults) -> `handle.submit`; + envelope-level errors (invalid JSON, oversize frame, unknown type, + event acks) are answered locally; +4. writer pump (session-scoped, survives reconnects): `async for ev in + handle.events(): attachment.send_event(ev.to_realtime())`, journaling + for replay until the journal overflows (`session.resync_required`); +5. disconnect: a resumable session is detached (`attachment.detach` + + `omni.detach_session`, engine-owned grace); a superseded socket's + disconnect is a no-op; a non-resumable session is closed. Takeover sends + `session.replaced` to the old socket and closes it. + +## Validation + +- CPU suites: `tests/engine/duplex/**` (session, lease, manager, runner, + plugin, events, commands), `tests/engine/test_duplex_orchestrator.py`, + `tests/engine/test_duplex_omni_engine.py`, + `tests/entrypoints/test_duplex_omni.py`, + `tests/entrypoints/duplex/test_duplex_serving.py`, + `tests/entrypoints/openai/test_duplex_session_attachment.py`, + `tests/entrypoints/openai_api/test_duplex_api_server.py` (the duplex + server: pipeline probe, app state, routes, warmup gate), + `tests/entrypoints/duplex/test_chat_completions_adapter.py`, + `tests/clients/**`, `tests/engine/test_duplex_import_boundary.py`, + `tests/model_executor/models/minicpmo_4_5/duplex/**`, + `tests/worker/test_native_duplex_hooks.py`. +- GPU: `examples/online_serving/minicpmo/realtime_duplex_demo.py` (one turn + over the websocket), `examples/online_serving/barge_in_client.py` (cancel + during playback; `--inline` drives the same scenario through + `InlineDuplexClient`), `tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py` + (admission, resume, takeover, expiry). diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index d4dbf51a2..df22f3ce1 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -16,6 +16,7 @@ required_reviewers: - "@yenuo26" - "@NickCao" primary_code_paths: + - vllm_omni/engine/omni_engine_base.py - vllm_omni/engine/async_omni_engine.py - vllm_omni/engine/async_engine_utils.py - vllm_omni/engine/orchestrator.py @@ -28,8 +29,12 @@ related_code_paths: - vllm_omni/engine/stage_pool.py - vllm_omni/outputs/output_processor.py - vllm_omni/distributed/omni_coordinator/** - # Duplex kernel design is owned by ../fullduplex.md; listed here because the - # orchestrator hosts the duplex control plane. + # The duplex siblings (DuplexOmniEngine / DuplexOrchestrator) and the + # engine-resident session framework are owned by ../fullduplex.md; listed + # here because they fill the template seams of OmniEngineBase and + # OrchestratorBase. + - vllm_omni/engine/duplex_omni_engine.py + - vllm_omni/engine/duplex_orchestrator.py - vllm_omni/engine/duplex/** depends_on: - input_output_modality_contracts.md @@ -48,6 +53,8 @@ validation_paths: - tests/engine/test_cfg_companion_lifecycle.py - tests/engine/test_rpc_result_router.py - tests/engine/duplex/** + - tests/engine/test_duplex_orchestrator.py + - tests/engine/test_duplex_omni_engine.py - tests/engine/test_duplex_import_boundary.py upstream_refs: - vllm.v1.engine.EngineCoreRequest diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md index a0dee90b9..db62f9cb6 100644 --- a/docs/design/module/entrypoints.md +++ b/docs/design/module/entrypoints.md @@ -19,7 +19,9 @@ required_reviewers: primary_code_paths: - vllm_omni/entrypoints/omni.py - vllm_omni/entrypoints/async_omni.py + - vllm_omni/entrypoints/async_omni_base.py - vllm_omni/entrypoints/omni_base.py + - vllm_omni/entrypoints/duplex_omni.py - vllm_omni/entrypoints/cli/** - vllm_omni/entrypoints/openai/** - vllm_omni/entrypoints/openpi/** @@ -27,14 +29,13 @@ primary_code_paths: - vllm_omni/entrypoints/stage_utils.py - vllm_omni/entrypoints/utils.py - vllm_omni/entrypoints/duplex/** - - vllm_omni/entrypoints/duplex_request_client.py - vllm_omni/clients/** primary_path_exceptions: - path: vllm_omni/entrypoints/openai/errors.py owner: error_contracts.md - path: vllm_omni/entrypoints/duplex/** owner: ../fullduplex.md - - path: vllm_omni/entrypoints/duplex_request_client.py + - path: vllm_omni/entrypoints/duplex_omni.py owner: ../fullduplex.md - path: vllm_omni/clients/** owner: ../fullduplex.md @@ -42,7 +43,9 @@ related_code_paths: - vllm_omni/errors.py - vllm_omni/inputs/** - vllm_omni/outputs/** + - vllm_omni/engine/omni_engine_base.py - vllm_omni/engine/async_omni_engine.py + - vllm_omni/engine/duplex_omni_engine.py - vllm_omni/config/** - vllm_omni/deploy/** depends_on: @@ -53,7 +56,9 @@ validation_paths: - tests/entrypoints/test_omni_entrypoints.py - tests/entrypoints/test_async_omni.py - tests/entrypoints/test_async_omni_pause_sleep_routing.py - - tests/entrypoints/test_async_omni_duplex.py + - tests/entrypoints/test_duplex_omni.py + - tests/entrypoints/duplex/** + - tests/entrypoints/openai_api/test_duplex_api_server.py - tests/entrypoints/test_serve.py - tests/entrypoints/test_stream_finish_reason.py - tests/entrypoints/openai/** diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index 3130e688f..d236c1129 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/docs/mkdocs/hooks/generate_api_readme.py @@ -25,8 +25,8 @@ CATEGORIES = { "clients": { "name": "Clients", "description": ( - "Client libraries that talk to a running vLLM-Omni server " - "(the full-duplex Realtime client and its per-model session presets)." + "Client libraries for full-duplex sessions: the Realtime WebSocket client, " + "the in-process client over `DuplexOmni`, and the per-model session presets." ), }, "inputs": { diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f8db9eb92..42971ed97 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -107,7 +107,7 @@ th { | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — | | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — | | `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — | -| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — | +| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) | | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — | | `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — | diff --git a/docs/serving/README.md b/docs/serving/README.md index 8167f4c9b..472ab2a49 100644 --- a/docs/serving/README.md +++ b/docs/serving/README.md @@ -145,15 +145,22 @@ interchangeably. | Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) | | Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) | | Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) | -| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) | +| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` (alias `WS /v1/duplex`) | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) | | Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) | | Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) | All six routes are model- or configuration-dependent. In particular, -`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the -deployment explicitly enables duplex sessions. Clients should also verify the -duplex capability payload because the query-parameter form falls back to the -ordinary realtime handler when duplex is unavailable. +`/v1/realtime` is full duplex when the model is a duplex model (its pipeline +declares a `duplex_plugin` and the deploy configuration sets +`session_mode: duplex`); a stock Realtime client needs no vendor query +parameter, and `duplex=0` is the explicit opt-out. Such a server serves the +websocket route plus `POST /v1/chat/completions`, and no other turn-based +HTTP route. On a duplex server that chat route is not the turn-based path: +each request runs on a short-lived duplex session, so it holds one of +`duplex_session.max_sessions` for its lifetime and answers at the model's +real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should +also verify the duplex capability payload because the query-parameter form +falls back to the ordinary realtime handler when duplex is unavailable. ## Related Endpoints diff --git a/docs/serving/full_duplex_api.md b/docs/serving/full_duplex_api.md index 4e8bf9748..92511e84f 100644 --- a/docs/serving/full_duplex_api.md +++ b/docs/serving/full_duplex_api.md @@ -15,40 +15,52 @@ capability gates, see the [Realtime Duplex API guide](realtime_duplex_api.md). | Endpoint | Protocol | Recommended use | | --- | --- | --- | -| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients | -| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing | - -Both endpoints use the same duplex engine and require the same model-side -adapter. Prefer the Realtime projection unless the native lifecycle events are -specifically required. +| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style events (the normative contract) | Applications and browser clients | +| `WS /v1/duplex` | Alias of `/v1/realtime?duplex=1` (same protocol, same handler) | Clients that prefer a dedicated path | +| Python: `DuplexOmni` / `InlineDuplexClient` | Typed commands and events in-process | Embedding the model without a server | ## Enable Full Duplex -The route becomes usable only when the deployment configuration explicitly -sets: +A model is served full duplex when its registered pipeline declares a +`duplex_plugin` (the model's `DuplexModelPlugin`). That declaration alone +decides it: `vllm-omni serve` constructs `DuplexOmni` instead of `AsyncOmni`, +and every surface the server exposes is backed by a duplex session. It serves +`/v1/realtime?duplex=1` (and its alias `/v1/duplex`), +`POST /v1/chat/completions`, `/v1/models` and `/health`; every other +turn-based HTTP route (speech, batch, embeddings, video, ...) answers "not +available". Turn-based use of such a model stays available offline through the +Python API (`Omni` / `AsyncOmni`). + +`/v1/chat/completions` is the ordinary chat service running on the duplex +engine: a request is a turn-based generation on the same stages, served +alongside the live websocket sessions, with the request options the chat +service supports. It is wired only when the model's plugin declares +`DuplexCapabilities.supports_chat_completions`; on any other duplex model the +route reports "not available". A model that should not serve the route at all +lists it in the deploy configuration's `endpoint_restrictions`. + +The deploy configuration of such a model must agree: ```yaml session_mode: duplex ``` -The selected model pipeline must also provide a duplex serving adapter. -Model-native deployments configure the corresponding engine runtime extension -and control plane as part of their registered pipeline. +A duplex model started with a deploy configuration that does not set it fails +at startup rather than falling back to turn-based serving. !!! warning - `WS /v1/duplex` fails with `Duplex API is not available` when duplex is not - enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary - turn-based realtime handler when no duplex handler exists. Confirm that + On a deployment that is *not* duplex, `WS /v1/duplex` fails with + `Duplex API is not available` and `/v1/realtime?duplex=1` falls back to the + ordinary turn-based realtime handler. Confirm that `session.created.session.capabilities` is present before treating the connection as full duplex. -The current unified-runtime integrations are: - -- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5.yaml` - (`session_mode: duplex`); -- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex; -- Nemotron VoiceChat, via its registered duplex plugin package. +**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`) is the only model +served over this endpoint today. PersonaPlex and Nemotron VoiceChat still carry +their pre-framework duplex code: their pipelines declare no `duplex_plugin`, so +they run turn-based until the follow-up PRs port them to the plugin contract +(RFC [vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)). JoyVL is a separate HTTP interaction orchestrator and does not use these WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md). @@ -101,22 +113,23 @@ deferred, treated as a short acknowledgement, or used to interrupt output. The `session.created` payload includes capability fields such as `supports_barge_in`, `supports_playback_ack`, `supports_multi_session`, `supports_session_resume`, and `chunk_period_ms`. Treat this payload as the -runtime contract. +runtime contract and branch on the flags, never on the model name: a model +that supports native overlapping speech may still advertise +`supports_barge_in=false` when destructive output interruption and model-state +rewind are not validated for it. Capacity and session-resume behavior also +depend on the selected deployment configuration. The per-model table lives in +the [Realtime Duplex API guide](realtime_duplex_api.md#capability-negotiation-by-model). -For example, PersonaPlex supports native overlapping speech but currently -advertises `supports_barge_in=false`; destructive output interruption and -model-state rewind have not been validated for that integration. Capacity and -session-resume behavior also depend on the selected deployment configuration. +## Python API -## Native Protocol +`vllm_omni.entrypoints.duplex_omni.DuplexOmni` runs the same engine-resident +sessions in-process: `open_session()` returns a `DuplexSessionHandle` whose +`submit()` takes typed `DuplexCommand` objects and whose `events()` yields +typed `DuplexEvent` objects (each with a `to_realtime()` wire rendering). +`vllm_omni.clients.inline_duplex.InlineDuplexClient` exposes that handle +behind the `DuplexClient` API. See the +[Realtime Duplex API guide](realtime_duplex_api.md#using-the-python-api). -`WS /v1/duplex` exposes lower-level lifecycle names including -`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`, -and `session.close`. It returns native session, input, response, overlap, and -error events. This protocol is experimental and may evolve with the runtime; -applications should use the provided Realtime client where possible. - -See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo), -[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex), +See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo) and the [full-duplex runtime design](../design/fullduplex.md) for model-specific validation and architecture details. diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index 288c934e8..cc10d99c9 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -1,17 +1,18 @@ # Realtime Duplex API vLLM-Omni serves full-duplex speech models over a WebSocket endpoint, -`/v1/realtime`, that speaks the OpenAI Realtime vocabulary plus a +`/v1/realtime?duplex=1`, that speaks the OpenAI Realtime vocabulary plus a set of duplex extensions: the model decides when to listen and when to speak, the user can talk over the assistant, playback progress is reported back so history stays honest, and a dropped connection can resume the same session. This page covers how to run a duplex deployment, how to drive it from Python with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract. -The handler is enabled by either a model-native `session_mode: duplex` -deployment or an explicit `duplex_session` configuration for turn-based -Server VAD. The runtime architecture is described in -[Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). +The endpoint is served for models that ship a duplex plugin — currently +MiniCPM-o 4.5 only — and just for deploy configurations that declare +`session_mode: duplex`. PersonaPlex and Nemotron VoiceChat are ported to the +plugin contract in follow-up PRs and are not served over this endpoint yet. The runtime architecture +is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). ## Quick Start @@ -26,11 +27,38 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \ ``` `vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and -`duplex_session.max_sessions: 4`, which mounts the WebSocket routes -`ws://:8099/v1/realtime` (this page) and `ws://:8099/v1/duplex` -(the native dialect for raw-protocol clients) next to the usual HTTP API. -PersonaPlex and Nemotron VoiceChat use `vllm_omni/deploy/personaplex.yaml` and -`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml` in the same way. +`duplex_session.max_sessions: 4`. Because the MiniCPM-o 4.5 pipeline declares a +`duplex_plugin`, `vllm-omni serve` runs it through `DuplexOmni`: the server +mounts `ws://:8099/v1/realtime?duplex=1` (this page; `ws://:8099/v1/duplex` +is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and +`/health`, and nothing else. The remaining turn-based HTTP routes (speech, +batch, embeddings, video, ...) are not served by a duplex model. + +`/v1/chat/completions` is not a turn-based path on a duplex server: each +request opens a short-lived duplex session and reads the answer off it. Audio +content becomes an ordinary committed turn; text is seeded as the session's +opening turn, because a model-native model takes a turn only when it hears +speech. A model that cannot be seeded that way declares so, and a text request +to it is refused with 400. Two consequences are worth sizing for: + +- **Every request holds an admission slot** for its lifetime, so + `duplex_session.max_sessions` caps HTTP concurrency as well as websocket + sessions. Beyond it a request is refused with HTTP 503. +- **Latency follows the model's own clock.** A model-native session generates + per audio unit, so even a short answer costs what that many units cost -- + noticeably more than the same model would take turn-based. + +`n > 1`, `logprobs`, `tools`, and image or video content parts have no +representation in one duplex turn and are refused with HTTP 400 -- a Realtime +conversation item carries text and audio only. `temperature`, +`max_completion_tokens`, `modalities` and +`chat_template_kwargs.use_tts_template` apply; the other +`chat_template_kwargs` keys are logged as ignored. `extra_body` is passed to +the session unchanged, exactly as a websocket client's session `extra_body` +is, which is how a request asking for `modalities: ["text", "audio"]` supplies +the `ref_audio` MiniCPM-o requires for audio output. A model that should not +serve the route lists it in the deploy config's `endpoint_restrictions`, +exactly as a turn-based model does. ### Run the Example Client @@ -81,25 +109,24 @@ ships a preset: | Preset | Input / output audio | What it sets | | --- | --- | --- | -| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `extra_body.native_duplex=True` (opt in to the model-native lane), `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip | +| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip | | `vllm_omni.clients.personaplex.create_duplex_session_config(voice="NATF2.pt", persona="")` | 24 kHz `pcm_f32le` / 24 kHz `pcm16` | bundled `.pt` voice prompt and the persona as `instructions` | | `SessionConfig(...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | model-neutral defaults; pass `extra_body`, `turn_detection`, `overlap_policy`, `playback_commit_policy`, `instructions`, `voice`, `temperature` yourself | Keyword overrides on a preset replace the corresponding `SessionConfig` field, e.g. `create_duplex_session_config(ref_audio=..., temperature=0.6)`. -Constructor options: `session_id` (names the session the client creates; it -does not resume or take over an existing session — entering the client always -performs the `session.update` handshake), `reconnect` (a `ReconnectPolicy`, -default five attempts with 0.25–4 s jittered backoff; `None` disables -auto-resume), `heartbeat_interval_s` (default 30 s; `None` disables the lease -heartbeat), `handshake_timeout_s`, and `connect` (a custom transport factory, -used by tests and benchmarks). Resume is supported only as automatic reconnect -within the same `DuplexClient` instance (see -[Reconnect and resume](#reconnect-and-resume)); re-attaching from a new client -or process requires the wire-level `session.resume` handshake -(`resume_token`, `incarnation`, `last_received_server_event_seq`), which this -client does not expose. +Constructor options: `reconnect` (a `ReconnectPolicy`, default five +attempts with 0.25–4 s jittered backoff; `None` disables auto-resume), +`heartbeat_interval_s` (default 30 s; `None` disables the lease heartbeat), +`handshake_timeout_s`, and `connect` (a custom transport factory, used by +tests and benchmarks). The session id is allocated by the server and becomes +`client.session_id` once `session.created` arrives; a client cannot choose +it. Resume is supported only as automatic reconnect within the same +`DuplexClient` instance (see [Reconnect and resume](#reconnect-and-resume)); +re-attaching from a new client or process requires the wire-level +`session.resume` handshake (`resume_token`, +`last_received_server_event_seq`), which this client does not expose. ### Stream audio and commit a turn @@ -195,7 +222,7 @@ the stamped `event_id`. On a transport drop the client reconnects with the `ReconnectPolicy` backoff and sends `session.resume` from the last acknowledged `server_event_seq`. Replayed events are deduplicated, consumers see one `ConnectionResumed` -event, and `client.incarnation` / `client.resume_token` are refreshed. The +event, and `client.resume_token` is refreshed. The server keeps a detached session alive for `disconnect_grace_s` (30 s by default); after that a resume fails with `session_resume_expired`. Resume is gated by `capabilities.supports_session_resume` (MiniCPM-o 4.5 and Nemotron @@ -238,9 +265,67 @@ definitions out of the dependency-free client. `audio_bytes(response_id)`, `coun `openai-realtime-duplex` benchmark backend (`vllm-omni bench serve`) is built on the same helpers. +## Using the Python API + +`vllm_omni.entrypoints.duplex_omni.DuplexOmni` is the in-process Python API +of a duplex model: sessions run inside the engine, and the class only opens, +resumes and closes them and pipes typed commands in and typed events out +through a `DuplexSessionHandle`. + +```python +from vllm_omni.engine.duplex.commands import AppendAudio, Commit +from vllm_omni.engine.duplex.events import AudioDelta, SessionClosed +from vllm_omni.entrypoints.duplex_omni import DuplexOmni + +omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True) +async with await omni.open_session({"modalities": ["audio", "text"], "ref_audio": ref_audio_data_url}) as session: + async def consume(): + async for event in session.events(): # typed DuplexEvent objects + if isinstance(event, AudioDelta): + play(event.audio) + elif isinstance(event, ResponseDone): + break # one turn is enough here + task = asyncio.create_task(consume()) + await session.submit(AppendAudio(audio=pcm16_bytes, format="pcm16", sample_rate_hz=16000)) + await session.submit(Commit()) + await task +# leaving the block closes the session (the last event is SessionClosed) +``` + +`open_session` takes a `DuplexSessionConfig` or the same Realtime `session` +object the WebSocket handshake takes; the session id is always allocated by +`DuplexOmni`. `DuplexSessionHandle` offers `submit(command)` plus one +convenience wrapper per command (`append_audio`, `commit`, `create_response`, +`cancel_response`, `barge_in`, `ack_playback`, `update`, `create_item`, +`delete_item`, `truncate_item`, `heartbeat`, ...). `events()` is a +single-consumer async iterator that ends after `session.closed` / +`session.expired`; rejected commands come back as `ErrorEvent`s on it. +Every event renders the wire JSON with `to_realtime()`, so anything written +against the WebSocket protocol works unchanged on the typed stream. + +`vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni` +behind the `DuplexClient` API, so the same application code runs in-process +(no server) or over the network: + +```python +from vllm_omni.clients.inline_duplex import InlineDuplexClient +from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config + +config = create_duplex_session_config(ref_audio=audio_data_url(Path("reference_voice.wav"))) +async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=config) as client: + await client.stream_pcm(question) + await client.commit() + async for response in client.responses(): + ... +``` + +`examples/online_serving/barge_in_client.py --inline` drives the barge-in +scenario this way. The inline client has no transport, so it has no +reconnect, resume, heartbeat or event-acknowledgement machinery. + ## API Reference -This section is the complete wire contract of `/v1/realtime`: how +This section is the complete wire contract of `/v1/realtime?duplex=1`: how the dialect relates to the OpenAI Realtime protocol, which surfaces each model's capability flags gate, a catalogue of every client-to-server and server-to-client event, and a JSON example for each. `DuplexClient` speaks @@ -253,21 +338,26 @@ its protocol tests. ### Endpoint and transport -The endpoint is `ws(s):///v1/realtime`. Bare connections select the -session handler when the deployment declares either `session_mode: duplex` -or `duplex_session`. `?duplex=1` remains a compatibility alias, while -`?duplex=0` selects the legacy Realtime handler. Optional query parameters -are `model`, `session_id`, -`autostart` (`0` means resume-only), `resume`, and `native_duplex` -(`minicpmo45_native_duplex` is accepted as a deprecated alias and folded -into the canonical name). Every message is one JSON object per WebSocket +The endpoint is `ws(s):///v1/realtime`. On a deployment whose pipeline +declares a `duplex_plugin` and whose deploy configuration declares +`session_mode: duplex`, a bare connection is already a duplex session, so a +stock Realtime client needs no vendor query parameter. `?duplex=1` (`1`, +`true` or `on`) is the spelling this page uses and stays valid; +`?duplex=0` (`0`, `false`, `off`) explicitly opts out and selects the +turn-based Realtime handler, which a duplex server does not mount -- +that connection is answered with `Realtime API is not available`. +`ws(s):///v1/duplex` is an alias that always selects the duplex +handler. Optional query parameters are `model`, `autostart` (`0` means +resume-only) and `resume`. Every message is one JSON object per WebSocket text frame, discriminated by `type`. Inbound events are applied in arrival order through one per-session mailbox; outbound events preserve that order and carry a monotonically increasing `server_event_seq`, which is the replay cursor for `session.resume`. Client events may carry an `event_id`, echoed in `error.error.event_id` when the event is rejected. -Session identity is `(session_id, incarnation)`; `attachment_generation` +Session identity is the server-allocated `session_id` (`duplex-`, +announced in `session.created`; a `session_id` the client puts in the session +object is ignored, and ids are never reused); `attachment_generation` names the current socket; `epoch` is the interruption fence that advances on every cancel, clear, or barge-in; `turn_id` counts committed model turns; `response_id` (item id `item_`) names one assistant utterance. @@ -281,12 +371,13 @@ with no OpenAI counterpart. The event vocabulary is uniform, but several surfaces are gated by the `capabilities` object the server returns in `session.created`; a client must -branch on those flags rather than on the model name. The current model -plugins advertise: +branch on those flags rather than on the model name. MiniCPM-o 4.5 is the +only model on the plugin contract today; the other two columns record what +their integrations advertise once the follow-up PRs port them: | Capability | MiniCPM-o 4.5 | PersonaPlex | Nemotron VoiceChat | Gated surface | | --- | --- | --- | --- | --- | -| `implementation_level` | `model_native_duplex` (when `extra_body.native_duplex=true`) | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak`; the chat-fallback lane otherwise | +| `implementation_level` | `model_native_duplex` | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak` (constant: every duplex model is model-native) | | `chunk_period_ms` | 1000 | 80 | 80 | the model unit that `response.listen` decisions and camera frames align to | | `supports_session_resume` | yes | no | yes | `session.resume`, `session.resumed`, `session.replaced`, replay after a transport drop | | `supports_barge_in` | yes | no | no | `barge_in`, `turn.signal{event:"barge_in"}`, `overlap_policy=barge_in_on_speech`, `turn_detection.server_vad` | @@ -316,14 +407,8 @@ work unmodified. The overlap falls into the three tiers used by the event catalo | Object shapes | `realtime.session`, `realtime.item`, `realtime.response`; item content parts `input_text` / `input_audio` / `output_text` / `output_audio`; `function_call` and `function_call_output` items; `previous_item_id` chaining; `output_index` / `content_index` addressing | | Sequencing | `response.created → output_item.added → content_part.added → deltas → *.done → content_part.done → output_item.done → response.done → rate_limits.updated`; `speech_started → speech_stopped → committed` | -Audio/transcript output uses only the current (non-beta) OpenAI event names -(`response.output_audio.delta`, not the deprecated `response.audio.delta`), -matching `openai.types.realtime` in the official `openai` Python SDK and what -clients built against it (e.g. `livekit-plugins-openai`'s non-Azure code path) -expect; a client still on the old beta SDK types would not recognize these -events. Session fields and `conversation.item.*` events, in contrast, are -deliberately dual: the server emits and accepts **both** the beta and the GA -spellings at once (`output_modalities` alongside `modalities`, +The server deliberately emits and accepts **both** the beta and the GA +spellings at once (`response.output_audio.delta` alongside `output_modalities`, `conversation.item.added` alongside `conversation.item.created`) so either generation of OpenAI client parses the stream. @@ -333,7 +418,7 @@ A stock client ignores the extra keys; the extensions are additive. | Message | Extension | | --- | --- | -| `session.created` / `session.updated` | top-level `incarnation`, `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` | +| `session.created` / `session.updated` | top-level `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` | | `response.created` / `response.done` / `response.listen` | `response_id` at top level; the raw duplex event under `response.metadata` (`duplex_event`); `status_details.reason` uses vLLM reasons (`barge_in`, `client_cancelled`, `new_response`, …) | | `response.output_audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` | | `response.speak` (inserted before the first delta) | not an OpenAI event, but rides the OpenAI response envelope (`response_id`, `item_id`, `output_index`, `content_index`) | @@ -349,14 +434,9 @@ Semantic divergences hidden behind shared names: before any commit; a commit may end in `response.listen` and no response, and the model may open a response with no commit at all. OpenAI: commit ⇒ item, `response.create` ⇒ exactly one response. -- With `turn_detection={"type":"server_vad"}`, omitted `interrupt_response` - defaults to `false` on turn-based models such as Qwen3-Omni, enabling - endpointing and automatic responses without interruption. Explicit `true` - is rejected there. Model-native duplex keeps the `true` default and rejects - `false`; its VAD implies `overlap_policy=barge_in_on_speech`. - This capability-dependent default differs from OpenAI's `true` default; - the effective configuration is returned in `session.created` / `session.updated`. - `semantic_vad` is unsupported. +- `turn_detection.interrupt_response=false` is rejected; + `create_response` is ignored; `semantic_vad` is unsupported; the VAD + runs per session (Silero) and implies `overlap_policy=barge_in_on_speech`. - `rate_limits.updated` is always an empty list (compatibility only). - Cancellation never reuses a `response_id`; truncation is driven by `playback.ack` as well as `conversation.item.truncate`. @@ -365,11 +445,11 @@ Semantic divergences hidden behind shared names: | Area | Messages / fields | | --- | --- | -| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `incarnation`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume`, `native_duplex` | +| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume` | | Turn-taking | `response.speak`, `response.listen`, `overlap.decision`, `overlap_policy`, `barge_in`, `turn.signal`, `input.cancel`, `input.text.append`, `epoch` / `turn_id` fencing, `force_listen` | | Playback truth | `playback.ack` / `playback.acknowledged`, `playback_commit_policy` (`ack_only` \| `commit_all_on_done`), `playback{generated_ms, sent_ms, played_ms, committed_ms}`, `history_committed`, `audio_text_marks` | -| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `native_duplex` — deprecated alias `minicpmo45_native_duplex` is folded into it, `force_listen_count`, `duplex_initial_user_text`) | -| Diagnostics | `runtime.control`, `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/entrypoints/duplex/realtime_state.py`) | +| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `force_listen_count`, `duplex_initial_user_text`) | +| Diagnostics | `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/engine/duplex/events.py`) | #### Consequences for clients @@ -433,11 +513,15 @@ remains an approximation, not word-level audio alignment. ### Event catalogue Every event type the Realtime route accepts or emits, with a worked example in the next section: 21 -client→server and 42 server→client. Aliases (`push_chunk`, -`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`, -`audio.playback_ack`, `close_session`, `close`, `session_close`) share the -payload of their canonical event and are not listed separately; the alias -table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). +client→server and 42 server→client. Aliases (`input_text.append`, +`push_text`, `signal_turn`, `audio.playback_ack`, `close_session`, `close`) +share the payload of their canonical event and are not listed +separately; the alias table is `translate_realtime_command` in +`vllm_omni/engine/duplex/realtime_commands.py`. The pre-Realtime WAV-append +aliases `push_chunk` and `input.audio.append` are **not** accepted on this +route: send `input_audio_buffer.append` with one of the supported input +formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`, +`g711_alaw`). #### Client to server @@ -447,18 +531,18 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | `session.heartbeat` | 3 | Keepalive; refreshes the engine session lease. | | `session.event_ack` | 3 | Acknowledges received server events by `server_event_seq`; trims the resume replay journal. | | `session.close` | 3 | Graceful close; answered with `session.closed`. | -| `session.resume` | 3 | Re-attach a live session on a new socket (`resume_token`, `incarnation`, `last_received_server_event_seq`). | +| `session.resume` | 3 | Re-attach a live session on a new socket (`session_id`, `resume_token`, `last_received_server_event_seq`). | | `input_audio_buffer.append` | 2 | Append one audio chunk. OpenAI name; extended with `is_speech`, `video_frames`, `duration_ms`, `audio_end_ms`, per-event `format`/`sample_rate_hz`. | | `input_audio_buffer.commit` | 2 | Seal the buffered utterance into a user item. OpenAI name; extended with `final`, `response_create`, `is_speech:false`. | | `input.commit` | 3 | Alias of `input_audio_buffer.commit` accepted by the runner. | | `input_audio_buffer.clear` | 1 | Drop un-committed input audio. | | `input.cancel` | 3 | Cancel pending input; advances `epoch`. | -| `input.text.append` | 3 | Append text to the open input item (chat-fallback lane only). | +| `input.text.append` | 3 | Append text input; model-native sessions reject it with `native_text_append_unsupported`. | | `conversation.item.create` | 1 | Add a completed user item (text, audio, or `function_call_output`) to history. | | `conversation.item.retrieve` | 1 | Fetch a stored item by id. | | `conversation.item.delete` | 1 | Delete a stored item by id. | | `conversation.item.truncate` | 1 | Truncate an assistant item's audio/transcript at `audio_end_ms`. | -| `response.create` | 1 | Explicitly request a response with per-response overrides. | +| `response.create` | 1 | Explicitly request a response with per-response overrides. Requires committed audio input; otherwise `response_create_without_input`. | | `response.cancel` | 1 | Cancel the active (or named) response; advances `epoch`. | | `output_audio_buffer.clear` | 1 | Discard queued output audio; advances `epoch` (OpenAI: WebRTC-only, here also WebSocket). | | `barge_in` | 3 | Explicit hard interrupt; requires `capabilities.supports_barge_in`. | @@ -469,10 +553,10 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | Event type | Tier | Description | | --- | --- | --- | -| `session.created` | 2 | Session opened. OpenAI name; adds `incarnation`, `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`epoch`, `turn_id`, `playback`, `capabilities`, ...). | +| `session.created` | 2 | Session opened. OpenAI name; adds `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`id` is the server-allocated session id, `epoch`, `turn_id`, `playback`, `capabilities`, ...). | | `session.updated` | 2 | Echo of the effective session config after every `session.update` (same extended `session` object). | | `session.heartbeat_ack` | 3 | Reply to `session.heartbeat`. | -| `session.closed` | 3 | Last event on the socket; `reason` ∈ `session_close`, `timeout`, `disconnect`, ... | +| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... | | `session.resumed` | 3 | Resume accepted; carries the new `attachment_generation` and rotated `resume_token`; journaled events are replayed after it. | | `session.replaced` | 3 | Sent to the superseded socket when another socket resumes the session. | | `session.expired` | 3 | Engine lease reaped (`disconnect_grace_expired`, `idle_ttl_expired`); socket closes after it. | @@ -508,8 +592,8 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | `duplex.function_call.done` | 3 | Raw fallback when a model tool call is malformed. | | `output_audio_buffer.cleared` | 1 | Reply to `output_audio_buffer.clear`. | | `playback.acknowledged` | 3 | Reply to `playback.ack`; returns the playback ledger and `history_committed`. | -| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`). | -| `runtime.control` | 3 | Diagnostic: an engine control signal was partially unsupported. | +| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`); details under `metadata`. | +| `turn.event` | 3 | A turn-taking signal the session accepted (`event`, resulting `turn_state`). | | `error` | 2 | Error envelope. OpenAI shape and classes; vLLM-Omni error codes. | ### Message examples @@ -528,7 +612,6 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). "event_id": "evt_0001", "session": { "model": "openbmb/MiniCPM-o-4_5", - "session_id": "sess_demo_01", "modalities": ["audio", "text"], "instructions": "You are a concise voice assistant.", "input_audio_format": "pcm16", @@ -541,7 +624,6 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). "idle_timeout_s": 300, "extra_body": { "auto_response": true, - "native_duplex": true, "force_listen_count": 0 } } @@ -553,13 +635,12 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ```json { "type": "session.created", - "incarnation": 0, "attachment_generation": 1, "resume_token": "rt_3f9c…", "session": { "object": "realtime.session", "type": "realtime", - "id": "sess_demo_01", + "id": "duplex-5f2a…", "model": "openbmb/MiniCPM-o-4_5", "state": "open", "turn_state": "idle", @@ -612,7 +693,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). `session.created` and every later `session.update`) ```json -{"type": "session.updated", "session": {"object": "realtime.session", "id": "sess_demo_01", "…": "…"}} +{"type": "session.updated", "session": {"object": "realtime.session", "id": "duplex-5f2a…", "…": "…"}} ``` `C→S session.heartbeat` / `S→C session.heartbeat_ack` @@ -622,7 +703,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ``` ```json -{"type": "session.heartbeat_ack", "session_id": "sess_demo_01"} +{"type": "session.heartbeat_ack", "session_id": "duplex-5f2a…"} ``` `C→S session.event_ack` @@ -638,7 +719,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ``` ```json -{"type": "session.closed", "session_id": "sess_demo_01", "reason": "session_close"} +{"type": "session.closed", "session_id": "duplex-5f2a…", "epoch": 3, "reason": "client_close", "event": {"type": "session.closed", "reason": "client_close"}} ``` `C→S session.resume` (new socket, `?duplex=1&autostart=0`) @@ -646,8 +727,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ```json { "type": "session.resume", - "session_id": "sess_demo_01", - "incarnation": 0, + "session_id": "duplex-5f2a…", "resume_token": "rt_3f9c…", "last_received_server_event_seq": 40 } @@ -659,8 +739,8 @@ events with `server_event_seq > 40` are replayed in order) ```json { "type": "session.resumed", - "session_id": "sess_demo_01", - "incarnation": 0, + "session_id": "duplex-5f2a…", + "session": {"id": "duplex-5f2a…", "state": "open", "epoch": 3, "turn_id": 2, "...": "the engine's current session object"}, "attachment_generation": 2, "resume_token": "rt_a71e…" } @@ -669,19 +749,19 @@ events with `server_event_seq > 40` are replayed in order) `S→C session.replaced` (sent to the superseded socket, which is then closed) ```json -{"type": "session.replaced", "session_id": "sess_demo_01", "attachment_generation": 1} +{"type": "session.replaced", "session_id": "duplex-5f2a…", "attachment_generation": 1} ``` `S→C session.expired` ```json -{"type": "session.expired", "session_id": "sess_demo_01", "incarnation": 0, "reason": "disconnect_grace_expired"} +{"type": "session.expired", "session_id": "duplex-5f2a…", "reason": "disconnect_grace_expired"} ``` `S→C session.resync_required` ```json -{"type": "session.resync_required", "session_id": "sess_demo_01", "reason": "journal_gap"} +{"type": "session.resync_required", "session_id": "duplex-5f2a…", "reason": "journal_gap"} ``` #### Input @@ -772,7 +852,7 @@ silence/noise declaration (no runtime append; answered with `response.listen`) "item_id": "item_9b2d…", "event": { "type": "input.committed", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "turn_id": 1, "epoch": 0, "history_len": 1, @@ -802,7 +882,7 @@ deferred commit during an active response (`event.response_create_deferred`) "type": "input_audio_buffer.committed", "previous_item_id": "item_9b2d…", "item_id": "item_c044…", - "event": {"type": "input.committed", "session_id": "sess_demo_01", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}} + "event": {"type": "input.committed", "session_id": "duplex-5f2a…", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}} } ``` @@ -911,7 +991,7 @@ deferred commit during an active response (`event.response_create_deferred`) "status_details": null, "output": [], "modalities": ["audio", "text"], - "metadata": {"duplex_event": {"type": "response.created", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "turn_id": 1}} + "metadata": {"duplex_event": {"type": "response.created", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "turn_id": 1}} } } ``` @@ -947,7 +1027,7 @@ deferred commit during an active response (`event.response_create_deferred`) "item_id": "item_resp_01", "output_index": 0, "content_index": 0, - "metadata": {"session_id": "sess_demo_01", "epoch": 0, "model_speak": true} + "metadata": {"session_id": "duplex-5f2a…", "epoch": 0, "model_speak": true} } ``` @@ -964,7 +1044,7 @@ deferred commit during an active response (`event.response_create_deferred`) "format": "pcm16", "sample_rate_hz": 24000, "metadata": { - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, "model_speak": true, "end_of_turn": false, @@ -1029,7 +1109,7 @@ deferred commit during an active response (`event.response_create_deferred`) "status": "completed", "status_details": null, "output": [{"id": "item_resp_01", "object": "realtime.item", "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_audio", "transcript": "It is sunny, 24 degrees."}]}], - "metadata": {"type": "response.done", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}} + "metadata": {"type": "response.done", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}} } } ``` @@ -1055,9 +1135,9 @@ deferred commit during an active response (`event.response_create_deferred`) ```json { "type": "response.listen", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, - "response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "sess_demo_01", "epoch": 0, "reason": "model_listen"}} + "response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "duplex-5f2a…", "epoch": 0, "reason": "model_listen"}} } ``` @@ -1104,7 +1184,7 @@ function-call response (Nemotron VoiceChat tools) raw fallback when a model tool call is malformed (`call_id`/`name` missing) ```json -{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "sess_demo_01", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}} +{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "duplex-5f2a…", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}} ``` #### Playback @@ -1120,7 +1200,7 @@ raw fallback when a model tool call is malformed (`call_id`/`name` missing) "type": "playback.acknowledged", "event": { "type": "playback.acknowledged", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, "item_id": "item_resp_01", "played_ms": 1000, @@ -1145,20 +1225,17 @@ truncating ack after an interruption (what was actually heard) ```json { "type": "overlap.decision", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, "policy": "listen_only", "action": "listen", "reason": "long_overlap_speech", - "overlap_ms": 1400, - "buffer_audio": true, - "defer_runtime_append": true, - "force_listen": true + "metadata": {"overlap_ms": 1400, "buffer_audio": true, "defer_runtime_append": true, "force_listen": true} } ``` ```json -{"type": "overlap.decision", "session_id": "sess_demo_01", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"} +{"type": "overlap.decision", "session_id": "duplex-5f2a…", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"} ``` server-VAD session (`session.update` fragment) and the annotated append the @@ -1172,12 +1249,6 @@ translator forwards internally (shown for reference; clients never see it) {"type": "input_audio_buffer.append", "audio": "…", "format": "pcm16", "sample_rate_hz": 16000, "is_speech": true, "force_listen": true, "vad": {"backend": "silero", "is_speech": true, "speech_active": true, "speech_started": true, "speech_stopped": false, "speech_probability": 0.93}} ``` -`S→C runtime.control` (only when an engine signal was partially unsupported) - -```json -{"type": "runtime.control", "session_id": "sess_demo_01", "epoch": 1, "result": {"…": "redacted engine control result"}} -``` - #### Errors `S→C error` — Realtime shape @@ -1203,71 +1274,79 @@ translator forwards internally (shown for reference; clients never see it) ``` ```json -{"type": "error", "error": {"type": "rate_limit_error", "code": "duplex_session_busy", "message": "duplex session capacity reached (max_sessions=4)"}} +{"type": "error", "error": {"type": "rate_limit_error", "code": "resource_exhausted", "message": "duplex_session_capacity_exhausted: limit=4"}} ``` Handshake-stage errors (before the Realtime projector is bound) use the flat native shape: ```json -{"type": "error", "error": "Unknown or expired duplex session: sess_demo_01", "code": "session_resume_expired"} +{"type": "error", "error": "Unknown or expired duplex session: duplex-5f2a…", "code": "session_resume_expired"} ``` #### Coverage The examples above cover every event type the Realtime route emits or accepts (verified by diffing the `type` literals in -`vllm_omni/entrypoints/duplex/` against this document). Not shown, because -they never reach a `/v1/realtime?duplex=1` client unprojected: the native -`/v1/duplex` dialect (`session.create` / `open_session` / `session.config`, -`input.committed`, `input.cancelled`, `audio.cancelled`, -`response.output_audio.delta`, `response.text.delta`, `response.message`, -`function_call.done`) — their Realtime projections are the examples above -(see the name map at the end of this section) — and the runner-internal markers `__timeout__`, -`__disconnect__`, `__replaced_attachment__`. Input aliases (`push_chunk`, -`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`, -`audio.playback_ack`, `close_session`, `close`, `session_close`) share the -payload of their canonical event (see the alias table in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md)). +`vllm_omni/engine/duplex/events.py` and `vllm_omni/engine/duplex/commands.py` +against this document). Not shown, because they never reach a client +unprojected: the session-internal events the runner produces +(`input.committed`, `input.cancelled`, `audio.cancelled`, +`response.output_audio.delta`, `response.text.delta`, `function_call.done`) +— their Realtime projections are the examples above (see the name map at the +end of this section). Every public event is a typed `DuplexEvent` whose +`to_realtime()` is the wire object; every client event is parsed into a +typed `DuplexCommand` by `command_from_realtime()`. -#### Native to Realtime name map +#### Internal to Realtime name map -The native `/v1/duplex` dialect is what the session runner produces; the -Realtime projector renames and fans out these events before they reach a -`/v1/realtime?duplex=1` client. +The session runner emits these session-internal events; the Realtime +projector (`vllm_omni/engine/duplex/realtime_events.py`) renames and fans +them out into typed events before they reach a client. -| Native / internal event | Realtime projection | +| Internal event | Realtime projection | | --- | --- | -| `session.create` (`open_session`, `session.config`) | first `session.update` | | `input.committed` | `conversation.item.added` / `conversation.item.created`, `input_audio_buffer.committed`, `conversation.item.input_audio_transcription.completed`, `conversation.item.done` | | `input.cancelled` | `input_audio_buffer.cleared` | | `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.output_audio.delta` and `response.output_audio_transcript.delta` | | `response.output_audio.done` / `response.output_text.done` | `response.output_audio.done` / `response.output_text.done` | | `response.text.delta` | `response.output_text.delta` | -| `response.message` (chat-fallback raw chunk) | passed through | | `audio.cancelled` | optional `output_audio_buffer.cleared`, then the cancelled terminal events and `response.done` with `status: "cancelled"` | | `function_call.done` | `response.created`, `response.output_item.added` (`function_call` item), `response.function_call_arguments.delta` / `.done`, `response.output_item.done`, `response.done` | | `playback.acknowledged` | `playback.acknowledged` (wrapped verbatim under `event`) | -| `runtime.control`, `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged | +| `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged | ## Known Limitations +- Only MiniCPM-o 4.5 is served over this endpoint today; PersonaPlex and + Nemotron VoiceChat arrive with the follow-up PRs that port them to the + plugin contract. - Several surfaces are capability-gated per model (see *Capability negotiation by model* above): PersonaPlex does not support session resume, barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in or audio truncation; camera frames are consumed only by MiniCPM-o 4.5; tool calls are produced only by Nemotron VoiceChat. -- `turn_detection` supports `server_vad` or `null`, not `semantic_vad`. - Turn-based models support endpointing without interruption, including - `create_response=false` for explicit response creation. Model-native duplex - requires `interrupt_response=true` and `create_response=true`. +- A response is generated from committed **audio**, so the OpenAI text-prompt + shape does not drive one: `conversation.item.create` with an `input_text` + part adds the item to history, but a following `response.create` with no + committed audio is rejected with `response_create_without_input`, and + `input.text.append` is rejected with `native_text_append_unsupported`. + Text-to-speech has its own shape here — put the text in + `extra_body.duplex_initial_user_text` on `session.update`, then stream audio + units (silence is enough). The seeded turn is what the model answers. +- `turn_detection` supports only `server_vad` with `interrupt_response=true`; + `semantic_vad`, `interrupt_response=false`, and `create_response` are not + supported. - `input_audio_transcription` and `input_audio_noise_reduction` are accepted and echoed but no separate transcription or noise-reduction stage runs; transcripts come from the model. - `rate_limits.updated` is emitted for compatibility only and always carries an empty list. - Session capacity is bounded by `duplex_session.max_sessions` in the deploy - configuration; admission beyond it fails with `duplex_session_busy` or - `resource_exhausted`. -- The native `/v1/duplex` route speaks the internal dialect and does not - support `session.resume`; use `/v1/realtime?duplex=1` for the contract on - this page. + configuration; admission beyond it fails with `resource_exhausted` + (`rate_limit_error`, retryable). +- A duplex deployment serves the duplex session surfaces only: the websocket + route and `/v1/chat/completions`, the latter on a session per request (see + [Start the Server](#start-the-server)). The other turn-based HTTP routes + report "not available" on a MiniCPM-o 4.5 server; turn-based use of the + model stays available offline through `Omni` / `AsyncOmni`. diff --git a/examples/online_serving/barge_in_client.py b/examples/online_serving/barge_in_client.py index b144da9f4..b9d95c790 100644 --- a/examples/online_serving/barge_in_client.py +++ b/examples/online_serving/barge_in_client.py @@ -14,7 +14,10 @@ Scenario (see ``barge_in_client_flow.md`` for the diagram) The demo is model-neutral: it drives any duplex model through the public :class:`vllm_omni.clients.duplex.DuplexClient` over -``/v1/realtime?duplex=1``. Model-specific session shape comes from the +``/v1/realtime?duplex=1``, or — with ``--inline`` — through +:class:`vllm_omni.clients.inline_duplex.InlineDuplexClient` on an in-process +:class:`vllm_omni.entrypoints.duplex_omni.DuplexOmni` (no server needed; the +model loads in this process). Model-specific session shape comes from the per-model presets under ``vllm_omni.clients.`` selected with ``--preset`` (default ``minicpmo_4_5``, whose model decides listen/speak on its own — no VAD). @@ -31,6 +34,15 @@ then: --interrupt-wav follow_up_16k.wav \ --output-dir ./duplex_out +or, without a server, in-process (``--deploy-config`` is optional and defaults +to the model's bundled duplex profile): + + python barge_in_client.py --inline \ + --model /path/to/MiniCPM-o-4_5 \ + --ref-audio /path/to/reference_voice.wav \ + --question-wav question_16k.wav \ + --interrupt-wav follow_up_16k.wav + Input WAVs must be mono 16 kHz PCM16; when a preset's session takes a different capture format (the PersonaPlex preset streams 24 kHz float32), the audio is converted before streaming. ``--ref-audio`` is required by the @@ -55,6 +67,7 @@ if str(REPO_ROOT) not in sys.path: from vllm_omni.clients.duplex import ( # noqa: E402 AudioFormat, DuplexClient, + DuplexClientBase, EventCollector, SessionConfig, acknowledge_collected_playback, @@ -139,6 +152,39 @@ def _fold_responses(collector: EventCollector) -> dict[str, dict[str, object]]: return out +def _make_client(args: argparse.Namespace, config: SessionConfig) -> DuplexClientBase: + """Websocket client by default; ``--inline`` drives an in-process DuplexOmni.""" + if not args.inline: + return DuplexClient( + args.url, + model=args.model, + config=config, + reconnect=None, + heartbeat_interval_s=None, + handshake_timeout_s=args.timeout_s, + ) + # Imported lazily: the in-process path loads the model (and its engine + # dependencies) in this process, which the websocket path never needs. + from vllm_omni.clients.inline_duplex import InlineDuplexClient + from vllm_omni.entrypoints.duplex_omni import DuplexOmni + + # Models shipped with custom HF code (MiniCPM-o) need trust_remote_code to + # resolve their pipeline; harmless for the others. + omni_kwargs: dict[str, object] = {"model": args.model, "trust_remote_code": True} + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config + omni = DuplexOmni(**omni_kwargs) + client = InlineDuplexClient( + omni, + model=args.model, + config=config, + handshake_timeout_s=args.timeout_s, + ) + # The client owns no engine; stop the in-process one when the run ends. + client.owned_omni = omni # type: ignore[attr-defined] + return client + + async def run(args: argparse.Namespace) -> int: config = _session_config(args.preset, ref_audio_path=args.ref_audio) question = _convert_to_session_format(read_pcm16_wav(Path(args.question_wav)), config.input_audio) @@ -148,105 +194,117 @@ async def run(args: argparse.Namespace) -> int: output_dir.mkdir(parents=True, exist_ok=True) collector = EventCollector() - client = DuplexClient( - args.url, - model=args.model, - config=config, - reconnect=None, - heartbeat_interval_s=None, - handshake_timeout_s=args.timeout_s, - ) + client = _make_client(args, config) # GREET: open the voice session (handshake happens on enter). - async with client: - consume_task = asyncio.create_task(collector.consume(client)) + try: + async with client: + consume_task = asyncio.create_task(collector.consume(client)) - # ASK: stream the question at realtime pace, then commit the turn. - # The model decides on its own to answer (no VAD). - print("[1] streaming question ...", file=sys.stderr) - await client.stream_pcm(question, chunk_ms=chunk_ms) - await client.commit(create_response=False) - await wait_for_condition( - lambda: collector.count("response.created") > 0, - timeout_s=args.timeout_s, - label="first response.created", - ) - print("[1] assistant is answering", file=sys.stderr) + # ASK: stream the question at realtime pace, then commit the turn. + # The model decides on its own to answer (no VAD). + print("[1] streaming question ...", file=sys.stderr) + await client.stream_pcm(question, chunk_ms=chunk_ms) + await client.commit(create_response=False) + await wait_for_condition( + lambda: collector.count("response.created") > 0, + timeout_s=args.timeout_s, + label="first response.created", + ) + print("[1] assistant is answering", file=sys.stderr) - # Let the answer play for a moment, as a listening user would. - await asyncio.sleep(args.listen_s) + # Let the answer play for a moment, as a listening user would. + await asyncio.sleep(args.listen_s) - # INTERRUPT: talk over the assistant while its audio still streams. - # The serving overlap policy decides: enough overlapped speech barges - # in (the first response is cancelled) or a short remark is deferred - # until the model yields the turn. - print("[2] interrupting mid-answer ...", file=sys.stderr) - ids_before_interrupt = set(collector.response_ids) - await client.stream_pcm(interrupt, chunk_ms=chunk_ms) - await client.commit(create_response=False) + # INTERRUPT: talk over the assistant while its audio still streams. + # The serving overlap policy decides: enough overlapped speech barges + # in (the first response is cancelled) or a short remark is deferred + # until the model yields the turn. + print("[2] interrupting mid-answer ...", file=sys.stderr) + ids_before_interrupt = set(collector.response_ids) + await client.stream_pcm(interrupt, chunk_ms=chunk_ms) + await client.commit(create_response=False) - # ANSWER AGAIN: wait for a response created *after* the interruption - # to finish. The first answer may still be streaming here; matching on - # any response.done would let its completion satisfy the wait and cut - # off the actual follow-up answer at close. - def follow_up_answer_done() -> bool: - return any( - event.get("type") == "response.done" - and collector.response_id(event) is not None - and collector.response_id(event) not in ids_before_interrupt - for event in collector.events + # ANSWER AGAIN: wait for a response created *after* the interruption + # to finish. The first answer may still be streaming here; matching on + # any response.done would let its completion satisfy the wait and cut + # off the actual follow-up answer at close. + def follow_up_answer_done() -> bool: + return any( + event.get("type") == "response.done" + and collector.response_id(event) is not None + and collector.response_id(event) not in ids_before_interrupt + for event in collector.events + ) + + await wait_for_condition( + follow_up_answer_done, + timeout_s=args.timeout_s, + label="response.done for the response answering the interruption", ) - await wait_for_condition( - follow_up_answer_done, - timeout_s=args.timeout_s, - label="response.done for the response answering the interruption", - ) + # HANG UP: report playback, close the session cleanly. + await acknowledge_collected_playback(client, collector) + await client.close(timeout_s=args.timeout_s) + try: + await asyncio.wait_for(consume_task, timeout=5.0) + except (asyncio.TimeoutError, asyncio.CancelledError): + consume_task.cancel() - # HANG UP: report playback, close the session cleanly. - await acknowledge_collected_playback(client, collector) - await client.close(timeout_s=args.timeout_s) - try: - await asyncio.wait_for(consume_task, timeout=5.0) - except (asyncio.TimeoutError, asyncio.CancelledError): - consume_task.cancel() - - # Report: one WAV per response; statuses show what the overlap did. - responses = _fold_responses(collector) - summary: list[dict[str, object]] = [] - for index, (response_id, item) in enumerate(responses.items(), start=1): - audio = item["audio"] - assert isinstance(audio, bytes) - sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type] - wav_path = output_dir / f"response_{index}_{item['status']}.wav" - if audio: - write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz) - summary.append( - { - "response_id": response_id, - "status": item["status"], - "audio_s": round(len(audio) / (sample_rate_hz * 2), 2), - "text": item["text"], - "wav": wav_path.name if audio else None, - } + # Report: one WAV per response; statuses show what the overlap did. + responses = _fold_responses(collector) + summary: list[dict[str, object]] = [] + for index, (response_id, item) in enumerate(responses.items(), start=1): + audio = item["audio"] + assert isinstance(audio, bytes) + sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type] + wav_path = output_dir / f"response_{index}_{item['status']}.wav" + if audio: + write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz) + summary.append( + { + "response_id": response_id, + "status": item["status"], + "audio_s": round(len(audio) / (sample_rate_hz * 2), 2), + "text": item["text"], + "wav": wav_path.name if audio else None, + } + ) + (output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False)) + (output_dir / "events.jsonl").write_text( + "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events), + encoding="utf-8", ) - (output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False)) - (output_dir / "events.jsonl").write_text( - "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events), - encoding="utf-8", - ) - print(json.dumps(summary, indent=2, ensure_ascii=False)) - cancelled = sum(1 for item in summary if item["status"] == "cancelled") - print( - f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}", - file=sys.stderr, - ) + print(json.dumps(summary, indent=2, ensure_ascii=False)) + cancelled = sum(1 for item in summary if item["status"] == "cancelled") + print( + f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}", + file=sys.stderr, + ) + finally: + owned_omni = getattr(client, "owned_omni", None) + if owned_omni is not None: + owned_omni.shutdown() return 0 def main() -> int: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime") - parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument( + "--model", + default="openbmb/MiniCPM-o-4_5", + help="served model name, or the model path with --inline", + ) + parser.add_argument( + "--inline", + action="store_true", + help="run the model in this process through DuplexOmni + InlineDuplexClient instead of connecting to --url", + ) + parser.add_argument( + "--deploy-config", + default=None, + help="deploy YAML for --inline (default: the model's bundled profile)", + ) parser.add_argument( "--preset", choices=["minicpmo_4_5", "personaplex", "none"], diff --git a/examples/online_serving/barge_in_client_flow.md b/examples/online_serving/barge_in_client_flow.md index d32a00e5a..9562d4c9c 100644 --- a/examples/online_serving/barge_in_client_flow.md +++ b/examples/online_serving/barge_in_client_flow.md @@ -1,8 +1,10 @@ # `barge_in_client.py` — what happens, at a glance -The demo plays one everyday conversation pattern against the full-duplex -server: ask something, get talked-over mid-answer, and see the assistant -cope. High-level flow below — for the exact wire events, read the +The demo plays one everyday conversation pattern against a full-duplex +model: ask something, get talked-over mid-answer, and see the assistant +cope. It runs either against a server over `/v1/realtime?duplex=1` or, with +`--inline`, against an in-process `DuplexOmni` — the conversation below is the +same in both modes. High-level flow — for the exact wire events, read the docstring in `barge_in_client.py`. ```mermaid @@ -48,7 +50,8 @@ answer is saved as `response_1_cancelled.wav`, a completed one as duration per answer. Companion: the runtime architecture lives in -[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client -API is `vllm_omni.clients.duplex.DuplexClient`, documented with the wire -protocol in +[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client API is +`vllm_omni.clients.duplex.DuplexClient` over the websocket and +`vllm_omni.clients.inline_duplex.InlineDuplexClient` in-process — both are +documented with the wire protocol in [`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md). diff --git a/examples/online_serving/minicpmo/barge_in_serve.sh b/examples/online_serving/minicpmo/barge_in_serve.sh index 6dade8c5c..1aabd3042 100644 --- a/examples/online_serving/minicpmo/barge_in_serve.sh +++ b/examples/online_serving/minicpmo/barge_in_serve.sh @@ -5,10 +5,10 @@ # Server side for barge_in_client.py: start MiniCPM-o 4.5 in native # full-duplex mode. # -# The duplex deploy profile registers the WS /v1/duplex endpoint (with -# /v1/realtime?duplex=1 projecting onto the same handler), enables the -# engine duplex control plane, and bounds the Thinker/Talker/Code2Wav -# stages to four live duplex sessions on one GPU. See +# MiniCPM-o 4.5 declares a duplex plugin, so vllm-omni serve runs it through +# DuplexOmni: the server is duplex-only (WS /v1/realtime?duplex=1, alias +# /v1/duplex, plus /v1/models and /health) and bounds the +# Thinker/Talker/Code2Wav stages to four live duplex sessions on one GPU. See # vllm_omni/deploy/minicpmo_4_5.yaml (session_mode: duplex) for the session # limits; every other duplex_session knob keeps its runtime default # (idle TTL 300s, disconnect grace 30s, 16MiB pending input per session). @@ -22,4 +22,4 @@ exec vllm-omni serve "$MODEL" \ --deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \ --trust-remote-code \ --host 0.0.0.0 --port "$PORT" -# duplex endpoints: ws://:$PORT/v1/duplex and /v1/realtime?duplex=1 +# duplex endpoint: ws://:$PORT/v1/realtime?duplex=1 (alias /v1/duplex) diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py index 9fc75621e..6dc3fd6e7 100644 --- a/examples/online_serving/minicpmo/realtime_duplex_demo.py +++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py @@ -445,7 +445,6 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: ref_audio=_ref_audio_data_url(args.ref_audio), temperature=args.temperature, ), - session_id=args.session_id, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=args.timeout_s, @@ -638,7 +637,6 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") - parser.add_argument("--session-id") parser.add_argument( "--input-wav", help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.", diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/app.js b/examples/online_serving/minicpmo/realtime_web/app/static/app.js index 1ea2ba093..4d2b78247 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/app.js +++ b/examples/online_serving/minicpmo/realtime_web/app/static/app.js @@ -86,7 +86,6 @@ url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; url.searchParams.set('duplex', '1'); url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5'); - url.searchParams.set('native_duplex', '1'); url.searchParams.set('autostart', '0'); return url.toString(); } @@ -458,7 +457,6 @@ settled = true; const extraBody = { auto_response: true, - native_duplex: true, }; const session = { modalities: ['audio', 'text'], diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index fb0934d4f..01580843f 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -1,6 +1,6 @@ # MiniCPM-o 4.5 -> Online serving and offline inference for omni multimodal chat +> Online full-duplex and chat-completions serving, and offline inference, for omni multimodal chat > (text / image / audio / video → text + 24 kHz speech) ## Summary @@ -9,8 +9,12 @@ - Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5) - Task: Omni multimodal chat — accepts text / image / audio / video input; emits text and 24 kHz mono speech in the same response -- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions` - API (plus Gradio demo), and offline inference via `Omni.generate` +- Mode: Online full-duplex serving over the OpenAI Realtime protocol + (`/v1/realtime?duplex=1`, alias `/v1/duplex`), the OpenAI-compatible + `/v1/chat/completions` API (plus Gradio demo) on the same server, and + offline turn-based inference via `Omni.generate`. Besides those the server + exposes only `/v1/models` and `/health`; the other turn-based HTTP routes + are not available on a duplex deployment. - Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team) ## When to use this recipe @@ -18,9 +22,10 @@ Use this recipe as a known-good starting point for serving `openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member of the MiniCPM-o family — it runs a multimodal thinker, a streaming -MiniCPMTTS codec talker, and a separate batched Code2Wav stage so a single -`/v1/chat/completions` call can return text and 24 kHz speech in one -shot. The recommended batching deploy isolates the Thinker on GPU 0 and +MiniCPMTTS codec talker, and a separate batched Code2Wav stage, so a live +duplex session streams text and 24 kHz speech while it keeps listening, and a +single `/v1/chat/completions` call returns both in one shot. The +recommended batching deploy isolates the Thinker on GPU 0 and co-locates Talker and Code2Wav on GPU 1; 1-GPU, 3-GPU, and 8x4090 layouts are also provided. @@ -36,8 +41,11 @@ also provided. [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) - 8x RTX 4090 layout: [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml) -- Online example + Gradio demo: +- Online examples (Realtime CLI demo, browser client, barge-in client, + chat-completions client, curl script, Gradio demo): [`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/) +- Duplex API and Python client docs: + [`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md) - Offline end-to-end example: [`examples/offline_inference/minicpmo/`](../../examples/offline_inference/minicpmo/) - Pipeline / talker source: @@ -179,6 +187,13 @@ as one aggregated chunk, so the table reports Stage 0 metrics. #### Verification +**Health and model listing**: + +```bash +curl http://localhost:8099/health +curl http://localhost:8099/v1/models +``` + **Quick smoke test (text-only output)**: ```bash @@ -191,8 +206,31 @@ curl http://localhost:8099/v1/chat/completions \ }' ``` -**Text + speech in one response** (the headline 4.5 feature). The model -bridge conditions the Talker from the generated assistant span, so the +**One duplex turn over the Realtime WebSocket** (text + speech stream back +while the session keeps listening): + +```bash +python examples/online_serving/minicpmo/realtime_duplex_demo.py \ + --url ws://localhost:8099/v1/realtime?duplex=1 \ + --model openbmb/MiniCPM-o-4_5 \ + --input-wav /path/to/input_16k_mono_pcm16.wav \ + --ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \ + --output-dir /tmp/minicpmo_realtime_duplex_demo +``` + +The demo prints transcript deltas as they arrive and saves the streamed +24 kHz audio to WAV files. The same session shape is available from Python +through `vllm_omni.clients.duplex.DuplexClient` (over the server) or +`vllm_omni.clients.inline_duplex.InlineDuplexClient` (in-process, no server; +`examples/online_serving/barge_in_client.py --inline`). + +**Barge-in**: `examples/online_serving/barge_in_client.py` plays a question, +interrupts the answer with a second WAV and checks that the first response is +cancelled and the second one is answered. + +**Text + speech in one response over `/v1/chat/completions`** (the headline +4.5 feature, served by the ordinary chat service on the duplex engine). The +model bridge conditions the Talker from the generated assistant span, so the generic serving layer does not inject MiniCPM-specific template defaults. `use_tts_template=true` remains supported when explicitly requested: @@ -243,6 +281,12 @@ python examples/online_serving/minicpmo/gradio_demo.py \ Open `http://:7862` and try a text prompt with the **"Generate speech output (TTS)"** checkbox on / off. +**Turn-based generation** is also available offline +(`examples/offline_inference/minicpmo/`, `Omni.generate`). + +The Daily-Omni numbers above were measured on the `/v1/chat/completions` +route. + #### Notes - Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and @@ -326,9 +370,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ and defaults to deterministic seed 42. Stage-1 deploy sampling parameters control only vLLM's binary continue/stop token. -- **Output audio**: 24 kHz mono WAV inside the OpenAI-style - `message.audio.data` (base64). The Gradio demo's WAV player decodes - this automatically. +- **Output audio**: 24 kHz mono. On `/v1/chat/completions` it is base64 WAV + inside the OpenAI-style `message.audio.data` (the Gradio demo's WAV player + decodes this automatically); on a duplex session it streams as base64 + `response.audio.delta` events; offline `Omni.generate` returns WAV. +- **Response choices**: on `/v1/chat/completions` text and audio are separate + choices. SDK clients should select the choice whose `message.audio.data` is + populated rather than assuming `choices[0]` contains audio. - **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures= ["MiniCPMO"]` in HF config; routing is disambiguated by @@ -340,6 +388,3 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ - **Async chunking**: enabled in all deploy configs. Talker sends 25-code chunks with three-code left context to Code2Wav through `SharedMemoryConnector`; terminal chunks flush held lookahead state. -- **Response choices**: text and audio are separate choices. SDK clients - should select the choice whose `message.audio.data` is populated rather - than assuming `choices[0]` contains audio. diff --git a/recipes/README.md b/recipes/README.md index 33e29c086..078466e80 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -71,7 +71,7 @@ recipes/ | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | | [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) | | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | -| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | +| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online full-duplex serving and offline inference for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB | | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB | | [`OpenMOSS/MOSS-TTS-Nano-NPU.md`](./OpenMOSS/MOSS-TTS-Nano-NPU.md) | Online multilingual voice-cloned TTS | 1x Ascend 910B3 64GB | diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index 60e071ad5..d2552bbff 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -135,6 +135,7 @@ class MockResponse: async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch): class FakeRealtimeClient: last_instance = None + instances: list = [] def __init__(self, url): assert url == "ws://localhost:8000/v1/realtime?duplex=1" @@ -143,6 +144,9 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) self.sent = [] self.response_count = 0 self.ack_count = 0 + self.silence_seconds = [] + self.closed = 0 + type(self).instances.append(self) type(self).last_instance = self async def __aenter__(self): @@ -155,10 +159,22 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) assert model == "openbmb/MiniCPM-o-4_5" self.configure_kwargs = kwargs + async def stream_silence(self, *, seconds, chunk_ms=200, until=None): + # A model-native session speaks off its seeded context once audio + # units arrive; the silence itself carries no content. The probe + # stops the silence as soon as the turn settles. + self.silence_seconds.append(seconds) + self._emit_response() + assert until is not None and until() + return 1.0 + async def send(self, event): self.sent.append(event) if event["type"] != "response.create": return + self._emit_response() + + def _emit_response(self): self.response_count += 1 response_id = f"resp-{self.response_count}" now = time.monotonic() @@ -210,7 +226,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) self.ack_count += 1 async def close_session(self, **_kwargs): - return None + self.closed += 1 monkeypatch.setattr( "vllm_omni.benchmarks.patch.patch._RealtimeTTSProbe", @@ -239,30 +255,27 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) session=None, ) - client = FakeRealtimeClient.last_instance - assert client.configure_kwargs["native_duplex"] is False - assert client.configure_kwargs["extra_body"] == { - "ref_audio": "data:audio/wav;base64,AAAA", - "return_stage_metrics": True, - } - assert client.sent == [ - event - for index in range(4) - for event in ( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": f"text {index}"}], - }, - }, - {"type": "response.create"}, - ) - ] - assert client.ack_count == 4 + # One session per utterance: a model-native duplex session takes its text + # once, in the session context, so it cannot be re-seeded for a second one. + sessions = FakeRealtimeClient.instances + assert len(sessions) == 4 + for index, client in enumerate(sessions): + assert "native_duplex" not in client.configure_kwargs + assert client.configure_kwargs["extra_body"] == { + "ref_audio": "data:audio/wav;base64,AAAA", + "return_stage_metrics": True, + "duplex_initial_user_text": f"text {index}", + "force_listen_count": 0, + } + # The target text rides the session context, never a conversation item: + # a text-only response.create is rejected by a model-native session. + assert client.sent == [] + assert client.silence_seconds and all(seconds > 0 for seconds in client.silence_seconds) + assert client.ack_count == 1 + assert client.closed == 1 + assert output.success is True - assert output.generated_text == "turn 1 turn 2 turn 3 turn 4" + assert output.generated_text == "turn 1 turn 1 turn 1 turn 1" assert output.audio_duration == pytest.approx(0.4) assert output.ttft > 0 assert output.audio_ttfp > output.ttft @@ -281,11 +294,12 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) "session_id": session_id, "request_index": index, "utterance_id": f"utt-{index}", - "response_id": f"resp-{index + 1}", + # Response ids are per session, and each utterance now gets its own. + "response_id": "resp-1", "source": "client_monotonic_receive", "measurement_origin": { - "ttft": "conversation.item.create client send to first non-empty text delta", - "ttfp": "conversation.item.create client send to first audio packet", + "ttft": "first silence append client send to first non-empty text delta", + "ttfp": "first silence append client send to first audio packet", "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", "tpot": "Stage-0 engine mean time per output token", }, diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py index 73cedfb32..1a6acaca5 100644 --- a/tests/benchmarks/test_omniinteract.py +++ b/tests/benchmarks/test_omniinteract.py @@ -467,6 +467,27 @@ def test_tolerated_playback_ack_rejection_is_a_warning_not_a_failure(): oi._raise_if_session_terminated(collector, 0, warnings=warnings) +def test_our_own_close_is_expected_even_though_the_server_stamps_a_reason(): + """``session.close`` is answered with ``session.closed`` carrying ``client_close``. + + The guard exists to catch a session that ended for a reason we did not ask + for. A close we requested is not that, whether or not the server names it. + """ + collector = _collector(({"type": "session.closed", "reason": "client_close"}, 1.0)) + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + collector = _collector(({"type": "session.closed", "event": {"reason": "client_close"}}, 1.0)) + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + collector = _collector(({"type": "session.closed", "reason": "disconnect"}, 1.0)) + with pytest.raises(RuntimeError, match="Unexpected session.closed: disconnect"): + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + collector = _collector(({"type": "session.expired", "reason": "timeout"}, 1.0)) + with pytest.raises(RuntimeError, match="session.expired: timeout"): + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + @pytest.mark.parametrize( ("event", "match"), [ @@ -561,8 +582,8 @@ class _RealtimeClient(oi._RealtimeSession): instances: list[_RealtimeClient] = [] - def __init__(self, config: oi.OmniInteractBenchmarkConfig, session_id: str, reference_audio: str): - super().__init__(config, session_id, reference_audio) + def __init__(self, config: oi.OmniInteractBenchmarkConfig, reference_audio: str): + super().__init__(config, reference_audio) self.acks: list[tuple[str, int]] = [] self.instances.append(self) @@ -621,7 +642,7 @@ async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monke assert "autostart=0" in _RealtimeClient.instances[-1].url session_config = _RealtimeClient.instances[-1].session_config assert session_config.extra_body["custom"] == "value" - assert session_config.extra_body["native_duplex"] is True + assert "native_duplex" not in session_config.extra_body assert session_config.ref_audio == "data:audio/wav;base64,ref" acks = _RealtimeClient.instances[-1].acks # Cumulative acks for the one response: an optional 0 ms checkpoint the diff --git a/tests/benchmarks/test_realtime_tts_probe.py b/tests/benchmarks/test_realtime_tts_probe.py new file mode 100644 index 000000000..c07a09b62 --- /dev/null +++ b/tests/benchmarks/test_realtime_tts_probe.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The Seed-TTS Realtime probe queues for a duplex session slot instead of failing.""" + +from __future__ import annotations + +import pytest + +from vllm_omni.benchmarks.patch import patch as bench_patch +from vllm_omni.clients import duplex as duplex_client + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _AdmissionGate: + """Stands in for ``DuplexClient``: refuses admission ``refusals`` times, then admits.""" + + instances: list[_AdmissionGate] = [] + refusals = 0 + refusal_code = "resource_exhausted" + + def __init__(self, url: str, **kwargs: object) -> None: + del url, kwargs + self.entered = False + type(self).instances.append(self) + + async def __aenter__(self) -> _AdmissionGate: + cls = type(self) + if cls.refusals > 0: + cls.refusals -= 1 + raise duplex_client.DuplexProtocolError("duplex_session_capacity_exhausted: limit=4", code=cls.refusal_code) + self.entered = True + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + del exc_type, exc, tb + + async def events(self): # the collector's consume loop ends at once + return + yield # pragma: no cover + + +@pytest.fixture +def gate(monkeypatch: pytest.MonkeyPatch) -> type[_AdmissionGate]: + _AdmissionGate.instances = [] + _AdmissionGate.refusals = 0 + _AdmissionGate.refusal_code = "resource_exhausted" + monkeypatch.setattr(duplex_client, "DuplexClient", _AdmissionGate) + return _AdmissionGate + + +async def _configure(probe: bench_patch._RealtimeTTSProbe) -> None: + try: + await probe.configure("test-model", auto_response=True) + finally: + await probe.__aexit__(None, None, None) + + +@pytest.mark.asyncio +async def test_configure_waits_for_a_free_session_slot(gate) -> None: + """Two refusals, then a slot: the request is served, not counted as failed.""" + gate.refusals = 2 + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + await _configure(probe) + assert len(gate.instances) == 3 + assert gate.instances[-1].entered is True + assert probe._client is gate.instances[-1] + + +@pytest.mark.asyncio +async def test_configure_raises_other_rejections_at_once(gate) -> None: + gate.refusals = 1 + gate.refusal_code = "unsupported_turn_detection" + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + with pytest.raises(duplex_client.DuplexProtocolError): + await _configure(probe) + assert len(gate.instances) == 1 + + +@pytest.mark.asyncio +async def test_configure_gives_up_when_the_slot_wait_budget_is_spent(gate, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bench_patch._RealtimeTTSProbe, "_SESSION_SLOT_WAIT_S", 0.0) + gate.refusals = 5 + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + with pytest.raises(duplex_client.DuplexProtocolError) as excinfo: + await _configure(probe) + assert excinfo.value.code == "resource_exhausted" + assert len(gate.instances) == 1 + + +class _SilenceSink: + """Stands in for a configured ``DuplexClient``: records every appended chunk.""" + + def __init__(self) -> None: + self.config = duplex_client.SessionConfig() + self.chunks: list[bytes] = [] + + async def append_audio(self, pcm: bytes, *, is_speech: bool | None = None, video_frames=None) -> None: + del video_frames + assert is_speech is False + self.chunks.append(pcm) + + +@pytest.mark.asyncio +async def test_stream_silence_stops_at_the_turn_end(monkeypatch: pytest.MonkeyPatch) -> None: + """Silence is fed until the turn settles, not for the whole cap.""" + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + sink = _SilenceSink() + probe._client = sink + + async def no_wait(_: float) -> None: + return + + monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait) + streamed = await probe.stream_silence(seconds=12.0, chunk_ms=200, until=lambda: len(sink.chunks) >= 3) + assert len(sink.chunks) == 3 + assert streamed == pytest.approx(0.6) + assert all(len(chunk) == sink.config.input_audio.byte_count(200) for chunk in sink.chunks) + + +@pytest.mark.asyncio +async def test_stream_silence_runs_to_the_cap_when_the_turn_never_settles(monkeypatch: pytest.MonkeyPatch) -> None: + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + sink = _SilenceSink() + probe._client = sink + + async def no_wait(_: float) -> None: + return + + monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait) + streamed = await probe.stream_silence(seconds=1.0, chunk_ms=200, until=lambda: False) + assert len(sink.chunks) == 5 + assert streamed == pytest.approx(1.0) + + +class _TurnEvents: + def __init__(self, response_ids: list[str], audio: dict[str, bytes], text: dict[str, str] | None = None) -> None: + self.response_ids = response_ids + self._audio = audio + self._text = text or {} + + def audio_bytes(self, response_id: str) -> bytes: + return self._audio.get(response_id, b"") + + def response_text(self, response_id: str) -> str: + return self._text.get(response_id, "") + + +def test_turn_response_is_the_first_one_with_audio_even_if_the_model_spoke_again(caplog) -> None: + events = _TurnEvents(["r0", "r1", "r2"], {"r1": b"\x01\x00", "r2": b"\x02\x00"}) + with caplog.at_level("WARNING", logger=bench_patch.logger.name): + assert bench_patch._seed_tts_turn_response_id(events, 0, 0) == "r1" + assert "spoke again" in caplog.text + + +def test_turn_response_ignores_responses_from_before_the_turn() -> None: + events = _TurnEvents(["r0", "r1"], {"r0": b"\x01\x00", "r1": b"\x02\x00"}) + assert bench_patch._seed_tts_turn_response_id(events, 1, 0) == "r1" + + +def test_turn_without_audio_is_an_error() -> None: + events = _TurnEvents(["r0"], {}) + with pytest.raises(RuntimeError, match="no audio response"): + bench_patch._seed_tts_turn_response_id(events, 0, 0) + + +def test_stall_report_says_the_model_never_answered() -> None: + events = _TurnEvents(["r0"], {"r0": b"\x01\x00"}) + report = bench_patch._seed_tts_turn_stall_report(events, 1, 3, 30.0) + assert "turn 3 never started a response" in report + assert "30.0s of silence" in report + + +def test_stall_report_lists_a_response_that_never_finished() -> None: + events = _TurnEvents(["r1"], {"r1": b"\x01\x00\x02\x00"}, {"r1": "hello"}) + report = bench_patch._seed_tts_turn_stall_report(events, 0, 0, 4.2) + assert "started 1 response(s) after 4.2s of silence" in report + assert "r1 (4 audio bytes, text 'hello')" in report diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index ab2a38951..43b7547cb 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for the public duplex client (fake transport, no server).""" +"""Unit tests for the public duplex clients (fake transports, no server). + +``DuplexClient`` runs over a fake websocket; ``InlineDuplexClient`` over a +fake ``DuplexOmni`` whose handle records typed commands and replays typed +events. The demux / handshake tests that only need "feed a server event" are +parameterized over both clients. +""" from __future__ import annotations @@ -8,6 +14,8 @@ import asyncio import base64 import json import wave +from collections.abc import Callable +from typing import Any from urllib.parse import parse_qs, urlsplit import pytest @@ -16,8 +24,11 @@ from vllm_omni.clients.duplex import ( AudioFormat, ConnectionResumed, DuplexClient, + DuplexClientBase, + DuplexConnectionError, DuplexProtocolError, DuplexSessionClosedError, + ErrorEvent, EventCollector, ReconnectPolicy, SessionConfig, @@ -31,17 +42,26 @@ from vllm_omni.clients.duplex import ( summarize_session_request_metrics, write_pcm16_wav, ) +from vllm_omni.clients.inline_duplex import InlineDuplexClient +from vllm_omni.engine.duplex import commands as duplex_commands +from vllm_omni.engine.duplex import events as duplex_events +from vllm_omni.engine.duplex.messages import DuplexSessionError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +SESSION_ID = "duplex-0123456789abcdef" +# session.created is not journaled by the server, so it carries no server_event_seq. SESSION_CREATED = { "type": "session.created", - "session": {"session_id": "sess-1"}, - "incarnation": 0, + "session": {"id": SESSION_ID, "model": "test-model"}, + "attachment_generation": 1, "resume_token": "tok-1", - "server_event_seq": 1, } -SESSION_CLOSED = {"type": "session.closed", "session_id": "sess-1", "server_event_seq": 99} +SESSION_CLOSED = {"type": "session.closed", "session_id": SESSION_ID, "reason": "client_close", "server_event_seq": 99} + + +# --------------------------------------------------------------------------- +# Fake websocket transport class FakeSocket: @@ -88,14 +108,139 @@ def make_client(*sockets: FakeSocket, **kwargs) -> tuple[DuplexClient, list[str] return client, calls +# --------------------------------------------------------------------------- +# Fake DuplexOmni for the inline client + + +class _WireEvent: + """A typed-event stand-in that renders a fixed wire payload.""" + + def __init__(self, raw: dict[str, object]) -> None: + self.raw = raw + + @property + def is_terminal(self) -> bool: + return self.raw.get("type") in {"session.closed", "session.expired"} + + def to_realtime(self) -> dict[str, object]: + return dict(self.raw) + + +class FakeInlineHandle: + def __init__(self, session_id: str) -> None: + self.session_id = session_id + self.closed = False + self.close_calls = 0 + self.commands: list[duplex_commands.DuplexCommand] = [] + self.consumer_active = False + self._outbox: asyncio.Queue = asyncio.Queue() + + def deliver(self, event: Any) -> None: + self._outbox.put_nowait(event) + + def feed(self, raw: dict[str, object]) -> None: + self.deliver(_WireEvent(raw)) + + async def submit(self, command: duplex_commands.DuplexCommand) -> None: + if self.closed: + raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id) + self.commands.append(command) + + async def events(self): + self.consumer_active = True + try: + while True: + event = await self._outbox.get() + yield event + if event.is_terminal: + return + finally: + self.consumer_active = False + + async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None: + self.close_calls += 1 + if self.closed: + return + self.closed = True + self.deliver(duplex_events.SessionClosed(session_id=self.session_id, reason=reason)) + + +class FakeInlineOmni: + def __init__(self) -> None: + self.opened: list[dict[str, object]] = [] + self.handles: list[FakeInlineHandle] = [] + + async def open_session(self, config: dict[str, object]) -> FakeInlineHandle: + self.opened.append(dict(config)) + handle = FakeInlineHandle(SESSION_ID) + self.handles.append(handle) + handle.deliver( + duplex_events.SessionCreated( + session_id=SESSION_ID, + session={"id": SESSION_ID, "model": config.get("model")}, + ) + ) + return handle + + +def make_inline_client(**kwargs) -> tuple[InlineDuplexClient, FakeInlineOmni]: + omni = FakeInlineOmni() + kwargs.setdefault("handshake_timeout_s", 5.0) + return InlineDuplexClient(omni, model="test-model", **kwargs), omni + + +# --------------------------------------------------------------------------- +# One rig over both clients: ``feed`` a server event, drive ``client`` + + +class Rig: + def __init__(self, client: DuplexClientBase, feed: Callable[[dict[str, object]], None]) -> None: + self.client = client + self.feed = feed + + +def _websocket_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig: + sock = FakeSocket() + if handshake is not None: + sock.feed(handshake) + client, _ = make_client(sock) + return Rig(client, sock.feed) + + +def _inline_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig: + client, omni = make_inline_client() + original_open = omni.open_session + + async def open_session(config: dict[str, object]) -> FakeInlineHandle: + handle = await original_open(config) + if handshake is not SESSION_CREATED: + # Replace the default session.created with the requested first event. + handle._outbox = asyncio.Queue() + if handshake is not None: + handle.feed(handshake) + return handle + + omni.open_session = open_session # type: ignore[method-assign] + + def feed(raw: dict[str, object]) -> None: + omni.handles[-1].feed(raw) + + return Rig(client, feed) + + +@pytest.fixture(params=["websocket", "inline"]) +def make_rig(request: pytest.FixtureRequest) -> Callable[..., Rig]: + return _websocket_rig if request.param == "websocket" else _inline_rig + + # --------------------------------------------------------------------------- # SessionConfig def test_session_config_payload_defaults(): - payload = SessionConfig().to_session_payload(model="m", session_id="s") + payload = SessionConfig().to_session_payload(model="m") assert payload["model"] == "m" - assert payload["session_id"] == "s" + assert "session_id" not in payload assert payload["modalities"] == ["audio", "text"] assert payload["input_audio_format"] == "pcm16" assert payload["output_audio_format"] == "pcm16" @@ -137,19 +282,59 @@ def test_audio_format_math(): # --------------------------------------------------------------------------- -# Handshake and lifecycle +# Handshake and lifecycle (both clients) -async def test_handshake_adopts_session_and_acks(): +@pytest.mark.asyncio +async def test_handshake_adopts_server_allocated_session_id(make_rig): + rig = make_rig() + async with rig.client as client: + assert client.session_id == SESSION_ID + assert client.session_info == {"id": SESSION_ID, "model": "test-model"} + rig.feed(SESSION_CLOSED) + + +@pytest.mark.asyncio +async def test_handshake_error_raises_protocol_error(make_rig): + rig = make_rig(handshake={"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}}) + with pytest.raises(DuplexProtocolError) as excinfo: + async with rig.client: + pass + assert excinfo.value.code == "unsupported_audio_format" + + +@pytest.mark.asyncio +async def test_session_expired_raises_from_event_stream(make_rig): + rig = make_rig() + async with rig.client as client: + rig.feed( + {"type": "session.expired", "session_id": SESSION_ID, "reason": "lease_expired", "server_event_seq": 2} + ) + with pytest.raises(DuplexSessionClosedError) as excinfo: + async for _ in client.events(): + pass + assert "lease_expired" in excinfo.value.reason + + +# --------------------------------------------------------------------------- +# Websocket transport specifics + + +@pytest.mark.asyncio +async def test_websocket_handshake_sends_session_update_and_acks_sequenced_events(): sock = FakeSocket() sock.feed(SESSION_CREATED) client, calls = make_client(sock) async with client: - assert client.session_id == "sess-1" + assert client.session_id == SESSION_ID assert client.resume_token == "tok-1" assert calls == ["ws://test-host:8099/v1/realtime?duplex=1&model=test-model&autostart=0"] assert sock.sent[0]["type"] == "session.update" assert sock.sent[0]["session"]["model"] == "test-model" + assert "session_id" not in sock.sent[0]["session"] + # session.created is unjournaled (no seq): nothing to ack yet. + assert _acks(sock) == [] + sock.feed({"type": "response.listen", "server_event_seq": 1}) await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 1} in _acks(sock)) close_task = asyncio.create_task(client.close()) await _drain(lambda: "session.close" in sock.sent_types()) @@ -158,28 +343,7 @@ async def test_handshake_adopts_session_and_acks(): assert "session.close" in sock.sent_types() -async def test_handshake_error_raises_protocol_error(): - sock = FakeSocket() - sock.feed({"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}}) - client, _ = make_client(sock) - with pytest.raises(DuplexProtocolError) as excinfo: - async with client: - pass - assert excinfo.value.code == "unsupported_audio_format" - - -async def test_session_expired_raises_from_event_stream(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "session.expired", "reason": "lease_expired", "server_event_seq": 2}) - with pytest.raises(DuplexSessionClosedError) as excinfo: - async for _ in client.events(): - pass - assert "lease_expired" in excinfo.value.reason - - +@pytest.mark.asyncio async def test_resync_required_drops_resume_credential(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -187,7 +351,7 @@ async def test_resync_required_drops_resume_credential(): async with client: assert client.resume_token == "tok-1" stream = client.events() - sock.feed({"type": "session.resync_required", "session_id": "sess-1", "server_event_seq": 2}) + sock.feed({"type": "session.resync_required", "session_id": SESSION_ID, "server_event_seq": 2}) async for event in stream: if event.type == "session.resync_required": break @@ -202,9 +366,10 @@ async def test_resync_required_drops_resume_credential(): # --------------------------------------------------------------------------- -# Input events +# Input events (websocket wire shape) +@pytest.mark.asyncio async def test_append_audio_tracks_cumulative_end_ms(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -224,6 +389,7 @@ async def test_append_audio_tracks_cumulative_end_ms(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_append_audio_strips_video_frame_data_url_prefix(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -238,6 +404,7 @@ async def test_append_audio_strips_video_frame_data_url_prefix(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_stream_pcm_chunking(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -249,6 +416,7 @@ async def test_stream_pcm_chunking(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_interruption_primitives_send_documented_events(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -262,22 +430,21 @@ async def test_interruption_primitives_send_documented_events(): # --------------------------------------------------------------------------- -# Response demultiplexing +# Response demultiplexing (both clients) def _b64(data: bytes) -> str: return base64.b64encode(data).decode("ascii") -async def test_response_handle_flow(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) +@pytest.mark.asyncio +async def test_response_handle_flow(make_rig): + rig = make_rig() chunk = b"\x00\x01" * 2400 # 100 ms of pcm16 @ 24 kHz - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) - sock.feed( + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) + rig.feed( { "type": "response.output_audio.delta", "response_id": "resp-1", @@ -286,7 +453,7 @@ async def test_response_handle_flow(): "server_event_seq": 4, } ) - sock.feed( + rig.feed( { "type": "response.output_audio_transcript.delta", "response_id": "resp-1", @@ -294,7 +461,7 @@ async def test_response_handle_flow(): "server_event_seq": 5, } ) - sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6}) + rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6}) async for response in client.responses(): chunks = [piece async for piece in response.audio()] @@ -304,49 +471,46 @@ async def test_response_handle_flow(): assert response.played_ms == pytest.approx(100.0) assert response.finished break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_listen_decision_yields_finished_silent_handle(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.listen", "server_event_seq": 2}) +@pytest.mark.asyncio +async def test_listen_decision_yields_finished_silent_handle(make_rig): + rig = make_rig() + async with rig.client as client: + rig.feed({"type": "response.listen", "server_event_seq": 2}) async for response in client.responses(): assert response.decision == "listen" assert response.finished assert [piece async for piece in response.audio()] == [] break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_listen_terminates_active_response(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3}) +@pytest.mark.asyncio +async def test_listen_terminates_active_response(make_rig): + rig = make_rig() + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3}) async for response in client.responses(): await response.wait(timeout_s=5.0) assert response.decision == "listen" break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_id_less_listen_never_closes_an_in_flight_response(): +@pytest.mark.asyncio +async def test_id_less_listen_never_closes_an_in_flight_response(make_rig): # An id-less listen is a standalone decision beat (e.g. a silence-skip # while a response streams); it must surface as its own finished handle # and leave the in-flight response to its real terminal. - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) - sock.feed({"type": "response.listen", "server_event_seq": 4}) - sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5}) + rig = make_rig() + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) + rig.feed({"type": "response.listen", "server_event_seq": 4}) + rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5}) seen: list[tuple[str | None, str | None, bool]] = [] async for response in client.responses(): await response.wait(timeout_s=5.0) @@ -355,20 +519,19 @@ async def test_id_less_listen_never_closes_an_in_flight_response(): break assert ("resp-1", "speak", True) in seen assert (None, "listen", True) in seen - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_terminal_listen_keeps_spoken_decision(): +@pytest.mark.asyncio +async def test_terminal_listen_keeps_spoken_decision(make_rig): # A terminal listen after the response spoke is the model yielding the # turn; the handle must stay decision="speak" with its audio intact. - sock = FakeSocket() - sock.feed(SESSION_CREATED) + rig = make_rig() chunk = b"\x00\x01" * 2400 - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) - sock.feed( + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) + rig.feed( { "type": "response.output_audio.delta", "response_id": "resp-1", @@ -377,24 +540,23 @@ async def test_terminal_listen_keeps_spoken_decision(): "server_event_seq": 4, } ) - sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5}) + rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5}) async for response in client.responses(): chunks = [piece async for piece in response.audio()] assert chunks == [chunk] assert response.decision == "speak" assert response.finished break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_error_event_surfaces_on_response_path(): +@pytest.mark.asyncio +async def test_error_event_surfaces_on_response_path(make_rig): # A rejected send produces no response; a consumer waiting in responses() # must see the rejection instead of waiting forever. - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed( + rig = make_rig() + async with rig.client as client: + rig.feed( { "type": "error", "error": {"code": "input_audio_buffer_empty", "message": "empty commit"}, @@ -405,9 +567,10 @@ async def test_error_event_surfaces_on_response_path(): async for _ in client.responses(): pass assert excinfo.value.code == "input_audio_buffer_empty" - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): from vllm_omni.clients.duplex import AudioDelta, ResponseHandle @@ -425,24 +588,24 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): # --------------------------------------------------------------------------- -# Resume +# Resume (websocket only) +@pytest.mark.asyncio async def test_resume_after_transport_drop(): first = FakeSocket() first.feed(SESSION_CREATED) second = FakeSocket() # Mirror the real resume activation payload # (entrypoints/duplex/serving.py, activation_payload_factory): it carries - # no nested "session" object, only the identity fields. + # the identity fields plus the (empty) public session object. second.feed( { "type": "session.resumed", - "session_id": "sess-1", - "incarnation": 0, - "attachment_generation": 1, + "session_id": SESSION_ID, + "session": {}, + "attachment_generation": 2, "resume_token": "tok-2", - "server_event_seq": 5, } ) client, calls = make_client( @@ -453,34 +616,39 @@ async def test_resume_after_transport_drop(): async with client: stream = client.events() - async def take_two(): + async def collect_until_resumed(): seen = [] async for event in stream: seen.append(event) - if len(seen) == 2: + if isinstance(event, SessionResumed): return seen - consumer = asyncio.create_task(take_two()) + consumer = asyncio.create_task(collect_until_resumed()) await asyncio.sleep(0) # let the consumer subscribe before the drop await asyncio.sleep(0) + first.feed({"type": "response.listen", "server_event_seq": 7}) + await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 7} in _acks(first)) first.feed(RuntimeError("transport dropped")) seen = await asyncio.wait_for(consumer, timeout=5.0) - assert isinstance(seen[0], ConnectionResumed) - assert isinstance(seen[1], SessionResumed) + assert [type(event) for event in seen[-2:]] == [ConnectionResumed, SessionResumed] assert client.resume_token == "tok-2" + # The superseded transport is closed, not leaked. + assert first.closed # The activation payload has no session object; the info captured at # the original handshake must survive the resume. - assert client.session_info == {"session_id": "sess-1"} + assert client.session_info == {"id": SESSION_ID, "model": "test-model"} assert len(calls) == 2 resume = second.sent[0] - assert resume["type"] == "session.resume" - assert resume["session_id"] == "sess-1" - assert resume["incarnation"] == 0 - assert resume["resume_token"] == "tok-1" - assert resume["last_received_server_event_seq"] == 1 + assert resume == { + "type": "session.resume", + "session_id": SESSION_ID, + "resume_token": "tok-1", + "last_received_server_event_seq": 7, + } second.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_no_reconnect_policy_surfaces_closed(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -493,6 +661,7 @@ async def test_no_reconnect_policy_surfaces_closed(): pass +@pytest.mark.asyncio async def test_heartbeat_survives_transient_send_failure(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -514,6 +683,7 @@ async def test_heartbeat_survives_transient_send_failure(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_aexit_sends_session_close_on_error_path(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -534,6 +704,7 @@ def test_target_url_forces_autostart_off(): assert query["autostart"] == ["0"] +@pytest.mark.asyncio async def test_fatal_resume_error_gives_up(): first = FakeSocket() first.feed(SESSION_CREATED) @@ -553,6 +724,121 @@ async def test_fatal_resume_error_gives_up(): assert len(calls) == 2 # no retry after a fatal resume error +# --------------------------------------------------------------------------- +# Inline client over a fake DuplexOmni + + +@pytest.mark.asyncio +async def test_inline_handshake_opens_session_with_payload_and_adopts_id(): + client, omni = make_inline_client(config=SessionConfig(ref_audio="data:audio/wav;base64,AAA=")) + async with client: + assert omni.opened == [client.config.to_session_payload(model="test-model")] + assert omni.opened[0]["ref_audio"] == "data:audio/wav;base64,AAA=" + assert client.session_id == SESSION_ID + handle = omni.handles[0] + # The session was opened with the full session object already; the + # base handshake's session.update must not be re-submitted as a patch + # (a ref_audio patch would be rejected by the runner). + assert handle.commands == [] + await client.send({"type": "session.update", "session": {"instructions": "later"}}) + assert [type(command) for command in handle.commands] == [duplex_commands.UpdateSession] + assert handle.commands[0].patch == {"instructions": "later"} + + +@pytest.mark.asyncio +async def test_inline_client_submits_typed_commands_in_order(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + pcm = b"\x01\x00" * 1600 + await client.append_audio(pcm, is_speech=True) + await client.commit(create_response=True) + await client.cancel_response("resp-1") + await client.ack_playback(120.9, response_id="resp-1") + assert [type(command) for command in handle.commands] == [ + duplex_commands.AppendAudio, + duplex_commands.Commit, + duplex_commands.CancelResponse, + duplex_commands.AckPlayback, + ] + append, commit, cancel, ack = handle.commands + assert append.is_speech is True + assert append.audio_end_ms == 100 + assert append.event_id + assert commit.final is True and commit.create_response is True + assert cancel.response_id == "resp-1" + assert ack.played_ms == 120 and ack.response_id == "resp-1" + + +@pytest.mark.asyncio +async def test_inline_session_close_maps_to_handle_close(): + client, omni = make_inline_client() + async with client: + pass + handle = omni.handles[0] + assert handle.closed + assert handle.close_calls >= 1 + assert not any(isinstance(command, duplex_commands.CloseSession) for command in handle.commands) + # The typed session.closed reached the client as the expected end. + with pytest.raises(DuplexSessionClosedError): + await client.send({"type": "session.heartbeat"}) + + +@pytest.mark.asyncio +async def test_inline_client_forwards_heartbeats_and_drops_event_acks(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + await client.send({"type": "session.heartbeat"}) + await client.send({"type": "session.event_ack", "server_event_seq": 3}) + # A heartbeat is a lease touch the engine must see; the ack is a journal concern. + assert [type(command).__name__ for command in handle.commands] == ["Heartbeat"] + + +@pytest.mark.asyncio +async def test_inline_malformed_payload_surfaces_as_error_event(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + # Subscribe before sending: the translation error is dispatched inline. + first_event = asyncio.create_task(client.wait_for("error", timeout_s=5.0)) + await asyncio.sleep(0) + event_id = await client.send({"type": "playback.ack"}) # played_ms missing + event = await first_event + assert isinstance(event, ErrorEvent) + assert event.code == "bad_event" + assert event.related_event_id == event_id + assert handle.commands == [] + with pytest.raises(DuplexProtocolError) as excinfo: + async for _ in client.responses(): + pass + assert excinfo.value.code == "bad_event" + + +@pytest.mark.asyncio +async def test_inline_teardown_cancels_pump_and_closes_handle(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + await _drain(lambda: handle.consumer_active) + pump = client._pump_task + assert pump is not None and not pump.done() + assert pump.done() + assert not handle.consumer_active + assert handle.closed + + +@pytest.mark.asyncio +async def test_inline_send_after_handle_closed_raises_connection_error(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + handle.closed = True # the engine ended the session without an event yet + with pytest.raises(DuplexConnectionError, match="not open"): + await client.send({"type": "input_audio_buffer.clear"}) + handle.closed = False + + # --------------------------------------------------------------------------- # Collector @@ -600,8 +886,7 @@ def test_build_realtime_url_with_model_extra_query(): url = build_realtime_url( "ws://localhost:8099/v1/realtime?custom=1", "openbmb/MiniCPM-o-4_5", - session_id="session-a", - extra_query={"native_duplex": "1"}, + extra_query={"probe": "1"}, ) query = parse_qs(urlsplit(url).query) @@ -609,8 +894,7 @@ def test_build_realtime_url_with_model_extra_query(): "custom": ["1"], "duplex": ["1"], "model": ["openbmb/MiniCPM-o-4_5"], - "native_duplex": ["1"], - "session_id": ["session-a"], + "probe": ["1"], } @@ -619,12 +903,11 @@ def test_build_realtime_url_resume_only_when_autostart_disabled(): "ws://localhost:8099/v1/realtime?duplex=1", "openbmb/MiniCPM-o-4_5", autostart=False, - extra_query={"native_duplex": "1"}, ) query = parse_qs(urlsplit(url).query) assert query["autostart"] == ["0"] - assert query["native_duplex"] == ["1"] + assert "session_id" not in query def test_event_collector_partitions_audio_by_response(): @@ -989,6 +1272,7 @@ def test_duplex_unit_boundary_and_residual_math(): assert chunk_period_ms([{"type": "session.created", "session": {}}]) == 1000 +@pytest.mark.asyncio async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -1011,16 +1295,16 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame(): sock.feed(SESSION_CLOSED) -def test_build_realtime_url_native_duplex_flag_and_http_scheme(): - url = build_realtime_url("http://localhost:8099/v1/realtime", None, native_duplex=None) +def test_build_realtime_url_rewrites_http_scheme_and_rejects_others(): + url = build_realtime_url("http://localhost:8099/v1/realtime", None) parts = urlsplit(url) assert parts.scheme == "ws" assert parse_qs(parts.query) == {"duplex": ["1"]} - url = build_realtime_url("https://host/v1/realtime", "m", native_duplex=False) + url = build_realtime_url("https://host/v1/realtime", "m", autostart=True) parts = urlsplit(url) assert parts.scheme == "wss" - assert parse_qs(parts.query)["native_duplex"] == ["0"] + assert parse_qs(parts.query) == {"duplex": ["1"], "model": ["m"], "autostart": ["1"]} with pytest.raises(ValueError): build_realtime_url("ftp://host/v1/realtime", "m") diff --git a/tests/clients/test_model_session_configs.py b/tests/clients/test_model_session_configs.py index 00741d3c6..080fea879 100644 --- a/tests/clients/test_model_session_configs.py +++ b/tests/clients/test_model_session_configs.py @@ -18,13 +18,13 @@ from vllm_omni.clients.personaplex import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def test_minicpmo45_session_config_matches_native_duplex_deployment(): +def test_minicpmo45_session_config_matches_duplex_deployment(): config = create_minicpmo45_session_config(ref_audio="data:audio/wav;base64,AAA=", temperature=0.0) payload = config.to_session_payload(model="openbmb/MiniCPM-o-4_5") assert payload["ref_audio"] == "data:audio/wav;base64,AAA=" assert payload["overlap_policy"] == "listen_only" assert payload["playback_commit_policy"] == "ack_only" - assert payload["extra_body"]["native_duplex"] is True + assert "native_duplex" not in payload["extra_body"] assert payload["extra_body"]["auto_response"] is True assert payload["temperature"] == 0.0 @@ -41,8 +41,9 @@ def test_personaplex_session_config_matches_deployment(): @pytest.mark.parametrize( "duplex_package", [ + # PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up + # PRs that port them to the duplex plugin framework. "vllm_omni.model_executor.models.minicpmo_4_5.duplex", - "vllm_omni.model_executor.models.personaplex.duplex", ], ) def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str): diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5.json b/tests/dfx/perf/tests/test_minicpmo_4_5.json index 7d7ca08ad..b384d61a0 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5.json @@ -26,8 +26,8 @@ { "dataset_name": "seed-tts", "dataset_path": "zhaochenyang20/seed-tts-eval", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", + "backend": "openai-realtime-tts", + "endpoint": "/v1/realtime", "num_prompts": [ 32, 64, @@ -42,73 +42,7 @@ "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", - "extra_body": { - "modalities": [ - "text", - "audio" - ], - "chat_template_kwargs": { - "enable_thinking": false, - "use_tts_template": true - } - }, - "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration", - "baseline": { - "H100": { - "request_throughput": [ - 0.9731, - 1.1986, - 1.686 - ], - "mean_ttft_ms": [ - 152.5665, - 270.2559, - 301.455 - ], - "mean_e2el_ms": [ - 1027.2953, - 3306.3735, - 4686.5366 - ], - "mean_audio_ttfp_ms": [ - 458.1553, - 2030.2335, - 2068.6233 - ], - "mean_audio_rtf": [ - 0.2357, - 0.791, - 1.0324 - ] - }, - "A3": { - "request_throughput": [ - 0.5383, - 0.6042, - 0.7547 - ], - "mean_ttft_ms": [ - 333.2633, - 533.8656, - 547.3388 - ], - "mean_e2el_ms": [ - 1857.2154, - 6600.5101, - 10450.6811 - ], - "mean_audio_ttfp_ms": [ - 986.4666, - 3411.0754, - 3352.7515 - ], - "mean_audio_rtf": [ - 0.4423, - 1.5734, - 2.3024 - ] - } - } + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" } ] }, @@ -138,8 +72,8 @@ { "dataset_name": "seed-tts", "dataset_path": "zhaochenyang20/seed-tts-eval", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", + "backend": "openai-realtime-tts", + "endpoint": "/v1/realtime", "num_prompts": [ 32, 64, @@ -154,16 +88,6 @@ "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", - "extra_body": { - "modalities": [ - "text", - "audio" - ], - "chat_template_kwargs": { - "enable_thinking": false, - "use_tts_template": true - } - }, "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" } ] diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json index bb0599419..a145459ed 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json @@ -43,43 +43,7 @@ "extra_body": { "save_duplex_request_metrics": true }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration", - "baseline": { - "H100": { - "request_throughput": [ - 0.3038 - ], - "mean_ttft_ms": [ - 139.997 - ], - "mean_e2el_ms": [ - 3269.8807 - ], - "mean_audio_ttfp_ms": [ - 456.073 - ], - "mean_audio_rtf": [ - 0.2254 - ] - }, - "A3": { - "request_throughput": [ - 0.1594 - ], - "mean_ttft_ms": [ - 274.164 - ], - "mean_e2el_ms": [ - 6250.7696 - ], - "mean_audio_ttfp_ms": [ - 994.422 - ], - "mean_audio_rtf": [ - 0.46163 - ] - } - } + "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration" } ] } diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json index 26dc54aab..4fd78d681 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json @@ -26,8 +26,8 @@ { "dataset_name": "seed-tts", "dataset_path": "zhaochenyang20/seed-tts-eval", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", + "backend": "openai-realtime-tts", + "endpoint": "/v1/realtime", "num_prompts": [ 10, 40 @@ -40,41 +40,7 @@ "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", - "extra_body": { - "modalities": [ - "text", - "audio" - ], - "chat_template_kwargs": { - "enable_thinking": false, - "use_tts_template": true - } - }, - "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration", - "baseline": { - "H100": { - "request_throughput": [ - 0.9731, - 1.1986 - ], - "mean_ttft_ms": [ - 152.5665, - 270.2559 - ], - "mean_e2el_ms": [ - 1027.2953, - 3306.3735 - ], - "mean_audio_ttfp_ms": [ - 458.1553, - 2030.2335 - ], - "mean_audio_rtf": [ - 0.2357, - 0.791 - ] - } - } + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" } ] } diff --git a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py index c833ae7a1..e0b0a6c53 100644 --- a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py +++ b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py @@ -39,7 +39,8 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage from vllm_omni.engine.stage_init_utils import StageMetadata -from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py index 2e5ceb8e2..7780f0475 100644 --- a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py +++ b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py @@ -42,7 +42,8 @@ from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_client import StagePoolClient from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.engine.stage_pool import StagePool -from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index 1dad9face..ec222bcd3 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -90,6 +90,7 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) -> Args: source: Full file source code. class_name: Enclosing class name, or None for module-level functions. + The method may be inherited from a base class defined in the same file. func_name: Function/method name. Returns: @@ -97,13 +98,24 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) -> """ tree = ast.parse(source) if class_name is not None: - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and node.name == class_name: - for item in node.body: - if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name: - result = ast.get_source_segment(source, item) - assert result is not None, f"{class_name}.{func_name} source not found" - return result + # Follow the bases: ``Orchestrator`` inherits its loops and helpers from + # ``OrchestratorBase`` in the same module, and the lookup stays anchored + # on the public class rather than on where the split happens to fall. + classes = {node.name: node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)} + pending = [class_name] + seen: set[str] = set() + while pending: + name = pending.pop(0) + if name in seen or name not in classes: + continue + seen.add(name) + node = classes[name] + for item in node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name: + result = ast.get_source_segment(source, item) + assert result is not None, f"{class_name}.{func_name} source not found" + return result + pending.extend(base.id for base in node.bases if isinstance(base, ast.Name)) else: for node in ast.walk(tree): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: diff --git a/tests/e2e/features/fullduplex/test_core_contracts.py b/tests/e2e/features/fullduplex/test_core_contracts.py index 12a87640f..45493745c 100644 --- a/tests/e2e/features/fullduplex/test_core_contracts.py +++ b/tests/e2e/features/fullduplex/test_core_contracts.py @@ -14,7 +14,7 @@ import dataclasses import pytest -from vllm_omni.engine.duplex.messages import DuplexFence +from vllm_omni.engine.duplex.contracts import DuplexFence from vllm_omni.experimental.fullduplex.core import protocol as ev from vllm_omni.experimental.fullduplex.core.adapter import ( DuplexAdapter, @@ -32,12 +32,15 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_fence_is_frozen_value_identity(): + # The fence carries (session_id, epoch, turn_id). response_seq was + # dropped with the incarnation concept: response identity is the session's + # active_response_id, not a field of the fence. f = DuplexFence(session_id="s1") - assert (f.epoch, f.turn_id, f.response_seq) == (0, 0, 0) + assert (f.epoch, f.turn_id) == (0, 0) with pytest.raises(dataclasses.FrozenInstanceError): f.epoch = 1 # type: ignore[misc] - assert DuplexFence("s1", 1, 2, 3) == DuplexFence("s1", 1, 2, 3) - assert DuplexFence("s1", 1, 2, 3) != DuplexFence("s1", 2, 2, 3) + assert DuplexFence("s1", 1, 2) == DuplexFence("s1", 1, 2) + assert DuplexFence("s1", 1, 2) != DuplexFence("s1", 2, 2) def test_session_identity_rules(): diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py index f5bba11e4..0f9b4af51 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py @@ -71,15 +71,8 @@ from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402 from vllm_omni.metrics.definitions import compute_audio_rtf # noqa: E402 -def _url_with_model(url, model, *, autostart=None, session_id=None): - # This driver speaks the MiniCPM-o native duplex wire contract. - return build_realtime_url( - url, - model, - autostart=autostart, - session_id=session_id, - extra_query={"native_duplex": "1"}, - ) +def _url_with_model(url, model, *, autostart=None): + return build_realtime_url(url, model, autostart=autostart) _read_wav_pcm16 = read_pcm16_wav @@ -630,7 +623,6 @@ def _session_update_event(args: DemoArgs) -> dict[str, object]: "playback_commit_policy": "ack_only", "extra_body": { "auto_response": True, - "native_duplex": True, "force_listen_count": 0, }, } @@ -646,9 +638,6 @@ def _session_update_event(args: DemoArgs) -> dict[str, object]: "type": "session.update", "session": session_payload, } - session_id = getattr(args, "session_id", None) - if session_id: - session_payload["session_id"] = session_id return event @@ -1462,7 +1451,6 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=getattr(args, "session_id", None), ) state = DemoState() stop = asyncio.Event() @@ -1674,6 +1662,10 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: ) unexpected_error_events = _unexpected_error_events(state) continuous_input_ok = not continuous_input or state.count("input_audio_buffer.committed") == 0 + # The server allocates every session id; label metrics with the one it announced. + created = state.events[0] if state.events else {} + created_session = created.get("session") if isinstance(created, dict) else None + server_session_id = created_session.get("id") if isinstance(created_session, dict) else None result = { "ok": terminal_activity_ok and state.count("session.closed") > 0 @@ -1711,8 +1703,8 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: "input_transcription_ok": input_transcription_ok, "completed_response_ids": completed_response_ids, "response_timings": state.response_timing_summaries(), - "request_metrics": state.session_request_metrics(session_id=args.session_id), - "session_metrics": state.session_metric_summary(session_id=args.session_id), + "request_metrics": state.session_request_metrics(session_id=server_session_id), + "session_metrics": state.session_metric_summary(session_id=server_session_id), "lifecycle_counts_ok": lifecycle_counts_ok, "validation_mode": validation_mode, "scenario": scenario, @@ -1773,7 +1765,6 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") - parser.add_argument("--session-id", help="Use an explicit public session ID, including for close/reopen tests.") parser.add_argument("--input-wav", required=True) parser.add_argument("--ref-audio", help="Optional WAV used as the MiniCPM-o voice prompt.") parser.add_argument( diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py index b5d04097b..19b4b8107 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py @@ -17,7 +17,6 @@ import hashlib import io import json import sys -import uuid import wave from pathlib import Path @@ -266,8 +265,6 @@ async def _run_demo_process( str(args.chunk_ms), "--timeout-s", str(args.timeout_s), - "--session-id", - f"duplex-pair-{label}-{uuid.uuid4().hex}", ] if args.no_realtime_pacing: command.append("--no-realtime-pacing") diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py index 32e6e01da..c13f8ca7f 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py @@ -21,7 +21,7 @@ from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed SCRIPT_DIR = Path(__file__).resolve().parent -REPO_ROOT = SCRIPT_DIR.parents[3] +REPO_ROOT = SCRIPT_DIR.parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -42,6 +42,15 @@ def _url_with_model(*args, **kwargs) -> str: return _scenario_module()._url_with_model(*args, **kwargs) +def _created_session_id(created: dict[str, object]) -> str: + """The server-allocated session id announced by ``session.created``.""" + session = created.get("session") + session_id = session.get("id") if isinstance(session, dict) else None + if not isinstance(session_id, str) or not session_id: + raise RuntimeError("session.created omitted the server-allocated session id") + return session_id + + async def run_demo(args): return await _scenario_module().run_demo(args) @@ -217,7 +226,6 @@ async def _open_admission_session( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=session_id, ) ws = await websockets.connect(url, max_size=64 * 1024 * 1024) await ws.send( @@ -225,10 +233,8 @@ async def _open_admission_session( { "type": "session.update", "session": { - "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **({"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) else {}), }, } @@ -255,12 +261,10 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, for index in range(limit): accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}")) - overflow_id = f"{prefix}-overflow" overflow_url = _url_with_model( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=overflow_id, ) async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow: await overflow.send( @@ -268,10 +272,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, { "type": "session.update", "session": { - "session_id": overflow_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **( {"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) @@ -326,7 +328,6 @@ async def _resume_probe( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=session_id, ) async with websockets.connect(url, max_size=64 * 1024 * 1024) as first: await first.send( @@ -334,10 +335,8 @@ async def _resume_probe( { "type": "session.update", "session": { - "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **( {"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) @@ -348,10 +347,10 @@ async def _resume_probe( ) ) created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s) + session_id = _created_session_id(created) token = created.get("resume_token") - incarnation = created.get("incarnation") generation = created.get("attachment_generation") - if not isinstance(token, str) or not isinstance(incarnation, int): + if not isinstance(token, str): raise RuntimeError("session.created omitted resumable credentials") last_seq = max(_server_event_sequences(first_events), default=0) @@ -366,7 +365,6 @@ async def _resume_probe( { "type": "session.resume", "session_id": session_id, - "incarnation": incarnation, "resume_token": token, "last_received_server_event_seq": last_seq, } @@ -429,7 +427,6 @@ async def _takeover_probe( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=session_id, ) first = await websockets.connect(url, max_size=64 * 1024 * 1024) second = None @@ -439,10 +436,8 @@ async def _takeover_probe( { "type": "session.update", "session": { - "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **( {"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) @@ -453,10 +448,10 @@ async def _takeover_probe( ) ) created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s) + session_id = _created_session_id(created) token = created.get("resume_token") - incarnation = created.get("incarnation") generation = created.get("attachment_generation") - if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int): + if not isinstance(token, str) or not isinstance(generation, int): raise RuntimeError("session.created omitted takeover credentials") last_seq = max(_server_event_sequences(first_events), default=0) @@ -466,7 +461,6 @@ async def _takeover_probe( { "type": "session.resume", "session_id": session_id, - "incarnation": incarnation, "resume_token": token, "last_received_server_event_seq": last_seq, } diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py index 5e0cfff3a..2eee43fae 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py @@ -18,7 +18,6 @@ import hashlib import io import json import sys -import uuid import wave from pathlib import Path @@ -412,8 +411,6 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]: str(args.chunk_ms), "--timeout-s", str(args.timeout_s), - "--session-id", - f"duplex-soft-interrupt-{uuid.uuid4().hex}", ] command.extend(["--ref-audio", str(_canonical_path(args.ref_audio))]) temperature = getattr(args, "temperature", None) diff --git a/tests/e2e/online_serving/test_duplex_client_live.py b/tests/e2e/online_serving/test_duplex_client_live.py index d90a8a683..28044cdd1 100644 --- a/tests/e2e/online_serving/test_duplex_client_live.py +++ b/tests/e2e/online_serving/test_duplex_client_live.py @@ -12,7 +12,6 @@ incremental playback acks, a forced transport drop recovered by automatic from __future__ import annotations import asyncio -import uuid from pathlib import Path import pytest @@ -47,7 +46,6 @@ async def _run_live_session(*, url: str, model: str, ref_audio: Path, input_wav: url, model=model, config=config, - session_id=f"duplex-client-live-{uuid.uuid4().hex}", reconnect=ReconnectPolicy(max_attempts=5, backoff_s=(0.25, 1.0)), # Short interval so the periodic heartbeat path runs during the # session — including across the forced transport drop below. diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py index c04d3acec..59b40f5e5 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py @@ -6,8 +6,8 @@ from __future__ import annotations import asyncio +import base64 import json -import uuid from pathlib import Path import pytest @@ -90,25 +90,16 @@ async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: f async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[dict[str, object]]: - session_id = f"duplex-ci-protocol-{uuid.uuid4().hex}" - websocket_url = build_realtime_url( - url, - model, - autostart=False, - session_id=session_id, - extra_query={"native_duplex": "1"}, - ) + websocket_url = build_realtime_url(url, model, autostart=False) async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: await ws.send( json.dumps( { "type": "session.update", "session": { - "session_id": session_id, "model": model, "modalities": ["audio", "text"], "ref_audio": _ref_audio_data_url(str(ref_audio)), - "extra_body": {"native_duplex": True}, }, } ) @@ -123,6 +114,170 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[ return events +async def _run_text_only_response_create( + *, + url: str, + model: str, + ref_audio: Path, + text: str, + timeout_s: float = 120.0, +) -> dict[str, object]: + """Drive a text prompt with no audio and return the server's answer. + + A Realtime client asks for speech from text with a user message item + followed by ``response.create``. A model-native duplex session generates + from audio units, so this returns whatever the server answers rather than + assuming a response. + """ + websocket_url = build_realtime_url(url, model, autostart=False) + async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: + await ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "model": model, + "modalities": ["audio", "text"], + "output_audio_format": "pcm16", + "turn_detection": None, + "temperature": 0.0, + "ref_audio": _ref_audio_data_url(str(ref_audio)), + "extra_body": {"auto_response": False}, + }, + } + ) + ) + await _receive_protocol_events(ws, {"session.created", "session.updated"}, timeout_s=60) + await ws.send( + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": text}], + }, + } + ) + ) + await ws.send(json.dumps({"type": "response.create"})) + + async def receive_outcome() -> dict[str, object]: + while True: + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if isinstance(event, dict) and event.get("type") in {"error", "response.done"}: + return event + + return await asyncio.wait_for(receive_outcome(), timeout=timeout_s) + + +async def _run_seeded_text_to_audio( + *, + url: str, + model: str, + ref_audio: Path | None, + text: str, + modalities: tuple[str, ...] = ("audio", "text"), + silence_seconds: float = 12.0, + timeout_s: float = 180.0, +) -> dict[str, object]: + """Speak a seeded text: the duplex route's text-to-speech shape. + + A model-native session takes its text once, in the session context + (``duplex_initial_user_text``), and then generates per audio unit. Silence + carries no content of its own, so it only advances the clock and lets the + model answer the seeded turn. + """ + websocket_url = build_realtime_url(url, model, autostart=False) + audio_bytes = 0 + transcript: list[str] = [] + output_text: list[str] = [] + seen: list[str] = [] + session_payload: dict[str, object] = { + "model": model, + "modalities": list(modalities), + "input_audio_format": "pcm16", + "output_audio_format": "pcm16", + "turn_detection": None, + "temperature": 0.0, + "extra_body": { + "auto_response": True, + "force_listen_count": 0, + "duplex_initial_user_text": text, + }, + } + if ref_audio is not None: + session_payload["ref_audio"] = _ref_audio_data_url(str(ref_audio)) + async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: + await ws.send(json.dumps({"type": "session.update", "session": session_payload})) + + async def reader() -> None: + nonlocal audio_bytes + while True: + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if not isinstance(event, dict): + continue + event_type = event.get("type") + if isinstance(event_type, str): + seen.append(event_type) + delta = event.get("delta") + if event_type == "response.output_audio.delta" and isinstance(delta, str): + audio_bytes += len(base64.b64decode(delta)) + elif event_type == "response.output_audio_transcript.delta" and isinstance(delta, str): + transcript.append(delta) + elif event_type == "response.output_text.delta" and isinstance(delta, str): + output_text.append(delta) + elif event_type == "error": + raise AssertionError(f"seeded text turn received an error: {event}") + + reader_task = asyncio.create_task(reader()) + silence = bytes(2 * 16_000 * 200 // 1000) + sent_ms = 0 + try: + while sent_ms < silence_seconds * 1000 and "response.done" not in seen: + sent_ms += 200 + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(silence).decode("ascii"), + "input_audio_format": "pcm16", + "sample_rate_hz": 16_000, + "duration_ms": 200, + "audio_end_ms": sent_ms, + } + ) + ) + await asyncio.sleep(0.2) + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + while loop.time() < deadline and "response.done" not in seen: + await asyncio.sleep(0.5) + if reader_task.done(): + await reader_task + # Close explicitly. Dropping the socket parks the session in its + # disconnect grace, where it keeps holding an admission slot and + # starves the tests that follow. + await ws.send(json.dumps({"type": "session.close"})) + deadline = loop.time() + 30 + while loop.time() < deadline and "session.closed" not in seen: + await asyncio.sleep(0.25) + finally: + reader_task.cancel() + return { + "audio_bytes": audio_bytes, + "transcript": "".join(transcript), + "output_text": "".join(output_text), + "event_types": seen, + } + + @pytest.mark.core_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) @@ -205,6 +360,215 @@ def test_duplex_single_session_video_input(omni_server, tmp_path: Path) -> None: _assert_session_metrics(result["session_metrics"], expected_count=2) +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +@pytest.mark.parametrize( + ("locale", "text"), + [ + ("en", "Please say exactly: the quick brown fox jumps over the lazy dog."), + ("zh", "请朗读：今天天气很好，我们一起去公园散步。"), + ], +) +def test_duplex_seeded_text_to_audio(omni_server, locale: str, text: str) -> None: + """text -> audio over the duplex route, in both locales. + + This is the case the deleted turn-based ``text -> audio`` tests covered. + It is also the shape the Seed-TTS Realtime backend uses, so it guards the + benchmark path as well as the modality. + """ + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text=text, + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + assert int(result["audio_bytes"]) > 0, f"{locale} seeded text produced no audio" + assert str(result["transcript"]).strip(), f"{locale} seeded text produced audio with no transcript" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_seeded_text_to_text_needs_no_reference_voice(omni_server) -> None: + """text -> text, the one duplex session that opens without a reference voice. + + ``ref_audio`` is required only when the session asks for audio output, so a + ``modalities: ["text"]`` session is the duplex equivalent of the deleted + turn-based ``text -> text`` case. The model is model-native and still + speaks its answer, so this asserts the text side and the absence of the + ``ref_audio_required`` rejection, not the absence of audio. + """ + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=None, + text="What is the capital of France? Answer in one short sentence.", + modalities=("text",), + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + produced_text = str(result["output_text"]) or str(result["transcript"]) + assert produced_text.strip(), f"text-only session produced nothing: {result['event_types']}" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_seeded_text_to_long_audio_output(omni_server) -> None: + """A longer answer, so Code2Wav runs across many frames rather than one. + + Replaces the deleted ``text_to_audio_long_output`` case. The assertion is + on the amount of audio rather than its content: the point is that a long + answer streams to completion, not that the model says any given word. + """ + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="What is the capital of China? Answer in about 40 words.", + silence_seconds=30.0, + timeout_s=240.0, + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + # 24 kHz mono pcm16: 2 s of speech is 96000 bytes, comfortably more than a + # single Code2Wav frame and well under a 40-word answer. + assert int(result["audio_bytes"]) > 96_000, f"expected a long answer, got {result['audio_bytes']} bytes" + assert str(result["transcript"]).strip() + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_seeded_long_form_generation(omni_server) -> None: + """Long-form Chinese generation, replacing the deleted long-form case.""" + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="帮我讲一个100字的故事。", + silence_seconds=30.0, + timeout_s=240.0, + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + assert int(result["audio_bytes"]) > 96_000, f"expected long-form audio, got {result['audio_bytes']} bytes" + assert str(result["transcript"]).strip() + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_sequential_sessions_are_independent(omni_server) -> None: + """A second session must answer its own prompt, not replay the first one. + + Replaces the deleted ``sequential_requests_independent`` case. Sessions are + engine-resident now, so this is the check that one session's context does + not leak into the next one on the same stages. + """ + first = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="What is the capital of France? Answer in one short sentence.", + ) + ) + second = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="What is the capital of China? Answer in one short sentence.", + ) + ) + + for result in (first, second): + assert "response.done" in result["event_types"], result["event_types"] + assert int(result["audio_bytes"]) > 0 + assert str(result["transcript"]).strip() + + # Different prompts must not produce the same answer: that would mean the + # second session inherited the first one's context. + assert str(first["transcript"]).strip() != str(second["transcript"]).strip() + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_text_only_response_create_is_rejected(omni_server) -> None: + """A duplex session generates from audio: a text-only prompt is refused, not ignored. + + Every session on this route is model-native, and such a model decides per + audio unit whether to speak: conversation items are context, not a turn. + Refusing is the point -- opening a response the model never fills costs the + caller the session's whole idle timeout to discover. A text prompt reaches + the model through the session's seeded opening turn + (``initial_user_text``) instead, which is what ``/v1/chat/completions`` + uses and what the Seed-TTS Realtime backend has to be ported to. + """ + outcome = asyncio.run( + _run_text_only_response_create( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="Please say: the quick brown fox jumps over the lazy dog.", + ) + ) + + assert outcome.get("type") == "error", outcome + error = outcome.get("error") + assert isinstance(error, dict), outcome + assert error.get("code") == "text_only_turn_unsupported", outcome + assert error.get("type") == "invalid_request_error", outcome + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_single_session_still_image_input(omni_server, tmp_path: Path) -> None: + """A spoken turn with one still image attached. + + The still-image case is a one-element frame list, which the client repeats + on every unit of the turn: unlike the camera track it never advances, so + this covers the image path the turn-based ``image -> text + audio`` request + used to cover. + """ + args = demo_args( + omni_server=omni_server, + input_wav=validated_input_wav(), + ref_audio=resolve_ref_audio(), + output_dir=tmp_path / "still_image", + ) + args.turns = 2 + args.turn_duration_ms = [args.first_turn_ms] * args.turns + args.video_frames_b64 = duplex_camera_frames(seconds=1, cache_dir=tmp_path / "still")[:1] + + result = asyncio.run(run_demo(args)) + + assert result["ok"] is True + assert result["video_frame_count"] == 1 + assert result["video_stacked_frame_count"] == 0 + _assert_positive_int(result["audio_delta_count"]) + assert result["done_count"] == 2 + assert result["error_count"] == 0 + assert result["all_audio_responses_have_transcript"] is True + assert result["transcript_delta_done_ok"] is True + _assert_request_metrics(result["request_metrics"], expected_count=2) + _assert_session_metrics(result["session_metrics"], expected_count=2) + + @pytest.mark.advanced_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py index 628040001..f931fb95d 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -30,7 +30,7 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import run_soft_interrupt, ) from tests.helpers.mark import hardware_test -from vllm_omni.entrypoints.duplex.server_vad import ( +from vllm_omni.engine.duplex.vad import ( SILERO_VAD_FILENAME, SILERO_VAD_REPO_ID, SILERO_VAD_REVISION, diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index ecd885dd6..a23ce73e8 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -152,7 +152,6 @@ def test_realtime_duplex_demo_pair_launches_demo_processes_concurrently(tmp_path "parser.add_argument('--output-dir')", "parser.add_argument('--chunk-ms')", "parser.add_argument('--timeout-s')", - "parser.add_argument('--session-id')", "parser.add_argument('--ref-audio')", "parser.add_argument('--require-audio', action='store_true')", "parser.add_argument('--no-realtime-pacing', action='store_true')", @@ -1027,7 +1026,7 @@ def test_realtime_duplex_demo_resolves_distinct_turn_inputs(): demo._turn_input_paths(primary, ["second.wav"], turns=3) -def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect(): +def test_realtime_duplex_demo_url_carries_no_model_opt_in_or_session_id(): demo = _load_demo_module() url = demo._url_with_model( @@ -1036,23 +1035,12 @@ def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect() ) query = parse_qs(urlsplit(url).query) - assert query["native_duplex"] == ["1"] + assert query["duplex"] == ["1"] + assert "native_duplex" not in query + assert "session_id" not in query -def test_realtime_duplex_demo_explicit_session_id_reaches_autostart_query(): - demo = _load_demo_module() - - url = demo._url_with_model( - "ws://localhost:8099/v1/realtime?duplex=1", - "openbmb/MiniCPM-o-4_5", - session_id="reopen-e2e", - ) - - query = parse_qs(urlsplit(url).query) - assert query["session_id"] == ["reopen-e2e"] - - -def test_realtime_duplex_demo_session_update_uses_explicit_session_id(): +def test_realtime_duplex_demo_session_update_leaves_session_id_to_the_server(): demo = _load_demo_module() event = demo._session_update_event( @@ -1066,8 +1054,8 @@ def test_realtime_duplex_demo_session_update_uses_explicit_session_id(): assert event["type"] == "session.update" assert "session_id" not in event - assert event["session"]["session_id"] == "reopen-e2e" - assert event["session"]["extra_body"]["native_duplex"] is True + assert "session_id" not in event["session"] + assert "native_duplex" not in event["session"]["extra_body"] def test_realtime_duplex_demo_response_required_uses_deterministic_sampling(): diff --git a/tests/entrypoints/duplex/test_duplex_audio.py b/tests/engine/duplex/test_duplex_audio.py similarity index 95% rename from tests/entrypoints/duplex/test_duplex_audio.py rename to tests/engine/duplex/test_duplex_audio.py index d0be02072..ca358abb0 100644 --- a/tests/entrypoints/duplex/test_duplex_audio.py +++ b/tests/engine/duplex/test_duplex_audio.py @@ -6,7 +6,7 @@ import base64 import numpy as np import pytest -from vllm_omni.entrypoints.duplex.audio import convert_input_audio_with_rate +from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate pytestmark = [pytest.mark.core_model, pytest.mark.cpu] diff --git a/tests/engine/duplex/test_duplex_commands.py b/tests/engine/duplex/test_duplex_commands.py new file mode 100644 index 000000000..ad45c8355 --- /dev/null +++ b/tests/engine/duplex/test_duplex_commands.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI Realtime client events -> typed ``DuplexCommand`` objects and their mailbox payloads.""" + +from __future__ import annotations + +import base64 + +import numpy as np +import pytest + +from vllm_omni.engine.duplex.commands import ( + REALTIME_COMMAND_TYPES, + AckPlayback, + AppendAudio, + AppendText, + BargeIn, + CancelInput, + CancelResponse, + ClearInput, + ClearOutputAudio, + CloseSession, + Commit, + CreateItem, + CreateResponse, + DeleteItem, + DuplexCommand, + DuplexCommandError, + Heartbeat, + SignalTurn, + TruncateItem, + UpdateSession, + command_from_realtime, +) +from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + translate_realtime_command, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_LOUD_PCM16 = base64.b64encode(b"\x00\x10" * 8).decode("ascii") +_SILENT_PCM16 = base64.b64encode(b"\x00\x00" * 8).decode("ascii") +_JPEG_FRAME = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 12).decode("ascii") + +_MINIMAL_PAYLOADS: dict[str, tuple[dict[str, object], type[DuplexCommand]]] = { + "input_audio_buffer.append": ({"audio": _LOUD_PCM16}, AppendAudio), + "input_audio_buffer.commit": ({}, Commit), + "input_audio_buffer.clear": ({}, ClearInput), + "output_audio_buffer.clear": ({}, ClearOutputAudio), + "response.create": ({}, CreateResponse), + "response.cancel": ({}, CancelResponse), + "conversation.item.create": ({"item": {"type": "message", "role": "user", "content": []}}, CreateItem), + "conversation.item.delete": ({"item_id": "item_1"}, DeleteItem), + "conversation.item.truncate": ({"item_id": "item_1", "audio_end_ms": 100}, TruncateItem), + "session.update": ({"session": {"instructions": "hi"}}, UpdateSession), + "playback.ack": ({"played_ms": 10}, AckPlayback), + "session.heartbeat": ({}, Heartbeat), + "session.close": ({}, CloseSession), + "turn.signal": ({"event": "user_started"}, SignalTurn), + "input.text.append": ({"text": "hello"}, AppendText), + "input.cancel": ({}, CancelInput), + "barge_in": ({}, BargeIn), +} + + +def test_every_realtime_command_type_has_a_mapping_case(): + assert set(_MINIMAL_PAYLOADS) == set(REALTIME_COMMAND_TYPES) + + +@pytest.mark.parametrize("event_type", sorted(REALTIME_COMMAND_TYPES)) +def test_command_from_realtime_maps_each_type_to_its_dataclass_and_propagates_event_id(event_type: str): + body, expected_cls = _MINIMAL_PAYLOADS[event_type] + + command = command_from_realtime({"type": event_type, "event_id": f"evt-{event_type}", **body}) + + assert type(command) is expected_cls + assert command.event_id == f"evt-{event_type}" + payload = command.payload() + assert payload["type"] == expected_cls.type + assert payload["realtime_event_id"] == f"evt-{event_type}" + assert "event_id" not in payload + + anonymous = command_from_realtime({"type": event_type, **body}) + assert anonymous.event_id is None + assert "realtime_event_id" not in anonymous.payload() + + +@pytest.mark.parametrize( + ("alias", "expected_cls", "body"), + [ + ("push_text", AppendText, {"text": "hello"}), + ("input_text.append", AppendText, {"text": "hello"}), + ("signal_turn", SignalTurn, {"event": "user_started"}), + ("close_session", CloseSession, {}), + ("close", CloseSession, {}), + ("audio.playback_ack", AckPlayback, {"played_ms": 10}), + ("input.commit", Commit, {}), + ], +) +def test_pre_realtime_client_event_aliases_map_to_their_canonical_command( + alias: str, expected_cls: type[DuplexCommand], body: dict[str, object] +): + assert type(command_from_realtime({"type": alias, **body})) is expected_cls + + +def test_append_audio_decodes_and_converts_pcm16_to_16k_float32(): + command = command_from_realtime( + {"type": "input_audio_buffer.append", "audio": _LOUD_PCM16, "format": "pcm16", "sample_rate_hz": 16000} + ) + + assert isinstance(command, AppendAudio) + assert command.format == "pcm_f32le" + assert command.sample_rate_hz == 16000 + samples = np.frombuffer(command.audio, dtype=" 8 * 4 + + passthrough = translate_realtime_command( + {"type": "input_audio_buffer.append", "audio": base64.b64encode(np.zeros(4, dtype=" None: - self.calls = [] - - def execute(self, key, message, **kwargs): - self.calls.append((key, message, kwargs)) - return DuplexControlResultMessage( - control_id=message.control_id, - fence=message.fence, - operation="touch" if isinstance(message, TouchDuplexSessionMessage) else "resume", - session_id=message.session_id, - ok=True, - stage_results=[], - ) - - -class _UnsupportedTransport: - def execute(self, key, message, **kwargs): - del key, kwargs - return DuplexControlResultMessage( - control_id=message.control_id, - fence=message.fence, - operation="touch", - session_id=message.session_id, - ok=False, - stage_results=[{"result": {"supported": False}}], - unsupported_count=1, - error=DuplexControlError(code="invalid_capability", message="unsupported", retryable=False), - accepted_fence=DuplexFence(message.session_id, epoch=2), - lease_generation=7, - ) - - -def test_control_client_routes_touch_and_resume_by_control_id() -> None: - transport = _Transport() - control_ids = iter(("touch-id", "resume-id")) - client = DuplexControlClient(transport, control_id_factory=lambda: next(control_ids)) - fence = DuplexFence("sid-client") - - assert ( - client.touch( - fence.session_id, - fence=fence, - activity=DuplexLeaseActivity.PLAYBACK_ACK, - timeout=2.0, - )["ok"] - is True - ) - assert ( - client.resume( - fence.session_id, - fence=fence, - expected_lease_generation=7, - timeout=3.0, - )["ok"] - is True - ) - - touch_key, touch_message, touch_kwargs = transport.calls[0] - assert touch_key == ("duplex", "touch-id") - assert isinstance(touch_message, TouchDuplexSessionMessage) - assert touch_message.activity == DuplexLeaseActivity.PLAYBACK_ACK.value - assert touch_kwargs["timeout"] == 2.0 - - resume_key, resume_message, resume_kwargs = transport.calls[1] - assert resume_key == ("duplex", "resume-id") - assert isinstance(resume_message, ResumeDuplexSessionMessage) - assert resume_message.expected_lease_generation == 7 - assert resume_kwargs["timeout"] == 3.0 - - -def test_control_client_treats_ok_false_as_authoritative_failure() -> None: - client = DuplexControlClient(_UnsupportedTransport(), control_id_factory=lambda: "unsupported") - fence = DuplexFence("sid-unsupported") - - with pytest.raises(DuplexControlRequestError, match="duplex touch failed") as exc_info: - client.touch( - fence.session_id, - fence=fence, - activity=DuplexLeaseActivity.HEARTBEAT, - timeout=1.0, - ) - - assert exc_info.value.code == "invalid_capability" - assert exc_info.value.retryable is False - assert exc_info.value.accepted_fence == DuplexFence(fence.session_id, epoch=2) - assert exc_info.value.lease_generation == 7 diff --git a/tests/engine/duplex/test_duplex_control_plane.py b/tests/engine/duplex/test_duplex_control_plane.py deleted file mode 100644 index 14da4d21c..000000000 --- a/tests/engine/duplex/test_duplex_control_plane.py +++ /dev/null @@ -1,1270 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import logging -from dataclasses import FrozenInstanceError - -import msgspec -import pytest - -from vllm_omni.engine.duplex.control_plane import ( - DuplexControlPlane, - DuplexOutputContext, - DuplexRequestIdentity, - DuplexStageRequestContext, - DuplexStageSubmission, - DuplexStageSubmissionResult, -) -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity, DuplexLeaseConfig -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexFence, - DuplexSessionLifecycleMessage, - OpenDuplexSessionMessage, - ResumeDuplexSessionMessage, - SignalDuplexTurnMessage, - TouchDuplexSessionMessage, -) -from vllm_omni.engine.duplex.runtime import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexRuntimeCapabilities, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _Extension: - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return tuple(f"configured-{stage_id}" for stage_id, _ in enumerate(defaults)) - - def plan_append( - self, - *, - request_id, - fence, - session_config, - runtime_config, - seq, - turn_seq, - mode, - payload, - final, - sampling_params, - ): - del request_id, fence, session_config, runtime_config, seq, turn_seq, mode, payload, final - assert sampling_params == "configured-0" - return DuplexAppendPlan(prompt={"prompt_token_ids": [1, 2, 3]}) - - def decide_output(self, **kwargs): - del kwargs - return None - - -class _TypedStagePort: - stage_count = 2 - - def __init__(self) -> None: - self.ensure_calls: list[DuplexStageRequestContext] = [] - self.submit_calls: list[DuplexStageSubmission] = [] - self.cleanup_calls: list[tuple[list[str], bool]] = [] - - def sampling_defaults(self) -> tuple[object, ...]: - return ("default-0", "default-1") - - def ensure_request(self, context: DuplexStageRequestContext) -> None: - self.ensure_calls.append(context) - - async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: - self.submit_calls.append(submission) - return DuplexStageSubmissionResult( - request_id=submission.context.request_id, - stage_id=submission.context.stage_id, - replica_id=3, - ) - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_calls.append((request_ids, abort)) - - -class _Clock: - def __init__(self) -> None: - self.value = 0.0 - - def __call__(self) -> float: - return self.value - - def advance(self, seconds: float) -> None: - self.value += seconds - - -@pytest.mark.asyncio -async def test_control_plane_uses_frozen_typed_stage_context_without_request_state() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=_Extension(), - stage_port=stage_port, - result_sink=result_sink, - ) - fence = DuplexFence("typed-port") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open", - fence=fence, - session_id=fence.session_id, - capabilities={ - "input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value], - }, - session_config={"voice": "test"}, - runtime_config={"runtime": "test"}, - ) - ) - assert (await result_sink.get()).ok is True - - context = stage_port.ensure_calls[-1] - assert context.session_id == fence.session_id - assert context.fence == fence - assert context.stage_id == 0 - assert context.final_stage_id == 1 - assert context.sampling_params == ("configured-0", "configured-1") - assert context.session_config == {"voice": "test"} - assert context.runtime_config == {"runtime": "test"} - with pytest.raises(FrozenInstanceError): - context.stage_id = 1 # type: ignore[misc] - - await plane.handle( - AppendDuplexInputMessage( - control_id="append", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - final=True, - ) - ) - append_result = await result_sink.get() - assert append_result.ok is True - assert append_result.stage_results[0]["replica_id"] == 3 - assert append_result.stage_results[0]["result"]["response_stage_id"] == 1 - - submission = stage_port.submit_calls[-1] - assert submission.context == stage_port.ensure_calls[-1] - assert submission.prompt == {"prompt_token_ids": [1, 2, 3]} - assert submission.already_submitted is False - assert not hasattr(submission, "request_state") - - -def test_control_plane_accepts_only_typed_duplex_messages() -> None: - plane = DuplexControlPlane( - extension=None, - stage_port=_TypedStagePort(), - result_sink=asyncio.Queue(), - ) - - assert plane.accepts( - OpenDuplexSessionMessage( - control_id="open", - fence=DuplexFence("typed-message"), - session_id="typed-message", - capabilities=DuplexRuntimeCapabilities().__dict__, - ) - ) - assert plane.accepts(type("Lookalike", (), {"type": "open_duplex_session"})()) is False - - -@pytest.mark.parametrize( - "message", - [ - TouchDuplexSessionMessage( - control_id="touch-1", - fence=DuplexFence("sid-message"), - session_id="sid-message", - activity=DuplexLeaseActivity.HEARTBEAT.value, - ), - ResumeDuplexSessionMessage( - control_id="resume-1", - fence=DuplexFence("sid-message"), - session_id="sid-message", - expected_lease_generation=3, - ), - DuplexSessionLifecycleMessage( - fence=DuplexFence("sid-message"), - session_id="sid-message", - event="expired", - reason="idle_ttl_expired", - lease_generation=4, - submitted_request_ids=["req-submitted"], - reserved_request_ids=["req-reserved"], - ), - ], -) -def test_duplex_lease_messages_round_trip(message) -> None: - encoded = msgspec.json.encode(message) - decoded = msgspec.json.decode(encoded, type=type(message)) - - assert decoded == message - - -@pytest.mark.asyncio -async def test_control_plane_touches_append_signal_and_explicit_activity() -> None: - clock = _Clock() - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=_Extension(), - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=30.0, disconnect_grace_s=5.0), - clock=clock, - ) - fence = DuplexFence("sid-touch") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open", - fence=fence, - session_id=fence.session_id, - capabilities={"input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value]}, - ) - ) - await result_sink.get() - session = plane.sessions.require(fence.session_id) - - clock.advance(1.0) - await plane.handle( - AppendDuplexInputMessage( - control_id="append", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - ) - assert (await result_sink.get()).ok is True - assert session.lease.last_activity == 1.0 - - clock.advance(1.0) - await plane.handle( - SignalDuplexTurnMessage( - control_id="signal", - fence=fence, - session_id=fence.session_id, - event="session.update", - ) - ) - assert (await result_sink.get()).ok is True - assert session.lease.last_activity == 2.0 - - submit_count = len(stage_port.submit_calls) - clock.advance(1.0) - await plane.handle( - TouchDuplexSessionMessage( - control_id="heartbeat", - fence=fence, - session_id=fence.session_id, - activity=DuplexLeaseActivity.HEARTBEAT.value, - ) - ) - touch_result = await result_sink.get() - assert touch_result.ok is True - assert session.lease.last_activity == 3.0 - assert len(stage_port.submit_calls) == submit_count - assert lifecycle_sink.empty() - - -@pytest.mark.asyncio -async def test_control_plane_resume_requires_expected_lease_generation() -> None: - clock = _Clock() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=_TypedStagePort(), - result_sink=result_sink, - lifecycle_sink=asyncio.Queue(), - lease_config=DuplexLeaseConfig(), - clock=clock, - ) - fence = DuplexFence("sid-resume-control") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open", - fence=fence, - session_id=fence.session_id, - capabilities={}, - ) - ) - await result_sink.get() - await plane.handle( - TouchDuplexSessionMessage( - control_id="detach", - fence=fence, - session_id=fence.session_id, - activity=DuplexLeaseActivity.DETACH.value, - ) - ) - await result_sink.get() - - await plane.handle( - ResumeDuplexSessionMessage( - control_id="resume", - fence=fence, - session_id=fence.session_id, - expected_lease_generation=0, - ) - ) - result = await result_sink.get() - - assert result.ok is True - assert result.stage_results[0]["result"]["lease_generation"] == 1 - - await plane.handle( - ResumeDuplexSessionMessage( - control_id="stale-resume", - fence=fence, - session_id=fence.session_id, - expected_lease_generation=0, - ) - ) - stale_result = await result_sink.get() - assert stale_result.ok is False - assert stale_result.accepted_fence == fence - assert stale_result.lease_generation == 1 - - -@pytest.mark.asyncio -async def test_control_plane_reaps_one_session_through_cleanup_and_lifecycle_sink() -> None: - clock = _Clock() - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=2.0, disconnect_grace_s=1.0), - clock=clock, - ) - for session_id in ("sid-a", "sid-b"): - fence = DuplexFence(session_id) - await plane.handle( - OpenDuplexSessionMessage( - control_id=f"open-{session_id}", - fence=fence, - session_id=session_id, - capabilities={}, - ) - ) - await result_sink.get() - clock.advance(1.0) - plane.sessions.require("sid-b").touch(DuplexFence("sid-b"), DuplexLeaseActivity.HEARTBEAT) - clock.advance(1.1) - - expired_count = await plane.reap_expired() - - assert expired_count == 1 - assert plane.sessions.get("sid-a") is None - assert plane.sessions.get("sid-b") is not None - assert stage_port.cleanup_calls == [] - lifecycle = await lifecycle_sink.get() - assert isinstance(lifecycle, DuplexSessionLifecycleMessage) - assert lifecycle.session_id == "sid-a" - assert lifecycle.reason == "idle_ttl_expired" - assert result_sink.empty() - - -@pytest.mark.asyncio -async def test_control_plane_retries_expired_cleanup_before_publishing_lifecycle() -> None: - class _FailOnceStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.failures_remaining = 1 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if self.failures_remaining: - self.failures_remaining -= 1 - raise RuntimeError("transient cleanup failure") - await super().cleanup(request_ids, abort=abort) - - clock = _Clock() - stage_port = _FailOnceStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - ) - fence = DuplexFence("sid-retry-expiry") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-retry", - fence=fence, - session_id=fence.session_id, - capabilities={}, - ) - ) - await result_sink.get() - request_id = plane.stage_request_id(fence, stage_id=0) - plane.sessions.require(fence.session_id).reserve_stage_request(0, request_id, fence=fence) - clock.advance(2.0) - - assert await plane.reap_expired() == 0 - assert lifecycle_sink.empty() - - assert await plane.reap_expired() == 1 - assert (await lifecycle_sink.get()).session_id == fence.session_id - - -@pytest.mark.asyncio -async def test_expired_session_holds_admission_slot_while_cleanup_is_blocked() -> None: - class _BlockedCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.cleanup_started = asyncio.Event() - self.release_cleanup = asyncio.Event() - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_started.set() - await self.release_cleanup.wait() - await super().cleanup(request_ids, abort=abort) - - clock = _Clock() - stage_port = _BlockedCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - max_sessions=1, - ) - old_fence = DuplexFence("sid-expired-blocked") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-expired-blocked", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-expired-blocked", fence=old_fence) - clock.advance(2.0) - - reap_task = asyncio.create_task(plane.reap_expired()) - await asyncio.wait_for(stage_port.cleanup_started.wait(), timeout=1.0) - replacement_fence = DuplexFence("sid-replacement-blocked") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-blocked", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - - blocked_result = await result_sink.get() - assert blocked_result.ok is False - assert blocked_result.error is not None - assert blocked_result.error.code == "resource_exhausted" - assert plane.sessions.get(old_fence.session_id) is old_session - - stage_port.release_cleanup.set() - assert await reap_task == 1 - assert plane.sessions.get(old_fence.session_id) is None - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-after-cleanup", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - - -@pytest.mark.asyncio -async def test_failed_expiry_cleanup_retains_admission_slot_for_retry() -> None: - class _FailingCleanupStagePort(_TypedStagePort): - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - del request_ids, abort - raise RuntimeError("cleanup failed") - - clock = _Clock() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=_FailingCleanupStagePort(), - result_sink=result_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - max_sessions=1, - ) - old_fence = DuplexFence("sid-expired-failed") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-expired-failed", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-expired-failed", fence=old_fence) - clock.advance(2.0) - - assert await plane.reap_expired() == 0 - replacement_fence = DuplexFence("sid-replacement-failed") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-failed", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - - blocked_result = await result_sink.get() - assert blocked_result.ok is False - assert blocked_result.error is not None - assert blocked_result.error.code == "resource_exhausted" - assert plane.sessions.get(old_fence.session_id) is old_session - - -@pytest.mark.asyncio -async def test_capacity_rejection_is_an_info_event(caplog: pytest.LogCaptureFixture) -> None: - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=_TypedStagePort(), - result_sink=result_sink, - max_sessions=1, - ) - first_fence = DuplexFence("sid-capacity-first") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-capacity-first", - fence=first_fence, - session_id=first_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - caplog.clear() - - rejected_fence = DuplexFence("sid-capacity-rejected") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-capacity-rejected", - fence=rejected_fence, - session_id=rejected_fence.session_id, - capabilities={}, - ) - ) - - result = await result_sink.get() - assert result.error.code == "resource_exhausted" - assert not [record for record in caplog.records if record.levelno >= logging.ERROR] - - -@pytest.mark.asyncio -async def test_failed_request_cleanup_retries_before_releasing_admission() -> None: - class _FailOnceCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.attempts = 0 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_calls.append((list(request_ids), abort)) - self.attempts += 1 - if self.attempts == 1: - raise RuntimeError("transient request cleanup failure") - - stage_port = _FailOnceCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - max_sessions=1, - ) - old_fence = DuplexFence("sid-request-cleanup") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-request-cleanup", fence=old_fence) - - closed = plane.close_sessions_for_request_ids( - ["req-request-cleanup"], - abort=True, - ) - assert closed == {old_fence.session_id: ["req-request-cleanup"]} - assert plane.sessions.get(old_fence.session_id) is old_session - with pytest.raises(RuntimeError, match="transient request cleanup failure"): - await stage_port.cleanup(closed[old_fence.session_id], abort=True) - - replacement_fence = DuplexFence("sid-request-cleanup-replacement") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-replacement-blocked", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).error.code == "resource_exhausted" - - assert await plane.reap_expired() == 1 - assert plane.sessions.get(old_fence.session_id) is None - assert stage_port.cleanup_calls == [ - (["req-request-cleanup"], True), - (["req-request-cleanup"], True), - ] - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-replacement", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - - -@pytest.mark.asyncio -async def test_stale_request_cleanup_does_not_finalize_reopened_incarnation() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - max_sessions=1, - ) - old_fence = DuplexFence("sid-request-cleanup-reopen") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-old", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-request-cleanup-old", fence=old_fence) - closed = plane.close_sessions_for_request_ids( - ["req-request-cleanup-old"], - abort=True, - cleanup_in_progress=True, - ) - assert closed == {old_fence.session_id: ["req-request-cleanup-old"]} - plane.defer_request_cleanups([old_fence.session_id]) - - await plane.handle( - CloseDuplexSessionMessage( - control_id="close-request-cleanup-old", - fence=old_fence, - session_id=old_fence.session_id, - reason="client_close", - ) - ) - assert (await result_sink.get()).ok is True - assert plane.sessions.get(old_fence.session_id) is None - - new_fence = DuplexFence(old_fence.session_id, incarnation=1) - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-new", - fence=new_fence, - session_id=new_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - new_session = plane.sessions.require(new_fence.session_id) - assert new_session.resume(new_fence, expected_lease_generation=0) == 1 - - assert await plane.reap_expired() == 1 - assert plane.sessions.get(new_fence.session_id) is new_session - - -@pytest.mark.asyncio -async def test_failed_open_cleanup_is_retried_before_releasing_admission() -> None: - class _FailOpenAndFirstCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.ensure_attempts = 0 - self.cleanup_attempts = 0 - - def ensure_request(self, context: DuplexStageRequestContext) -> None: - self.ensure_attempts += 1 - if self.ensure_attempts == 1: - raise RuntimeError("initial stage setup failed") - super().ensure_request(context) - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_calls.append((list(request_ids), abort)) - self.cleanup_attempts += 1 - if self.cleanup_attempts == 1: - raise RuntimeError("transient open rollback failure") - - stage_port = _FailOpenAndFirstCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=_Extension(), - stage_port=stage_port, - result_sink=result_sink, - max_sessions=1, - ) - failed_fence = DuplexFence("sid-open-rollback") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-fails-during-stage-setup", - fence=failed_fence, - session_id=failed_fence.session_id, - capabilities={"input_modes": ["append_audio_chunk"]}, - ) - ) - - failed_result = await result_sink.get() - assert failed_result.ok is False - assert plane.sessions.get(failed_fence.session_id) is not None - - replacement_fence = DuplexFence("sid-open-rollback-replacement") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-before-rollback", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={"input_modes": ["append_audio_chunk"]}, - ) - ) - assert (await result_sink.get()).error.code == "resource_exhausted" - - assert await plane.reap_expired() == 0 - assert plane.sessions.get(failed_fence.session_id) is None - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-after-rollback", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={"input_modes": ["append_audio_chunk"]}, - ) - ) - assert (await result_sink.get()).ok is True - - -@pytest.mark.asyncio -async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: - class _OneStuckExpiryStagePort(_TypedStagePort): - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if request_ids == [stuck_request_id]: - raise RuntimeError("stuck expiry cleanup") - await super().cleanup(request_ids, abort=abort) - - clock = _Clock() - stage_port = _OneStuckExpiryStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - ) - stuck_fence = DuplexFence("sid-expiry-stuck") - ready_fence = DuplexFence("sid-expiry-ready") - stuck_request_id = plane.stage_request_id(stuck_fence, stage_id=0) - ready_request_id = plane.stage_request_id(ready_fence, stage_id=0) - for fence in (stuck_fence, ready_fence): - await plane.handle( - OpenDuplexSessionMessage( - control_id=f"open-{fence.session_id}", - fence=fence, - session_id=fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - plane.sessions.require(fence.session_id).reserve_stage_request( - 0, - plane.stage_request_id(fence, stage_id=0), - fence=fence, - ) - clock.advance(2.0) - - assert await plane.reap_expired() == 1 - - assert (await lifecycle_sink.get()).session_id == ready_fence.session_id - assert stage_port.cleanup_calls == [([ready_request_id], False)] - - -@pytest.mark.asyncio -async def test_cancel_retains_request_resources_until_abort_succeeds() -> None: - class _FailOnceStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.failures_remaining = 1 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if self.failures_remaining: - self.failures_remaining -= 1 - raise RuntimeError("transient abort failure") - await super().cleanup(request_ids, abort=abort) - - stage_port = _FailOnceStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - cancelled_fence = DuplexFence("sid-cancel-retry") - next_fence = DuplexFence("sid-cancel-retry", epoch=1) - session = plane.sessions.open_session(cancelled_fence) - session.bind_stage_request(0, "req-cancel-retry", fence=cancelled_fence) - - def signal(control_id: str) -> SignalDuplexTurnMessage: - return SignalDuplexTurnMessage( - control_id=control_id, - fence=cancelled_fence, - next_fence=next_fence, - session_id=cancelled_fence.session_id, - event="input.cancel", - ) - - await plane.handle(signal("cancel-first")) - assert (await result_sink.get()).ok is False - assert session.fence == next_fence - assert session.resource_request_ids(cancelled_fence) == ["req-cancel-retry"] - - await plane.handle(signal("cancel-retry")) - assert (await result_sink.get()).ok is True - assert session.resource_request_ids(cancelled_fence) == [] - assert stage_port.cleanup_calls == [(["req-cancel-retry"], True)] - - -@pytest.mark.asyncio -async def test_close_retains_session_resources_until_cleanup_succeeds() -> None: - class _FailOnceStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.failures_remaining = 1 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if self.failures_remaining: - self.failures_remaining -= 1 - raise RuntimeError("transient close cleanup failure") - await super().cleanup(request_ids, abort=abort) - - stage_port = _FailOnceStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-close-retry") - session = plane.sessions.open_session(fence) - session.bind_stage_request(0, "req-close-retry", fence=fence) - - def close(control_id: str) -> CloseDuplexSessionMessage: - return CloseDuplexSessionMessage( - control_id=control_id, - fence=fence, - session_id=fence.session_id, - reason="client_close", - ) - - await plane.handle(close("close-first")) - assert (await result_sink.get()).ok is False - assert plane.sessions.get(fence.session_id) is session - assert session.resource_request_ids() == ["req-close-retry"] - - await plane.handle(close("close-retry")) - assert (await result_sink.get()).ok is True - assert plane.sessions.get(fence.session_id) is None - assert stage_port.cleanup_calls == [(["req-close-retry"], True)] - - -@pytest.mark.asyncio -async def test_control_cleanup_failure_does_not_block_other_sessions() -> None: - class _IndependentFailureStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.attempts: dict[str, int] = {} - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - request_id = request_ids[0] - self.attempts[request_id] = self.attempts.get(request_id, 0) + 1 - if request_id == "req-stuck" or self.attempts[request_id] == 1: - raise RuntimeError(f"cleanup failed for {request_id}") - await super().cleanup(request_ids, abort=abort) - - stage_port = _IndependentFailureStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - - for session_id, request_id in (("sid-stuck", "req-stuck"), ("sid-ready", "req-ready")): - fence = DuplexFence(session_id) - session = plane.sessions.open_session(fence) - session.bind_stage_request(0, request_id, fence=fence) - await plane.handle( - CloseDuplexSessionMessage( - control_id=f"close-{session_id}", - fence=fence, - session_id=session_id, - reason="client_close", - ) - ) - assert (await result_sink.get()).ok is False - - await plane.reap_expired() - - assert plane.sessions.get("sid-stuck") is not None - assert plane.sessions.get("sid-ready") is None - assert stage_port.cleanup_calls == [(["req-ready"], True)] - - -@pytest.mark.asyncio -async def test_control_cleanup_is_single_flight_with_concurrent_reaper() -> None: - class _SlowCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.cleanup_started = asyncio.Event() - self.release_cleanup = asyncio.Event() - self.cleanup_attempts = 0 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_attempts += 1 - self.cleanup_started.set() - await self.release_cleanup.wait() - await super().cleanup(request_ids, abort=abort) - - stage_port = _SlowCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-concurrent-cleanup") - session = plane.sessions.open_session(fence) - session.bind_stage_request(0, "req-concurrent-cleanup", fence=fence) - - close_task = asyncio.create_task( - plane.handle( - CloseDuplexSessionMessage( - control_id="close-concurrent-cleanup", - fence=fence, - session_id=fence.session_id, - reason="client_close", - ) - ) - ) - await asyncio.wait_for(stage_port.cleanup_started.wait(), timeout=1) - - reaper_task = asyncio.create_task(plane.reap_expired()) - await asyncio.sleep(0) - await asyncio.sleep(0) - assert stage_port.cleanup_attempts == 1 - - stage_port.release_cleanup.set() - await asyncio.gather(close_task, reaper_task) - - assert (await result_sink.get()).ok is True - assert plane.sessions.get(fence.session_id) is None - assert stage_port.cleanup_calls == [(["req-concurrent-cleanup"], True)] - - -@pytest.mark.asyncio -async def test_unknown_input_mode_rejects_open_without_registering_session() -> None: - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=_TypedStagePort(), result_sink=result_sink) - fence = DuplexFence("sid-unknown-mode") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-unknown-mode", - fence=fence, - session_id=fence.session_id, - capabilities={"input_modes": ["future_mode"]}, - ) - ) - - result = await result_sink.get() - assert result.ok is False - assert result.error is not None - assert result.error.code == "invalid_capability" - assert plane.sessions.get(fence.session_id) is None - - -@pytest.mark.asyncio -async def test_extension_free_turn_commit_does_not_allocate_or_submit_stage_request() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-turn-commit") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-turn-commit", - fence=fence, - session_id=fence.session_id, - capabilities={"input_modes": [DuplexInputMode.TURN_COMMIT_ONLY.value]}, - ) - ) - assert (await result_sink.get()).ok is True - - await plane.handle( - AppendDuplexInputMessage( - control_id="append-turn-commit", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.TURN_COMMIT_ONLY.value, - payload=None, - final=True, - ) - ) - result = await result_sink.get() - - assert result.ok is True - assert result.stage_results[0]["result"]["data_plane_append"] is False - assert stage_port.ensure_calls == [] - assert stage_port.submit_calls == [] - - -@pytest.mark.asyncio -async def test_session_update_replaces_both_configs_with_one_generation() -> None: - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=_TypedStagePort(), result_sink=result_sink) - fence = DuplexFence("sid-atomic-config") - session = plane.sessions.open_session(fence) - - await plane.handle( - SignalDuplexTurnMessage( - control_id="update-config", - fence=fence, - session_id=fence.session_id, - event="session.update", - session_config={"voice": "new"}, - runtime_config={"temperature": 0.5}, - ) - ) - - assert (await result_sink.get()).ok is True - assert session.config_generation == 1 - assert session.session_config == {"voice": "new"} - assert session.runtime_config == {"temperature": 0.5} - - -@pytest.mark.asyncio -async def test_stage_submission_is_compensated_when_local_commit_fails() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-submit-compensation") - session = plane.sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), - ) - request_id = plane.stage_request_id(fence, stage_id=0) - - def fail_commit(_reservation): - raise RuntimeError("local commit failed") - - session.commit_append = fail_commit # type: ignore[method-assign] - await plane.handle( - AppendDuplexInputMessage( - control_id="append-compensation", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - ) - - assert (await result_sink.get()).ok is False - assert stage_port.cleanup_calls == [([request_id], True)] - - -@pytest.mark.asyncio -async def test_failed_submission_compensation_blocks_append_until_reaper_cleans_it() -> None: - class _FailingCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.cleanup_failures = 2 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - await super().cleanup(request_ids, abort=abort) - if self.cleanup_failures > 0: - self.cleanup_failures -= 1 - raise RuntimeError("stage abort failed") - - stage_port = _FailingCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-durable-submit-compensation") - session = plane.sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), - ) - original_commit = session.commit_append - - def fail_commit(_reservation): - raise RuntimeError("local commit failed") - - session.commit_append = fail_commit # type: ignore[method-assign] - first = AppendDuplexInputMessage( - control_id="append-compensation-1", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - await plane.handle(first) - assert (await result_sink.get()).ok is False - assert plane.pending_submission_cleanup_count == 1 - assert len(stage_port.submit_calls) == 1 - - await plane.handle( - AppendDuplexInputMessage( - control_id="append-compensation-2", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - ) - assert (await result_sink.get()).ok is False - assert plane.pending_submission_cleanup_count == 1 - assert len(stage_port.submit_calls) == 1 - - session.commit_append = original_commit # type: ignore[method-assign] - stage_port.cleanup_failures = 0 - await plane.reap_expired() - - assert plane.pending_submission_cleanup_count == 0 - assert session.resource_request_ids() == [] - - -def test_stale_output_is_rejected_before_extension_decision() -> None: - class _CountingExtension(_Extension): - def __init__(self) -> None: - self.decision_calls = 0 - - def decide_output(self, **kwargs): - del kwargs - self.decision_calls += 1 - return None - - extension = _CountingExtension() - plane = DuplexControlPlane(extension=extension, stage_port=_TypedStagePort(), result_sink=asyncio.Queue()) - current = DuplexFence("sid-stale-output", epoch=1) - plane.sessions.open_session(current) - context = DuplexOutputContext( - identity=DuplexRequestIdentity( - session_id=current.session_id, - fence=DuplexFence(current.session_id, epoch=0), - ), - final_stage_id=1, - segment_finished=True, - ) - - assert plane.decide_output(0, object(), context) is None - assert extension.decision_calls == 0 - - -@pytest.mark.asyncio -async def test_control_dispatch_is_ordered_per_session_without_blocking_other_sessions() -> None: - class _BlockingStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.started = asyncio.Event() - self.release = asyncio.Event() - - async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: - if submission.context.session_id == "sid-blocked": - self.started.set() - await self.release.wait() - return await super().submit(submission) - - stage_port = _BlockingStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) - blocked = DuplexFence("sid-blocked") - independent = DuplexFence("sid-independent") - capabilities = DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}) - plane.sessions.open_session(blocked, capabilities=capabilities) - plane.sessions.open_session(independent, capabilities=capabilities) - - plane.dispatch( - AppendDuplexInputMessage( - control_id="blocked-append", - fence=blocked, - session_id=blocked.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - operation_id="blocked-operation", - ) - ) - await asyncio.wait_for(stage_port.started.wait(), timeout=1) - plane.dispatch( - AppendDuplexInputMessage( - control_id="blocked-append-retry", - fence=blocked, - session_id=blocked.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - operation_id="blocked-operation", - ) - ) - plane.dispatch( - TouchDuplexSessionMessage( - control_id="blocked-touch", - fence=blocked, - session_id=blocked.session_id, - activity=DuplexLeaseActivity.HEARTBEAT.value, - ) - ) - plane.dispatch( - TouchDuplexSessionMessage( - control_id="independent-touch", - fence=independent, - session_id=independent.session_id, - activity=DuplexLeaseActivity.HEARTBEAT.value, - ) - ) - - independent_result = await asyncio.wait_for(result_sink.get(), timeout=1) - assert independent_result.control_id == "independent-touch" - - stage_port.release.set() - await plane.drain() - assert [result_sink.get_nowait().control_id for _ in range(3)] == [ - "blocked-append", - "blocked-append-retry", - "blocked-touch", - ] - assert [ - submission.context.session_id - for submission in stage_port.submit_calls - if submission.context.session_id == blocked.session_id - ] == [blocked.session_id] diff --git a/tests/engine/duplex/test_duplex_events.py b/tests/engine/duplex/test_duplex_events.py new file mode 100644 index 000000000..a4ebfbf5a --- /dev/null +++ b/tests/engine/duplex/test_duplex_events.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Typed duplex events: pure ``to_realtime()`` rendering and the stateful Realtime projection.""" + +from __future__ import annotations + +import base64 +import inspect + +import pytest + +from vllm_omni.engine.duplex import events as events_module +from vllm_omni.engine.duplex.events import ( + REALTIME_ERROR_TYPES_BY_CODE, + AudioDelta, + DuplexEvent, + DuplexRawEvent, + ErrorEvent, + InputCommitted, + ItemAdded, + ItemCreated, + ItemDone, + Listen, + OutputItemAdded, + OverlapDecision, + ResponseDone, + SessionClosed, + SessionCreated, + SessionExpired, + SessionHeartbeatAck, + SessionReplaced, + SessionResumed, + SessionResyncRequired, + TextDelta, + TranscriptDone, + TurnEvent, + error_event, +) +from vllm_omni.engine.duplex.realtime_events import ( + RealtimeProjectionState, + project_internal_event, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: Classes whose wire object carries ``session_id`` by contract. +_SESSION_ID_ON_WIRE = { + SessionClosed, + SessionExpired, + SessionHeartbeatAck, + Listen, + TurnEvent, + OverlapDecision, + SessionResumed, + SessionReplaced, + SessionResyncRequired, +} +#: Classes whose wire object carries ``epoch`` by contract. +_EPOCH_ON_WIRE = {Listen, TurnEvent, OverlapDecision} + + +def _public_event_classes() -> list[type[DuplexEvent]]: + classes = [ + getattr(events_module, name) + for name in events_module.__all__ + if inspect.isclass(getattr(events_module, name)) and issubclass(getattr(events_module, name), DuplexEvent) + ] + assert len(classes) > 40 + return [cls for cls in classes if cls is not DuplexEvent] + + +@pytest.mark.parametrize("event_cls", _public_event_classes(), ids=lambda cls: cls.__name__) +def test_to_realtime_renders_wire_type_event_id_and_documented_fields(event_cls: type[DuplexEvent]): + event = event_cls(session_id="duplex-hidden", epoch=7) + + wire = event.to_realtime() + + assert wire["type"] == event.type + assert wire["event_id"] == event.event_id + assert wire["event_id"].startswith("event_") + if event_cls is DuplexRawEvent: + assert wire["type"] == "duplex." + else: + assert wire["type"] == event_cls.wire_type + if event_cls in _SESSION_ID_ON_WIRE: + assert wire["session_id"] == "duplex-hidden" + else: + assert "session_id" not in wire + if event_cls in _EPOCH_ON_WIRE: + assert wire["epoch"] == 7 + else: + assert "epoch" not in wire + # Every value is plain JSON (mappings / sequences are copied into dict / list). + for value in wire.values(): + assert isinstance(value, str | int | float | bool | dict | list) or value is None + + +def test_optional_wire_fields_are_omitted_when_none(): + created = SessionCreated(session={"id": "sid"}) + assert created.optional_wire_fields == frozenset({"attachment_generation", "resume_token"}) + wire = created.to_realtime() + assert "attachment_generation" not in wire + assert "resume_token" not in wire + assert wire["session"] == {"id": "sid"} + + with_resume = SessionCreated(session={"id": "sid"}, attachment_generation=2, resume_token="tok").to_realtime() + assert with_resume["attachment_generation"] == 2 + assert with_resume["resume_token"] == "tok" + + delta = AudioDelta(response_id="resp_1", item_id="item_resp_1", delta="AAAA").to_realtime() + assert "sample_rate_hz" not in delta + assert "metadata" not in delta + assert delta["format"] == "pcm16" + assert delta["output_index"] == 0 + assert delta["content_index"] == 0 + + anonymous_listen = Listen(session_id="sid", epoch=0, details={"reason": "silence"}).to_realtime() + assert "response_id" not in anonymous_listen + assert "id" not in anonymous_listen["response"] + assert anonymous_listen["response"]["status"] == "listening" + assert anonymous_listen["response"]["metadata"] == {"reason": "silence"} + + bound_listen = Listen(session_id="sid", epoch=0, response_id="resp_1").to_realtime() + assert bound_listen["response_id"] == "resp_1" + assert bound_listen["response"]["id"] == "resp_1" + + +def test_error_event_maps_codes_to_openai_error_types_and_echoes_client_event_id(): + assert REALTIME_ERROR_TYPES_BY_CODE["resource_exhausted"] == "rate_limit_error" + assert REALTIME_ERROR_TYPES_BY_CODE["internal_error"] == "server_error" + assert REALTIME_ERROR_TYPES_BY_CODE["bad_event"] == "invalid_request_error" + + error = ErrorEvent( + code="resource_exhausted", + message="too many sessions", + related_event_id="evt_client", + param="session", + extra={"retryable": True}, + ) + + assert error.error_type == "rate_limit_error" + assert error.text == "too many sessions" + assert error.error == { + "type": "rate_limit_error", + "code": "resource_exhausted", + "message": "too many sessions", + "event_id": "evt_client", + "param": "session", + "retryable": True, + } + wire = error.to_realtime() + assert wire["type"] == "error" + assert wire["error"] == error.error + # The server event id and the echoed client event id are distinct. + assert wire["event_id"] != "evt_client" + assert ErrorEvent(code="never-seen-code").error_type == "invalid_request_error" + + +def test_error_event_helper_normalizes_optional_client_event_id_and_param(): + event = error_event("bad_event", "nope", event_id="evt_1", param="item", extra={"retryable": False}) + assert isinstance(event, ErrorEvent) + assert event.related_event_id == "evt_1" + assert event.param == "item" + assert event.error["retryable"] is False + + anonymous = error_event("bad_event", "nope", event_id=42, param="") + assert anonymous.related_event_id is None + assert anonymous.param is None + assert "event_id" not in anonymous.error + assert "param" not in anonymous.error + + +def test_audio_delta_decodes_base64_and_invalid_payloads_yield_none(): + raw = bytes(range(16)) + delta = AudioDelta(delta=base64.b64encode(raw).decode("ascii"), format="pcm16", sample_rate_hz=24000) + + assert delta.audio == raw + assert delta.to_realtime()["sample_rate_hz"] == 24000 + assert AudioDelta(delta="@@not base64@@").audio is None + + +def test_session_closed_and_expired_are_the_only_terminal_events(): + closed = SessionClosed(session_id="sid", reason="client_close", details={"type": "session.closed"}) + expired = SessionExpired(session_id="sid", reason="idle_timeout") + + assert closed.is_terminal is True + assert expired.is_terminal is True + assert closed.to_realtime() == { + "type": "session.closed", + "event_id": closed.event_id, + "session_id": "sid", + "reason": "client_close", + "event": {"type": "session.closed"}, + } + assert expired.to_realtime() == { + "type": "session.expired", + "event_id": expired.event_id, + "session_id": "sid", + "reason": "idle_timeout", + } + non_terminal = [cls for cls in _public_event_classes() if cls().is_terminal] + assert non_terminal == [SessionClosed, SessionExpired] + + +def test_generic_accessors_expose_response_item_text_and_status(): + assert TextDelta(response_id="resp_1", item_id="item_resp_1", delta="hi").text == "hi" + assert TranscriptDone(transcript="done").text == "done" + assert OutputItemAdded(item={"id": "item_1"}).item_id == "item_1" + assert ItemAdded(item={"id": 3}).item_id is None + assert ResponseDone(response={"status": "cancelled"}).status == "cancelled" + assert ResponseDone().status is None + assert TurnEvent(event="user_started", turn_state="user_speaking").response_id is None + assert TurnEvent().audio is None + + +def test_raw_event_prefixes_internal_type_and_wraps_details(): + raw = DuplexRawEvent(internal_type="model.custom", details={"type": "model.custom", "value": (1, 2)}) + + assert raw.type == "duplex.model.custom" + assert raw.to_realtime() == { + "type": "duplex.model.custom", + "event_id": raw.event_id, + "event": {"type": "model.custom", "value": [1, 2]}, + } + + +def test_turn_event_and_overlap_decision_carry_session_identity_on_the_wire(): + turn = TurnEvent(session_id="sid", epoch=2, event="barge_in", turn_state="barge_in").to_realtime() + assert turn == { + "type": "turn.event", + "event_id": turn["event_id"], + "session_id": "sid", + "event": "barge_in", + "turn_state": "barge_in", + "epoch": 2, + } + decision = OverlapDecision( + session_id="sid", epoch=2, policy="listen_only", action="ignore", reason="short", details={"ms": 300} + ).to_realtime() + assert decision["policy"] == "listen_only" + assert decision["action"] == "ignore" + assert decision["metadata"] == {"ms": 300} + assert decision["session_id"] == "sid" + + +# ---- RealtimeProjectionState / project_internal_event ---- + + +def _types(events: list[DuplexEvent]) -> list[str]: + return [event.type for event in events] + + +def test_projection_of_the_main_internal_event_sequence(): + state = RealtimeProjectionState(session_id="duplex-proj", model="test-model", initial_session_update=True) + pcm = base64.b64encode(b"\x00\x10" * 8).decode("ascii") + + created = project_internal_event(state, {"type": "session.created", "session": {"voice": "alloy"}}) + assert _types(created) == ["session.created", "session.updated"] + assert isinstance(created[0], SessionCreated) + session_payload = created[0].session + assert session_payload["id"] == "duplex-proj" + assert session_payload["model"] == "test-model" + assert session_payload["voice"] == "alloy" + assert session_payload["input_audio_format"] == "pcm16" + assert session_payload["output_audio_format"] == "pcm16" + assert state.initial_session_update is False + # The extra session.updated is a one-shot for the opening session.update. + assert _types(project_internal_event(state, {"type": "session.created", "session": {}})) == ["session.created"] + + response_created = project_internal_event( + state, {"type": "response.created", "response_id": "resp_1", "modalities": ["audio", "text"]} + ) + assert _types(response_created) == [ + "response.created", + "conversation.item.added", + "conversation.item.created", + "response.output_item.added", + "response.content_part.added", + ] + assert state.active_response_id == "resp_1" + assert response_created[0].response_id == "resp_1" + assert response_created[0].to_realtime()["response"]["status"] == "in_progress" + assert isinstance(response_created[1], ItemAdded) + assert isinstance(response_created[2], ItemCreated) + assert response_created[1].item_id == "item_resp_1" + assert response_created[3].item_id == "item_resp_1" + assert response_created[4].to_realtime()["part"] == {"type": "audio", "transcript": ""} + + audio_delta = project_internal_event( + state, + {"type": "response.output_audio.delta", "response_id": "resp_1", "audio": pcm, "text": "hi", "format": "pcm16"}, + ) + assert _types(audio_delta) == ["response.output_audio.delta", "response.output_audio_transcript.delta"] + assert isinstance(audio_delta[0], AudioDelta) + assert audio_delta[0].delta == pcm + assert audio_delta[0].format == "pcm16" + assert audio_delta[0].item_id == "item_resp_1" + assert audio_delta[1].text == "hi" + assert state.conversation_items["item_resp_1"]["content"][0]["transcript"] == "hi" + + done = project_internal_event(state, {"type": "response.done", "response_id": "resp_1"}) + assert _types(done) == [ + "response.output_audio.done", + "response.output_audio_transcript.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + "response.done", + "rate_limits.updated", + ] + assert done[1].text == "hi" + assert isinstance(done[4], ItemDone) + assert done[4].item["status"] == "completed" + assert isinstance(done[5], ResponseDone) + assert done[5].status == "completed" + assert done[5].response["output"][0]["content"][0] == {"type": "output_audio", "transcript": "hi"} + assert state.active_response_id is None + # A duplicate terminal for the same response is idempotent. + assert project_internal_event(state, {"type": "response.done", "response_id": "resp_1"}) == [] + + committed = project_internal_event( + state, {"type": "input.committed", "message": {"role": "user", "content": "hello"}, "input_commit_seq": 1} + ) + assert _types(committed) == [ + "conversation.item.added", + "conversation.item.created", + "input_audio_buffer.committed", + "conversation.item.done", + ] + assert isinstance(committed[2], InputCommitted) + assert committed[2].previous_item_id == "item_resp_1" + assert committed[2].item_id == committed[0].item_id + assert committed[2].to_realtime()["event"]["input_commit_seq"] == 1 + assert committed[0].item["content"] == [{"type": "input_text", "text": "hello"}] + + project_internal_event(state, {"type": "response.created", "response_id": "resp_2"}) + project_internal_event(state, {"type": "response.output_audio.delta", "response_id": "resp_2", "audio": pcm}) + cancelled = project_internal_event( + state, {"type": "audio.cancelled", "response_id": "resp_2", "reason": "barge_in", "committed_ms": 0} + ) + assert _types(cancelled) == [ + "response.output_audio.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + "response.done", + "rate_limits.updated", + ] + assert cancelled[4].status == "cancelled" + assert cancelled[4].response["status_details"] == {"type": "cancelled", "reason": "barge_in"} + assert state.active_response_id is None + + closed = project_internal_event(state, {"type": "session.closed", "reason": "client_close"}) + assert _types(closed) == ["session.closed"] + assert closed[0].is_terminal is True + assert closed[0].to_realtime()["reason"] == "client_close" + + +def test_projection_of_output_audio_buffer_clear_emits_cleared_before_terminals(): + state = RealtimeProjectionState(session_id="duplex-clear") + pcm = base64.b64encode(b"\x00\x10" * 8).decode("ascii") + project_internal_event(state, {"type": "response.created", "response_id": "resp_1"}) + project_internal_event(state, {"type": "response.output_audio.delta", "response_id": "resp_1", "audio": pcm}) + + cleared = project_internal_event( + state, {"type": "audio.cancelled", "reason": "output_audio_buffer_clear", "committed_ms": 250} + ) + + assert _types(cleared)[:2] == ["output_audio_buffer.cleared", "response.output_audio.done"] + assert cleared[0].response_id == "resp_1" + assert _types(cleared)[-2:] == ["response.done", "rate_limits.updated"] + assert state.item_truncation_cursors["item_resp_1"] == (0, 250) + + +def test_projection_leaves_error_to_typed_emit_sites(): + # Errors are constructed typed (``error_event``) at the emit site; the projector + # has no ``error`` branch and falls back to the raw wrapper for such a dict. + state = RealtimeProjectionState(session_id="duplex-err") + + projected = project_internal_event(state, {"type": "error", "code": "bad_event", "error": "nope"}) + + assert len(projected) == 1 + assert isinstance(projected[0], DuplexRawEvent) + assert projected[0].type == "duplex.error" + assert projected[0].to_realtime()["event"]["code"] == "bad_event" diff --git a/tests/engine/duplex/test_duplex_lease.py b/tests/engine/duplex/test_duplex_lease.py index 62fc8c70c..bd26b5740 100644 --- a/tests/engine/duplex/test_duplex_lease.py +++ b/tests/engine/duplex/test_duplex_lease.py @@ -1,19 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lease semantics of the engine-resident duplex session (idle TTL, disconnect grace, resume CAS).""" + from __future__ import annotations from dataclasses import FrozenInstanceError import pytest -from vllm_omni.engine.duplex.lease import ( +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.lease import ( DuplexLeaseActivity, DuplexLeaseConfig, + DuplexLeaseState, ) -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import DuplexInputMode -from vllm_omni.engine.duplex.session import DuplexSessionRuntimeManager pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -30,195 +33,148 @@ class FakeMonotonicClock: def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig: - return DuplexLeaseConfig( - idle_ttl_s=idle_ttl_s, - disconnect_grace_s=30.0, + return DuplexLeaseConfig(idle_ttl_s=idle_ttl_s, disconnect_grace_s=30.0) + + +def _session(session_id: str, clock: FakeMonotonicClock, *, idle_ttl_s: float | None = 300.0) -> DuplexEngineSession: + return DuplexEngineSession( + session_id=session_id, + config=DuplexSessionConfig(), + lease=DuplexLeaseState(config=_lease_config(idle_ttl_s=idle_ttl_s), generation=0, last_activity=clock()), + _clock=clock, ) -def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None: +def test_open_touch_detach_and_expiry_use_monotonic_time() -> None: clock = FakeMonotonicClock(100.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-expiry") - session = manager.open_session(fence, lease_config=_lease_config()) + session = _session("sid-expiry", clock) assert session.lease.last_activity == 100.0 clock.advance(10.0) - session.touch(fence, DuplexLeaseActivity.HEARTBEAT) - session.detach(fence) + session.touch_lease(DuplexLeaseActivity.HEARTBEAT) + assert session.lease.last_activity == 110.0 + session.detach_lease() clock.advance(29.0) assert session.lease.disconnect_grace_expired(clock()) is False - assert manager.collect_expired() == [] + assert session.lease.idle_expired(clock()) is False - clock.advance(272.0) - expired = manager.collect_expired() + clock.advance(1.0) + assert session.lease.disconnect_grace_expired(clock()) is True - assert [item.session_id for item in expired] == ["sid-expiry"] - assert expired[0].reason == "disconnect_grace_expired" - assert manager.get("sid-expiry") is session - manager.finalize_close_session(session) - assert manager.get("sid-expiry") is None + clock.advance(271.0) + assert session.lease.idle_expired(clock()) is True -def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None: +def test_detach_grace_and_resume_advance_the_lease_generation() -> None: clock = FakeMonotonicClock(10.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-resume") - session = manager.open_session(fence, lease_config=_lease_config()) + session = _session("sid-resume", clock) - session.detach(fence) + session.detach_lease() clock.advance(29.0) assert session.lease.disconnect_grace_expired(clock()) is False clock.advance(2.0) assert session.lease.disconnect_grace_expired(clock()) is True - generation = session.resume(fence, expected_lease_generation=0) + generation = session.resume_lease(expected_lease_generation=0) assert generation == 1 + assert session.lease_generation == 1 assert session.lease.detached_at is None assert session.lease.disconnect_grace_expired(clock()) is False with pytest.raises(ValueError, match="lease generation mismatch"): - session.resume(fence, expected_lease_generation=0) + session.resume_lease(expected_lease_generation=0) def test_active_operation_prevents_mid_transaction_expiry() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-operation") - session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0)) + session = _session("sid-operation", clock, idle_ttl_s=5.0) - session.begin_operation(fence, "append-1") + session.begin_lease_operation(session.fence, "append-1") clock.advance(10.0) - assert manager.collect_expired() == [] + assert session.lease.idle_expired(clock()) is False - session.end_operation(fence, "append-1") - assert manager.collect_expired() == [] + session.end_lease_operation("append-1") + assert session.lease.idle_expired(clock()) is False clock.advance(6.0) - assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"] + assert session.lease.idle_expired(clock()) is True def test_sessions_have_independent_activity_deadlines_and_resources() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence_a = DuplexFence("sid-a") - fence_b = DuplexFence("sid-b") - session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0)) - session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0)) - session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a) - session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a) - session_b.bind_stage_request(0, "req-b", fence=fence_b) + session_a = _session("sid-a", clock, idle_ttl_s=10.0) + session_b = _session("sid-b", clock, idle_ttl_s=10.0) + session_a.reserve_stage_request(0, "req-a-reserved", fence=session_a.fence) + session_a.bind_stage_request(1, "req-a-submitted", fence=session_a.fence) + session_b.bind_stage_request(0, "req-b", fence=session_b.fence) clock.advance(6.0) - session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT) + session_b.touch_lease(DuplexLeaseActivity.MODEL_OUTPUT) clock.advance(5.0) - expired = manager.collect_expired() - assert len(expired) == 1 - assert expired[0].session_id == "sid-a" - assert expired[0].reserved_request_ids == ("req-a-reserved",) - assert expired[0].submitted_request_ids == ("req-a-submitted",) - assert manager.get("sid-b") is session_b + assert session_a.lease.idle_expired(clock()) is True + assert session_b.lease.idle_expired(clock()) is False + assert session_a.resource_request_ids(submitted=False) == ["req-a-reserved"] + assert session_a.resource_request_ids(submitted=True) == ["req-a-submitted"] assert session_b.resource_request_ids() == ["req-b"] -def test_close_and_reaper_select_exactly_one_terminal_transition() -> None: +def test_close_is_a_single_terminal_transition() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-race") - manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0)) + session = _session("sid-race", clock, idle_ttl_s=1.0) clock.advance(2.0) + assert session.lease.idle_expired(clock()) is True - assert len(manager.collect_expired()) == 1 - assert manager.close_session(fence, reason="explicit_close") is None - assert manager.collect_expired() == [] - - fence_2 = DuplexFence("sid-race-2") - manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0)) - assert manager.close_session(fence_2, reason="explicit_close") is not None - clock.advance(2.0) - assert manager.collect_expired() == [] + assert session.begin_close(reason="explicit_close") is True + assert session.lease.terminal_reason == "explicit_close" + # A terminal lease is never reaped again and a second close is a no-op. + assert session.lease.idle_expired(clock()) is False + assert session.lease.disconnect_grace_expired(clock()) is False + assert session.begin_close(reason="idle_ttl_expired") is True + assert session.lease.terminal_reason == "explicit_close" + with pytest.raises(RuntimeError, match="terminal"): + session.touch_lease(DuplexLeaseActivity.HEARTBEAT) -def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None: +def test_stale_fence_cannot_start_a_lease_operation() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - current = DuplexFence("sid-fence", epoch=1) + session = _session("sid-fence", clock) + session.epoch = 1 + session.sync_fence() stale = DuplexFence("sid-fence", epoch=0) - session = manager.open_session(current, lease_config=_lease_config()) - with pytest.raises(RuntimeError, match="fence mismatch"): - session.touch(stale, DuplexLeaseActivity.APPEND) - with pytest.raises(RuntimeError, match="fence mismatch"): - session.detach(stale) - with pytest.raises(RuntimeError, match="fence mismatch"): - session.resume(stale, expected_lease_generation=0) + with pytest.raises(DuplexFenceMismatchError, match="fence mismatch"): + session.begin_lease_operation(stale, "append-1") + with pytest.raises(DuplexFenceMismatchError, match="fence mismatch"): + session.begin_lease_operation(DuplexFence("other-session", epoch=1), "append-2") -def test_disabled_idle_expiry_never_collects_session() -> None: +def test_disabled_idle_expiry_never_expires_session() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None)) + session = _session("sid-no-expiry", clock, idle_ttl_s=None) clock.advance(1_000_000.0) - assert manager.collect_expired() == [] + assert session.lease.expires_at is None + assert session.lease.idle_expired(clock()) is False def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-disconnect-grace") - session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None)) + session = _session("sid-disconnect-grace", clock, idle_ttl_s=None) - session.detach(fence) + session.detach_lease() clock.advance(29.0) - assert manager.collect_expired() == [] + assert session.lease.disconnect_grace_expired(clock()) is False clock.advance(1.0) - expired = manager.collect_expired() - - assert [item.session_id for item in expired] == [fence.session_id] - assert expired[0].reason == "disconnect_grace_expired" + assert session.lease.disconnect_grace_expired(clock()) is True -def test_runtime_manager_enforces_server_owned_session_admission() -> None: - manager = DuplexSessionRuntimeManager(max_sessions=2) - manager.open_session(DuplexFence("sid-admission-a")) - manager.open_session(DuplexFence("sid-admission-b")) - - with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"): - manager.open_session(DuplexFence("sid-admission-c")) - - -def test_completed_append_cache_is_bounded() -> None: - manager = DuplexSessionRuntimeManager(completed_append_limit=2) - fence = DuplexFence("sid-completed-cache") - session = manager.open_session(fence) - - for index in range(3): - session.record_completed_append( - f"operation-{index}", - fence=fence, - mode=DuplexInputMode.TURN_COMMIT_ONLY, - final=True, - stage_results=[{"index": index}], - ) - - assert list(session.completed_appends) == ["operation-1", "operation-2"] - - session.accept_fence(DuplexFence(fence.session_id, epoch=1)) - assert session.completed_appends == {} - - -def test_lease_config_and_expiry_record_are_immutable() -> None: +def test_lease_config_is_immutable_and_validated() -> None: config = _lease_config() with pytest.raises(FrozenInstanceError): config.idle_ttl_s = 1.0 # type: ignore[misc] - - clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0)) - clock.advance(2.0) - record = manager.collect_expired()[0] - with pytest.raises(FrozenInstanceError): - record.reason = "changed" # type: ignore[misc] + with pytest.raises(ValueError, match="idle_ttl_s"): + DuplexLeaseConfig(idle_ttl_s=0.0) + with pytest.raises(ValueError, match="disconnect_grace_s"): + DuplexLeaseConfig(disconnect_grace_s=0.0) diff --git a/tests/engine/duplex/test_duplex_plugin.py b/tests/engine/duplex/test_duplex_plugin.py new file mode 100644 index 000000000..f45b625e7 --- /dev/null +++ b/tests/engine/duplex/test_duplex_plugin.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexModelPlugin`` loading/validation helpers and the ``PipelineConfig.duplex_plugin`` binding.""" + +from __future__ import annotations + +from dataclasses import fields +from pathlib import Path +from typing import Any + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import DuplexAppendPlan +from vllm_omni.engine.duplex.plugin import ( + DuplexDataPlane, + DuplexModelPlugin, + DuplexModelSessionState, + DuplexRuntimeConfigError, + coerce_int, + load_duplex_plugin, + payload_turn_id, + reject_changed_runtime_value, + validate_duplex_plugin_sampling, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _encode_audio(audio: object, sample_rate_hz: int, fmt: str, speed: float | None) -> str | None: + del audio, sample_rate_hz, fmt, speed + return None + + +class FakeDataPlane(DuplexDataPlane): + def begin_request(self, request_id: str) -> None: + pass + + def is_terminal(self, request_id: str | None) -> bool: + return False + + def mark_terminal(self, request_id: str) -> None: + pass + + def close_stream(self, request_id: str) -> None: + pass + + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: + pass + + def project(self, result, *, context=None): + del result, context + return [] + + +class FakePlugin(DuplexModelPlugin): + """Minimal plugin; referenced by dotted path from ``load_duplex_plugin`` tests.""" + + plugin_id = "fake-duplex" + + def __init__(self, encode_audio) -> None: + super().__init__(encode_audio) + self.data_plane = FakeDataPlane() + + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(defaults) + + def plan_append(self, **kwargs) -> DuplexAppendPlan: + del kwargs + return DuplexAppendPlan(prompt={"prompt_token_ids": [1]}) + + def decide_output(self, **kwargs): + del kwargs + return None + + def create_session_state(self) -> DuplexModelSessionState: + raise NotImplementedError + + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: + del max_sessions + return DuplexCapabilities() + + def validate_client_extra_body(self, extra_body: object) -> None: + pass + + async def prepare_runtime_config(self, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]: + del config, model_config + return {} + + def runtime_config_for_update(self, config, current): + del config + return dict(current) + + def data_plane_context(self, **kwargs) -> object: + return dict(kwargs) + + +class PluginWithoutId(FakePlugin): + plugin_id = "" + + +class PluginWithoutDataPlane(FakePlugin): + def __init__(self, encode_audio) -> None: + DuplexModelPlugin.__init__(self, encode_audio) + + +class PluginWithWrongDataPlane(FakePlugin): + def __init__(self, encode_audio) -> None: + DuplexModelPlugin.__init__(self, encode_audio) + self.data_plane = object() # type: ignore[assignment] + + +class NotAPlugin: + plugin_id = "lookalike" + data_plane = FakeDataPlane() + + def __init__(self, encode_audio) -> None: + self.encode_audio = encode_audio + + +def _path(cls: type) -> str: + return f"{__name__}.{cls.__name__}" + + +# --------------------------------------------------------------------------- # +# load_duplex_plugin # +# --------------------------------------------------------------------------- # + + +def test_load_duplex_plugin_instantiates_the_class_at_the_dotted_path() -> None: + plugin = load_duplex_plugin(_path(FakePlugin), _encode_audio) + + assert type(plugin) is FakePlugin + assert plugin.plugin_id == "fake-duplex" + assert isinstance(plugin.data_plane, FakeDataPlane) + # Defaults of the optional class-level policy. + assert plugin.private_runtime_config_keys == frozenset() + assert plugin.silence_continuation_samples == 16000 + assert plugin.runtime_config_for_function_output(DuplexSessionConfig(), {}, {"name": "tool"}) is None + + +def test_load_duplex_plugin_rejects_a_non_plugin_class() -> None: + with pytest.raises(TypeError, match="not a DuplexModelPlugin"): + load_duplex_plugin(_path(NotAPlugin), _encode_audio) + + +def test_load_duplex_plugin_requires_a_plugin_id() -> None: + with pytest.raises(TypeError, match="plugin_id"): + load_duplex_plugin(_path(PluginWithoutId), _encode_audio) + + +@pytest.mark.parametrize("plugin_cls", [PluginWithoutDataPlane, PluginWithWrongDataPlane]) +def test_load_duplex_plugin_requires_a_data_plane(plugin_cls: type) -> None: + with pytest.raises(TypeError, match="DuplexDataPlane as data_plane"): + load_duplex_plugin(_path(plugin_cls), _encode_audio) + + +def test_load_duplex_plugin_rejects_paths_without_a_module() -> None: + with pytest.raises(ValueError, match="Invalid duplex plugin path"): + load_duplex_plugin("FakePlugin", _encode_audio) + + +def test_load_duplex_plugin_propagates_missing_attributes() -> None: + with pytest.raises(AttributeError): + load_duplex_plugin(f"{__name__}.NoSuchPlugin", _encode_audio) + + +# --------------------------------------------------------------------------- # +# validate_duplex_plugin_sampling # +# --------------------------------------------------------------------------- # + + +def test_validate_duplex_plugin_sampling_accepts_one_parameter_per_stage() -> None: + plugin = FakePlugin(_encode_audio) + + validate_duplex_plugin_sampling(plugin, sampling_defaults=(SamplingParams(), SamplingParams())) + validate_duplex_plugin_sampling(plugin, sampling_defaults=()) + + +def test_validate_duplex_plugin_sampling_rejects_stage_count_mismatch() -> None: + class WrongStageCountPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return defaults[:1] + + with pytest.raises(ValueError, match="one sampling parameter per stage"): + validate_duplex_plugin_sampling( + WrongStageCountPlugin(_encode_audio), + sampling_defaults=(SamplingParams(), SamplingParams()), + ) + + +def test_validate_duplex_plugin_sampling_rejects_type_mismatch() -> None: + class WrongSamplingTypePlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(object() for _ in defaults) + + with pytest.raises(TypeError, match="sampling parameter type mismatch for stage 0"): + validate_duplex_plugin_sampling( + WrongSamplingTypePlugin(_encode_audio), + sampling_defaults=(SamplingParams(), SamplingParams()), + ) + + +def test_validate_duplex_plugin_sampling_skips_type_check_for_stages_without_defaults() -> None: + class ObjectSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(object() for _ in defaults) + + validate_duplex_plugin_sampling(ObjectSamplingPlugin(_encode_audio), sampling_defaults=(None, None)) + + +def test_validate_duplex_plugin_sampling_rejects_non_tuple_results() -> None: + class ListSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return list(defaults) + + with pytest.raises(TypeError, match="as a tuple"): + validate_duplex_plugin_sampling(ListSamplingPlugin(_encode_audio), sampling_defaults=(SamplingParams(),)) + + +# --------------------------------------------------------------------------- # +# Runtime-config helpers # +# --------------------------------------------------------------------------- # + + +def test_reject_changed_runtime_value_only_raises_on_change() -> None: + reject_changed_runtime_value("voice-a", "voice-a", message="voice is fixed", code="voice_changed") + + with pytest.raises(DuplexRuntimeConfigError, match="voice is fixed") as excinfo: + reject_changed_runtime_value("voice-b", "voice-a", message="voice is fixed", code="voice_changed") + + assert excinfo.value.code == "voice_changed" + assert isinstance(excinfo.value, ValueError) + + +def test_reject_changed_runtime_value_uses_the_requested_error_class() -> None: + class CustomRuntimeConfigError(DuplexRuntimeConfigError): + pass + + with pytest.raises(CustomRuntimeConfigError) as excinfo: + reject_changed_runtime_value(1, 2, message="changed", code="custom", error_cls=CustomRuntimeConfigError) + + assert excinfo.value.code == "custom" + + +def test_duplex_runtime_config_error_defaults_its_code() -> None: + assert DuplexRuntimeConfigError("bad").code == "invalid_duplex_runtime_config" + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"duplex_turn_id": 3}, 3), + ({"duplex_turn_id": "7"}, 7), + ({"model_turn_id": 2}, 2), + ({"duplex_turn_id": 5, "model_turn_id": 9}, 5), + ({"duplex_turn_id": "not-a-number"}, None), + ({}, None), + ("turn", None), + (None, None), + ], +) +def test_payload_turn_id_prefers_duplex_turn_id_and_coerces(payload: object, expected: int | None) -> None: + assert payload_turn_id(payload) == expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [(4, 4), ("11", 11), (2.9, 2), (True, 1), (None, None), ("x", None), (object(), None), ([1], None)], +) +def test_coerce_int(value: object, expected: int | None) -> None: + assert coerce_int(value) == expected + + +# --------------------------------------------------------------------------- # +# PipelineConfig binding # +# --------------------------------------------------------------------------- # + + +def test_pipeline_config_binds_one_duplex_plugin_path() -> None: + field_defaults = {f.name: f.default for f in fields(PipelineConfig)} + + assert "duplex_plugin" in field_defaults + assert field_defaults["duplex_plugin"] is None + + # The pre-framework fields survive only so that the pipelines of the models + # that are not ported yet still construct (RFC vllm-omni#7181 ports them). + # Nothing in the duplex framework may read them: a pipeline is a duplex + # model iff it declares ``duplex_plugin``. + legacy_fields = ("duplex_runtime_extension", "duplex_serving_adapter", "duplex_control_enabled") + for legacy_field in legacy_fields: + assert legacy_field in field_defaults + + framework_sources = [ + *(Path("vllm_omni/engine/duplex").glob("*.py")), + *(Path("vllm_omni/entrypoints/duplex").glob("*.py")), + Path("vllm_omni/engine/duplex_omni_engine.py"), + Path("vllm_omni/engine/duplex_orchestrator.py"), + Path("vllm_omni/entrypoints/duplex_omni.py"), + ] + readers = [ + f"{source}:{legacy_field}" + for source in framework_sources + for legacy_field in legacy_fields + if legacy_field in source.read_text(encoding="utf-8") + ] + assert readers == [] + + # Same rule for the runtime-config knob the framework stopped reading: the + # append-retry table it bounded went away with the correlated append RPC. + assert "completed_append_cache_size" in {f.name for f in fields(DuplexSessionRuntimeConfig)} + assert [ + str(source) + for source in framework_sources + if "completed_append_cache_size" in source.read_text(encoding="utf-8") + ] == [] diff --git a/tests/engine/duplex/test_duplex_runtime.py b/tests/engine/duplex/test_duplex_runtime.py deleted file mode 100644 index 475cf79d0..000000000 --- a/tests/engine/duplex/test_duplex_runtime.py +++ /dev/null @@ -1,563 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from dataclasses import FrozenInstanceError -from types import SimpleNamespace - -import pytest -from vllm.sampling_params import SamplingParams - -from vllm_omni.engine.duplex import runtime as duplex_runtime -from vllm_omni.engine.duplex.contracts import ( - duplex_data_plane_request_info, - duplex_resource_request_belongs_to_session, - duplex_resource_request_id, -) -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexInputMode, - DuplexOutputAction, - DuplexOutputDecision, - DuplexRuntimeCapabilities, - DuplexSessionRuntimeManager, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import ( - MiniCPMO45DuplexRuntimeExtension, - _duplex_hd_slice_count, - build_duplex_data_plane_prompt, - duplex_scheduler_token_budget, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_duplex_fence_is_immutable(): - fence = DuplexFence("session") - - with pytest.raises(FrozenInstanceError): - fence.epoch = 1 # type: ignore[misc] - - assert not hasattr(fence, "__dict__") - - -def test_duplex_runtime_tracks_stage_bindings_and_barge_in_epoch(): - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - DuplexFence("sid-1"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_TOKENS}, - ), - ) - session.bind_stage_request(stage_id=0, request_id="req-stage0", fence=session.fence) - session.bind_stage_request(stage_id=1, request_id="req-stage1", fence=session.fence) - - update = session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence) - next_fence = DuplexFence("sid-1", epoch=1) - stale_request_ids = session.release_fence(session.fence) - session.accept_fence(next_fence) - - assert update.seq == 1 - assert session.fence == next_fence - assert stale_request_ids == ["req-stage0", "req-stage1"] - assert session.stage_bindings == {} - - -def test_duplex_runtime_tracks_same_request_id_for_each_pipeline_stage(): - manager = DuplexSessionRuntimeManager() - fence = DuplexFence("sid-shared-pipeline-request") - session = manager.open_session(fence) - - session.reserve_stage_request(0, "req-shared", fence=fence) - session.bind_stage_request(0, "req-shared", fence=fence) - session.bind_stage_request(1, "req-shared", fence=fence) - - assert session.stage_bindings[0].request_id == "req-shared" - assert session.stage_bindings[1].request_id == "req-shared" - assert session.resource_request_ids() == ["req-shared"] - assert session.input_seq == 0 - - -def test_duplex_runtime_extension_validation_rejects_missing_methods(): - class IncompleteExtension: - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return defaults - - with pytest.raises(TypeError, match="plan_append"): - duplex_runtime.validate_duplex_runtime_extension(IncompleteExtension()) - - -def test_duplex_runtime_extension_validation_rejects_stage_count_mismatch(): - class WrongStageCountExtension(MiniCPMO45DuplexRuntimeExtension): - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return defaults[:1] - - with pytest.raises(ValueError, match="one sampling parameter per stage"): - duplex_runtime.validate_duplex_runtime_extension( - WrongStageCountExtension(), - sampling_defaults=(SamplingParams(), SamplingParams()), - ) - - -def test_duplex_runtime_extension_validation_rejects_sampling_type_mismatch(): - class WrongSamplingTypeExtension(MiniCPMO45DuplexRuntimeExtension): - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return tuple(object() for _ in defaults) - - with pytest.raises(TypeError, match="sampling parameter type"): - duplex_runtime.validate_duplex_runtime_extension( - WrongSamplingTypeExtension(), - sampling_defaults=(SamplingParams(), SamplingParams()), - ) - - -def test_duplex_output_decision_metadata_is_immutable(): - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={"model_listen": True}, - ) - - with pytest.raises(TypeError): - decision.metadata["model_listen"] = False - - -def _decide_minicpmo_output( - output: object, - *, - segment_token_ids: tuple[int, ...] = (), - segment_output_metadata: dict | None = None, -): - return MiniCPMO45DuplexRuntimeExtension().decide_output( - stage_id=0, - final_stage_id=1, - segment_finished=True, - segment_token_ids=segment_token_ids, - segment_output_metadata=segment_output_metadata or {}, - output=output, - ) - - -def test_minicpmo_extension_owns_stage_sampling_overrides(): - defaults = ( - SamplingParams(max_tokens=4), - SamplingParams(max_tokens=8), - ) - - configured = MiniCPMO45DuplexRuntimeExtension().configure_sampling_params( - runtime_config={ - "duplex_stage_max_tokens": {"0": 20}, - "duplex_stage_sampling_params": {"1": {"stop_token_ids": [151645]}}, - }, - defaults=defaults, - ) - - assert configured[0].max_tokens == 20 - assert configured[1].stop_token_ids == [151645] - assert defaults[0].max_tokens == 4 - assert 151645 not in (defaults[1].stop_token_ids or []) - - -def test_minicpmo_output_decision_uses_raw_streaming_token_snapshot(): - decision = _decide_minicpmo_output( - SimpleNamespace(outputs=[SimpleNamespace()]), - segment_token_ids=(151705,), - segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, - ) - - assert decision is not None - assert decision.action is DuplexOutputAction.DIRECT_RESPONSE - assert decision.metadata["duplex_native_decision"] == "listen" - assert decision.metadata["model_listen"] is True - - -@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"]) -def test_minicpmo_output_decision_ignores_output_level_token_history(attr): - output = SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - outputs=[SimpleNamespace()], - **{attr: [42, 151705]}, - ) - - assert _decide_minicpmo_output(output) is None - - -@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"]) -def test_minicpmo_output_decision_uses_completion_token_ids(attr): - output = SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - outputs=[SimpleNamespace(**{attr: [42, 151705]})], - ) - - assert _decide_minicpmo_output(output) is not None - - -def test_minicpmo_output_decision_uses_completion_stop_reason(): - output = SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - outputs=[SimpleNamespace(stop_reason=151705)], - ) - - assert _decide_minicpmo_output(output) is not None - - -def test_duplex_runtime_cancel_fence_rejects_late_append_and_accepts_next_epoch(): - manager = DuplexSessionRuntimeManager() - cancelled_fence = DuplexFence("sid-cancel-race") - next_fence = DuplexFence("sid-cancel-race", epoch=1) - session = manager.open_session( - cancelled_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request( - stage_id=0, - request_id="req-cancelled", - fence=cancelled_fence, - ) - - stale_request_ids = session.cancel_fence(cancelled_fence, next_fence) - - assert stale_request_ids == ["req-cancelled"] - assert session.fence == next_fence - assert session.stage_bindings == {} - with pytest.raises(RuntimeError, match="fence mismatch"): - session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=cancelled_fence, - ) - update = session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=next_fence, - ) - assert update.seq == 1 - - -def test_duplex_runtime_stale_close_preserves_live_session_and_bindings(): - manager = DuplexSessionRuntimeManager() - current_fence = DuplexFence("sid-stale-close", epoch=1) - session = manager.open_session( - current_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request(0, "req-live", fence=current_fence) - - with pytest.raises(RuntimeError, match="fence mismatch"): - manager.close_session(DuplexFence("sid-stale-close")) - - assert manager.get("sid-stale-close") is session - assert session.stage_request_ids() == ["req-live"] - - -def test_duplex_runtime_reopen_rejects_late_append_from_old_incarnation(): - manager = DuplexSessionRuntimeManager() - old_fence = DuplexFence("sid-reopen", incarnation=0) - old_session = manager.open_session( - old_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - manager.close_session(old_fence) - - new_fence = DuplexFence("sid-reopen", incarnation=1) - new_session = manager.open_session( - new_fence, - capabilities=old_session.capabilities, - ) - - with pytest.raises(RuntimeError, match="fence mismatch"): - new_session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=old_fence, - ) - assert ( - new_session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=new_fence, - ).seq - == 1 - ) - - -def test_duplex_prompt_expands_incarnation_metadata(): - fence = DuplexFence("sid-incarnation", incarnation=3) - - prompt = build_duplex_data_plane_prompt( - request_id=duplex_resource_request_id(fence, "stage0"), - fence=fence, - session_config={}, - runtime_config={}, - seq=1, - turn_seq=1, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - payload={"is_speech": True}, - final=False, - ) - - assert prompt["model_intermediate_buffer"]["duplex"]["incarnation"] == 3 - - -def test_duplex_runtime_tracks_turn_local_append_sequence(): - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - DuplexFence("sid-turn-seq"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - first = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence) - second = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence) - next_turn = DuplexFence("sid-turn-seq", turn_id=1, response_seq=1) - third = session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=next_turn, - ) - fourth = session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=next_turn, - ) - - assert [first.seq, second.seq, third.seq, fourth.seq] == [1, 2, 3, 4] - assert [first.turn_seq, second.turn_seq, third.turn_seq, fourth.turn_seq] == [1, 2, 1, 2] - assert [first.turn_id, second.turn_id, third.turn_id, fourth.turn_id] == [0, 0, 1, 1] - - -def test_duplex_runtime_rejects_unsupported_append_mode(): - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - DuplexFence("sid-2"), - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.TURN_COMMIT_ONLY}), - ) - - with pytest.raises(ValueError, match="not supported"): - session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence) - - -def test_duplex_data_plane_request_info_extracts_structured_stage_result(): - request_id, response_stage_id = duplex_data_plane_request_info( - { - "stage_results": [ - {"result": {"supported": True}}, - { - "result": { - "data_plane_append": True, - "request_id": "duplex-sid-e0-stage0-s1", - "response_stage_id": 1, - } - }, - ] - } - ) - - assert request_id == "duplex-sid-e0-stage0-s1" - assert response_stage_id == 1 - - -def test_duplex_data_plane_request_info_rejects_missing_request_id(): - assert duplex_data_plane_request_info( - { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": "", - "response_stage_id": 1, - } - } - ] - } - ) == (None, None) - - -def test_duplex_scheduler_token_budget_estimates_pcm_slots(): - assert ( - duplex_scheduler_token_budget( - { - "audio": "AAAAAA==", - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - == 16 - ) - - -def test_duplex_scheduler_token_budget_ignores_client_budget_fields(): - assert ( - duplex_scheduler_token_budget( - { - "audio": "AAAAAA==", - "format": "pcm_f32le", - "duplex_num_input_tokens": 999, - "num_input_tokens": 999, - } - ) - == 16 - ) - - -def test_resource_state_rejects_fence_regression_and_requires_explicit_fence(): - current = DuplexFence("sid", epoch=2, turn_id=3, response_seq=4) - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - current, - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), - ) - - for stale in ( - DuplexFence("sid", epoch=1, turn_id=99, response_seq=99), - DuplexFence("sid", epoch=2, turn_id=2, response_seq=4), - DuplexFence("sid", epoch=2, turn_id=3, response_seq=3), - ): - with pytest.raises(RuntimeError, match="fence mismatch"): - session.accept_fence(stale) - assert session.fence == current - - with pytest.raises(TypeError, match="fence"): - session.bind_stage_request(0, "request") - with pytest.raises(TypeError, match="fence"): - session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK) - with pytest.raises(TypeError, match="DuplexFence"): - manager.open_session("legacy-session") - with pytest.raises(TypeError, match="DuplexFence"): - manager.close_session("legacy-session") - - -def test_resource_request_id_is_derived_from_fence_and_role(): - fence = DuplexFence("sid-with-dashes", epoch=7, turn_id=11, response_seq=13) - - assert duplex_resource_request_id(fence, "stage0") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage0" - assert duplex_resource_request_id(fence, "stage1") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage1" - - -def test_resource_request_id_codec_separates_session_id_from_incarnation(): - embedded_incarnation = duplex_resource_request_id( - DuplexFence("foo-i1", incarnation=0), - "stage0", - ) - actual_incarnation = duplex_resource_request_id( - DuplexFence("foo", incarnation=1), - "stage0", - ) - - assert embedded_incarnation != actual_incarnation - - -def test_resource_request_id_session_membership_uses_encoded_identity(): - request_id = duplex_resource_request_id( - DuplexFence("sid-with-dashes", incarnation=2, epoch=7), - "stage0", - ) - - assert duplex_resource_request_belongs_to_session(request_id, "sid-with-dashes") is True - assert duplex_resource_request_belongs_to_session(request_id, "sid") is False - assert duplex_resource_request_belongs_to_session("duplex-s.invalid.i.x.e.7.r.stage0", "sid") is False - - -def test_placeholder_budget_is_planned_inside_omni_engine_boundary(): - fence = DuplexFence("sid", turn_id=1, response_seq=1) - prompt = build_duplex_data_plane_prompt( - request_id=duplex_resource_request_id(fence, "stage0"), - fence=fence, - session_config={}, - runtime_config={}, - seq=2, - turn_seq=1, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - payload={ - "audio": "AAAAAA==", - "format": "pcm_f32le", - "duplex_num_input_tokens": 999, - "num_input_tokens": 999, - }, - final=False, - ) - - assert len(prompt["prompt_token_ids"]) == 16 - assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence - assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16 - - -def _jpeg_b64(width: int, height: int) -> str: - import base64 - from io import BytesIO - - from PIL import Image - - buffer = BytesIO() - Image.new("RGB", (width, height), (32, 64, 96)).save(buffer, format="JPEG") - return base64.b64encode(buffer.getvalue()).decode("ascii") - - -def _pcm_unit_payload(frames: list[str]) -> dict[str, object]: - import base64 - - return { - "audio": base64.b64encode(b"\x00" * (16000 * 4)).decode("ascii"), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "video_frames": frames, - } - - -def test_duplex_scheduler_token_budget_one_frame_is_one_block(): - audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) - one_frame = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(960, 540)])) - - assert one_frame - audio_only == 66 - - -@pytest.mark.parametrize( - ("base_size", "expected_blocks"), - [ - ((960, 540), 4), # official camera frame: source + 2 HD patches + composite - ((540, 960), 4), # portrait phone capture - ((640, 360), 4), - ((448, 448), 2), # fits one tile: the processor does not slice it - ((320, 240), 2), - ], -) -def test_duplex_scheduler_token_budget_matches_stage0_hd_slicing(base_size, expected_blocks): - """A stacked pair is processed with max_slice_nums=[2, 1]; the budget must - equal the blocks Stage0 builds for that base-frame size.""" - audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) - pair = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(*base_size), _jpeg_b64(448, 448)])) - - assert pair - audio_only == expected_blocks * 66 - - -@pytest.mark.parametrize( - ("size", "expected_patches"), - [ - # multiple = min(ceil(w * h / 448**2), max_slice_nums); <= 1 is never sliced. - ((448, 448), 0), # exactly one tile - ((320, 240), 0), # smaller than one tile - ((447, 449), 0), # just under one tile - ((449, 449), 2), # just over: the processor cuts a 2-tile grid - ((640, 360), 2), - ((960, 540), 2), # official camera frame - ((540, 960), 2), # same frame rotated - ((1920, 1080), 2), # capped by max_slice_nums=2, not by area - ], -) -def test_duplex_hd_slice_count_matches_official_grid(size, expected_patches): - """Lock the port of ``MiniCPMVImageProcessor.get_sliced_grid``. - - Stage0 slices a stacked unit with ``max_slice_nums=[2, 1]``, so the patch - count for the base frame decides the scheduler budget. A frame that fits one - ``scale_resolution=448`` tile is not sliced at all. - """ - assert _duplex_hd_slice_count(size, 2) == expected_patches - - -def test_duplex_scheduler_token_budget_unreadable_base_frame_keeps_hd_fallback(): - audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) - pair = duplex_scheduler_token_budget(_pcm_unit_payload(["not-an-image", _jpeg_b64(448, 448)])) - - assert pair - audio_only == 4 * 66 diff --git a/tests/engine/duplex/test_engine_session.py b/tests/engine/duplex/test_engine_session.py new file mode 100644 index 000000000..fccaa59e8 --- /dev/null +++ b/tests/engine/duplex/test_engine_session.py @@ -0,0 +1,612 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session-state semantics of the engine-resident duplex session (``DuplexEngineSession``).""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest + +from vllm_omni.engine.duplex.config import ( + DuplexCapabilities, + DuplexConfigError, + DuplexOverlapPolicy, + DuplexSessionConfig, + DuplexTurnEventType, + DuplexTurnState, + ResponseCreateOptions, +) +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.engine.duplex.events import TurnEvent +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( + minicpmo45_native_capabilities, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _session(session_id: str = "duplex-test", config: DuplexSessionConfig | None = None) -> DuplexEngineSession: + return DuplexEngineSession(session_id=session_id, config=config or DuplexSessionConfig(model="test-model")) + + +def test_commit_audio_input_does_not_advance_model_turn_identity(): + session = _session() + + first = session.commit_audio_input(transcript="first chunk") + second = session.commit_audio_input(transcript="second chunk") + + assert session.input_commit_seq == 2 + assert first.input_commit_seq == 1 + assert second.input_commit_seq == 2 + assert first.turn_id == second.turn_id == session.turn_id == 0 + assert first.message["transcript"] == "first chunk" + assert first.message["content"] == [ + {"type": "audio_url", "audio_url": {"url": "native-duplex:input-audio"}, "transcript": "first chunk"} + ] + assert len(session.history) == 2 + assert session.turn_state == DuplexTurnState.USER_COMMITTED + + +def test_commit_audio_input_accepts_explicit_turn_id(): + session = _session() + + committed = session.commit_audio_input(turn_id=4) + + assert committed.turn_id == 4 + assert session.turn_id == 0 + assert "transcript" not in committed.message + + +# ---- response options ---- + + +def test_response_options_apply_to_one_response_without_mutating_session_defaults(): + session = _session(config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64)) + session.reserve_response_options( + ResponseCreateOptions(instructions="one response", voice="override-voice", max_tokens=8) + ) + + assert session.config.instructions == "base" + assert session.config.voice == "base-voice" + assert session.response_config is session.config + session.begin_response() + assert session.config.instructions == "base" + assert session.config.voice == "base-voice" + assert session.response_config.instructions == "one response" + assert session.response_config.voice == "override-voice" + assert session.response_config.max_tokens == 8 + + session.end_response() + assert session.response_config is session.config + assert session.config.instructions == "base" + assert session.config.voice == "base-voice" + assert session.config.max_tokens == 64 + + +def test_response_options_cannot_overwrite_an_unconsumed_reservation(): + session = _session(config=DuplexSessionConfig(instructions="base")) + session.reserve_response_options(ResponseCreateOptions(instructions="first")) + + with pytest.raises(RuntimeError, match="already reserved"): + session.reserve_response_options(ResponseCreateOptions(instructions="second")) + + session.begin_response() + assert session.response_config.instructions == "first" + + +def test_response_options_cannot_be_reserved_while_response_is_active(): + session = _session(config=DuplexSessionConfig(instructions="base")) + session.begin_response() + + with pytest.raises(RuntimeError, match="active"): + session.reserve_response_options(ResponseCreateOptions(instructions="too late")) + + +def test_discarded_response_options_do_not_apply(): + session = _session(config=DuplexSessionConfig(instructions="base")) + session.reserve_response_options(ResponseCreateOptions(instructions="dropped")) + session.discard_response_options() + + session.begin_response() + + assert session.response_config.instructions == "base" + + +# ---- response / overlap identity ---- + + +def test_session_owns_response_and_overlap_identity(): + session = _session() + + response_id = session.begin_response() + assert session.turn_state == DuplexTurnState.ASSISTANT_GENERATING + session.accumulate_overlap_speech(320) + session.accumulate_overlap_speech(180) + session.end_response() + + assert session.active_response_id is None + assert session.last_response_id == response_id + assert session.overlap_speech_ms == 500 + assert session.reset_overlap_speech() == 500 + assert session.overlap_speech_ms == 0 + assert session.turn_state == DuplexTurnState.IDLE + + +def test_session_composes_single_owner_ledgers_with_immutable_views(): + session = _session() + + session.bind_request("req-1") + response_id = session.begin_response(turn_id=3) + session.mark_audio_sent(duration_ms=240) + + assert session.active_request_id == "req-1" + assert session.active_response_id == response_id + assert session.active_response_turn_id == 3 + assert session.playback.sent_ms == 240 + assert session.turn_state == DuplexTurnState.ASSISTANT_PLAYING + with pytest.raises(FrozenInstanceError): + session.playback.sent_ms = 480 # type: ignore[misc] + assert session.playback.sent_ms == 240 + + +def test_barge_in_advances_epoch_and_drops_uncommitted_assistant_text(): + session = _session() + response_id = session.begin_response() + session.bind_request("chatcmpl-duplex-test") + session.append_assistant_text("unplayed answer") + + new_epoch = session.barge_in() + + assert response_id is not None + assert new_epoch == 1 + assert session.epoch == 1 + assert session.accepted_fence == DuplexFence(session.session_id, epoch=1, turn_id=0) + assert session.active_request_id is None + assert session.active_response_id is None + assert session.assistant_text_buffer == () + assert session.history == () + assert session.turn_state == DuplexTurnState.BARGE_IN + + +# ---- playback ledger ---- + + +def test_playback_ack_tracks_committed_cursor_separately(): + session = _session() + session.mark_audio_sent(duration_ms=10_000) + + session.acknowledge_playback(played_ms=2_000) + + assert session.playback.generated_ms == 10_000 + assert session.playback.sent_ms == 10_000 + assert session.playback.played_ms == 2_000 + assert session.playback.committed_ms == 2_000 + assert session.playback.as_dict() == { + "generated_ms": 10_000, + "sent_ms": 10_000, + "played_ms": 2_000, + "committed_ms": 2_000, + } + + +def test_playback_ack_is_scoped_per_response(): + session = _session() + first = session.begin_response() + session.mark_audio_sent(duration_ms=3_000) + session.end_response() + second = session.begin_response() + session.mark_audio_sent(duration_ms=1_000) + + session.acknowledge_playback(played_ms=2_500, response_id=first) + + assert session.playback_for_response(first).played_ms == 2_500 + assert session.playback_for_response(second).played_ms == 0 + assert session.playback.played_ms == 0 + + +def test_history_commit_uses_audio_text_alignment_marks(): + session = _session() + session.begin_response() + + session.append_assistant_text("hello ") + session.mark_audio_sent(duration_ms=1_000, text_chars=6) + session.append_assistant_text("world") + session.mark_audio_sent(duration_ms=2_000, text_chars=11) + assert [(mark.text_chars, mark.audio_end_ms) for mark in session.assistant_audio_text_marks] == [ + (6, 1_000), + (11, 2_000), + ] + session.acknowledge_playback(played_ms=1_200, committed_ms=1_200) + + committed = session.end_response(commit_text=True) + + assert committed == {"role": "assistant", "content": "hello w"} + assert session.history[-1] == committed + assert session.last_assistant_full_message == {"role": "assistant", "content": "hello world"} + + +def test_history_commit_with_ack_only_policy_defers_unacknowledged_text(): + session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only")) + response_id = session.begin_response() + session.append_assistant_text("never played") + session.mark_audio_sent(duration_ms=1_000, text_chars=12) + + committed = session.end_response(commit_text=True) + + assert committed is None + assert session.history == () + assert f"item_{response_id}" in session.pending_history_item_ids + + +def test_truncate_history_item_uses_response_alignment_marks(): + session = _session() + response_id = session.begin_response() + session.append_assistant_text("hello ") + session.mark_audio_sent(duration_ms=1_000, text_chars=6) + session.append_assistant_text("world") + session.mark_audio_sent(duration_ms=2_000, text_chars=11) + session.acknowledge_playback(played_ms=2_000) + committed = session.end_response(commit_text=True) + assert committed == {"role": "assistant", "content": "hello world"} + item_id = f"item_{response_id}" + session.register_history_item(item_id, committed) + + assert session.truncate_history_item(item_id, audio_end_ms=1_000) is True + + assert session.history[-1] == {"role": "assistant", "content": "hello"} + assert session.history_item_ids[item_id] == {"role": "assistant", "content": "hello"} + assert session.delete_history_item(item_id) is True + assert session.history == () + + +# ---- turn signals / model turns ---- + + +@pytest.mark.parametrize( + ("event_type", "turn_state"), + [ + (DuplexTurnEventType.USER_STARTED, DuplexTurnState.USER_SPEAKING), + (DuplexTurnEventType.USER_COMMITTED, DuplexTurnState.USER_COMMITTED), + (DuplexTurnEventType.ASSISTANT_STARTED, DuplexTurnState.ASSISTANT_GENERATING), + (DuplexTurnEventType.ASSISTANT_DONE, DuplexTurnState.IDLE), + (DuplexTurnEventType.BARGE_IN, DuplexTurnState.BARGE_IN), + ], +) +def test_signal_turn_transitions_and_returns_typed_turn_event(event_type, turn_state): + session = _session() + + event = session.signal_turn(event_type.value) + + assert isinstance(event, TurnEvent) + assert session.turn_state == turn_state + wire = event.to_realtime() + assert wire["type"] == "turn.event" + assert wire["event"] == event_type.value + assert wire["turn_state"] == turn_state.value + + +def test_signal_turn_playback_ack_updates_cursor_and_close_marks_session_closing(): + session = _session() + session.mark_audio_sent(duration_ms=5_000) + + ack = session.signal_turn("playback_ack", {"played_ms": 1_500, "committed_ms": 1_000}) + assert ack.event == "playback_ack" + assert session.playback.played_ms == 1_500 + assert session.playback.committed_ms == 1_000 + + close = session.signal_turn(DuplexTurnEventType.CLOSE.value) + assert close.event == "close" + assert session.state.value == "closing" + + +def test_complete_model_turn_advances_turn_id_and_fence(): + session = _session() + assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=0) + + session.complete_model_turn(0) + + assert session.turn_id == 1 + assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=1) + assert session.accepted_fence == session.fence + + # A stale terminal for an already-completed turn does not move identity backwards. + session.complete_model_turn(0) + assert session.turn_id == 1 + session.complete_model_turn(3) + assert session.turn_id == 4 + assert session.accepted_fence.turn_id == 4 + + +# ---- fence validation ---- + + +def test_accept_fence_rejects_stale_or_foreign_fences(): + session = _session("sid-fence") + session.complete_model_turn(0) + + with pytest.raises(DuplexFenceMismatchError): + session.accept_fence(DuplexFence("sid-fence", epoch=0, turn_id=0)) + with pytest.raises(DuplexFenceMismatchError): + session.accept_fence(DuplexFence("sid-other", epoch=0, turn_id=1)) + + session.accept_fence(DuplexFence("sid-fence", epoch=1, turn_id=0)) + assert session.accepted_fence == DuplexFence("sid-fence", epoch=1, turn_id=0) + + +def test_prepare_and_commit_append_sequence_chunks_per_turn_and_epoch(): + session = _session("sid-append") + fence = session.fence + + first = session.commit_append(session.prepare_append(fence)) + second = session.commit_append(session.prepare_append(fence)) + assert (first.seq, first.turn_seq, first.turn_id) == (1, 1, 0) + assert (second.seq, second.turn_seq, second.turn_id) == (2, 2, 0) + + session.complete_model_turn(0) + third = session.commit_append(session.prepare_append(session.fence)) + assert (third.seq, third.turn_seq, third.turn_id) == (3, 1, 1) + + # Barge-in advances the epoch (turn_id is kept) and restarts the append sequence. + session.barge_in() + assert session.fence == DuplexFence("sid-append", epoch=1, turn_id=1) + fourth = session.commit_append(session.prepare_append(session.fence)) + assert (fourth.seq, fourth.turn_seq, fourth.turn_id) == (1, 1, 1) + assert session.input_seq == 1 + + +def test_commit_append_rejects_stale_reservation(): + session = _session("sid-append-stale") + fence = session.fence + reservation = session.prepare_append(fence) + session.commit_append(session.prepare_append(fence)) + + with pytest.raises(RuntimeError, match="stale"): + session.commit_append(reservation) + + with pytest.raises(DuplexFenceMismatchError): + session.prepare_append(DuplexFence("sid-other", epoch=0, turn_id=0)) + + +def test_cancel_fence_releases_stage_requests_and_advances_identity(): + session = _session("sid-cancel") + cancelled = session.fence + session.reserve_stage_request(0, "req-a", fence=cancelled) + session.bind_stage_request(1, "req-b", fence=cancelled) + assert session.stage_request_submitted(1, "req-b") is True + assert session.stage_request_submitted(0, "req-a") is False + next_fence = DuplexFence("sid-cancel", epoch=1, turn_id=0) + + stale = session.cancel_fence(cancelled, next_fence) + + assert stale == ["req-a", "req-b"] + assert session.resource_request_ids() == [] + assert session.accepted_fence == next_fence + + with pytest.raises(DuplexFenceMismatchError): + session.cancel_fence(next_fence, DuplexFence("sid-cancel", epoch=1, turn_id=1)) + with pytest.raises(DuplexFenceMismatchError): + session.cancel_fence(DuplexFence("sid-other", epoch=1, turn_id=0), DuplexFence("sid-cancel", epoch=2)) + + +# ---- public view ---- + + +def test_as_public_dict_exposes_identity_capabilities_and_playback(): + session = _session("duplex-public", config=DuplexSessionConfig(model="m", voice="alloy")) + session.mark_audio_sent(duration_ms=100) + + payload = session.as_public_dict() + + assert payload["id"] == "duplex-public" + assert payload["model"] == "m" + assert payload["voice"] == "alloy" + assert payload["state"] == "open" + assert payload["turn_state"] == "assistant_playing" + assert payload["epoch"] == 0 + assert payload["turn_id"] == 0 + assert payload["capabilities"] == DuplexCapabilities().as_dict() + assert payload["playback"] == {"generated_ms": 100, "sent_ms": 100, "played_ms": 0, "committed_ms": 0} + + +# ---- DuplexSessionConfig.from_realtime ---- + + +def test_from_realtime_rejects_unsupported_audio_formats(): + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"input_audio_format": "mp3"}) + assert excinfo.value.code == "unsupported_audio_format" + + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"audio": {"output": {"format": {"type": "audio/opus"}}}}) + assert excinfo.value.code == "unsupported_audio_format" + + +def test_from_realtime_validates_turn_detection(): + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"turn_detection": {"type": "semantic_vad"}}) + assert excinfo.value.code == "unsupported_turn_detection" + assert excinfo.value.param == "turn_detection" + + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"overlap_policy": "barge_in_on_speech"}) + assert excinfo.value.code == "unsupported_turn_detection" + + server_vad = DuplexSessionConfig.from_realtime({"turn_detection": {"type": "server_vad"}}) + assert server_vad.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + assert server_vad.extra_body["realtime_turn_detection"]["threshold"] == 0.5 + + model_owned = DuplexSessionConfig.from_realtime({"turn_detection": None}) + assert model_owned.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert model_owned.extra_body["realtime_turn_detection"] is None + + +def test_from_realtime_maps_wire_fields_and_stores_realtime_keys_in_extra_body(): + config = DuplexSessionConfig.from_realtime( + { + "model": "openbmb/MiniCPM-o-4_5", + "instructions": "be brief", + "output_audio_format": "pcm16", + "max_response_output_tokens": "inf", + "audio": {"output": {"voice": "alloy", "speed": 1.25}}, + "tools": [{"type": "function", "name": "lookup"}], + "tool_choice": "auto", + "metadata": {"tenant": "t1"}, + "include": ["item.input_audio_transcription.logprobs"], + "input_audio_transcription": {"model": "whisper-1"}, + "extra_body": {"custom": True}, + } + ) + + assert config.model == "openbmb/MiniCPM-o-4_5" + assert config.instructions == "be brief" + assert config.voice == "alloy" + assert config.speed == 1.25 + assert config.max_tokens is None + assert config.response_format == "pcm" + assert config.modalities == ["text", "audio"] + assert config.extra_body["custom"] is True + assert config.extra_body["realtime_tools"] == [{"type": "function", "name": "lookup"}] + assert config.extra_body["realtime_tool_choice"] == "auto" + assert config.extra_body["realtime_metadata"] == {"tenant": "t1"} + assert config.extra_body["realtime_include"] == ["item.input_audio_transcription.logprobs"] + assert config.extra_body["realtime_input_audio_transcription"] == {"model": "whisper-1"} + assert config.extra_body["realtime_audio"] == {"output": {"voice": "alloy", "speed": 1.25}} + assert config.extra_body["realtime_output_audio_format"] == "pcm16" + assert "extra_body" not in config.extra_body["realtime_session_payload"] + assert config.extra_body["realtime_session_payload"]["model"] == "openbmb/MiniCPM-o-4_5" + + +def test_from_realtime_uses_served_model_when_payload_has_none(): + config = DuplexSessionConfig.from_realtime({}, model="served-model") + + assert config.model == "served-model" + assert config.response_format == "pcm" + assert config.idle_timeout_s == 300.0 + + +def test_from_realtime_ignores_client_chosen_session_ids(): + config = DuplexSessionConfig.from_realtime({"id": "client-chosen", "session_id": "also-client-chosen"}) + + assert not hasattr(config, "session_id") + assert "id" not in config.as_dict() + session = _session("duplex-engine-allocated", config=config) + assert session.as_public_dict()["id"] == "duplex-engine-allocated" + + +# ---- DuplexSessionConfig.apply_realtime_update ---- + + +def test_apply_realtime_update_rejects_model_change(): + config = DuplexSessionConfig(model="served-model") + + with pytest.raises(DuplexConfigError) as excinfo: + config.apply_realtime_update({"model": "other-model"}, session_id="sid") + + assert excinfo.value.code == "model_update_unsupported" + assert config.model == "served-model" + + +def test_apply_realtime_update_rejects_voice_change_after_audio_started(): + config = DuplexSessionConfig(model="m", voice="alloy") + + config.apply_realtime_update({"voice": "verse"}, audio_started=False) + assert config.voice == "verse" + with pytest.raises(DuplexConfigError) as excinfo: + config.apply_realtime_update({"audio": {"output": {"voice": "alloy"}}}, audio_started=True) + assert excinfo.value.code == "voice_update_after_audio_unsupported" + assert config.voice == "verse" + + +def test_apply_realtime_update_rejects_ref_audio_change(): + config = DuplexSessionConfig(model="m") + + with pytest.raises(DuplexConfigError) as excinfo: + config.apply_realtime_update({"ref_audio": "/tmp/voice.wav"}) + + assert excinfo.value.code == "ref_audio_update_unsupported" + assert config.ref_audio is None + + +def test_apply_realtime_update_patches_fields_and_realtime_keys(): + config = DuplexSessionConfig(model="m", instructions="old", extra_body={"realtime_tools": [{"name": "x"}]}) + + config.apply_realtime_update( + { + "instructions": "new", + "output_audio_format": "g711_ulaw", + "max_output_tokens": 32, + "tools": None, + "metadata": {"k": "v"}, + "overlap_policy": "not-a-policy", + } + ) + + assert config.instructions == "new" + assert config.response_format == "pcm" + assert config.max_tokens == 32 + assert "realtime_tools" not in config.extra_body + assert config.extra_body["realtime_metadata"] == {"k": "v"} + assert config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert config.extra_body["realtime_session_payload"]["instructions"] == "new" + + +# ---- overlap policy / capabilities ---- + + +def test_overlap_policy_defaults_and_invalid_values_to_listen_only(): + assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value + assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value + assert ( + DuplexSessionConfig._normalize_overlap_policy(" Barge_In_On_Speech ") + == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + ) + assert DuplexSessionConfig.from_event({"session": {"overlap_policy": "bogus"}}).overlap_policy == "listen_only" + + +def test_capabilities_as_dict_reports_model_native_duplex_constants(): + caps = DuplexCapabilities().as_dict() + + assert caps["implementation_level"] == "model_native_duplex" + assert caps["input_modes"] == ["append_audio_chunk"] + assert caps["supports_kv_lease"] is False + assert caps["supports_core_kv_lease"] is False + assert "client_event" in caps["signal_sources"] + assert set(caps) == set(minicpmo45_native_capabilities().as_dict()) + + +def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): + caps = minicpmo45_native_capabilities(max_sessions=2).as_dict() + + assert caps["implementation_level"] == "model_native_duplex" + assert caps["supports_input_append"] is True + assert caps["input_modes"] == ["append_audio_chunk"] + assert caps["adapter_patterns"] == ["scheduler_data_plane"] + assert caps["supports_model_internal_state"] is True + assert caps["requires_model_runner_kv"] is True + assert caps["requires_native_stage_role"] is True + assert caps["supports_kv_lease"] is False + assert caps["supports_core_kv_lease"] is False + assert caps["supports_stage_resumption"] is True + assert caps["supports_scheduler_native_append"] is False + assert caps["supports_core_resumable_request"] is True + assert caps["supports_stage_connector_handoff"] is True + assert caps["supports_audio_truncate"] is True + assert caps["supports_barge_in"] is True + assert caps["target_barge_in_latency_ms"] is None + assert caps["supports_multi_session"] is True + assert caps["supports_multi_session_same_replica"] is True + assert caps["supports_session_lease"] is True + assert caps["supports_session_resume"] is True + assert caps["session_admission_mode"] == "engine_managed" + assert caps["stage_handoff_transport"] == "scheduler_data_plane" + + +def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment(): + caps = minicpmo45_native_capabilities(max_sessions=1).as_dict() + + assert caps["supports_multi_session"] is False + assert caps["supports_multi_session_same_replica"] is False diff --git a/tests/engine/duplex/test_overlap_policy.py b/tests/engine/duplex/test_overlap_policy.py new file mode 100644 index 000000000..096b7c7c9 --- /dev/null +++ b/tests/engine/duplex/test_overlap_policy.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The overlap decision, reached directly. + +These rules decide what happens when input arrives while the model is speaking. +They were unreachable except by driving a whole append through the runner, so +none of them had a test that named the rule it was checking. +""" + +from __future__ import annotations + +import base64 + +import numpy as np +import pytest + +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.engine.duplex.session import overlap_policy +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _session(**config: object) -> DuplexEngineSession: + values: dict[str, object] = {"model": "fake-model"} + values.update(config) + return DuplexEngineSession(session_id="sid", config=DuplexSessionConfig(**values)) + + +def _pcm16(rms: float, samples: int = 1600) -> str: + data = np.full(samples, int(rms * 32768), dtype=np.int16) + return base64.b64encode(data.tobytes()).decode("ascii") + + +# --------------------------------------------------------------------------- # +# Speech detection # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("flag", [True, False]) +def test_an_explicit_is_speech_flag_wins_over_the_audio(flag: bool) -> None: + session = _session() + payload = {"format": "pcm16", "audio": _pcm16(0.9 if not flag else 0.0)} + assert overlap_policy.input_looks_like_speech(session, {"is_speech": flag}, payload) is flag + + +def test_vad_probability_decides_when_no_flag_is_given() -> None: + session = _session() + payload = {"format": "pcm16", "audio": _pcm16(0.0)} + assert overlap_policy.input_looks_like_speech(session, {"vad": {"speech_probability": 0.9}}, payload) is True + assert overlap_policy.input_looks_like_speech(session, {"vad": {"speech_probability": 0.1}}, payload) is False + + +def test_quiet_audio_below_the_configured_rms_is_not_speech() -> None: + session = _session(overlap_silence_rms=0.05) + loud = {"format": "pcm16", "audio": _pcm16(0.5)} + quiet = {"format": "pcm16", "audio": _pcm16(0.001)} + assert overlap_policy.input_looks_like_speech(session, {}, loud) is True + assert overlap_policy.input_looks_like_speech(session, {}, quiet) is False + + +def test_undecodable_audio_is_treated_as_speech() -> None: + """Fail toward listening: dropping real speech is worse than a spurious turn.""" + session = _session() + assert overlap_policy.input_looks_like_speech(session, {}, {"format": "pcm16", "audio": "!!not base64!!"}) is True + + +# --------------------------------------------------------------------------- # +# Barge-in # +# --------------------------------------------------------------------------- # + + +def test_an_explicit_client_overlap_action_requests_barge_in() -> None: + session = _session() + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + decision = overlap_policy.decide(session, {"overlap_action": "barge_in"}, payload, auto_responds=True) + assert decision["action"] == "barge_in" + assert decision["reason"] == "client_overlap_action" + + +def test_barge_in_is_deferred_when_the_session_cannot_support_it() -> None: + session = _session() + session.capabilities = type(session.capabilities)(supports_barge_in=False) + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + decision = overlap_policy.decide(session, {"force_barge_in": True}, payload, auto_responds=True) + assert decision["action"] != "barge_in" + + +# --------------------------------------------------------------------------- # +# Force-listen # +# --------------------------------------------------------------------------- # + + +def test_a_short_response_create_is_forced_to_listen() -> None: + """A sub-threshold ack should not take the turn from the model.""" + session = _session(overlap_short_ack_ms=700) + short = {"response_create": True, "duration_ms": 200} + long = {"response_create": True, "duration_ms": 1500} + assert overlap_policy.should_force_listen_for_short_commit(session, short, {}) is True + assert overlap_policy.should_force_listen_for_short_commit(session, long, {}) is False + + +def test_force_barge_in_overrides_the_short_ack_rule() -> None: + session = _session(overlap_short_ack_ms=700) + event = {"response_create": True, "duration_ms": 200, "force_barge_in": True} + assert overlap_policy.should_force_listen_for_short_commit(session, event, {}) is False + + +def test_auto_response_overlap_forces_listen_only_when_auto_responding() -> None: + event = {"force_listen": True} + assert overlap_policy.should_force_listen_for_auto_response_overlap(event, {}, auto_responds=True) is True + assert overlap_policy.should_force_listen_for_auto_response_overlap(event, {}, auto_responds=False) is False + + +# --------------------------------------------------------------------------- # +# Payload merging # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("speech_first", [True, False], ids=["speech-then-silence", "silence-then-speech"]) +def test_merging_two_appends_keeps_speech_if_either_half_was_speech(speech_first: bool) -> None: + """Matching rates and decodable audio, so the halves really are concatenated. + + Without a ``sample_rate_hz`` on both, ``merge_audio_payloads`` falls back + to the newer chunk before merging, and the speech flag it returned was the + second chunk's own -- the test passed without exercising the merge. + """ + speech = { + "format": "pcm_f32le", + "audio": base64.b64encode(b"\x01" * 8).decode("ascii"), + "sample_rate_hz": 16000, + "is_speech": True, + } + silence = { + "format": "pcm_f32le", + "audio": base64.b64encode(b"\x00" * 8).decode("ascii"), + "sample_rate_hz": 16000, + "is_speech": False, + } + first, second = (speech, silence) if speech_first else (silence, speech) + merged = overlap_policy.merge_audio_payloads(first, second) + assert base64.b64decode(merged["audio"]) == base64.b64decode(first["audio"]) + base64.b64decode(second["audio"]) + assert merged["sample_rate_hz"] == 16000 + assert merged["is_speech"] is True + + +@pytest.mark.parametrize( + ("first", "second"), + [ + ({"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, {"format": "pcm16", "audio": ""}), + ( + {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, + {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 24000}, + ), + ( + {"format": "pcm_f32le", "audio": "!!not base64!!", "sample_rate_hz": 16000}, + {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, + ), + ], + ids=["format-mismatch", "rate-mismatch", "undecodable"], +) +def test_unmergeable_payloads_fall_back_to_the_newer_one(first: dict, second: dict) -> None: + """Whatever cannot be concatenated keeps the newer chunk and drops the older. + + Worth pinning: the fallback discards ``first`` silently, so a bug that makes + two mergeable payloads look mismatched loses buffered audio rather than + raising. + """ + assert overlap_policy.merge_audio_payloads(first, second) == second + + +def test_mergeable_payloads_are_concatenated_in_order() -> None: + head = bytes([1, 2, 3, 4]) + tail = bytes([5, 6, 7, 8]) + first = {"format": "pcm_f32le", "audio": base64.b64encode(head).decode(), "sample_rate_hz": 16000} + second = {"format": "pcm_f32le", "audio": base64.b64encode(tail).decode(), "sample_rate_hz": 16000} + merged = overlap_policy.merge_audio_payloads(first, second) + assert base64.b64decode(merged["audio"]) == head + tail + assert merged["sample_rate_hz"] == 16000 + + +# --------------------------------------------------------------------------- # +# Server VAD + barge_in_on_speech # +# --------------------------------------------------------------------------- # + + +def test_server_vad_speech_started_reaches_a_barge_in_decision() -> None: + """The configuration this PR introduces must decide, not raise. + + Both call sites read the VAD verdict through the module-level + ``vad_speech_started`` helper while binding its result to a local of the + same name, which made the name local for the whole scope and raised + ``UnboundLocalError`` on the call itself. Every ``server_vad`` + + ``barge_in_on_speech`` append hit it, and the client saw ``internal_error`` + instead of a barge-in. The MiniCPM-o preset is ``listen_only``, so no demo + exercised this branch. + """ + session = _session(overlap_policy="barge_in_on_speech") + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + event = {"vad": {"speech_started": True}} + + decision = overlap_policy.decide(session, event, payload, auto_responds=True) + + assert decision["action"] == "barge_in" + assert decision["reason"] == "server_vad_speech_started" + assert decision["cancel_reason"] == "turn_detected" + + +def test_server_vad_mid_utterance_keeps_listening_instead_of_barging_in() -> None: + """``speech_started`` false means the utterance is already running: not a new turn.""" + session = _session(overlap_policy="barge_in_on_speech") + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + event = {"vad": {"speech_started": False}} + + decision = overlap_policy.decide(session, event, payload, auto_responds=True) + + assert decision["action"] != "barge_in" + + +def test_barge_in_on_speech_without_vad_still_barges_in_on_speech() -> None: + """No VAD verdict: the policy falls back to its own speech classification.""" + session = _session(overlap_policy="barge_in_on_speech") + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + + decision = overlap_policy.decide(session, {}, payload, auto_responds=True) + + assert decision["action"] == "barge_in" + assert decision["reason"] == "barge_in_on_speech" diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py new file mode 100644 index 000000000..0323f16d8 --- /dev/null +++ b/tests/engine/duplex/test_session_manager.py @@ -0,0 +1,1569 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexSessionManager``: admission, control RPC, command backpressure, reaper and request cleanup.""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +from dataclasses import FrozenInstanceError, dataclass, field +from typing import Any + +import pytest + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.commands import AppendAudio, CloseSession, Commit, DuplexCommand, Heartbeat +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig, DuplexSessionState +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, + duplex_resource_request_belongs_to_session, +) +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + ErrorEvent, + SessionClosed, + SessionCreated, + SessionExpired, + SessionHeartbeatAck, +) +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionEventMessage, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexDataPlane, + DuplexModelPlugin, + DuplexModelSessionState, + DuplexRuntimeConfigError, + PcmAppendBuffer, + PcmAppendReservation, +) +from vllm_omni.engine.duplex.session import manager as session_manager_module +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.duplex.session.runner import _Internal + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- # +# Fakes # +# --------------------------------------------------------------------------- # + + +class FakeClock: + def __init__(self, value: float = 0.0) -> None: + self.value = value + + def __call__(self) -> float: + return self.value + + def advance(self, seconds: float) -> None: + self.value += seconds + + +class FakePcmAppendBuffer(PcmAppendBuffer): + @property + def pending_byte_count(self) -> int: + return 0 + + def clear(self) -> None: + pass + + def clear_force_listen(self) -> None: + pass + + def has_pending(self) -> bool: + return False + + def has_reserved(self) -> bool: + return False + + def prepare_append(self, payload, *, operation_id, chunk_period_ms, allow_emit) -> PcmAppendReservation | None: + del payload, operation_id, chunk_period_ms, allow_emit + return None + + def prepare_commit(self, *, operation_id, chunk_period_ms) -> PcmAppendReservation: + del operation_id, chunk_period_ms + raise NotImplementedError("fake buffer never commits") + + def flush(self, *, chunk_period_ms) -> dict[str, object] | None: + del chunk_period_ms + return None + + +class FakeModelSessionState(DuplexModelSessionState): + def __init__(self) -> None: + self.audio_buffer = FakePcmAppendBuffer() + self.input_since_commit = False + self.speech_since_commit = False + self.context_locked = False + self.committed_audio_payload = None + self.committed_audio_operation_id = None + self.committed_audio_reserved_bytes = 0 + self.deferred_response_create = False + self.deferred_precreate_response = False + self.continuation_owner_id = None + self.continuation_units = 0 + self.pending_silence_task = None + self.pending_silence_owner_id = None + + def retain_committed_audio(self, payload, *, operation_id, reserved_bytes=0) -> None: + self.committed_audio_payload = payload + self.committed_audio_operation_id = operation_id + self.committed_audio_reserved_bytes = reserved_bytes + + def clear_committed_audio(self) -> int: + reserved = self.committed_audio_reserved_bytes + self.committed_audio_payload = None + self.committed_audio_operation_id = None + self.committed_audio_reserved_bytes = 0 + return reserved + + def clear_continuation(self) -> None: + self.continuation_owner_id = None + self.continuation_units = 0 + + +class FakeDataPlane(DuplexDataPlane): + def __init__(self) -> None: + self.closed_sessions: list[str] = [] + + def begin_request(self, request_id: str) -> None: + pass + + def is_terminal(self, request_id: str | None) -> bool: + return False + + def mark_terminal(self, request_id: str) -> None: + pass + + def close_stream(self, request_id: str) -> None: + pass + + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: + self.closed_sessions.append(session_id) + + def project(self, result, *, context=None): + del result, context + return [] + + +class FakePlugin(DuplexModelPlugin): + plugin_id = "fake-duplex" + + def __init__(self, encode_audio=None) -> None: + super().__init__(encode_audio or (lambda *args: None)) + self.data_plane = FakeDataPlane() + self.runtime_config: dict[str, object] = {"runtime": "test"} + self.runtime_config_error: Exception | None = None + #: Session ids whose ``prepare_runtime_config`` waits for ``runtime_config_gate``. + self.blocked_session_ids: set[str] = set() + self.runtime_config_gate = asyncio.Event() + self.runtime_config_started = asyncio.Event() + + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(f"configured-{stage_id}" for stage_id, _ in enumerate(defaults)) + + def plan_append(self, **kwargs) -> DuplexAppendPlan: + del kwargs + return DuplexAppendPlan(prompt={"prompt_token_ids": [1, 2, 3]}) + + def decide_output(self, **kwargs): + del kwargs + return None + + def create_session_state(self) -> DuplexModelSessionState: + return FakeModelSessionState() + + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: + del max_sessions + return DuplexCapabilities(supports_input_append=True) + + def validate_client_extra_body(self, extra_body: object) -> None: + pass + + async def prepare_runtime_config(self, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]: + del model_config + if self.runtime_config_error is not None: + raise self.runtime_config_error + if config.instructions in self.blocked_session_ids: + self.runtime_config_started.set() + await self.runtime_config_gate.wait() + return dict(self.runtime_config) + + def runtime_config_for_update(self, config, current): + del config + return dict(current) + + def data_plane_context(self, **kwargs) -> object: + return dict(kwargs) + + +class FakeStagePort(DuplexStagePort): + def __init__(self, stage_count: int = 2) -> None: + self._stage_count = stage_count + self.ensure_calls: list[DuplexStageRequestContext] = [] + self.submit_calls: list[DuplexStageSubmission] = [] + #: Every cleanup attempt, including the ones that were made to fail. + self.cleanup_calls: list[tuple[list[str], bool]] = [] + self.abort_calls: list[list[str]] = [] + self.cleanup_failures_remaining = 0 + self.cleanup_gate: asyncio.Event | None = None + self.cleanup_started = asyncio.Event() + + @property + def stage_count(self) -> int: + return self._stage_count + + def sampling_defaults(self) -> tuple[object, ...]: + return tuple(f"default-{stage_id}" for stage_id in range(self._stage_count)) + + def ensure_request(self, context: DuplexStageRequestContext) -> None: + self.ensure_calls.append(context) + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: + self.submit_calls.append(submission) + return DuplexStageSubmissionResult( + request_id=submission.context.request_id, + stage_id=submission.context.stage_id, + replica_id=3, + ) + + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + self.cleanup_calls.append((list(request_ids), abort)) + self.cleanup_started.set() + if self.cleanup_gate is not None: + await self.cleanup_gate.wait() + if self.cleanup_failures_remaining: + self.cleanup_failures_remaining -= 1 + raise RuntimeError("transient cleanup failure") + + async def abort_requests(self, request_ids: list[str]) -> None: + self.abort_calls.append(list(request_ids)) + + +# --------------------------------------------------------------------------- # +# Harness # +# --------------------------------------------------------------------------- # + + +def _runtime_config(**overrides: Any) -> DuplexSessionRuntimeConfig: + values: dict[str, Any] = { + "idle_ttl_s": 300.0, + "disconnect_grace_s": 30.0, + "max_sessions": 4, + "max_pending_input_bytes_per_session": 1024, + "max_pending_turns_per_session": 2, + } + values.update(overrides) + return DuplexSessionRuntimeConfig(**values) + + +def stage0_request_id(session_id: str, *, epoch: int = 0) -> str: + encoded = base64.urlsafe_b64encode(session_id.encode("utf-8")).decode("ascii").rstrip("=") + return f"duplex-s.{encoded}.e.{epoch}.r.stage0" + + +@dataclass +class Harness: + manager: DuplexSessionManager + plugin: FakePlugin + stage_port: FakeStagePort + clock: FakeClock + output_sink: asyncio.Queue = field(default_factory=asyncio.Queue) + result_sink: asyncio.Queue = field(default_factory=asyncio.Queue) + + @classmethod + def create( + cls, + *, + plugin: FakePlugin | None = None, + stage_port: FakeStagePort | None = None, + clock: FakeClock | None = None, + **runtime_overrides: Any, + ) -> Harness: + plugin = plugin or FakePlugin() + stage_port = stage_port or FakeStagePort() + clock = clock or FakeClock() + output_sink: asyncio.Queue = asyncio.Queue() + result_sink: asyncio.Queue = asyncio.Queue() + manager = DuplexSessionManager( + plugin=plugin, + stage_port=stage_port, + output_sink=output_sink, + result_sink=result_sink, + runtime_config=_runtime_config(**runtime_overrides), + model_config=None, + clock=clock, + ) + return cls(manager, plugin, stage_port, clock, output_sink, result_sink) + + async def __aenter__(self) -> Harness: + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self.manager.shutdown() + + async def open(self, session_id: str, config: DuplexSessionConfig | None = None) -> DuplexControlResultMessage: + await self.manager.handle( + OpenDuplexSessionMessage( + control_id=f"open-{session_id}", + session_id=session_id, + session_config=config or DuplexSessionConfig(model="fake-model"), + ) + ) + return await self.result() + + async def close(self, session_id: str, *, reason: str = "client_close") -> DuplexControlResultMessage: + await self.manager.handle( + CloseDuplexSessionMessage(control_id=f"close-{session_id}", session_id=session_id, reason=reason) + ) + return await self.result() + + async def resume(self, session_id: str, *, expected_lease_generation: int) -> DuplexControlResultMessage: + await self.manager.handle( + ResumeDuplexSessionMessage( + control_id=f"resume-{session_id}-{expected_lease_generation}", + session_id=session_id, + expected_lease_generation=expected_lease_generation, + ) + ) + return await self.result() + + async def touch(self, session_id: str, activity: str) -> DuplexControlResultMessage: + await self.manager.handle( + TouchDuplexSessionMessage( + control_id=f"touch-{session_id}-{activity}", session_id=session_id, activity=activity + ) + ) + return await self.result() + + def command(self, session_id: str, command: DuplexCommand) -> None: + self.manager.dispatch(DuplexSessionCommandMessage(session_id=session_id, command=command)) + + async def result(self) -> DuplexControlResultMessage: + result = await asyncio.wait_for(self.result_sink.get(), timeout=2.0) + assert isinstance(result, DuplexControlResultMessage) + return result + + def events(self, session_id: str | None = None) -> list[DuplexEvent]: + events: list[DuplexEvent] = [] + while not self.output_sink.empty(): + message = self.output_sink.get_nowait() + assert isinstance(message, DuplexSessionEventMessage) + assert message.event.session_id == message.session_id + if session_id is None or message.session_id == session_id: + events.append(message.event) + return events + + def session(self, session_id: str) -> DuplexEngineSession: + session = self.manager.get(session_id) + assert session is not None + return session + + def capture_submissions(self, session_id: str) -> list[DuplexCommand]: + """Replace the runner mailbox with a list so admitted commands can be inspected.""" + runner = self.manager.runners[session_id] + submitted: list[DuplexCommand] = [] + runner.submit = submitted.append # type: ignore[method-assign] + return submitted + + +async def _settle() -> None: + for _ in range(10): + await asyncio.sleep(0) + + +# --------------------------------------------------------------------------- # +# Open # +# --------------------------------------------------------------------------- # + + +async def test_open_answers_with_capabilities_and_emits_session_created() -> None: + async with Harness.create() as harness: + config = DuplexSessionConfig(model="fake-model", voice="test") + result = await harness.open("sid-open", config) + + assert result.ok is True + assert result.operation == "open" + assert result.control_id == "open-sid-open" + assert result.session_id == "sid-open" + assert result.lease_generation == 0 + assert result.capabilities == DuplexCapabilities(supports_input_append=True) + assert result.public_session is not None + assert result.public_session["id"] == "sid-open" + assert result.public_session["voice"] == "test" + assert result.error_code is None + + events = harness.events("sid-open") + assert isinstance(events[0], SessionCreated) + assert events[0].epoch == 0 + assert events[0].session["id"] == "sid-open" + wire = events[0].to_realtime() + assert wire["type"] == "session.created" + assert "incarnation" not in wire and "incarnation" not in wire["session"] + + session = harness.session("sid-open") + assert session.state == DuplexSessionState.OPEN + assert session.runtime_config == {"runtime": "test"} + assert isinstance(session.model_state, FakeModelSessionState) + assert harness.manager.active_count() == 1 + assert harness.manager.sessions() == {"sid-open": session} + + +async def test_open_reserves_the_fenced_stage0_request_with_a_frozen_context() -> None: + async with Harness.create() as harness: + await harness.open("sid-context", DuplexSessionConfig(model="fake-model", voice="test")) + + request_id = stage0_request_id("sid-context") + context = harness.stage_port.ensure_calls[-1] + assert isinstance(context, DuplexStageRequestContext) + assert context.request_id == request_id + assert context.session_id == "sid-context" + assert context.fence == DuplexFence("sid-context") + assert context.stage_id == 0 + assert context.final_stage_id == 1 + assert context.config_generation == 0 + assert context.sampling_params == ("configured-0", "configured-1") + assert context.stage_sampling_params == "configured-0" + assert context.session_config["voice"] == "test" + assert context.runtime_config == {"runtime": "test"} + with pytest.raises(FrozenInstanceError): + context.stage_id = 1 # type: ignore[misc] + + session = harness.session("sid-context") + assert session.resource_request_ids() == [request_id] + assert session.stage_request_submitted(0, request_id) is False + assert harness.manager.runner_for_request_id(request_id) is harness.manager.runners["sid-context"] + assert harness.stage_port.submit_calls == [] + + +def test_stage_request_id_is_derived_from_the_fence() -> None: + fence = DuplexFence("sid/with+special", epoch=2, turn_id=5) + + request_id = DuplexSessionManager.stage_request_id(fence, stage_id=0) + + assert request_id == stage0_request_id("sid/with+special", epoch=2) + assert duplex_resource_request_belongs_to_session(request_id, "sid/with+special") + assert not duplex_resource_request_belongs_to_session(request_id, "sid/other") + + +async def test_ensure_stage_request_ignores_stages_beyond_the_pipeline() -> None: + async with Harness.create() as harness: + await harness.open("sid-stages") + session = harness.session("sid-stages") + + assert harness.manager.ensure_stage_request(session, stage_id=2) is None + context = harness.manager.ensure_stage_request(session, stage_id=1) + + assert context is not None + assert context.request_id.endswith(".r.stage1") + assert harness.manager.runner_for_request_id(context.request_id) is harness.manager.runners["sid-stages"] + assert session.resource_request_ids() == [stage0_request_id("sid-stages"), context.request_id] + + +async def test_sampling_params_for_validates_the_plugin_policy() -> None: + class ShortSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + if runtime_config: + return defaults[:1] + return tuple(defaults) + + class ListSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + if runtime_config: + return list(defaults) + return tuple(defaults) + + async with Harness.create(plugin=ShortSamplingPlugin()) as harness: + result = await harness.open("sid-short") + assert result.ok is False + assert result.error_code == "invalid_argument" + assert "one sampling parameter per stage" in str(result.error_message) + assert harness.manager.get("sid-short") is None + assert harness.manager.active_count() == 0 + + async with Harness.create(plugin=ListSamplingPlugin()) as harness: + result = await harness.open("sid-list") + assert result.ok is False + assert result.error_code == "invalid_argument" + assert "as a tuple" in str(result.error_message) + + +async def test_duplicate_open_is_rejected_with_session_exists() -> None: + async with Harness.create() as harness: + assert (await harness.open("sid-dup")).ok is True + first_session = harness.session("sid-dup") + + result = await harness.open("sid-dup") + + assert result.ok is False + assert result.error_code == "session_exists" + assert result.error_retryable is False + assert harness.session("sid-dup") is first_session + assert harness.manager.active_count() == 1 + assert len(harness.stage_port.ensure_calls) == 1 + + +async def test_capacity_rejection_is_retryable_and_logged_at_info(caplog: pytest.LogCaptureFixture) -> None: + manager_logger = session_manager_module.logger + manager_logger.addHandler(caplog.handler) + try: + async with Harness.create(max_sessions=1) as harness: + assert (await harness.open("sid-first")).ok is True + caplog.clear() + + with caplog.at_level(logging.INFO, logger=manager_logger.name): + result = await harness.open("sid-rejected") + + assert result.ok is False + assert result.error_code == "resource_exhausted" + assert result.error_retryable is True + assert harness.manager.get("sid-rejected") is None + assert harness.manager.active_count() == 1 + rejections = [record for record in caplog.records if "rejected" in record.getMessage()] + assert rejections and all(record.levelno == logging.INFO for record in rejections) + assert not [record for record in caplog.records if record.levelno >= logging.ERROR] + finally: + manager_logger.removeHandler(caplog.handler) + + +@pytest.mark.parametrize( + ("error", "expected_code"), + [ + (ValueError("reference audio not found"), "invalid_duplex_runtime_config"), + (DuplexRuntimeConfigError("bad ref", code="unsupported_ref_audio_path"), "unsupported_ref_audio_path"), + (DuplexRuntimeConfigError("bad voice", code="invalid_voice"), "invalid_voice"), + ], +) +async def test_plugin_runtime_config_rejection_rolls_back_the_open(error: Exception, expected_code: str) -> None: + plugin = FakePlugin() + plugin.runtime_config_error = error + async with Harness.create(plugin=plugin) as harness: + result = await harness.open("sid-runtime") + + assert result.ok is False + assert result.error_code == expected_code + assert result.error_message == str(error) + assert result.capabilities is None + assert harness.manager.get("sid-runtime") is None + assert harness.manager.active_count() == 0 + assert harness.stage_port.ensure_calls == [] + assert harness.events() == [] + + # Nothing was registered under the rejected id and the slot is free again. + plugin.runtime_config_error = None + assert (await harness.open("sid-runtime")).ok is True + + +# --------------------------------------------------------------------------- # +# Command dispatch # +# --------------------------------------------------------------------------- # + + +async def test_command_for_unknown_session_emits_unknown_session_error() -> None: + async with Harness.create() as harness: + harness.command("sid-missing", Heartbeat(event_id="evt-hb")) + + events = harness.events("sid-missing") + assert len(events) == 1 + error = events[0] + assert isinstance(error, ErrorEvent) + assert error.code == "unknown_session" + assert error.related_event_id == "evt-hb" + assert error.session_id == "sid-missing" + assert error.epoch is None + assert error.to_realtime()["error"]["code"] == "unknown_session" + + +async def test_command_for_closed_session_emits_unknown_session_error() -> None: + async with Harness.create() as harness: + await harness.open("sid-closed") + assert (await harness.close("sid-closed")).ok is True + harness.events() + + harness.command("sid-closed", Heartbeat(event_id="evt-late")) + + events = harness.events("sid-closed") + assert [type(event) for event in events] == [ErrorEvent] + assert events[0].code == "unknown_session" + assert events[0].related_event_id == "evt-late" + + +async def test_admitted_command_reaches_the_session_runner_in_order() -> None: + async with Harness.create() as harness: + await harness.open("sid-runner") + harness.events() + harness.clock.advance(5.0) + + harness.command("sid-runner", Heartbeat(event_id="evt-1")) + harness.command("sid-runner", Heartbeat(event_id="evt-2")) + await _settle() + + events = harness.events("sid-runner") + assert [type(event) for event in events] == [SessionHeartbeatAck, SessionHeartbeatAck] + assert all(event.epoch == 0 for event in events) + assert harness.session("sid-runner").lease.last_activity == 5.0 + + +async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None: + async with Harness.create(max_pending_input_bytes_per_session=8) as harness: + await harness.open("sid-bytes") + harness.events() + session = harness.session("sid-bytes") + submitted = harness.capture_submissions("sid-bytes") + assert session.reserve_input_bytes(6, limit=8) + + harness.command("sid-bytes", AppendAudio(audio=b"pcm", event_id="evt-over")) + harness.command("sid-bytes", AppendAudio(audio=b"pc", event_id="evt-fits")) + + events = harness.events("sid-bytes") + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "input_backpressure" + assert events[0].related_event_id == "evt-over" + assert events[0].epoch == 0 + assert [command.event_id for command in submitted] == ["evt-fits"] + # Admission reserved the wire size of the admitted append; the runner + # released it on dequeue and re-reserved the decoded size (2 bytes). + assert session.pending_input_bytes == 8 + + +async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues_it() -> None: + async with Harness.create(max_pending_turns_per_session=1) as harness: + await harness.open("sid-turns") + harness.events() + session = harness.session("sid-turns") + submitted = harness.capture_submissions("sid-turns") + + harness.command("sid-turns", Commit(event_id="evt-commit-1")) + harness.command("sid-turns", Commit(event_id="evt-commit-2")) + + assert [command.event_id for command in submitted] == ["evt-commit-1"] + assert session.pending_input_turns == 1 + events = harness.events("sid-turns") + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "input_backpressure" + assert events[0].related_event_id == "evt-commit-2" + + # The runner releases the reservation when it dequeues the commit. + session.release_pending_turn() + harness.command("sid-turns", Commit(event_id="evt-commit-3")) + assert [command.event_id for command in submitted] == ["evt-commit-1", "evt-commit-3"] + assert harness.events("sid-turns") == [] + + +def test_accepts_only_duplex_message_types() -> None: + manager = DuplexSessionManager( + plugin=FakePlugin(), + stage_port=FakeStagePort(), + output_sink=asyncio.Queue(), + result_sink=asyncio.Queue(), + runtime_config=_runtime_config(), + model_config=None, + clock=FakeClock(), + ) + try: + config = DuplexSessionConfig() + assert manager.accepts(OpenDuplexSessionMessage(control_id="c", session_id="s", session_config=config)) + assert manager.accepts(CloseDuplexSessionMessage(control_id="c", session_id="s")) + assert manager.accepts(ResumeDuplexSessionMessage(control_id="c", session_id="s", expected_lease_generation=0)) + assert manager.accepts(TouchDuplexSessionMessage(control_id="c", session_id="s", activity="heartbeat")) + assert manager.accepts(DuplexSessionCommandMessage(session_id="s", command=Heartbeat())) + assert manager.accepts(type("Lookalike", (), {"type": "open_duplex_session", "session_id": "s"})()) is False + assert ( + manager.accepts(DuplexControlResultMessage(control_id="c", operation="open", session_id="s", ok=True)) + is False + ) + assert manager.accepts(Heartbeat()) is False + with pytest.raises(TypeError, match="Unsupported duplex control message"): + manager.dispatch(object()) + finally: + manager.executor.shutdown(wait=False) + + +# --------------------------------------------------------------------------- # +# Close / resume / touch # +# --------------------------------------------------------------------------- # + + +async def test_close_releases_stage_resources_and_is_idempotent() -> None: + async with Harness.create() as harness: + await harness.open("sid-close") + harness.events() + session = harness.session("sid-close") + request_id = stage0_request_id("sid-close") + session.bind_stage_request(0, request_id, fence=session.fence) + reserved_context = harness.manager.ensure_stage_request(session, stage_id=1) + assert reserved_context is not None + + result = await harness.close("sid-close", reason="client_close") + + assert result.ok is True + assert result.operation == "close" + assert result.lease_generation == 1 + assert result.public_session is not None and result.public_session["state"] == "closed" + assert session.state == DuplexSessionState.CLOSED + assert session.lease.terminal_reason == "client_close" + assert session.resource_request_ids() == [] + assert harness.manager.get("sid-close") is None + assert harness.manager.active_count() == 0 + assert harness.manager.runner_for_request_id(request_id) is None + assert harness.stage_port.cleanup_calls == [([request_id], True), ([reserved_context.request_id], False)] + assert harness.plugin.data_plane.closed_sessions == ["sid-close"] + closed = [event for event in harness.events("sid-close") if isinstance(event, SessionClosed)] + assert len(closed) == 1 + assert closed[0].reason == "client_close" + assert closed[0].is_terminal is True + + again = await harness.close("sid-close") + assert again.ok is True + assert again.lease_generation is None + assert again.public_session is None + assert harness.stage_port.cleanup_calls[2:] == [] + assert harness.events() == [] + + +async def test_close_retains_the_admission_slot_until_stage_cleanup_succeeds() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-close-retry") + session = harness.session("sid-close-retry") + session.bind_stage_request(0, stage0_request_id("sid-close-retry"), fence=session.fence) + harness.stage_port.cleanup_failures_remaining = 1 + + result = await harness.close("sid-close-retry") + + # The session is closed from the client's point of view; only the + # stage cleanup is outstanding and the reaper retries it. + assert result.ok is True + assert harness.manager.get("sid-close-retry") is None + assert session.state == DuplexSessionState.CLOSED + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-close-retry")], True)] * 2 + assert (await harness.open("sid-replacement")).ok is True + + +async def test_resume_requires_the_expected_lease_generation() -> None: + async with Harness.create() as harness: + await harness.open("sid-resume") + session = harness.session("sid-resume") + harness.clock.advance(3.0) + assert (await harness.touch("sid-resume", DuplexLeaseActivity.DETACH.value)).ok is True + assert session.lease.detached_at == 3.0 + + harness.clock.advance(1.0) + result = await harness.resume("sid-resume", expected_lease_generation=0) + + assert result.ok is True + assert result.operation == "resume" + assert result.lease_generation == 1 + assert result.public_session is not None and result.public_session["id"] == "sid-resume" + assert session.lease_generation == 1 + assert session.lease.detached_at is None + assert session.lease.last_activity == 4.0 + + stale = await harness.resume("sid-resume", expected_lease_generation=0) + assert stale.ok is False + assert stale.error_code == "session_resume_conflict" + assert stale.lease_generation == 1 + assert session.lease_generation == 1 + + unknown = await harness.resume("sid-unknown", expected_lease_generation=0) + assert unknown.ok is False + assert unknown.error_code == "unknown_session" + assert unknown.lease_generation is None + + +async def test_touch_records_activity_and_detach_marks_the_lease_detached() -> None: + async with Harness.create() as harness: + await harness.open("sid-touch") + session = harness.session("sid-touch") + assert session.lease.last_activity == 0.0 + + harness.clock.advance(2.0) + result = await harness.touch("sid-touch", DuplexLeaseActivity.HEARTBEAT.value) + assert result.ok is True + assert result.operation == "touch" + assert session.lease.last_activity == 2.0 + assert session.lease.detached_at is None + + harness.clock.advance(1.0) + assert (await harness.touch("sid-touch", DuplexLeaseActivity.DETACH.value)).ok is True + assert session.lease.last_activity == 3.0 + assert session.lease.detached_at == 3.0 + + harness.clock.advance(1.0) + assert (await harness.touch("sid-touch", DuplexLeaseActivity.ATTACH.value)).ok is True + assert session.lease.last_activity == 4.0 + # Only resume clears a detach. + assert session.lease.detached_at == 3.0 + + invalid = await harness.touch("sid-touch", "bogus") + assert invalid.ok is False + assert invalid.error_code == "invalid_argument" + + unknown = await harness.touch("sid-unknown", DuplexLeaseActivity.HEARTBEAT.value) + assert unknown.ok is False + assert unknown.error_code == "unknown_session" + assert harness.stage_port.submit_calls == [] + + +# --------------------------------------------------------------------------- # +# Reaper # +# --------------------------------------------------------------------------- # + + +async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None: + async with Harness.create(idle_ttl_s=2.0, disconnect_grace_s=1.0) as harness: + await harness.open("sid-idle") + await harness.open("sid-active") + harness.events() + harness.clock.advance(1.0) + assert (await harness.touch("sid-active", DuplexLeaseActivity.HEARTBEAT.value)).ok is True + harness.clock.advance(1.1) + + assert await harness.manager.reap_expired() == 1 + + assert harness.manager.get("sid-idle") is None + assert harness.manager.get("sid-active") is not None + assert harness.manager.active_count() == 1 + expired = harness.events("sid-idle") + assert len(expired) == 1 + assert isinstance(expired[0], SessionExpired) + assert expired[0].reason == "idle_ttl_expired" + assert expired[0].to_realtime() == { + "type": "session.expired", + "event_id": expired[0].event_id, + "session_id": "sid-idle", + "reason": "idle_ttl_expired", + } + assert harness.events("sid-active") == [] + assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-idle")], False)] + assert harness.result_sink.empty() + + assert await harness.manager.reap_expired() == 0 + + +async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() -> None: + async with Harness.create(idle_ttl_s=300.0, disconnect_grace_s=1.0) as harness: + await harness.open("sid-detached") + harness.events() + assert (await harness.touch("sid-detached", DuplexLeaseActivity.DETACH.value)).ok is True + harness.clock.advance(0.5) + assert await harness.manager.reap_expired() == 0 + harness.clock.advance(0.5) + + assert await harness.manager.reap_expired() == 1 + + assert harness.manager.get("sid-detached") is None + expired = harness.events("sid-detached") + assert [type(event) for event in expired] == [SessionExpired] + assert expired[0].reason == "disconnect_grace_expired" + assert (await harness.resume("sid-detached", expected_lease_generation=0)).error_code == "unknown_session" + + +async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is_admitted() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-closing") + harness.events() + session = harness.session("sid-closing") + session.bind_stage_request(0, stage0_request_id("sid-closing"), fence=session.fence) + harness.stage_port.cleanup_gate = asyncio.Event() + + close_task = asyncio.create_task(harness.close("sid-closing")) + await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) + # The runner is gone but the stage request is still being aborted: + # no session.closed yet, so a client cannot race the admission slot. + assert harness.events("sid-closing") == [] + + harness.stage_port.cleanup_gate.set() + assert (await close_task).ok is True + events = harness.events("sid-closing") + assert [type(event) for event in events] == [SessionClosed] + assert events[0].reason == "client_close" + assert (await harness.open("sid-replacement")).ok is True + + +async def test_wire_close_command_frees_the_slot_before_session_closed() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-wire-close") + harness.events() + session = harness.session("sid-wire-close") + request_id = stage0_request_id("sid-wire-close") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.stage_port.cleanup_gate = asyncio.Event() + + harness.command("sid-wire-close", CloseSession(reason="client_close")) + await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) + assert harness.manager.get("sid-wire-close") is None + assert harness.events("sid-wire-close") == [] + + harness.stage_port.cleanup_gate.set() + for _ in range(50): + if harness.manager._admission_count() == 0: + break + await asyncio.sleep(0.01) + assert harness.stage_port.cleanup_calls == [([request_id], True)] + events = harness.events("sid-wire-close") + assert [type(event) for event in events] == [SessionClosed] + assert events[0].reason == "client_close" + assert (await harness.open("sid-replacement")).ok is True + + +async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-wire") + harness.events() + runner = harness.manager.runners["sid-wire"] + worker = runner._worker + assert worker is not None + harness.stage_port.cleanup_gate = asyncio.Event() + + harness.command("sid-wire", CloseSession(reason="client_close")) + await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) + # Control operations issued while the close is in flight are refused, + # not silently accepted on a half-torn session. + resumed = await harness.resume("sid-wire", expected_lease_generation=0) + assert resumed.ok is False + assert resumed.error_code in {"session_closed", "unknown_session"} + harness.command("sid-wire", Heartbeat()) + + harness.stage_port.cleanup_gate.set() + for _ in range(50): + if worker.done() and harness.manager._admission_count() == 0: + break + await asyncio.sleep(0.01) + assert worker.done(), "the mailbox worker must exit after a wire close" + events = harness.events("sid-wire") + assert [type(event) for event in events] == [ErrorEvent, SessionClosed] or [ + type(event) for event in events + ] == [SessionClosed, ErrorEvent] + assert (await harness.open("sid-after")).ok is True + + +async def test_concurrent_opens_cannot_exceed_max_sessions() -> None: + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("first") + first = asyncio.create_task( + harness.open("first", DuplexSessionConfig(model="fake-model", instructions="first")) + ) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + # The first open holds the only slot while it awaits the plugin. + second = await harness.open("second", DuplexSessionConfig(model="fake-model", instructions="second")) + assert second.ok is False + assert second.error_code == "resource_exhausted" + + harness.plugin.runtime_config_gate.set() + assert (await first).ok is True + assert harness.manager.active_count() == 1 + + +async def test_runtime_close_without_stage_requests_frees_the_slot() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-runtime") + harness.events() + runner = harness.manager.runners["sid-runtime"] + runner.session.release_all_requests() + + await runner._close_from_runtime("model_failed") + for _ in range(50): + if harness.manager._admission_count() == 0: + break + await asyncio.sleep(0.01) + + assert harness.manager.get("sid-runtime") is None + assert [type(event) for event in harness.events("sid-runtime")] == [SessionClosed] + assert (await harness.open("sid-next")).ok is True + + +async def test_close_of_a_forgotten_session_is_idempotent() -> None: + async with Harness.create(max_sessions=1) as harness: + result = await harness.close("never-opened") + + assert result.ok is True + + +async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues() -> None: + async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=8) as harness: + await harness.open("sid-bytes") + harness.events() + session = harness.session("sid-bytes") + runner = harness.manager.runners["sid-bytes"] + # Park the worker so admitted commands stay in the mailbox. + gate = asyncio.Event() + runner._mailbox.put_nowait(_Internal("wait", {"gate": gate})) + runner._on_internal = lambda item: gate.wait() # type: ignore[method-assign] + await asyncio.sleep(0) + + harness.command("sid-bytes", AppendAudio(audio=b"12345")) + assert session.pending_input_bytes == 5 + harness.command("sid-bytes", AppendAudio(audio=b"12345")) + errors = [event for event in harness.events("sid-bytes") if isinstance(event, ErrorEvent)] + assert [error.code for error in errors] == ["input_backpressure"] + gate.set() + + +async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds() -> None: + async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness: + await harness.open("sid-expired") + harness.events() + session = harness.session("sid-expired") + request_id = stage0_request_id("sid-expired") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.stage_port.cleanup_failures_remaining = 1 + harness.clock.advance(2.0) + + assert await harness.manager.reap_expired() == 0 + + assert harness.manager.get("sid-expired") is None + assert session.state == DuplexSessionState.CLOSED + assert [type(event) for event in harness.events("sid-expired")] == [SessionExpired] + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + assert await harness.manager.reap_expired() == 1 + + assert harness.stage_port.cleanup_calls == [([request_id], True)] * 2 + assert session.resource_request_ids() == [] + assert harness.events("sid-expired") == [] + assert (await harness.open("sid-replacement")).ok is True + + +async def test_expired_session_holds_the_admission_slot_while_cleanup_is_blocked() -> None: + async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness: + await harness.open("sid-blocked") + session = harness.session("sid-blocked") + session.bind_stage_request(0, stage0_request_id("sid-blocked"), fence=session.fence) + harness.stage_port.cleanup_gate = asyncio.Event() + harness.clock.advance(2.0) + + reap_task = asyncio.create_task(harness.manager.reap_expired()) + await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + harness.stage_port.cleanup_gate.set() + assert await reap_task == 1 + assert (await harness.open("sid-replacement")).ok is True + + +async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: + class OneStuckStagePort(FakeStagePort): + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + await super().cleanup(request_ids, abort=abort) + if request_ids == [stage0_request_id("sid-stuck")]: + raise RuntimeError("stuck expiry cleanup") + + async with Harness.create(stage_port=OneStuckStagePort(), idle_ttl_s=1.0) as harness: + await harness.open("sid-stuck") + await harness.open("sid-fine") + harness.events() + harness.clock.advance(2.0) + + assert await harness.manager.reap_expired() == 1 + + assert harness.manager.active_count() == 0 + assert {type(event) for event in harness.events()} == {SessionExpired} + assert (await harness.open("sid-fine")).ok is True + assert (await harness.open("sid-stuck")).error_code == "session_exists" + + +# --------------------------------------------------------------------------- # +# Request-triggered cleanup (orchestrator error paths) # +# --------------------------------------------------------------------------- # + + +async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_cleanup() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-request") + harness.events() + session = harness.session("sid-request") + request_id = stage0_request_id("sid-request") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.stage_port.cleanup_failures_remaining = 1 + + closed = harness.manager.close_sessions_for_request_ids([request_id, "req-unrelated"], abort=True) + + assert closed == {"sid-request": [request_id]} + assert harness.manager.get("sid-request") is None + assert session.lease.terminal_reason == "request_cleanup" + await _settle() + expired = harness.events("sid-request") + assert [type(event) for event in expired] == [SessionExpired] + assert expired[0].reason == "request_cleanup" + assert session.state == DuplexSessionState.CLOSED + # The stage resources are still owned by the pending cleanup. + assert session.resource_request_ids() == [request_id] + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [([request_id], True)] + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([request_id], True)] * 2 + assert session.resource_request_ids() == [] + assert harness.manager.runner_for_request_id(request_id) is None + assert (await harness.open("sid-replacement")).ok is True + assert harness.manager.close_sessions_for_request_ids([request_id], abort=True) == {} + + +async def test_in_progress_request_cleanup_is_finalized_by_the_orchestrator_not_the_reaper() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-in-progress") + session = harness.session("sid-in-progress") + request_id = stage0_request_id("sid-in-progress") + session.bind_stage_request(0, request_id, fence=session.fence) + + closed = harness.manager.close_sessions_for_request_ids([request_id], abort=True, cleanup_in_progress=True) + assert closed == {"sid-in-progress": [request_id]} + + # The orchestrator owns this cleanup: the reaper neither retries nor releases it. + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [] + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + harness.manager.finalize_closed_sessions(["sid-in-progress"]) + + assert session.resource_request_ids() == [] + assert session.state == DuplexSessionState.CLOSED + assert harness.manager.runner_for_request_id(request_id) is None + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [] + assert (await harness.open("sid-replacement")).ok is True + + +async def test_deferred_request_cleanup_is_retried_by_the_reaper() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-deferred") + session = harness.session("sid-deferred") + request_id = stage0_request_id("sid-deferred") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.manager.close_sessions_for_request_ids([request_id], abort=True, cleanup_in_progress=True) + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [] + + # The orchestrator's own cleanup failed: hand the retry to the reaper. + harness.manager.defer_request_cleanups(["sid-deferred"]) + + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([request_id], True)] + assert session.resource_request_ids() == [] + assert (await harness.open("sid-replacement")).ok is True + + +async def test_request_cleanup_closes_the_owner_of_any_of_its_stage_requests() -> None: + async with Harness.create() as harness: + await harness.open("sid-merge") + session = harness.session("sid-merge") + stage0 = stage0_request_id("sid-merge") + session.bind_stage_request(0, stage0, fence=session.fence) + stage1 = harness.manager.ensure_stage_request(session, stage_id=1) + assert stage1 is not None + + closed = harness.manager.close_sessions_for_request_ids([stage1.request_id]) + + assert closed == {"sid-merge": [stage0, stage1.request_id]} + assert harness.manager.close_sessions_for_request_ids([stage0]) == {} + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([stage0, stage1.request_id], False)] + + +# --------------------------------------------------------------------------- # +# Dispatch ordering / shutdown # +# --------------------------------------------------------------------------- # + + +async def test_control_dispatch_is_ordered_per_session_without_blocking_other_sessions() -> None: + plugin = FakePlugin() + plugin.blocked_session_ids.add("sid-blocked") + async with Harness.create(plugin=plugin) as harness: + manager = harness.manager + blocked_config = DuplexSessionConfig(model="fake-model", instructions="sid-blocked") + + manager.dispatch( + OpenDuplexSessionMessage(control_id="blocked-open", session_id="sid-blocked", session_config=blocked_config) + ) + await asyncio.wait_for(plugin.runtime_config_started.wait(), timeout=1.0) + manager.dispatch( + TouchDuplexSessionMessage( + control_id="blocked-touch", + session_id="sid-blocked", + activity=DuplexLeaseActivity.HEARTBEAT.value, + ) + ) + manager.dispatch( + OpenDuplexSessionMessage( + control_id="independent-open", + session_id="sid-independent", + session_config=DuplexSessionConfig(model="fake-model"), + ) + ) + manager.dispatch( + TouchDuplexSessionMessage( + control_id="independent-touch", + session_id="sid-independent", + activity=DuplexLeaseActivity.HEARTBEAT.value, + ) + ) + + assert [(await harness.result()).control_id for _ in range(2)] == ["independent-open", "independent-touch"] + assert harness.result_sink.empty() + assert manager.get("sid-blocked") is None + + plugin.runtime_config_gate.set() + results = [await harness.result() for _ in range(2)] + assert [result.control_id for result in results] == ["blocked-open", "blocked-touch"] + assert all(result.ok for result in results) + assert manager.active_count() == 2 + + +async def test_shutdown_closes_every_runner_and_stops_dispatch_tasks() -> None: + harness = Harness.create() + await harness.open("sid-a") + await harness.open("sid-b") + sessions = [harness.session("sid-a"), harness.session("sid-b")] + harness.events() + + await harness.manager.shutdown() + + assert harness.manager.active_count() == 0 + assert harness.manager.runners == {} + assert all(session.state == DuplexSessionState.CLOSED for session in sessions) + assert harness.manager.runner_for_request_id(stage0_request_id("sid-b")) is None + # Shutdown is silent: no session.closed / session.expired is emitted. + assert harness.events() == [] + assert harness.stage_port.cleanup_calls == [] + + +def test_manager_rejects_a_plugin_whose_sampling_policy_mismatches_the_stages() -> None: + class WrongStageCountPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return defaults[:1] + + with pytest.raises(ValueError, match="one sampling parameter per stage"): + DuplexSessionManager( + plugin=WrongStageCountPlugin(), + stage_port=FakeStagePort(), + output_sink=asyncio.Queue(), + result_sink=asyncio.Queue(), + runtime_config=_runtime_config(), + model_config=None, + ) + + +# --------------------------------------------------------------------------- # +# Admission bookkeeping # +# --------------------------------------------------------------------------- # + + +async def test_duplicate_open_in_flight_does_not_release_the_first_opens_slot() -> None: + """A duplicate id arriving mid-flight must not free the slot the first open holds. + + The rejection runs the same ``finally`` as a successful open, so without a + guard it discards an ``_admitting`` entry it never added, and the next open + of a *different* id sails past ``max_sessions``. Server-allocated ids make + this a guard rather than a live path, which is exactly why a sequential + duplicate cannot exercise it. + """ + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("dup") + first = asyncio.create_task(harness.open("dup", DuplexSessionConfig(model="fake-model", instructions="dup"))) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + duplicate = await harness.open("dup", DuplexSessionConfig(model="fake-model", instructions="dup")) + assert duplicate.ok is False + assert duplicate.error_code == "session_exists" + + # The in-flight open still holds the only slot. + overflow = await harness.open("other", DuplexSessionConfig(model="fake-model", instructions="other")) + assert overflow.ok is False + assert overflow.error_code == "resource_exhausted" + + harness.plugin.runtime_config_gate.set() + assert (await first).ok is True + assert harness.manager.active_count() == 1 + + +async def test_expiry_emits_a_terminal_event_even_when_a_close_deferred_it() -> None: + """Every session ends with exactly one terminal event, whoever tears it down. + + A wire ``session.close`` defers its ``session.closed`` to the manager. If a + stage failure then takes the runner away from the manager (the orchestrator + cleanup path), nobody would emit the deferred event unless ``expire`` + upgrades it. + """ + async with Harness.create() as harness: + await harness.open("sid-deferred") + runner = harness.manager.runners["sid-deferred"] + # What the runner does for a wire close: tear down, defer the terminal. + await runner.close("client_close", emit_closed=False) + harness.events("sid-deferred") + + await runner.expire("request_cleanup", emit_expired=True) + + terminal = [ + event for event in harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired) + ] + assert len(terminal) == 1, terminal + assert isinstance(terminal[0], SessionExpired) + assert terminal[0].reason == "request_cleanup" + + +async def test_an_open_abandoned_by_its_caller_does_not_hold_the_admission_slot() -> None: + """A control-RPC timeout must not leave the engine holding a session. + + The caller's timeout unregisters its RPC waiter but cannot cancel the + manager's open, which is already past the capacity check. Without the + compensating close, that session lands, keeps its slot and its Stage0 + reservation, and with ``max_sessions=1`` the next open is refused until + idle expiry. + """ + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("blocked") + config = DuplexSessionConfig(model="fake-model", instructions="blocked") + + open_task = asyncio.create_task(harness.open("sid-abandoned", config)) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + # The caller gave up: DuplexOmni sends the compensating close while the + # open is still awaiting the plugin. + await harness.close("sid-abandoned", reason="open_abandoned") + + harness.plugin.runtime_config_gate.set() + await open_task + + assert "sid-abandoned" not in harness.manager.runners + # The slot came back, so a replacement is admitted. + assert (await harness.open("sid-replacement")).ok is True + + +async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal() -> None: + """Two teardown drivers, one terminal event. + + The window is inside ``runner.close()``: ``cancel_append_tasks()`` is awaited + *before* the terminal is deferred, so an append task that dies there finds + the session merely CLOSING with neither ``closed_emitted`` nor + ``closed_deferred`` set, and emits ``session.closed`` itself. The manager + decided to emit before any of this and must re-check, or the session ends + with two terminal events. + """ + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-race") + harness.events() + session = harness.session("sid-race") + session.bind_stage_request(0, stage0_request_id("sid-race"), fence=session.fence) + runner = harness.manager.runners["sid-race"] + + async def dying_append() -> bool: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + # What AppendAttempt does when its submit fails mid-teardown. + await runner._close_from_runtime("runtime_append_task_failed") + raise + return True + + task = asyncio.ensure_future(dying_append()) + await asyncio.sleep(0) + runner.tasks.track_append_task(task, epoch=session.epoch, final=True, response_bound=True) + + await harness.close("sid-race") + + terminal = [event for event in harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)] + assert len(terminal) == 1, terminal + + +async def test_a_terminal_event_is_emitted_once_even_if_expiry_runs_twice() -> None: + async with Harness.create() as harness: + await harness.open("sid-once") + runner = harness.manager.runners["sid-once"] + + await runner.expire("idle_ttl_expired", emit_expired=True) + await runner.expire("idle_ttl_expired", emit_expired=True) + + terminal = [event for event in harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired)] + assert len(terminal) == 1, terminal + + +# --------------------------------------------------------------------------- # +# Resume and takeover lifecycle # +# --------------------------------------------------------------------------- # + + +async def test_resume_inside_the_grace_window_stops_the_reaper_expiring_the_session() -> None: + """Resuming clears the disconnect deadline, not just the lease generation. + + The reaper decides expiry from ``lease.detached_at``. A resume that bumped + the generation without clearing it would let the reaper kill a session a + client had already come back to. + """ + async with Harness.create(disconnect_grace_s=30.0) as harness: + await harness.open("sid-grace") + session = harness.session("sid-grace") + assert (await harness.touch("sid-grace", DuplexLeaseActivity.DETACH.value)).ok is True + harness.clock.advance(10.0) + + assert (await harness.resume("sid-grace", expected_lease_generation=0)).ok is True + assert session.lease.detached_at is None + + # Past the original deadline: the session survives because it is attached. + harness.clock.advance(25.0) + assert await harness.manager.reap_expired() == 0 + assert harness.manager.get("sid-grace") is not None + + +async def test_resume_after_the_grace_window_reports_unknown_session() -> None: + """Once the reaper has expired a detached session, resume has nothing to attach to.""" + async with Harness.create(disconnect_grace_s=5.0) as harness: + await harness.open("sid-late") + assert (await harness.touch("sid-late", DuplexLeaseActivity.DETACH.value)).ok is True + harness.clock.advance(6.0) + assert await harness.manager.reap_expired() == 1 + + late = await harness.resume("sid-late", expected_lease_generation=0) + + assert late.ok is False + assert late.error_code == "unknown_session" + + +async def test_takeover_invalidates_the_generation_the_replaced_client_held() -> None: + """Two clients cannot both hold the session: the first one's generation goes stale. + + This is the engine half of a websocket takeover — the transport picks the + winning socket, and the lease generation is what stops the loser from + resuming behind it. + """ + async with Harness.create() as harness: + await harness.open("sid-takeover") + session = harness.session("sid-takeover") + assert (await harness.touch("sid-takeover", DuplexLeaseActivity.DETACH.value)).ok is True + + winner = await harness.resume("sid-takeover", expected_lease_generation=0) + assert winner.ok is True + assert winner.lease_generation == 1 + + # The replaced client still believes it holds generation 0. + loser = await harness.resume("sid-takeover", expected_lease_generation=0) + assert loser.ok is False + assert loser.error_code == "session_resume_conflict" + assert session.lease_generation == 1 + + # The winner can still act on the session it took over. + assert (await harness.touch("sid-takeover", DuplexLeaseActivity.HEARTBEAT.value)).ok is True + + +async def test_resume_is_refused_once_a_close_has_begun() -> None: + async with Harness.create() as harness: + await harness.open("sid-closing") + await harness.close("sid-closing") + + resumed = await harness.resume("sid-closing", expected_lease_generation=0) + + assert resumed.ok is False + assert resumed.error_code == "unknown_session" + + +async def test_wire_close_of_an_idle_session_still_emits_session_closed() -> None: + """The protocol smoke shape: open, send nothing, close. + + A session that never bound a stage request has nothing for the reaper to + clean up, so the terminal event is the only thing the client ever gets + back for its ``session.close``. Regression guard for a close path that + only emitted after a cleanup it never performed. + """ + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-idle-close") + harness.events() + + harness.command("sid-idle-close", CloseSession(reason="client_close")) + events: list[DuplexEvent] = [] + for _ in range(200): + events.extend(harness.events("sid-idle-close")) + if events: + break + await asyncio.sleep(0.01) + + assert [type(event) for event in events] == [SessionClosed] + assert events[0].reason == "client_close" + assert harness.manager.get("sid-idle-close") is None + assert (await harness.open("sid-idle-replacement")).ok is True + + +# --------------------------------------------------------------------------- # +# Reaper loop # +# --------------------------------------------------------------------------- # + + +async def test_reaper_loop_waits_between_ticks() -> None: + """The loop paces itself; it does not spin on the expiry check. + + Ported from the pre-framework ``Orchestrator._duplex_reaper_loop`` tests: + the loop moved to the session manager with the sessions it reaps. + """ + manager = object.__new__(DuplexSessionManager) + calls = 0 + + async def reap_expired(now: float | None = None) -> int: + nonlocal calls + calls += 1 + return 0 + + manager.reap_expired = reap_expired # type: ignore[method-assign] + manager.runtime_config = DuplexSessionRuntimeConfig(reaper_interval_s=0.01) + shutdown = asyncio.Event() + + task = asyncio.create_task(manager.reaper_loop(shutdown)) + await asyncio.sleep(0.035) + shutdown.set() + await asyncio.wait_for(task, timeout=5.0) + + assert 2 <= calls <= 5 + + +@pytest.mark.parametrize("first_cleanup_delay", [0.0, 0.05], ids=["immediate", "delayed"]) +async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: float) -> None: + """A failed expiry sweep is retried on the next tick, not fatal to the loop.""" + recovered = asyncio.Event() + calls = 0 + + manager = object.__new__(DuplexSessionManager) + + async def reap_expired(now: float | None = None) -> int: + nonlocal calls + calls += 1 + if calls == 1: + await asyncio.sleep(first_cleanup_delay) + raise RuntimeError("transient cleanup failure") + recovered.set() + return 0 + + manager.reap_expired = reap_expired # type: ignore[method-assign] + manager.runtime_config = DuplexSessionRuntimeConfig(reaper_interval_s=0.01) + shutdown = asyncio.Event() + + task = asyncio.create_task(manager.reaper_loop(shutdown)) + try: + await asyncio.wait_for(recovered.wait(), timeout=5.0) + assert calls >= 2 + finally: + shutdown.set() + await asyncio.wait_for(task, timeout=5.0) diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py new file mode 100644 index 000000000..51711f52a --- /dev/null +++ b/tests/engine/duplex/test_session_runner.py @@ -0,0 +1,1195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Engine-resident session runner scenarios (the old serving handler tests, re-homed). + +The runner is driven exactly the way ``DuplexOrchestrator`` drives it: typed +commands go through ``DuplexSessionManager.dispatch``, stage outputs are pushed +with ``runner.on_stage_output`` and everything the session says is read back +from the manager's output sink as typed events. The stage port is a recording +fake; the model plugin is the real MiniCPM-o 4.5 one so append planning and +output projection are exercised end to end. +""" + +from __future__ import annotations + +import asyncio +import base64 +import struct +from collections.abc import Sequence +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands +from vllm_omni.engine.duplex.commands import DuplexCommand +from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState +from vllm_omni.engine.duplex.contracts import ( + DuplexOutputContext, + DuplexRequestIdentity, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, + duplex_resource_request_id, +) +from vllm_omni.engine.duplex.events import DuplexEvent +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionEventMessage, + OpenDuplexSessionMessage, +) +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner +from vllm_omni.metrics.stats import StageRequestStats, StageStats +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import MiniCPMO45DuplexPlugin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SESSION_ID = "duplex-test" +LISTEN_TOKEN_ID = 7 + + +# --------------------------------------------------------------------------- # +# Fakes # +# --------------------------------------------------------------------------- # + + +class RecordingStagePort(DuplexStagePort): + """Records what the runner asks of the orchestrator; never talks to a stage.""" + + def __init__(self, *, stage_count: int = 2) -> None: + self._stage_count = stage_count + self.ensured: list[DuplexStageRequestContext] = [] + self.submissions: list[DuplexStageSubmission] = [] + self.cleanups: list[tuple[list[str], bool]] = [] + self.aborts: list[list[str]] = [] + self.fail_submit: Exception | None = None + + @property + def stage_count(self) -> int: + return self._stage_count + + def sampling_defaults(self) -> tuple[object, ...]: + return tuple(SamplingParams(max_tokens=8) for _ in range(self._stage_count)) + + def ensure_request(self, context: DuplexStageRequestContext) -> None: + self.ensured.append(context) + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: + if self.fail_submit is not None: + raise self.fail_submit + self.submissions.append(submission) + return DuplexStageSubmissionResult( + request_id=submission.context.request_id, + stage_id=submission.context.stage_id, + replica_id=0, + ) + + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + self.cleanups.append((list(request_ids), abort)) + + async def abort_requests(self, request_ids: list[str]) -> None: + self.aborts.append(list(request_ids)) + + +def _fake_encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None: + del sample_rate_hz, response_format, speed + if audio is None: + return None + samples = int(np.asarray(audio, dtype=np.float32).size) + return f"wav-{samples}" if samples > 0 else None + + +@dataclass +class Harness: + manager: DuplexSessionManager + port: RecordingStagePort + output: asyncio.Queue[Any] + results: asyncio.Queue[Any] + runner: DuplexSessionRunner + events: list[DuplexEvent] = field(default_factory=list) + + @property + def session(self): + return self.runner.session + + def submit(self, command: DuplexCommand) -> None: + self.manager.dispatch(DuplexSessionCommandMessage(session_id=SESSION_ID, command=command)) + + async def settle(self, *, idle_s: float = 0.05, timeout_s: float = 3.0) -> list[DuplexEvent]: + """Run the loop until the runner mailbox and append tasks are quiet; return new events.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + quiet_since: float | None = None + collected: list[DuplexEvent] = [] + while True: + drained = False + while not self.output.empty(): + message = self.output.get_nowait() + if isinstance(message, DuplexSessionEventMessage): + collected.append(message.event) + drained = True + busy = ( + drained + or not self.runner._mailbox.empty() + or any(not task.done() for task in self.runner.tasks.append_tasks) + or any(not task.done() for task in self.runner._background_tasks) + ) + now = loop.time() + if busy: + quiet_since = None + elif quiet_since is None: + quiet_since = now + elif now - quiet_since >= idle_s: + break + if now >= deadline: + break + await asyncio.sleep(0.005) + self.events.extend(collected) + return collected + + async def run(self, command: DuplexCommand) -> list[DuplexEvent]: + self.submit(command) + return await self.settle() + + def deliver( + self, + output: object, + *, + stage_id: int = 1, + segment_finished: bool = False, + segment_output_metadata: dict[str, Any] | None = None, + segment_token_ids: Sequence[int] = (), + epoch: int | None = None, + metrics: Any = None, + ) -> bool: + session = self.session + fence = session.fence if epoch is None else session.fence.__class__(SESSION_ID, epoch=epoch) + context = DuplexOutputContext( + identity=DuplexRequestIdentity(session_id=SESSION_ID, fence=fence), + final_stage_id=self.port.stage_count - 1, + segment_finished=segment_finished, + segment_token_ids=tuple(segment_token_ids), + segment_output_metadata=dict(segment_output_metadata or {}), + ) + return self.runner.on_stage_output(stage_id, output, metrics, request_id=output.request_id, context=context) + + async def deliver_and_settle(self, output: object, **kwargs: Any) -> list[DuplexEvent]: + self.deliver(output, **kwargs) + return await self.settle() + + def stage0_request_id(self, epoch: int | None = None) -> str: + fence = self.session.fence if epoch is None else self.session.fence.__class__(SESSION_ID, epoch=epoch) + return duplex_resource_request_id(fence, "stage0") + + +async def open_harness( + *, + auto_response: bool = True, + modalities: Sequence[str] = ("text",), + extra_body: dict[str, object] | None = None, + runtime_config: DuplexSessionRuntimeConfig | None = None, + stage_count: int = 2, + clock: Any = None, +) -> Harness: + plugin = MiniCPMO45DuplexPlugin(_fake_encode_audio) + port = RecordingStagePort(stage_count=stage_count) + output: asyncio.Queue[Any] = asyncio.Queue() + results: asyncio.Queue[Any] = asyncio.Queue() + manager = DuplexSessionManager( + plugin=plugin, + stage_port=port, + output_sink=output, + result_sink=results, + runtime_config=runtime_config or DuplexSessionRuntimeConfig(), + model_config=None, + clock=clock, + ) + body: dict[str, object] = {"auto_response": auto_response, **(extra_body or {})} + config = DuplexSessionConfig( + model="openbmb/MiniCPM-o-4_5", + modalities=list(modalities), + instructions="You are a concise assistant.", + extra_body=body, + ) + await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config)) + result = await asyncio.wait_for(results.get(), timeout=2.0) + assert isinstance(result, DuplexControlResultMessage) and result.ok, result + harness = Harness(manager=manager, port=port, output=output, results=results, runner=manager.runners[SESSION_ID]) + await harness.settle() + return harness + + +async def close_harness(harness: Harness) -> None: + await harness.manager.shutdown() + + +# --------------------------------------------------------------------------- # +# Payload / output builders # +# --------------------------------------------------------------------------- # + + +def pcm_f32(samples: int, *, value: float = 0.05) -> bytes: + return struct.pack(f"<{samples}f", *([value] * samples)) + + +def append_audio( + samples: int = 16000, + *, + value: float = 0.05, + is_speech: bool | None = True, + event_id: str | None = None, +) -> commands.AppendAudio: + return commands.AppendAudio( + audio=pcm_f32(samples, value=value), + format="pcm_f32le", + sample_rate_hz=16000, + is_speech=is_speech, + event_id=event_id, + ) + + +def stage_stats( + *, + stage_id: int, + request_id: str, + num_tokens_out: int, + itls_ms: list[float] | None = None, +) -> StageRequestStats: + """The per-request stats the orchestrator hands to the runner with an output.""" + return StageRequestStats( + batch_id=0, + batch_size=1, + num_tokens_in=7, + num_tokens_out=num_tokens_out, + stage_gen_time_ms=120.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + stage_id=stage_id, + request_id=request_id, + final_output_type="text", + vllm_itls_ms=list(itls_ms or []), + ) + + +def tts_output( + request_id: str, + *, + samples: int = 24000, + finished: bool = False, + text: str = "hello", + tts_is_last_chunk: bool = False, + turn_end: bool = False, + turn_id: int = 0, + epoch: int = 0, +) -> SimpleNamespace: + """A Stage1 (TTS) output the way the orchestrator hands it to the runner.""" + return SimpleNamespace( + request_id=request_id, + finished=finished, + outputs=[SimpleNamespace(text=text, token_ids=[], multimodal_output={})], + multimodal_output={ + "audio": np.zeros(samples, dtype=np.float32), + "sr": 24000, + "meta.tts_is_last_chunk": np.array([int(tts_is_last_chunk)], dtype=np.int32), + "meta.turn_end": np.array([int(turn_end)], dtype=np.int32), + "meta.duplex_turn_id": np.array([turn_id], dtype=np.int32), + "meta.duplex_epoch": np.array([epoch], dtype=np.int32), + }, + ) + + +def listen_output(request_id: str) -> SimpleNamespace: + """A finished Stage0 segment that ends with the listen token.""" + return SimpleNamespace( + request_id=request_id, + finished=True, + outputs=[ + SimpleNamespace( + text="", + token_ids=[11, 12, LISTEN_TOKEN_ID], + stop_reason=LISTEN_TOKEN_ID, + multimodal_output={}, + ) + ], + multimodal_output={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + +def types(events: Sequence[DuplexEvent]) -> list[str]: + return [event.type for event in events] + + +def find(events: Sequence[DuplexEvent], wire_type: str) -> DuplexEvent: + for event in events: + if event.type == wire_type: + return event + raise AssertionError(f"no {wire_type!r} in {types(events)}") + + +def b64(data: bytes) -> str: + return base64.b64encode(data).decode("ascii") + + +# --------------------------------------------------------------------------- # +# Session lifecycle # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_open_announces_the_session_and_reserves_the_stage0_request() -> None: + h = await open_harness() + try: + assert types(h.events) == ["session.created", "session.updated"] + created = h.events[0] + assert created.session_id == SESSION_ID + assert created.session["id"] == SESSION_ID + assert created.session["capabilities"]["supports_input_append"] is True + assert created.to_realtime()["type"] == "session.created" + assert "incarnation" not in created.to_realtime() + assert h.session.state == DuplexSessionState.OPEN + # Admission reserves the resumable Stage0 request atomically with the open. + assert [context.request_id for context in h.port.ensured] == [h.stage0_request_id(epoch=0)] + assert h.stage0_request_id(epoch=0) == "duplex-s.ZHVwbGV4LXRlc3Q.e.0.r.stage0" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_close_emits_session_closed_and_releases_stage_requests() -> None: + h = await open_harness() + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.manager.handle( + CloseDuplexSessionMessage(control_id="c-close", session_id=SESSION_ID, reason="client_close") + ) + result = await asyncio.wait_for(h.results.get(), timeout=2.0) + events = await h.settle() + + assert result.ok and result.operation == "close" + assert types(events) == ["session.closed"] + assert events[0].reason == "client_close" + assert events[0].is_terminal + assert h.session.state == DuplexSessionState.CLOSED + assert h.port.cleanups == [([request_id], True)] + assert SESSION_ID not in h.manager.runners + await close_harness(h) + + +@pytest.mark.asyncio +async def test_idle_expiry_emits_session_expired() -> None: + clock = {"now": 100.0} + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(idle_ttl_s=1.0), clock=lambda: clock["now"]) + clock["now"] += 2.0 + assert await h.manager.reap_expired() == 1 + events = await h.settle() + + assert types(events) == ["session.expired"] + assert events[0].reason == "idle_ttl_expired" + assert events[0].is_terminal + assert SESSION_ID not in h.manager.runners + assert not h.manager._closing + await close_harness(h) + + +@pytest.mark.asyncio +async def test_commands_for_a_closed_session_are_answered_with_unknown_session() -> None: + h = await open_harness() + await h.manager.handle(CloseDuplexSessionMessage(control_id="c-close", session_id=SESSION_ID, reason="bye")) + await h.results.get() + await h.settle() + + h.submit(commands.Heartbeat(event_id="evt-late")) + events = await h.settle() + + assert types(events) == ["error"] + assert events[0].code == "unknown_session" + assert events[0].related_event_id == "evt-late" + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Input path # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_append_plans_and_submits_one_model_unit_in_wire_order() -> None: + h = await open_harness() + try: + events = await h.run(append_audio()) + assert types(events) == ["input_audio_buffer.speech_started"] + assert len(h.port.submissions) == 1 + first = h.port.submissions[0] + assert first.context.request_id == h.stage0_request_id(epoch=0) + assert first.context.stage_id == 0 + assert first.already_submitted is False + duplex = first.prompt["model_intermediate_buffer"]["duplex"] + assert duplex["session_id"] == SESSION_ID + assert (duplex["epoch"], duplex["seq"], duplex["turn_id"], duplex["turn_seq"]) == (0, 1, 0, 1) + assert duplex["final"] is False + assert "incarnation" not in duplex + assert len(base64.b64decode(duplex["payload"]["audio"])) == 16000 * 4 + assert duplex["payload"]["is_speech"] is True + assert h.session.input_seq == 1 + assert h.session.stage_request_submitted(0, first.context.request_id) + + await h.run(append_audio()) + second = h.port.submissions[1] + assert second.already_submitted is True + assert second.prompt["model_intermediate_buffer"]["duplex"]["seq"] == 2 + assert h.session.input_seq == 2 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_partial_units_are_buffered_until_a_whole_chunk_is_available() -> None: + h = await open_harness() + try: + await h.run(append_audio(samples=8000)) + assert h.port.submissions == [] + assert h.runner.model_state.audio_buffer.has_pending() + + await h.run(append_audio(samples=8000)) + assert len(h.port.submissions) == 1 + assert not h.runner.model_state.audio_buffer.has_pending() + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_input_backpressure_is_rejected_before_the_buffer_is_touched() -> None: + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(max_pending_input_bytes_per_session=1000)) + try: + events = await h.run(append_audio(event_id="evt-big")) + assert types(events) == ["error"] + assert events[0].code == "input_backpressure" + assert events[0].related_event_id == "evt-big" + assert h.port.submissions == [] + assert not h.runner.model_state.audio_buffer.has_pending() + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_text_append_is_not_supported_by_the_native_runtime() -> None: + h = await open_harness() + try: + events = await h.run(commands.AppendText(text="hello", event_id="evt-text")) + assert types(events) == ["error"] + assert events[0].code == "native_text_append_unsupported" + assert events[0].related_event_id == "evt-text" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_clear_input_drops_buffered_audio() -> None: + h = await open_harness() + try: + await h.run(append_audio(samples=8000)) + assert h.runner.model_state.audio_buffer.has_pending() + events = await h.run(commands.ClearInput()) + assert types(events) == ["input_audio_buffer.cleared"] + assert not h.runner.model_state.audio_buffer.has_pending() + assert h.session.pending_input_bytes == 0 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_commit_without_audio_is_rejected() -> None: + h = await open_harness() + try: + events = await h.run(commands.Commit(event_id="evt-commit")) + assert types(events) == ["error"] + assert events[0].code == "input_audio_buffer_empty" + assert events[0].related_event_id == "evt-commit" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_commit_projects_the_user_item_and_does_not_resubmit_consumed_audio() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + events = await h.run(commands.Commit()) + assert types(events) == [ + "input_audio_buffer.speech_stopped", + "conversation.item.added", + "conversation.item.created", + "input_audio_buffer.committed", + "conversation.item.done", + ] + committed = find(events, "input_audio_buffer.committed") + assert committed.item_id == find(events, "conversation.item.added").item_id + # The unit already reached Stage0; a commit is only a turn marker here. + assert len(h.port.submissions) == 1 + assert h.session.input_commit_seq == 1 + assert h.session.history[-1]["role"] == "user" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_heartbeat_touches_the_lease_and_acks() -> None: + clock = {"now": 100.0} + h = await open_harness(clock=lambda: clock["now"]) + try: + clock["now"] = 150.0 + events = await h.run(commands.Heartbeat()) + assert types(events) == ["session.heartbeat_ack"] + assert h.session.lease.last_activity == 150.0 + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Model output path # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_stage1_audio_opens_a_response_and_streams_deltas() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="he")) + assert types(events) == [ + "response.created", + "conversation.item.added", + "conversation.item.created", + "response.output_item.added", + "response.speak", + "response.content_part.added", + "response.output_audio.delta", + "response.output_audio_transcript.delta", + ] + response_id = find(events, "response.created").response_id + assert response_id == h.session.active_response_id + assert h.session.active_request_id == request_id + delta = find(events, "response.output_audio.delta") + assert (delta.delta, delta.format, delta.sample_rate_hz) == ("wav-24000", "wav", 24000) + assert delta.response_id == response_id + assert find(events, "response.output_audio_transcript.delta").delta == "he" + assert h.session.playback.sent_ms == 1000 + + # Cumulative Stage1 audio is sliced to the new samples only. + events = await h.deliver_and_settle(tts_output(request_id, samples=48000, text="hello")) + assert types(events) == ["response.output_audio.delta", "response.output_audio_transcript.delta"] + assert find(events, "response.output_audio.delta").delta == "wav-24000" + assert find(events, "response.output_audio_transcript.delta").delta == "llo" + assert h.session.playback.sent_ms == 2000 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_turn_end_completes_the_response_and_advances_the_model_turn() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello", turn_end=True)) + assert types(events) == [ + "response.output_audio.done", + "response.output_audio_transcript.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + "response.done", + "rate_limits.updated", + ] + done = find(events, "response.done") + assert done.status == "completed" + assert find(events, "response.output_audio_transcript.done").transcript == "hello" + assert h.session.active_response_id is None + assert h.session.turn_id == 1 + # Ack-only playback: the assistant text enters history on playback ack. + assert h.session.history == () + + ack = await h.run(commands.AckPlayback(played_ms=1000, response_id=done.response_id)) + assert types(ack) == ["playback.acknowledged"] + assert ack[0].details["history_committed"] is True + assert ack[0].details["playback"]["committed_ms"] == 1000 + assert h.session.history == ({"role": "assistant", "content": "hello"},) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_playback_ack_for_an_unknown_response_is_rejected() -> None: + h = await open_harness() + try: + events = await h.run(commands.AckPlayback(played_ms=10, response_id="resp-unknown", event_id="evt-ack")) + assert types(events) == ["error"] + assert events[0].code == "playback_item_not_found" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stale_epoch_output_is_dropped_after_barge_in() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="he")) + + events = await h.run(commands.BargeIn()) + assert h.session.epoch == 1 + assert h.session.active_response_id is None + assert h.port.aborts == [[request_id]] + assert h.port.cleanups == [([request_id], True)] + done = find(events, "response.done") + assert done.status == "cancelled" + assert types(events)[: types(events).index("response.done")] == [ + "response.output_audio.done", + "response.output_audio_transcript.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + ] + + late = tts_output(request_id, samples=48000, text="hello", epoch=0) + assert h.deliver(late, epoch=0) is True # consumed, never forwarded + assert await h.settle() == [] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_cancel_response_without_active_response_is_rejected() -> None: + h = await open_harness() + try: + events = await h.run(commands.CancelResponse(event_id="evt-cancel")) + assert types(events) == ["error"] + assert events[0].code == "response_not_active" + assert events[0].related_event_id == "evt-cancel" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_cancel_response_aborts_the_stage_request_and_reports_playback() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + await h.run(commands.AckPlayback(played_ms=400, response_id=h.session.active_response_id)) + + events = await h.run(commands.CancelResponse()) + done = find(events, "response.done") + assert done.status == "cancelled" + assert done.response["status_details"]["reason"] == "client_cancelled" + assert h.port.aborts == [[request_id]] + assert h.session.epoch == 1 + # Only the played prefix of the cancelled answer is kept in history. + assert h.session.history == ({"role": "assistant", "content": "he"},) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + # An unfinished Stage0 segment is plain text for the TTS stage: forwarded. + text_output = SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text="hi", token_ids=[11], multimodal_output={})], + multimodal_output={}, + ) + assert h.deliver(text_output, stage_id=0) is False + + consumed = h.deliver( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + events = await h.settle() + assert consumed is True + assert types(events) == ["response.listen"] + assert events[0].details["reason"] == "model_listen" + assert events[0].details["model_listen"] is True + assert events[0].to_realtime()["response"]["status"] == "listening" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_listen_decision_on_a_resumable_request_closes_the_bounded_response() -> None: + """A live stage-0 request is resumable: its output never says ``finished``, + only its segment does. The listen that answers the last continuation unit + must still close the response, or the session sits forever with no + terminal event (the live-client E2E hang).""" + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + response_id = h.session.active_response_id + assert response_id is not None + # The continuation budget is spent: the unit being answered was the forced listen. + model_state = h.runner.model_state + model_state.continuation_owner_id = f"response:{response_id}" + model_state.continuation_units = h.runner.model._AUTO_RESPONSE_MAX_CONTINUATION_UNITS + + listen = listen_output(request_id) + listen.finished = False + events = await h.deliver_and_settle( + listen, + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + assert types(events)[0] == "response.listen" + assert events[0].details["model_listen"] is True + assert find(events, "response.done").response_id == response_id + assert h.session.active_response_id is None + assert model_state.continuation_units == 0 + assert len(h.port.submissions) == 1 # no further silence unit was scheduled + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_listen_decision_on_a_resumable_request_keeps_the_turn_going() -> None: + """Same unfinished listen, budget left: it must schedule the next unit. + + A non-terminal auto-response listen is answered with the next silence + unit, not with a ``response.listen`` event, so the submission is the + observable. Before the fix the projector saw ``finished=False`` and + yielded nothing: no silence unit, the response left open with nothing in + flight -- the same hang, one unit earlier. + """ + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + response_id = h.session.active_response_id + assert response_id is not None + submissions_before = len(h.port.submissions) + + listen = listen_output(request_id) + listen.finished = False + events = await h.deliver_and_settle( + listen, + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + assert not [event for event in events if event.type == "error"], types(events) + assert len(h.port.submissions) == submissions_before + 1, "the next silence unit was not scheduled" + assert h.runner.model_state.continuation_units == 1 + assert h.session.active_response_id == response_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_tts_segment_end_schedules_a_silence_continuation_unit() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + await h.deliver_and_settle( + tts_output(request_id, samples=48000, text="hello", tts_is_last_chunk=True, finished=True) + ) + + assert len(h.port.submissions) == 2 + silence = h.port.submissions[1] + assert silence.already_submitted is True + duplex = silence.prompt["model_intermediate_buffer"]["duplex"] + assert duplex["seq"] == 2 + assert duplex["payload"]["duplex_turn_id"] == 0 + assert set(base64.b64decode(duplex["payload"]["audio"])) == {0} + assert h.runner.model_state.continuation_units == 1 + # The response stays open across the segment boundary. + assert h.session.active_response_id is not None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stage_failure_fails_the_active_response() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + + h.runner.on_stage_failure(1, RuntimeError("boom")) + events = await h.settle() + assert types(events)[0] == "error" + assert events[0].code == "runtime_data_plane_stream_failed" + done = find(events, "response.done") + assert done.status == "failed" + assert done.response["status_details"]["reason"] == "runtime_data_plane_stream_failed" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Turn mode (no auto response) # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_turn_mode_skips_silent_chunks() -> None: + h = await open_harness(auto_response=False) + try: + events = await h.run(append_audio(is_speech=False, value=0.0)) + assert types(events) == ["response.listen"] + assert events[0].details["reason"] == "silence_or_noise" + assert h.port.submissions == [] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_turn_mode_commit_with_response_create_starts_one_response() -> None: + h = await open_harness(auto_response=False) + try: + events = await h.run(append_audio()) + assert types(events) == ["input_audio_buffer.speech_started"] + assert h.port.submissions == [] # buffered until the commit + + events = await h.run(commands.Commit(create_response=True)) + assert "input_audio_buffer.committed" in types(events) + assert "response.created" in types(events) + assert len(h.port.submissions) == 1 + final = h.port.submissions[0] + assert final.prompt["model_intermediate_buffer"]["duplex"]["final"] is True + request_id = final.context.request_id + response_id = h.session.active_response_id + assert response_id is not None + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="sure")) + assert types(events)[-2:] == ["response.output_audio.delta", "response.output_audio_transcript.delta"] + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="sure", finished=True)) + done = find(events, "response.done") + assert done.response_id == response_id + assert done.status == "completed" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_response_create_with_no_input_at_all_is_rejected() -> None: + """Nothing committed and no items: refuse rather than emit an empty turn.""" + h = await open_harness(auto_response=False) + try: + events = await h.run(commands.CreateResponse(event_id="evt-resp")) + assert types(events) == ["error"] + assert events[0].code == "response_create_without_input" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_an_item_only_response_create_is_refused_not_left_hanging() -> None: + """Items are context, not a turn, and the refusal must be immediate. + + A model-native model decides to speak from the audio it hears, so a turn + with no audio in it is one it never answers. Opening a response anyway + submits work that never completes, and the caller only finds out when the + session idles out -- observed as a 300-second HTTP 500 against a real + server. Refusing costs the same information and none of the wait. + """ + h = await open_harness(auto_response=False) + try: + await h.run( + commands.CreateItem( + item={ + "id": "item_prompt", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What is 2+2?"}], + } + ) + ) + events = await h.run(commands.CreateResponse(event_id="evt-resp")) + assert types(events) == ["error"], events + assert events[0].code == "text_only_turn_unsupported" + assert not h.port.submissions, "nothing may be submitted for a turn the model cannot answer" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_conversation_items_are_consumed_by_the_turn_they_start() -> None: + """The items belong to the turn that answered them, not to the next one. + + Without this, a session that once received an item would let every later + ``response.create`` start an empty turn off the same stale input. + """ + h = await open_harness(auto_response=False) + try: + await h.run( + commands.CreateItem( + item={ + "id": "item_hi", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hi"}], + } + ) + ) + assert h.session.unanswered_user_items() == 1 + # Consumed by the response.create that refused them: a later one must + # not see the same stale input and refuse for a second time. + await h.run(commands.CreateResponse(event_id="evt-first")) + assert h.session.unanswered_user_items() == 0 + events = await h.run(commands.CreateResponse(event_id="evt-second")) + assert events[0].code == "response_create_without_input" + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Session update and conversation items # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_session_update_rejects_instruction_changes_after_the_context_is_locked() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + assert h.runner.model_state.context_locked + events = await h.run(commands.UpdateSession(patch={"instructions": "changed"}, event_id="evt-upd")) + assert types(events) == ["error"] + assert events[0].code == "instructions_update_unsupported" + assert events[0].related_event_id == "evt-upd" + assert h.session.config.instructions == "You are a concise assistant." + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_session_update_replaces_config_and_runtime_config_together() -> None: + h = await open_harness() + try: + generation = h.session.config_generation + events = await h.run(commands.UpdateSession(patch={"temperature": 0.3})) + assert types(events) == ["session.updated"] + assert events[0].session["temperature"] == 0.3 + assert h.session.config.temperature == 0.3 + assert h.session.runtime_config["duplex_stage_sampling_params"]["0"]["temperature"] == 0.3 + assert h.session.config_generation == generation + 2 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_conversation_items_can_be_injected_and_deleted() -> None: + h = await open_harness() + try: + item = {"id": "item_hist", "type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + events = await h.run(commands.CreateItem(item=item)) + assert types(events) == ["conversation.item.added", "conversation.item.created", "conversation.item.done"] + assert h.session.history == ({"role": "user", "content": "hi"},) + + events = await h.run(commands.DeleteItem(item_id="item_hist")) + assert types(events) == ["conversation.item.deleted"] + assert h.session.history == () + + events = await h.run(commands.DeleteItem(item_id="item_missing", event_id="evt-del")) + assert types(events) == ["error"] + assert events[0].code == "item_not_found" + finally: + await close_harness(h) + + +def _stage_metrics_of(event: object) -> dict[str, dict[str, object]]: + """Per-stage engine metrics as the client reads them off one wire event.""" + payload = event.to_realtime() + metadata = payload.get("metadata") + assert isinstance(metadata, dict), payload + vllm_omni = metadata.get("vllm_omni") + assert isinstance(vllm_omni, dict), metadata + stage_metrics = vllm_omni.get("stage_metrics") + assert isinstance(stage_metrics, dict), vllm_omni + return stage_metrics + + +@pytest.mark.asyncio +async def test_stage0_metrics_reach_the_response_even_though_its_output_feeds_tts() -> None: + """A pass-through Stage0 output still has to report its tokens. + + Stage0 text with no decision is forwarded to the TTS stage rather than + consumed, and before sessions moved into the engine the orchestrator + published its metrics separately. Returning ``False`` here must not also + throw the metrics away, or every response reports no engine token count and + the benchmark cannot compute TPOT. + """ + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + text_output = SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text="hi", token_ids=[11], multimodal_output={})], + multimodal_output={}, + ) + forwarded = h.deliver( + text_output, + stage_id=0, + metrics=stage_stats(stage_id=0, request_id=request_id, num_tokens_out=3, itls_ms=[9.0, 11.0]), + ) + assert forwarded is False + assert await h.settle() == [] + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hi")) + stage_metrics = _stage_metrics_of(find(events, "response.output_audio.delta")) + assert stage_metrics["0"]["num_tokens_out"] == 3 + assert stage_metrics["0"]["vllm_itls_ms"] == [9.0, 11.0] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stage0_metrics_from_several_units_are_summed_into_one_response() -> None: + """Two pass-through segments before the first audio are one response's tokens.""" + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + for tokens in (3, 4): + h.deliver( + SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text="hi", token_ids=[11], multimodal_output={})], + multimodal_output={}, + ), + stage_id=0, + metrics=stage_stats(stage_id=0, request_id=request_id, num_tokens_out=tokens), + ) + assert await h.settle() == [] + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hi")) + stage_metrics = _stage_metrics_of(find(events, "response.output_audio.delta")) + assert stage_metrics["0"]["num_tokens_out"] == 7 + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Server VAD # +# --------------------------------------------------------------------------- # + + +class _ScriptedDetector: + """A turn detector that answers each appended chunk with a pre-decided result.""" + + def __init__(self, results: list[object]) -> None: + self._results = list(results) + self.resets = 0 + + def process(self, base64_audio: str, *, fmt: str, sample_rate_hz: int | None, audio_end_ms: int | None = None): + del base64_audio, fmt, sample_rate_hz, audio_end_ms + return self._results.pop(0) + + def reset(self) -> None: + self.resets += 1 + + +def _speech_then_stop() -> _ScriptedDetector: + from vllm_omni.engine.duplex.turn_detection import TurnDetectionResult + + return _ScriptedDetector( + [ + TurnDetectionResult(is_speech=True, speech_active=True, speech_started=True, speech_probability=0.9), + TurnDetectionResult( + is_speech=False, + speech_active=False, + speech_stopped=True, + speech_probability=0.05, + should_commit=True, + create_response=True, + ), + ] + ) + + +def _final_submissions(h: Harness) -> list[object]: + return [s for s in h.port.submissions if s.prompt["model_intermediate_buffer"]["duplex"].get("final") is True] + + +@pytest.mark.asyncio +async def test_server_vad_speech_stopped_does_not_end_the_turn_of_an_auto_response_session() -> None: + """A model-native session keeps its turn open past the detector's stop. + + Committing at ``speech_stopped`` cut the utterance short: the model + listened on that early final unit, the trailing silence became a second, + near-empty turn, and the client's own commit answered nothing (the + server-VAD hard-interrupt E2E timed out waiting for the follow-up + response). The old translator synthesized that commit for turn-based + server VAD only. + """ + h = await open_harness() + try: + h.runner.control._detector = _speech_then_stop() + first = await h.run(append_audio()) + assert "input_audio_buffer.speech_started" in types(first) + # A partial unit, like the trailing silence of a real utterance: it is + # what a commit here would submit as the turn's final append. + second = await h.run(append_audio(samples=8000, is_speech=False, value=0.0)) + assert "input_audio_buffer.speech_stopped" in types(second) + assert "input_audio_buffer.committed" not in types(first) + types(second) + assert _final_submissions(h) == [], "the detector must not close a model-native turn" + + events = await h.run(commands.Commit()) + assert types(events).count("input_audio_buffer.committed") == 1 + assert len(_final_submissions(h)) == 1, "the client's commit closes the turn exactly once" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_server_vad_speech_stopped_still_commits_a_turn_mode_session() -> None: + """Turn-based server VAD keeps the contract the old translator implemented.""" + h = await open_harness(auto_response=False) + try: + h.runner.control._detector = _speech_then_stop() + await h.run(append_audio()) + events = await h.run(append_audio(is_speech=False, value=0.0)) + assert "input_audio_buffer.speech_stopped" in types(events) + assert "input_audio_buffer.committed" in types(events) + assert len(_final_submissions(h)) == 1, "the detector's stop commits the turn and starts the response" + finally: + await close_harness(h) diff --git a/tests/engine/duplex/test_vad_backend.py b/tests/engine/duplex/test_vad_backend.py new file mode 100644 index 000000000..8a22259ac --- /dev/null +++ b/tests/engine/duplex/test_vad_backend.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The engine-side Silero VAD: backend selection and endpoint rules. + +The endpoint rules are the part worth pinning hardest. They were reimplemented +when turn detection moved into the engine, and drifted from upstream's +``ThresholdEndpointPolicy`` in two ways that only show up mid-utterance: a loud +frame used to cancel an in-progress silence timer, and ``audio_end_ms`` used to +exclude the trailing silence OpenAI says it should include. The parity cases +below hold upstream's own outputs, so neither can come back unnoticed. +""" + +from __future__ import annotations + +import base64 +import hashlib +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import numpy as np +import pytest + +from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection +from vllm_omni.engine.duplex.vad import ( + SILERO_VAD_SHA256, + ServerVADUnavailableError, + SileroStreamingVAD, + SileroVADBackend, + SileroVADBackendProvider, + SileroVADConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +FRAME = 512 +SAMPLE_RATE_HZ = 16_000 + + +def _drive(config: SileroVADConfig, probabilities: list[float]) -> list[tuple[str, int | None]]: + """Feed one probability per frame and collect the endpoint events.""" + scores = iter(probabilities) + vad = SileroStreamingVAD(config, frame_scorer=lambda _frame: next(scores)) + events: list[tuple[str, int | None]] = [] + for _ in probabilities: + result = vad.process(np.zeros(FRAME, dtype=np.float32)) + if result.speech_started: + events.append(("start", result.speech_start_ms)) + if result.speech_stopped: + events.append(("stop", result.speech_end_ms)) + return events + + +# --------------------------------------------------------------------------- # +# Parity with the serving-side policy this was ported from # +# --------------------------------------------------------------------------- # + +#: Probability sequences with the endpoint events upstream's +#: ``ThresholdEndpointPolicy`` produced for them, captured from +#: ``entrypoints/duplex/server_vad.py`` at `e2d2617f` before that module was +#: removed. Golden rather than computed because the oracle no longer ships: the +#: serving layer is transport only, so the VAD it used to carry is gone. Any +#: drift in the engine's rules still breaks these. +_PARITY_CASES = { + "one utterance": ( + [0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 100, "min_speech_duration_ms": 32}, + [("start", 32), ("stop", 256)], + ), + "prefix padding reaches back": ( + [0.0, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0], + {"threshold": 0.5, "prefix_padding_ms": 300, "silence_duration_ms": 64, "min_speech_duration_ms": 32}, + [("start", 0), ("stop", 160)], + ), + "loud frames inside the silence timer": ( + [0.9, 0.9, 0.1, 0.45, 0.1, 0.1, 0.1], + {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 96, "min_speech_duration_ms": 32}, + [("start", 0), ("stop", 160)], + ), + "min speech duration rejects a blip": ( + [0.9, 0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0], + {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 64, "min_speech_duration_ms": 96}, + [("start", 64), ("stop", 224)], + ), + "negative threshold floor at a low threshold": ( + [0.9, 0.9, 0.1, 0.1, 0.1, 0.1], + {"threshold": 0.16, "prefix_padding_ms": 0, "silence_duration_ms": 64, "min_speech_duration_ms": 32}, + [("start", 0)], + ), +} + + +@pytest.mark.parametrize(("probabilities", "options", "expected"), _PARITY_CASES.values(), ids=list(_PARITY_CASES)) +def test_endpoint_decisions_match_the_serving_side_policy( + probabilities: list[float], + options: dict[str, object], + expected: list[tuple[str, int | None]], +) -> None: + """The engine detector reproduces the policy server VAD shipped with upstream.""" + assert _drive(SileroVADConfig(**options), probabilities) == expected + + +# --------------------------------------------------------------------------- # +# The two rules that had drifted # +# --------------------------------------------------------------------------- # + + +def test_a_loud_frame_does_not_restart_the_silence_timer() -> None: + """Silero v6.2 hysteresis: once silence is running, a loud frame only delays it. + + The frame at 0.45 sits above the negative threshold (0.35) but below the + activation threshold. It must not cancel the pending endpoint, or a speaker + who trails off unevenly never gets their turn committed. + """ + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=96, min_speech_duration_ms=32) + assert _drive(config, [0.9, 0.9, 0.1, 0.45, 0.1, 0.1]) == [("start", 0), ("stop", 160)] + + +def test_audio_end_ms_includes_the_trailing_silence() -> None: + """OpenAI defines audio_end_ms as the end of the audio sent to the model. + + That includes the silence spent deciding the turn was over, so it is the + frame boundary the detector stopped on, not where speech last was. + """ + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=100, min_speech_duration_ms=32) + events = _drive(config, [0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + stop_ms = next(ms for kind, ms in events if kind == "stop") + # 8 frames consumed when the timer expires: 8 * 512 / 16 kHz = 256 ms. + assert stop_ms == 256 + + +def test_reset_keeps_the_session_clock_and_clamps_the_prefix() -> None: + """A barge-in resets the detector but not the session's timeline. + + Timestamps have to stay comparable with the rest of the session, and prefix + padding must not reach back into audio the reset discarded. + """ + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=300, silence_duration_ms=64, min_speech_duration_ms=32) + scores = iter([0.0] * 10 + [0.9, 0.9, 0.9]) + vad = SileroStreamingVAD(config, frame_scorer=lambda _frame: next(scores)) + for _ in range(10): + vad.process(np.zeros(FRAME, dtype=np.float32)) + + stream_start_ms = round(10 * FRAME * 1000 / SAMPLE_RATE_HZ) + vad.reset() + result = vad.process(np.zeros(FRAME * 3, dtype=np.float32)) + + assert result.speech_started + # Without the clock the start would land at 0; without the clamp, 300 ms of + # prefix would reach back before the reset. + assert result.speech_start_ms == stream_start_ms + + +# --------------------------------------------------------------------------- # +# Input handling # +# --------------------------------------------------------------------------- # + + +def _pcm16_b64(samples: np.ndarray) -> str: + return base64.b64encode((samples * 32767).astype(" None: + config = SileroVADConfig(threshold=0.5) + tone = np.zeros(FRAME * 2, dtype=np.float32) + seen: list[int] = [] + vad = SileroStreamingVAD(config, frame_scorer=lambda frame: seen.append(frame.size) or 0.0) + vad.process_base64(_pcm16_b64(tone), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ) + assert seen == [FRAME, FRAME] + + +def test_an_odd_pcm16_byte_count_is_rejected() -> None: + vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=lambda _frame: 0.0) + with pytest.raises(ValueError, match="incomplete pcm16"): + vad.process_base64(base64.b64encode(b"\x00\x01\x02").decode(), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ) + + +def test_resampling_keeps_the_frame_grid_from_drifting() -> None: + """The 24 kHz remainder carries across chunks, so frame boundaries cannot drift. + + This is a weaker claim than upstream's ``StreamingAudioResampler``: the + interpolation is per chunk, so values near a chunk edge differ slightly. + What must hold is the frame *count* and alignment, which is what the + endpoint timestamps are derived from. + """ + source = np.sin(np.linspace(0, 40 * np.pi, 24_000, dtype=np.float32)) + + def frames_for(chunk_sizes: list[int]) -> int: + frames = 0 + + def count(_frame: np.ndarray) -> float: + nonlocal frames + frames += 1 + return 0.0 + + vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=count) + offset = 0 + for size in chunk_sizes: + vad.process_base64(_pcm16_b64(source[offset : offset + size]), fmt="pcm16", sample_rate_hz=24_000) + offset += size + assert offset == source.size + return frames + + whole = frames_for([source.size]) + split = frames_for([7, 13, 511, 512, 513, source.size - 1_556]) + assert whole == split == 16_000 // FRAME + + +def test_a_sample_rate_change_mid_stream_is_rejected() -> None: + vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=lambda _frame: 0.0) + chunk = _pcm16_b64(np.zeros(1_024, dtype=np.float32)) + vad.process_base64(chunk, fmt="pcm16", sample_rate_hz=24_000) + with pytest.raises(ValueError, match="sample rate cannot change"): + vad.process_base64(chunk, fmt="pcm16", sample_rate_hz=48_000) + + +# --------------------------------------------------------------------------- # +# Backend selection # +# --------------------------------------------------------------------------- # + + +def test_a_configured_model_path_that_does_not_exist_is_an_error() -> None: + """An operator who names an artifact gets told it is missing, not a fallback.""" + provider = SileroVADBackendProvider(model_path="/nonexistent/silero_vad.onnx") + with pytest.raises(ServerVADUnavailableError, match="does not exist"): + provider.get() + + +def test_a_configured_model_with_the_wrong_checksum_is_refused(tmp_path) -> None: + """A different Silero revision changes endpointing, so the pin is enforced.""" + impostor = tmp_path / "silero_vad.onnx" + impostor.write_bytes(b"not the pinned model") + provider = SileroVADBackendProvider(model_path=str(impostor)) + with pytest.raises(ServerVADUnavailableError, match="checksum mismatch"): + provider.get() + assert hashlib.sha256(impostor.read_bytes()).hexdigest() != SILERO_VAD_SHA256 + + +def test_the_onnx_backend_is_shared_by_every_session() -> None: + """One ONNX session for the whole engine: its state is passed in, not held.""" + provider = SileroVADBackendProvider() + try: + backend = provider.get() + except ServerVADUnavailableError as exc: + pytest.skip(f"no Silero backend available here: {exc}") + if not isinstance(backend, SileroVADBackend): + pytest.skip("no local ONNX artifact; the torch fallback is per-session by design") + + assert provider.get() is backend + # State is the caller's, so two sessions on one backend cannot interfere. + first, second = backend.new_state(), backend.new_state() + assert first is not second + probability, next_state = backend.infer(np.zeros(FRAME, dtype=np.float32), first) + assert 0.0 <= probability <= 1.0 + assert next_state is not first + + +def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tmp_path) -> None: + """Pin the ONNX call shape without needing the real artifact or runtime. + + The context carry is the part worth pinning: Silero v6.2 is fed 64 samples + of the previous frame ahead of the current 512, and that tail travels in the + caller's state rather than inside the session -- which is exactly what lets + one session serve every duplex session at once. + """ + sessions: list[FakeInferenceSession] = [] + + class FakeSessionOptions: + inter_op_num_threads = 0 + intra_op_num_threads = 0 + + class FakeInferenceSession: + def __init__(self, path: str, *, providers: list[str], sess_options: object) -> None: + self.providers = providers + self.sess_options = sess_options + self.calls: list[dict[str, np.ndarray]] = [] + self.run_barrier: threading.Barrier | None = None + sessions.append(self) + + def get_inputs(self) -> list[SimpleNamespace]: + return [SimpleNamespace(name=name) for name in ("input", "state", "sr")] + + def run(self, _outputs: object, inputs: dict[str, np.ndarray]) -> list[np.ndarray]: + self.calls.append({name: np.array(value, copy=True) for name, value in inputs.items()}) + if self.run_barrier is not None: + self.run_barrier.wait(timeout=5) + return [ + np.asarray([[0.75]], dtype=np.float32), + np.full((2, 1, 128), len(self.calls), dtype=np.float32), + ] + + monkeypatch.setitem( + sys.modules, + "onnxruntime", + SimpleNamespace(SessionOptions=FakeSessionOptions, InferenceSession=FakeInferenceSession), + ) + + model_path = tmp_path / "silero_vad.onnx" + model_path.write_bytes(b"fake-model") + backend = SileroVADBackend(model_path) + session = sessions[0] + + # Single-threaded on purpose: many sessions share this one ORT session, so + # it must not fan out per call. + assert session.providers == ["CPUExecutionProvider"] + assert session.sess_options.inter_op_num_threads == 1 + assert session.sess_options.intra_op_num_threads == 1 + # Warm-up already ran: 64 context + 512 frame. + assert session.calls[0]["input"].shape == (1, 576) + assert session.calls[0]["state"].shape == (2, 1, 128) + assert session.calls[0]["sr"].item() == SAMPLE_RATE_HZ + + state = backend.new_state() + first_frame = np.arange(backend.frame_samples, dtype=np.float32) + probability, state = backend.infer(first_frame, state) + first_call = session.calls[1] + + assert probability == pytest.approx(0.75) + np.testing.assert_array_equal( + first_call["input"][:, : backend.context_samples], + np.zeros((1, backend.context_samples), dtype=np.float32), + ) + np.testing.assert_array_equal(first_call["input"][:, backend.context_samples :], first_frame[None, :]) + + second_frame = -first_frame + backend.infer(second_frame, state) + second_call = session.calls[2] + # The previous frame's tail becomes this frame's context. + np.testing.assert_array_equal( + second_call["input"][:, : backend.context_samples], + first_frame[None, -backend.context_samples :], + ) + np.testing.assert_array_equal(second_call["state"], np.full((2, 1, 128), 2, dtype=np.float32)) + + # Two sessions may enter the shared ORT session at the same time. + session.run_barrier = threading.Barrier(2) + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(backend.infer, first_frame, backend.new_state()), + executor.submit(backend.infer, second_frame, backend.new_state()), + ] + assert [future.result()[0] for future in futures] == pytest.approx([0.75, 0.75]) + + +# --------------------------------------------------------------------------- # +# Session payload validation # +# --------------------------------------------------------------------------- # + + +def test_an_unknown_server_vad_field_is_refused() -> None: + """A misspelled knob that silently does nothing is worse than a refusal. + + The endpointing would still look configured while running on defaults. + """ + error = validate_realtime_turn_detection({"turn_detection": {"type": "server_vad", "semantic_eagerness": "high"}}) + assert error is not None + assert "semantic_eagerness" in error + + +def test_the_documented_server_vad_fields_are_all_accepted() -> None: + assert ( + validate_realtime_turn_detection( + { + "turn_detection": { + "type": "server_vad", + "threshold": 0.6, + "prefix_padding_ms": 200, + "silence_duration_ms": 400, + "create_response": True, + "interrupt_response": True, + "min_speech_duration_ms": 120, + }, + "overlap_policy": "barge_in_on_speech", + } + ) + is None + ) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py index 86276ec66..43a29bf3a 100644 --- a/tests/engine/test_async_omni_engine_abort_ack.py +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -54,7 +54,7 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture): rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="abort-rpc-1"), ) @@ -81,7 +81,7 @@ def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture): rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="abort-rpc-err"), ) @@ -109,7 +109,7 @@ def test_abort_async_times_out_without_result(mocker: MockerFixture): rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="abort-rpc-timeout"), ) @@ -139,7 +139,7 @@ def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture) rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-abort"), ) diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py index 6b0d6a188..425a16fc1 100644 --- a/tests/engine/test_async_omni_engine_outputs.py +++ b/tests/engine/test_async_omni_engine_outputs.py @@ -7,7 +7,6 @@ Focuses on the critical behavior: when the orchestrator thread dies, subsequent attempts to collect output raise RuntimeError. """ -import asyncio import queue import threading from concurrent.futures import ThreadPoolExecutor @@ -18,16 +17,6 @@ from pytest_mock import MockerFixture from vllm_omni.engine.async_engine_utils import weak_shutdown_async_omni_engine from vllm_omni.engine.async_omni_engine import AsyncOmniEngine -from vllm_omni.engine.duplex.control_client import ( - DuplexControlClient, - DuplexControlRequestError, -) -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - DuplexControlResultMessage, - DuplexFence, - SignalDuplexTurnMessage, -) from vllm_omni.engine.messages import ( CollectiveRPCResultMessage, ErrorMessage, @@ -166,9 +155,9 @@ def test_output_remains_on_shared_output_path(mocker: MockerFixture): raw_queue = queue.Queue() raw_queue.put_nowait( OutputMessage( - request_id="legacy-duplex-request", + request_id="shared-path-request", stage_id=0, - engine_outputs=OmniRequestOutput(request_id="legacy-duplex-request"), + engine_outputs=OmniRequestOutput(request_id="shared-path-request"), finished=False, ) ) @@ -176,293 +165,7 @@ def test_output_remains_on_shared_output_path(mocker: MockerFixture): output = engine.try_get_output(timeout=0.01) - assert output.request_id == "legacy-duplex-request" - - -def test_duplex_control_client_uses_engine_owned_correlated_transport(): - engine = object.__new__(AsyncOmniEngine) - engine._duplex_control_client = None - engine._correlated_rpc_client = CorrelatedRpcClient(queue.Queue(), queue.Queue()) - - try: - client = engine._get_duplex_control_client() - - assert isinstance(client, DuplexControlClient) - assert client._transport is engine._correlated_rpc_client - assert not hasattr(engine, "_submit_rpc_and_wait") - finally: - engine._correlated_rpc_client.close() - - -def test_duplex_control_api_requires_explicit_fence(): - engine = object.__new__(AsyncOmniEngine) - - with pytest.raises(TypeError, match="fence"): - engine.open_duplex_session("sid") - with pytest.raises(TypeError, match="fence"): - engine.append_duplex_input("sid", mode="append_audio_chunk", payload={}) - with pytest.raises(TypeError, match="fence"): - engine.signal_duplex_turn("sid", event="turn.end") - with pytest.raises(TypeError, match="fence"): - engine.close_duplex_session("sid") - - -def test_open_duplex_session_waits_for_control_ack(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - control_result = DuplexControlResultMessage( - control_id="ctrl-1", - fence=DuplexFence("sid"), - operation="open", - session_id="sid", - ok=False, - stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False}}], - unsupported_count=1, - error_count=0, - ) - - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-1")) - - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1) - msg = request_q.get(timeout=1) - assert msg.type == "open_duplex_session" - assert msg.control_id == "ctrl-1" - rpc_q.put(control_result) - with pytest.raises(DuplexControlRequestError) as exc_info: - pending.result(timeout=1) - - assert exc_info.value.result["unsupported_count"] == 1 - assert exc_info.value.result["stage_results"][0]["result"]["supported"] is False - engine._correlated_rpc_client.close() - - -def test_signal_duplex_turn_message_carries_next_fence(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - cancelled_fence = DuplexFence("sid-cancel", epoch=4, turn_id=2) - next_fence = DuplexFence("sid-cancel", epoch=5, turn_id=2) - control_result = DuplexControlResultMessage( - control_id="ctrl-signal", - fence=cancelled_fence, - operation="signal", - session_id="sid-cancel", - ok=True, - stage_results=[], - ) - - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-signal")) - - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit( - engine.signal_duplex_turn, - "sid-cancel", - event="input.cancel", - fence=cancelled_fence, - next_fence=next_fence, - session_config={"temperature": 0.0}, - timeout=1, - ) - msg = request_q.get(timeout=1) - assert isinstance(msg, SignalDuplexTurnMessage) - assert msg.fence == cancelled_fence - assert msg.next_fence == next_fence - assert msg.session_config == {"temperature": 0.0} - rpc_q.put(control_result) - assert pending.result(timeout=1)["ok"] is True - - engine._correlated_rpc_client.close() - - -@pytest.mark.asyncio -async def test_cancelled_async_append_does_not_stop_executor_callable(mocker: MockerFixture): - entered = threading.Event() - release = threading.Event() - completed = threading.Event() - engine = object.__new__(AsyncOmniEngine) - - def blocking_append(*args, **kwargs): - del args, kwargs - entered.set() - release.wait(timeout=1) - completed.set() - return {} - - mocker.patch.object(engine, "append_duplex_input", side_effect=blocking_append) - task = asyncio.create_task( - engine.append_duplex_input_async( - "sid-executor-cancel", - mode="append_audio_chunk", - payload=b"audio", - fence=DuplexFence("sid-executor-cancel"), - ) - ) - assert await asyncio.to_thread(entered.wait, 1) - - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - release.set() - - assert await asyncio.to_thread(completed.wait, 1) - - -def test_duplex_controls_route_out_of_order_without_global_lock(): - request_q = queue.Queue() - rpc_q = queue.Queue() - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - first_fence = DuplexFence("sid", epoch=1) - second_fence = DuplexFence("sid", epoch=2) - first = AppendDuplexInputMessage( - control_id="first", - fence=first_fence, - session_id="sid", - expected_epoch=1, - mode="audio", - payload=b"first", - ) - second = AppendDuplexInputMessage( - control_id="second", - fence=second_fence, - session_id="sid", - expected_epoch=2, - mode="audio", - payload=b"second", - ) - - with ThreadPoolExecutor(max_workers=2) as executor: - client = engine._get_duplex_control_client() - first_result = executor.submit(client.execute, first, timeout=1) - second_result = executor.submit(client.execute, second, timeout=1) - submitted = {request_q.get(timeout=1).control_id, request_q.get(timeout=1).control_id} - assert submitted == {"first", "second"} - rpc_q.put( - DuplexControlResultMessage( - control_id="second", - fence=second_fence, - operation="append", - session_id="sid", - ok=True, - stage_results=[], - ) - ) - rpc_q.put( - DuplexControlResultMessage( - control_id="first", - fence=first_fence, - operation="append", - session_id="sid", - ok=True, - stage_results=[], - ) - ) - - assert first_result.result(timeout=1)["fence"] == first_fence - assert second_result.result(timeout=1)["fence"] == second_fence - - assert not hasattr(engine, "_rpc_lock") - engine._correlated_rpc_client.close() - - -def test_terminal_rpc_router_rejects_without_enqueuing_new_control(): - request_q = queue.Queue() - rpc_q = queue.Queue() - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - pending_message = AppendDuplexInputMessage( - control_id="pending", - fence=DuplexFence("sid"), - session_id="sid", - mode="audio", - payload=b"audio", - ) - message = AppendDuplexInputMessage( - control_id="after-fatal", - fence=DuplexFence("sid"), - session_id="sid", - mode="audio", - payload=b"audio", - ) - - try: - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit(engine._get_duplex_control_client().execute, pending_message, timeout=1) - assert request_q.get(timeout=1).control_id == "pending" - rpc_q.put(ErrorMessage(error="orchestrator failed", fatal=True)) - with pytest.raises(RuntimeError, match="orchestrator failed"): - pending.result(timeout=1) - - with pytest.raises(RuntimeError, match="orchestrator failed"): - engine._get_duplex_control_client().execute(message, timeout=1) - - assert request_q.empty() - finally: - engine._correlated_rpc_client.close() - - -def test_collective_and_duplex_results_share_router_without_swapping(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - fence = DuplexFence("sid") - duplex = AppendDuplexInputMessage( - control_id="duplex-control", - fence=fence, - session_id="sid", - mode="audio", - payload=b"audio", - ) - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="collective-rpc")) - - with ThreadPoolExecutor(max_workers=2) as executor: - duplex_result = executor.submit(engine._get_duplex_control_client().execute, duplex, timeout=1) - collective_result = executor.submit(engine.collective_rpc, "health", timeout=1) - submitted_types = {request_q.get(timeout=1).type, request_q.get(timeout=1).type} - assert submitted_types == {"append_duplex_input", "collective_rpc"} - rpc_q.put( - CollectiveRPCResultMessage( - rpc_id="collective-rpc", - method="health", - stage_ids=[0], - results=["healthy"], - ) - ) - rpc_q.put( - DuplexControlResultMessage( - control_id="duplex-control", - fence=fence, - operation="append", - session_id="sid", - ok=True, - stage_results=[], - ) - ) - - assert collective_result.result(timeout=1) == ["healthy"] - assert duplex_result.result(timeout=1)["fence"] == fence - - engine._correlated_rpc_client.close() + assert output.request_id == "shared-path-request" def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixture): @@ -483,8 +186,7 @@ def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixtu engine.request_queue = SimpleNamespace(sync_q=request_q) engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc")) + mocker.patch("vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc")) with ThreadPoolExecutor(max_workers=1) as executor: pending = executor.submit(engine.collective_rpc, "health", timeout=1) @@ -506,36 +208,6 @@ def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixtu engine._correlated_rpc_client.close() -def test_open_duplex_session_raises_on_stage_control_error(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - control_result = DuplexControlResultMessage( - control_id="ctrl-error", - fence=DuplexFence("sid"), - operation="open", - session_id="sid", - ok=False, - stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False, "error": "boom"}}], - unsupported_count=1, - error_count=1, - ) - - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-error")) - - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1) - request_q.get(timeout=1) - rpc_q.put(control_result) - with pytest.raises(RuntimeError, match="duplex open failed"): - pending.result(timeout=1) - engine._correlated_rpc_client.close() - - def test_shutdown_does_not_race_runtime_cleanup_with_live_orchestrator(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) engine._shutdown_called = False diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 27677bd54..f052f7e9d 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -12,7 +12,7 @@ import types import pytest from vllm_omni.diffusion.data import AttentionConfig -from vllm_omni.engine import async_omni_engine as async_omni_engine_module +from vllm_omni.engine import omni_engine_base as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, @@ -188,7 +188,7 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.omni_engine_base as engine_mod engine = object.__new__(AsyncOmniEngine) engine.stage_configs = [types.SimpleNamespace()] diff --git a/tests/engine/test_diffusion_streaming_output_kwarg.py b/tests/engine/test_diffusion_streaming_output_kwarg.py index a6f3dff02..8bc00f587 100644 --- a/tests/engine/test_diffusion_streaming_output_kwarg.py +++ b/tests/engine/test_diffusion_streaming_output_kwarg.py @@ -33,7 +33,7 @@ def _resolve(monkeypatch, kwargs: dict[str, Any]) -> dict[str, Any]: return "config.yaml", [], None monkeypatch.setattr( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + "vllm_omni.engine.omni_engine_base.load_and_resolve_stage_configs", _fake_load, ) engine = SimpleNamespace( diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py index 013ab8661..1da513182 100644 --- a/tests/engine/test_duplex_import_boundary.py +++ b/tests/engine/test_duplex_import_boundary.py @@ -45,27 +45,62 @@ def test_isolated_import_ignores_conflicting_gpu_visibility(monkeypatch: pytest. _assert_isolated_import_succeeds("import vllm_omni.outputs") -def test_stable_engine_imports_load_duplex_kernel_eagerly() -> None: - # The duplex kernel is imported eagerly by the stable engine modules. - # Model-specific duplex adapters must still load only via the - # dotted-string plugin paths, and the client-side vllm_omni.clients - # package must never enter the stable runtime import graph. +def test_turn_based_stack_has_no_duplex_vocabulary() -> None: + # The generic bases and the turn-based siblings (AsyncOmni / AsyncOmniEngine / + # Orchestrator) carry no duplex code: importing them must not load the + # engine-side duplex package, the model plugins or the client library. _assert_isolated_import_succeeds(""" import sys import vllm_omni.engine.async_omni_engine +import vllm_omni.engine.omni_engine_base import vllm_omni.engine.orchestrator import vllm_omni.entrypoints.async_omni +import vllm_omni.entrypoints.async_omni_base + +forbidden_prefixes = ( + "vllm_omni.engine.duplex", + "vllm_omni.engine.duplex_omni_engine", + "vllm_omni.engine.duplex_orchestrator", + "vllm_omni.entrypoints.duplex", + "vllm_omni.entrypoints.duplex_omni", + "vllm_omni.clients", + "vllm_omni.model_executor.models.minicpmo_4_5.duplex", + "vllm_omni.model_executor.models.nemotron_voicechat.duplex", + "vllm_omni.model_executor.models.personaplex.duplex", +) +loaded = sorted( + name + for name in sys.modules + if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes) +) +if loaded: + raise SystemExit("turn-based imports loaded duplex modules: " + ", ".join(loaded)) +""") + + +def test_duplex_stack_loads_the_framework_but_not_the_model_plugins() -> None: + # The duplex siblings import the engine-side duplex framework eagerly; + # model plugins still load only through PipelineConfig.duplex_plugin, and + # the client-side package never enters the server import graph. + _assert_isolated_import_succeeds(""" +import sys + +import vllm_omni.engine.duplex_omni_engine +import vllm_omni.engine.duplex_orchestrator +import vllm_omni.entrypoints.duplex_omni +import vllm_omni.entrypoints.duplex.serving expected_eager = ( - "vllm_omni.engine.duplex.contracts", - "vllm_omni.engine.duplex.control_plane", - "vllm_omni.entrypoints.duplex_request_client", - "vllm_omni.outputs.duplex", + "vllm_omni.engine.duplex.commands", + "vllm_omni.engine.duplex.events", + "vllm_omni.engine.duplex.session.engine_session", + "vllm_omni.engine.duplex.session.manager", + "vllm_omni.engine.duplex.plugin", ) missing = sorted(name for name in expected_eager if name not in sys.modules) if missing: - raise SystemExit("duplex kernel modules not imported eagerly: " + ", ".join(missing)) + raise SystemExit("duplex framework modules not imported eagerly: " + ", ".join(missing)) forbidden_prefixes = ( "vllm_omni.clients", @@ -79,7 +114,7 @@ loaded = sorted( if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes) ) if loaded: - raise SystemExit("stable imports loaded plugin-only duplex modules: " + ", ".join(loaded)) + raise SystemExit("duplex imports loaded plugin-only or client modules: " + ", ".join(loaded)) """) @@ -190,16 +225,32 @@ def test_engine_duplex_uses_canonical_contract_module_names() -> None: core_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "core" for name in ( + "commands.py", "contracts.py", - "lease.py", + "events.py", "messages.py", - "session.py", - "control_plane.py", - "control_client.py", - "runtime.py", + "plugin.py", "intermediate.py", ): assert (engine_dir / name).is_file() + # one session and everything that runs it lives in the session subpackage + for name in ( + "__init__.py", + "engine_session.py", + "manager.py", + "runner.py", + "context.py", + "emitter.py", + "model_channel.py", + "control.py", + "append_task.py", + "helpers.py", + "lease.py", + ): + assert (engine_dir / "session" / name).is_file() + assert not (engine_dir / "session.py").exists() + for removed in ("control_plane.py", "control_client.py", "runtime.py"): + assert not (engine_dir / removed).exists() # the duplex_ prefix is dropped inside the duplex package assert not (engine_dir / "duplex_session.py").exists() assert not (engine_dir / "duplex_lease.py").exists() diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py new file mode 100644 index 000000000..06af9e7af --- /dev/null +++ b/tests/engine/test_duplex_omni_engine.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexOmniEngine``: the session message surface over the generic engine base.""" + +from __future__ import annotations + +import asyncio +import queue +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.engine.duplex.commands import Heartbeat +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionError, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeRpcClient: + def __init__(self, result: Any) -> None: + self.result = result + self.calls: list[tuple[tuple[str, str], Any, float | None]] = [] + #: Control RPCs must block rather than raise a raw queue.Full. + self.block_on_submit_flags: list[bool] = [] + + def execute(self, key, message, *, timeout=None, timeout_message=None, block_on_submit=False): + self.calls.append((key, message, timeout)) + self.block_on_submit_flags.append(block_on_submit) + if isinstance(self.result, BaseException): + raise self.result + return self.result + + +class _FakeSyncQueue: + def __init__(self, *, full: bool = False) -> None: + self.items: list[Any] = [] + self.full = full + + def put(self, item: Any, timeout: float | None = None) -> None: + if self.full: + raise queue.Full + self.items.append(item) + + +def _engine(result: Any, *, alive: bool = True, full: bool = False) -> DuplexOmniEngine: + engine = object.__new__(DuplexOmniEngine) + engine._correlated_rpc_client = _FakeRpcClient(result) + engine.request_queue = SimpleNamespace(sync_q=_FakeSyncQueue(full=full)) + engine.is_alive = lambda: alive + engine.plugin = None + return engine + + +def _ok(operation: str, **fields: Any) -> DuplexControlResultMessage: + return DuplexControlResultMessage(control_id="c", operation=operation, session_id="sid", ok=True, **fields) + + +@pytest.mark.asyncio +async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> None: + result = _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0) + engine = _engine(result) + config = DuplexSessionConfig(model="m") + + returned = await engine.open_session_async("sid", config, timeout=3.0) + + assert returned is result + ((key, message, timeout),) = engine.rpc_client.calls + assert key == ("duplex", message.control_id) + assert isinstance(message, OpenDuplexSessionMessage) + assert message.session_id == "sid" and message.session_config is config + assert timeout == 3.0 + # A momentarily full request queue is backpressure, not a failed open: the + # control RPC blocks rather than raising an untyped queue.Full at the caller. + assert engine.rpc_client.block_on_submit_flags == [True] + + +@pytest.mark.asyncio +async def test_close_resume_and_touch_build_their_messages_without_incarnation() -> None: + engine = _engine(_ok("close")) + await engine.close_session_async("sid", reason="bye") + await engine.resume_session_async("sid", expected_lease_generation=4) + await engine.touch_session_async("sid", activity="detach") + + close, resume, touch = (call[1] for call in engine.rpc_client.calls) + assert isinstance(close, CloseDuplexSessionMessage) and close.reason == "bye" + assert isinstance(resume, ResumeDuplexSessionMessage) and resume.expected_lease_generation == 4 + assert isinstance(touch, TouchDuplexSessionMessage) and touch.activity == "detach" + for message in (close, resume, touch): + assert message.session_id == "sid" + assert not hasattr(message, "incarnation") + + +@pytest.mark.asyncio +async def test_failed_control_result_maps_to_a_typed_session_error() -> None: + engine = _engine( + DuplexControlResultMessage( + control_id="c", + operation="open", + session_id="sid", + ok=False, + error_code="resource_exhausted", + error_message="duplex_session_capacity_exhausted: limit=1", + error_retryable=True, + ) + ) + with pytest.raises(DuplexSessionError) as excinfo: + await engine.open_session_async("sid", DuplexSessionConfig()) + assert excinfo.value.code == "resource_exhausted" + assert excinfo.value.retryable is True + assert excinfo.value.session_id == "sid" + + +@pytest.mark.asyncio +async def test_engine_error_and_timeout_map_to_session_errors() -> None: + # The correlated RPC client raises RuntimeError for a closed router or a + # terminal orchestrator error; both become a typed engine_error. + engine = _engine(RuntimeError("orchestrator exploded")) + with pytest.raises(DuplexSessionError, match="orchestrator exploded") as excinfo: + await engine.close_session_async("sid") + assert excinfo.value.code == "engine_error" + + engine = _engine(TimeoutError("duplex open timed out for session sid")) + with pytest.raises(DuplexSessionError) as excinfo: + await engine.open_session_async("sid", DuplexSessionConfig()) + assert excinfo.value.code == "timeout" and excinfo.value.retryable is True + + +@pytest.mark.asyncio +async def test_submit_command_is_one_way_on_the_request_queue() -> None: + engine = _engine(None) + command = Heartbeat(event_id="evt-1") + + await engine.submit_command_async("sid", command) + + (message,) = engine.request_queue.sync_q.items + assert isinstance(message, DuplexSessionCommandMessage) + assert message.session_id == "sid" and message.command is command + assert engine.rpc_client.calls == [] + + +@pytest.mark.asyncio +async def test_submit_command_rejects_dead_engine_and_full_queue() -> None: + with pytest.raises(DuplexSessionError) as excinfo: + await _engine(None, alive=False).submit_command_async("sid", Heartbeat()) + assert excinfo.value.code == "engine_dead" + + with pytest.raises(DuplexSessionError) as excinfo: + await _engine(None, full=True).submit_command_async("sid", Heartbeat()) + assert excinfo.value.code == "engine_backpressure" and excinfo.value.retryable is True + + +def test_duplex_capabilities_come_from_the_loaded_plugin() -> None: + engine = object.__new__(DuplexOmniEngine) + engine.plugin = None + with pytest.raises(RuntimeError, match="not loaded"): + engine.duplex_capabilities + capabilities = DuplexCapabilities(supports_barge_in=False) + engine.plugin = SimpleNamespace(capabilities=lambda *, max_sessions: capabilities) + engine.duplex_session_config = SimpleNamespace(max_sessions=3) + assert engine.duplex_capabilities is capabilities + + +def test_validate_deployment_requires_a_plugin_and_duplex_session_mode() -> None: + engine = object.__new__(DuplexOmniEngine) + engine.model = "plain-model" + engine.pipeline_config = SimpleNamespace(duplex_plugin=None) + engine.deploy_config = SimpleNamespace(session_mode="duplex") + with pytest.raises(ValueError, match="not a duplex model"): + engine._validate_deployment() + + engine.pipeline_config = SimpleNamespace(duplex_plugin="pkg.mod.Plugin") + engine.deploy_config = SimpleNamespace(session_mode="turn") + with pytest.raises(ValueError, match="session_mode: duplex"): + engine._validate_deployment() + + engine.deploy_config = None + with pytest.raises(ValueError, match="none resolved"): + engine._validate_deployment() + + +@pytest.mark.asyncio +async def test_async_wrappers_run_the_blocking_calls_off_the_event_loop() -> None: + loop = asyncio.get_running_loop() + engine = _engine(_ok("touch")) + seen: list[bool] = [] + original = engine.rpc_client.execute + + def execute(*args, **kwargs): + # Called from the default executor thread, never on the event loop thread. + try: + seen.append(asyncio.get_running_loop() is loop) + except RuntimeError: + seen.append(False) + return original(*args, **kwargs) + + engine.rpc_client.execute = execute + await engine.touch_session_async("sid", activity="heartbeat") + assert seen == [False] diff --git a/tests/engine/test_duplex_orchestrator.py b/tests/engine/test_duplex_orchestrator.py new file mode 100644 index 000000000..5b8130c36 --- /dev/null +++ b/tests/engine/test_duplex_orchestrator.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexOrchestrator``: the template seams and the stage port around ``OrchestratorBase``.""" + +from __future__ import annotations + +import asyncio +import struct +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from vllm.v1.engine.exceptions import EngineDeadError + +from tests.engine.test_orchestrator import ( + FakeOutputProcessor, + FakeRunningCounter, + FakeStageClient, + _build_stage_pools, +) +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands +from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState +from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + OpenDuplexSessionMessage, +) +from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator, DuplexOrchestratorRequestState +from vllm_omni.engine.messages import AbortRequestMessage, ShutdownRequestMessage +from vllm_omni.engine.orchestrator import Orchestrator +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import MiniCPMO45DuplexPlugin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SESSION_ID = "duplex-orch" + + +def _encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None: + del sample_rate_hz, response_format, speed + samples = int(np.asarray(audio, dtype=np.float32).size) if audio is not None else 0 + return f"wav-{samples}" if samples > 0 else None + + +def _stage_configs(count: int) -> list[object]: + return [SimpleNamespace(model_config=SimpleNamespace(max_model_len=4096)) for _ in range(count)] + + +def _build( + *, + stages: int = 1, + running_counter: FakeRunningCounter | None = None, + runtime_config: DuplexSessionRuntimeConfig | None = None, +) -> tuple[DuplexOrchestrator, list[FakeStageClient], asyncio.Queue, asyncio.Queue]: + clients = [FakeStageClient(stage_type="llm", final_output=index == stages - 1) for index in range(stages)] + pools = _build_stage_pools( + [[client] for client in clients], + output_processors=[FakeOutputProcessor() for _ in clients], + stage_vllm_configs=_stage_configs(stages), + ) + rpc_q: asyncio.Queue = asyncio.Queue() + output_q: asyncio.Queue = asyncio.Queue() + orchestrator = DuplexOrchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=output_q, + rpc_async_queue=rpc_q, + stage_pools=pools, + running_counter=running_counter, + plugin=MiniCPMO45DuplexPlugin(_encode_audio), + duplex_session_config=runtime_config or DuplexSessionRuntimeConfig(reaper_interval_s=0.01), + model_config=None, + ) + return orchestrator, clients, rpc_q, output_q + + +def _open_message(extra_body: dict[str, object] | None = None) -> OpenDuplexSessionMessage: + return OpenDuplexSessionMessage( + control_id=f"open-{SESSION_ID}", + session_id=SESSION_ID, + session_config=DuplexSessionConfig( + model="openbmb/MiniCPM-o-4_5", + modalities=["text"], + instructions="public instructions", + extra_body={"auto_response": True, **(extra_body or {})}, + ), + ) + + +async def _open(orchestrator: DuplexOrchestrator, rpc_q: asyncio.Queue, **kwargs: Any) -> DuplexControlResultMessage: + await orchestrator._dispatch_message(_open_message(**kwargs)) + result = await asyncio.wait_for(rpc_q.get(), timeout=2.0) + assert isinstance(result, DuplexControlResultMessage) + return result + + +async def _close(orchestrator: DuplexOrchestrator, rpc_q: asyncio.Queue) -> DuplexControlResultMessage: + await orchestrator._dispatch_message( + CloseDuplexSessionMessage(control_id=f"close-{SESSION_ID}", session_id=SESSION_ID, reason="test") + ) + result = await asyncio.wait_for(rpc_q.get(), timeout=2.0) + await _settle(orchestrator) + return result + + +def _stage0_request_id(epoch: int = 0) -> str: + return duplex_resource_request_id(DuplexFence(SESSION_ID, epoch=epoch), "stage0") + + +def _append_audio(samples: int = 16000) -> commands.AppendAudio: + return commands.AppendAudio( + audio=struct.pack(f"<{samples}f", *([0.05] * samples)), + format="pcm_f32le", + sample_rate_hz=16000, + is_speech=True, + ) + + +async def _submit(orchestrator: DuplexOrchestrator, command: commands.DuplexCommand) -> None: + assert await orchestrator._dispatch_message(DuplexSessionCommandMessage(session_id=SESSION_ID, command=command)) + await _settle(orchestrator) + + +async def _settle(orchestrator: DuplexOrchestrator, *, timeout_s: float = 3.0) -> None: + """Let the session runner drain its mailbox and append tasks.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + quiet = 0 + while loop.time() < deadline: + runner = orchestrator.session_manager.runners.get(SESSION_ID) + busy = runner is not None and ( + not runner._mailbox.empty() + or any(not task.done() for task in runner.tasks.append_tasks) + or any(not task.done() for task in runner._background_tasks) + ) + busy = busy or any(not task.done() for task in orchestrator.session_manager._dispatched_control_tasks) + quiet = 0 if busy else quiet + 1 + if quiet >= 5: + return + await asyncio.sleep(0.01) + + +def _tts_output(request_id: str, *, samples: int = 24000, text: str = "hello") -> SimpleNamespace: + return SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text=text, token_ids=[], multimodal_output={})], + multimodal_output={ + "audio": np.zeros(samples, dtype=np.float32), + "sr": 24000, + "meta.duplex_turn_id": np.array([0], dtype=np.int32), + "meta.duplex_epoch": np.array([0], dtype=np.int32), + }, + ) + + +# --------------------------------------------------------------------------- # +# Message dispatch and session lifecycle # +# --------------------------------------------------------------------------- # + + +def test_turn_based_orchestrator_has_no_session_manager() -> None: + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[], + ) + assert not hasattr(orchestrator, "session_manager") + + +@pytest.mark.asyncio +async def test_duplex_messages_are_dispatched_and_generic_ones_are_not() -> None: + orchestrator, _, rpc_q, _ = _build() + assert await orchestrator._dispatch_message(AbortRequestMessage(request_ids=["x"])) is False + assert await orchestrator._dispatch_message(ShutdownRequestMessage()) is False + assert await orchestrator._dispatch_message(_open_message()) is True + assert (await asyncio.wait_for(rpc_q.get(), timeout=2.0)).ok is True + await orchestrator.session_manager.shutdown() + + +@pytest.mark.asyncio +async def test_open_preregisters_the_stage0_request_and_close_releases_it() -> None: + counter = FakeRunningCounter() + orchestrator, clients, rpc_q, _ = _build(running_counter=counter) + result = await _open(orchestrator, rpc_q) + assert result.ok is True + assert result.public_session["id"] == SESSION_ID + assert result.capabilities is not None and result.capabilities.supports_input_append + + request_id = _stage0_request_id() + request_state = orchestrator.request_states[request_id] + assert isinstance(request_state, DuplexOrchestratorRequestState) + assert request_state.session_owned is True + assert request_state.session_id == SESSION_ID + assert request_state.fence == DuplexFence(SESSION_ID) + assert request_state.streaming.enabled is True + # Preregistration reserves the id; nothing is running until an append submits. + assert counter.value == 0 + assert clients[0].add_request_calls == [] + assert orchestrator.session_manager.runner_for_request_id(request_id) is not None + + closed = await _close(orchestrator, rpc_q) + assert closed.ok is True + assert request_id not in orchestrator.request_states + assert orchestrator.session_manager.runner_for_request_id(request_id) is None + assert counter.value == 0 + assert orchestrator.session_manager.active_count() == 0 + + +@pytest.mark.asyncio +async def test_open_failure_rolls_back_session_and_reserved_request() -> None: + orchestrator, _, rpc_q, _ = _build() + # Server-owned runtime keys may not come from the client. + result = await _open(orchestrator, rpc_q, extra_body={"duplex_stage_max_tokens": {"0": 99}}) + + assert result.ok is False + assert result.error_code == "invalid_duplex_runtime_config" + assert orchestrator.request_states == {} + assert orchestrator.session_manager.active_count() == 0 + + +@pytest.mark.asyncio +async def test_open_fails_cleanly_when_a_stage_has_no_live_replica() -> None: + orchestrator, _, rpc_q, _ = _build(stages=2) + orchestrator.stage_pools[1].evict_replica(0) + + result = await _open(orchestrator, rpc_q) + + assert result.ok is False + assert "stage 1 has no live replica" in (result.error_message or "") + assert orchestrator.request_states == {} + assert orchestrator.session_manager.active_count() == 0 + + +@pytest.mark.asyncio +async def test_bridge_state_keeps_public_and_runtime_config_separate() -> None: + orchestrator, _, rpc_q, _ = _build() + await _open(orchestrator, rpc_q) + request_state = orchestrator.request_states[_stage0_request_id()] + bridge = request_state.streaming.bridge_states["duplex"] + + assert bridge["session_id"] == SESSION_ID + assert bridge["fence"] == DuplexFence(SESSION_ID) + assert (bridge["epoch"], bridge["turn_id"], bridge["model_turn_id"]) == (0, 0, 0) + assert "incarnation" not in bridge + assert bridge["session_config"]["instructions"] == "public instructions" + assert "duplex_stage_sampling_params" not in bridge["session_config"] + assert bridge["runtime_config"]["instructions"] == "public instructions" + assert bridge["runtime_config"]["duplex_stage_max_tokens"] == {"0": 20, "1": 8192} + # The plugin's sampling policy is applied to the preregistered request. + assert request_state.sampling_params_list[0].max_tokens == 20 + await orchestrator.session_manager.shutdown() + + +# --------------------------------------------------------------------------- # +# Stage port # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_append_submits_the_resumable_stage0_request_and_counts_it_running() -> None: + counter = FakeRunningCounter() + orchestrator, clients, rpc_q, _ = _build(running_counter=counter) + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + + await _submit(orchestrator, _append_audio()) + assert len(clients[0].add_request_calls) == 1 + submitted = clients[0].add_request_calls[0][0] + assert submitted.request_id == request_id + assert submitted.resumable is True + assert submitted.sampling_params.max_tokens == 20 + request_state = orchestrator.request_states[request_id] + assert request_state.stage_fences[0] == DuplexFence(SESSION_ID) + assert 0 in request_state.stage_submit_ts + assert counter.value == 1 + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None and session.stage_request_submitted(0, request_id) + + # Later units update the same resumable request; the counter stays at one. + await _submit(orchestrator, _append_audio()) + assert len(clients[0].add_request_calls) == 2 + assert clients[0].add_request_calls[1][0].request_id == request_id + assert counter.value == 1 + + await _close(orchestrator, rpc_q) + assert clients[0].abort_calls == [[request_id]] + assert counter.value == 0 + + +@pytest.mark.asyncio +async def test_session_update_refreshes_the_next_append_sampling_params() -> None: + orchestrator, clients, rpc_q, _ = _build() + await _open(orchestrator, rpc_q) + request_state = orchestrator.request_states[_stage0_request_id()] + assert request_state.sampling_params_list[0].max_tokens == 20 + + await _submit(orchestrator, commands.UpdateSession(patch={"max_output_tokens": 7})) + await _submit(orchestrator, _append_audio()) + + assert request_state.sampling_params_list[0].max_tokens == 7 + assert clients[0].add_request_calls[-1][0].sampling_params.max_tokens == 7 + await orchestrator.session_manager.shutdown() + + +@pytest.mark.asyncio +async def test_forwarded_stage_requests_are_bound_and_barge_in_aborts_them() -> None: + orchestrator, clients, rpc_q, _ = _build(stages=2) + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + request_state = orchestrator.request_states[request_id] + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None + + # The base forwards Stage0 text to the TTS stage and reports the submission. + forwarded = SimpleNamespace(request_id=request_id, prompt_token_ids=[1, 2], resumable=True) + replica_id = await orchestrator.stage_pools[1].submit_initial(request_id, request_state, forwarded) + orchestrator._on_stage_submitted(1, request_id, replica_id, request_state) + assert request_state.stage_fences[1] == DuplexFence(SESSION_ID) + assert session.stage_request_submitted(1, request_id) + + await _submit(orchestrator, commands.BargeIn()) + assert session.epoch == 1 + assert clients[0].abort_calls == [[request_id]] + assert clients[1].abort_calls == [[request_id]] + assert request_id not in orchestrator.request_states + assert session.resource_request_ids() == [] + await orchestrator.session_manager.shutdown() + + +@pytest.mark.asyncio +async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue() -> None: + orchestrator, _, rpc_q, output_q = _build(stages=2) + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + request_state = orchestrator.request_states[request_id] + while not output_q.empty(): + output_q.get_nowait() + + consumed = await orchestrator._intercept_stage_output(1, 0, _tts_output(request_id), request_state, None, None) + await _settle(orchestrator) + + assert consumed is True + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None and session.active_response_id is not None + types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]] + assert "response.created" in types and "response.output_audio.delta" in types + + orphan = DuplexOrchestratorRequestState( + request_id="duplex-s.b3RoZXI.e.0.r.stage0", + prompt=None, + sampling_params_list=[], + final_stage_id=1, + session_owned=True, + session_id="other", + fence=DuplexFence("other"), + ) + assert await orchestrator._intercept_stage_output(1, 0, _tts_output(orphan.request_id), orphan, None, None) + await orchestrator.session_manager.shutdown() + + +@pytest.mark.asyncio +async def test_forward_failure_closes_the_owning_session() -> None: + orchestrator, clients, rpc_q, output_q = _build(stages=2) + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + request_state = orchestrator.request_states[request_id] + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None + + absorbed = await orchestrator._handle_forward_failure(request_id, 1, request_state, ValueError("bad thinker")) + await _settle(orchestrator) + + assert absorbed is True + assert request_id not in orchestrator.request_states + assert clients[0].abort_calls == [[request_id]] + assert SESSION_ID not in orchestrator.session_manager.runners + assert session.state == DuplexSessionState.CLOSED + assert orchestrator.session_manager.active_count() == 0 + types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]] + # The stage failure is reported before the session expires, never after. + assert types.index("error") < types.index("session.expired") + assert types[-1] == "session.expired" + + +@pytest.mark.asyncio +async def test_request_cleanup_failure_is_deferred_and_retried_by_the_reaper(monkeypatch) -> None: + orchestrator, clients, rpc_q, _ = _build() + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + manager = orchestrator.session_manager + + async def failing_abort(request_ids: list[str]) -> None: + raise RuntimeError("stage abort failed") + + monkeypatch.setattr(clients[0], "abort_requests_async", failing_abort) + with pytest.raises(RuntimeError, match="stage abort failed"): + await orchestrator._cleanup_request_ids([request_id], abort=True, release_owners=True) + await _settle(orchestrator) + + assert SESSION_ID not in manager.runners + assert SESSION_ID in manager._closing # admission slot retained + assert manager._pending_request_cleanups and not manager._request_cleanups_in_progress + assert manager.active_count() == 0 and manager._admission_count() == 1 + + monkeypatch.undo() + assert await manager.reap_expired() >= 1 + assert manager._pending_request_cleanups == {} + assert manager._admission_count() == 0 + assert request_id not in orchestrator.request_states + + +@pytest.mark.asyncio +async def test_reaper_runs_as_a_background_task_and_survives_one_failure(monkeypatch) -> None: + orchestrator, _, _, _ = _build(runtime_config=DuplexSessionRuntimeConfig(reaper_interval_s=0.01)) + calls = {"count": 0} + + async def reap_expired(now: float | None = None) -> int: + calls["count"] += 1 + if calls["count"] == 1: + raise RuntimeError("transient cleanup failure") + return 0 + + monkeypatch.setattr(orchestrator.session_manager, "reap_expired", reap_expired) + tasks = orchestrator._background_tasks() + assert len(tasks) == 1 + task = asyncio.create_task(tasks[0]) + await asyncio.sleep(0.05) + orchestrator._shutdown_event.set() + await task + assert calls["count"] >= 2 + await orchestrator._shutdown_extensions() + + +# --------------------------------------------------------------------------- # +# The hierarchy the two surfaces depend on # +# --------------------------------------------------------------------------- # + + +def test_the_duplex_stack_extends_the_turn_based_one() -> None: + """One engine serves a session and an ordinary request, so duplex must extend turn-based. + + Pinned because it is the kind of relationship a later tidy-up turns back + into siblings over a shared base. If it does, a duplex server silently + stops serving /v1/chat/completions: the chat service needs a real + EngineClient with generate(), which only the turn-based side provides. + """ + from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine + from vllm_omni.engine.orchestrator import Orchestrator + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.entrypoints.duplex_omni import DuplexOmni + + assert issubclass(DuplexOrchestrator, Orchestrator) + assert issubclass(DuplexOmniEngine, AsyncOmniEngine) + assert issubclass(DuplexOmni, AsyncOmni) + # What the chat route actually needs off the engine client. + assert hasattr(DuplexOmni, "generate") + + +@pytest.mark.asyncio +async def test_an_unrecognised_message_reaches_the_turn_based_handler() -> None: + """The duplex orchestrator consumes session messages and passes the rest down. + + Without the fall-through, an ordinary add_request on a duplex engine is + silently dropped and the chat request that sent it hangs forever. + """ + orchestrator, _clients, _rpc_q, _output_q = _build() + handled: list[str] = [] + + async def _turn_based(msg: object) -> None: + handled.append(getattr(msg, "type", "")) + + orchestrator._handle_add_request = _turn_based # type: ignore[method-assign] + msg = SimpleNamespace(type="add_request") + + assert await orchestrator._dispatch_message(msg) is True + assert handled == ["add_request"], "an add_request must reach Orchestrator, not be dropped" + + +@pytest.mark.asyncio +async def test_a_dead_replica_closes_the_sessions_it_was_serving() -> None: + """A session outlives its stage request, so replica death has to release the owner. + + ``_handle_dead_replica`` cleaned the stage request but not the runner that + owned it: the session stayed alive holding an admission slot, and the client + was told nothing — the request-scoped ``ErrorMessage`` it emits has no + frontend ``request_states`` entry to land in for a session-owned request, so + the frontend drops it. Heartbeats would then keep an unusable session + occupying capacity. + """ + orchestrator, clients, rpc_q, output_q = _build() + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None + assert orchestrator.session_manager.active_count() == 1 + + await orchestrator._handle_dead_replica(0, 0, EngineDeadError("stage 0 replica 0 died")) + await _settle(orchestrator) + + assert request_id not in orchestrator.request_states + assert SESSION_ID not in orchestrator.session_manager.runners + assert session.state == DuplexSessionState.CLOSED + assert orchestrator.session_manager.active_count() == 0, "the admission slot must come back" + messages = [output_q.get_nowait() for _ in range(output_q.qsize())] + types = [getattr(getattr(m, "event", None), "type", type(m).__name__) for m in messages] + assert types[-1] in {"session.expired", "session.closed"}, f"the client needs a terminal event, got {types}" diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index aa4f25ded..beddd2106 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -21,20 +21,6 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats -from vllm_omni.engine.duplex.control_plane import DuplexControlPlane -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexFence, - OpenDuplexSessionMessage, - SignalDuplexTurnMessage, -) -from vllm_omni.engine.duplex.runtime import ( - DuplexInputMode, - DuplexRuntimeCapabilities, - DuplexSessionRuntimeState, - duplex_resource_request_id, -) from vllm_omni.engine.messages import ( AbortRequestMessage, AbortResultMessage, @@ -54,7 +40,6 @@ from vllm_omni.engine.orchestrator import ( ) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import MiniCPMO45DuplexRuntimeExtension from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -1118,778 +1103,6 @@ async def test_handle_abort_emits_error_result_on_failure() -> None: assert "cleanup exploded" in (result.error or "") -def _duplex_open_message( - session_id: str, - *, - incarnation: int = 0, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, -) -> OpenDuplexSessionMessage: - return OpenDuplexSessionMessage( - control_id=f"open-{session_id}", - fence=DuplexFence(session_id, incarnation=incarnation), - session_id=session_id, - capabilities={ - "input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value], - "implementation_level": "model_native_duplex", - }, - session_config=session_config or {}, - runtime_config=runtime_config or {}, - ) - - -async def _handle_duplex(orchestrator: Orchestrator, message: object) -> None: - await orchestrator._require_duplex_control_plane().handle(message) - - -def _duplex_request_state( - orchestrator: Orchestrator, - session: DuplexSessionRuntimeState, - *, - stage_id: int, -) -> OrchestratorRequestState | None: - context = orchestrator._require_duplex_control_plane().ensure_stage_request( - session, - stage_id=stage_id, - ) - if context is None: - return None - return orchestrator.request_states.get(context.request_id) - - -@pytest.mark.asyncio -async def test_duplex_control_plane_keeps_public_and_runtime_config_separate() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - open_message = _duplex_open_message( - "sid-config-channels", - incarnation=3, - session_config={ - "instructions": "public instructions", - "extra_body": {"duplex_stage_max_tokens": {"0": 99}}, - }, - runtime_config={ - "duplex_stage_max_tokens": {"0": 3}, - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151705]}}, - }, - ) - - await _handle_duplex(orchestrator, open_message) - - assert rpc_q.get_nowait().ok is True - session = orchestrator.duplex_sessions.require(open_message.session_id) - request_state = _duplex_request_state(orchestrator, session, stage_id=0) - assert request_state.sampling_params_list[0].max_tokens == 3 - assert request_state.sampling_params_list[0].stop_token_ids == [151705] - bridge = request_state.streaming.bridge_states["duplex"] - assert bridge["incarnation"] == open_message.fence.incarnation - assert bridge["session_config"] == open_message.session_config - assert bridge["runtime_config"] == open_message.runtime_config - - -@pytest.mark.asyncio -async def test_duplex_control_plane_preserves_turn_commit_without_model_extension() -> None: - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=[], - duplex_runtime_extension=None, - enable_duplex_control=True, - ) - assert isinstance(orchestrator.duplex_control_plane, DuplexControlPlane) - - message = OpenDuplexSessionMessage( - control_id="open-turn-commit", - fence=DuplexFence("sid-turn-commit"), - session_id="sid-turn-commit", - capabilities={"input_modes": [DuplexInputMode.TURN_COMMIT_ONLY.value]}, - session_config={}, - ) - await _handle_duplex(orchestrator, message) - - result = rpc_q.get_nowait() - assert result.ok is True - assert result.stage_results[0]["result"]["scheduler_request_context"] is False - - -def test_ordinary_orchestrator_bypasses_duplex_control_plane() -> None: - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=asyncio.Queue(), - stage_pools=[], - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - ) - - assert orchestrator.duplex_control_plane is None - - -@pytest.mark.asyncio -async def test_duplex_close_cleans_preregistered_request_without_append() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - running_counter = FakeRunningCounter() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - running_counter=running_counter, - ) - open_message = _duplex_open_message("sid-preregister-close") - - await _handle_duplex(orchestrator, open_message) - open_result = rpc_q.get_nowait() - request_id = open_result.stage_results[0]["result"]["request_id"] - assert request_id in orchestrator.request_states - - await _handle_duplex( - orchestrator, - CloseDuplexSessionMessage( - control_id="close-preregistered", - fence=open_message.fence, - session_id=open_message.session_id, - ), - ) - - assert rpc_q.get_nowait().ok is True - assert request_id not in orchestrator.request_states - assert orchestrator.duplex_sessions.get(open_message.session_id) is None - assert running_counter.value == 0 - - -@pytest.mark.asyncio -async def test_duplex_open_failure_rolls_back_session_and_reserved_request() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - open_message = _duplex_open_message( - "sid-open-rollback", - runtime_config={ - "duplex_stage_sampling_params": { - "0": {"stop_token_ids": 123}, - } - }, - ) - - await _handle_duplex(orchestrator, open_message) - - assert rpc_q.get_nowait().ok is False - assert orchestrator.duplex_sessions.get(open_message.session_id) is None - assert orchestrator.request_states == {} - - -@pytest.mark.asyncio -async def test_duplex_running_counter_tracks_only_submitted_request() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - running_counter = FakeRunningCounter() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - running_counter=running_counter, - ) - open_message = _duplex_open_message("sid-duplex-counter") - await _handle_duplex(orchestrator, open_message) - assert rpc_q.get_nowait().ok is True - assert running_counter.value == 0 - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-duplex-counter", - fence=open_message.fence, - session_id=open_message.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - assert rpc_q.get_nowait().ok is True - assert running_counter.value == 1 - - await _handle_duplex( - orchestrator, - CloseDuplexSessionMessage( - control_id="close-duplex-counter", - fence=open_message.fence, - session_id=open_message.session_id, - ), - ) - assert rpc_q.get_nowait().ok is True - assert running_counter.value == 0 - - -@pytest.mark.asyncio -async def test_duplex_failed_append_does_not_advance_sequence_or_fence() -> None: - class FailingPlanExtension(MiniCPMO45DuplexRuntimeExtension): - def plan_append(self, **kwargs): - raise RuntimeError("planned append failure") - - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=FailingPlanExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-append-rollback", turn_id=1) - session = orchestrator.duplex_sessions.open_session( - DuplexFence("sid-append-rollback"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-rollback", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert session.fence == DuplexFence("sid-append-rollback") - assert session.input_seq == 0 - assert session.input_turn_seq == 0 - assert stage0.add_request_calls == [] - - -@pytest.mark.asyncio -async def test_duplex_duplicate_append_operation_is_submitted_once() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-idempotent-append") - session = orchestrator.duplex_sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - for control_id in ("append-first", "append-timeout-retry"): - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id=control_id, - operation_id="physical-input-1", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - first = rpc_q.get_nowait() - retry = rpc_q.get_nowait() - assert first.ok is True - assert retry.ok is True - assert first.stage_results == retry.stage_results - assert len(stage0.add_request_calls) == 1 - assert session.input_seq == 1 - - -@pytest.mark.asyncio -async def test_duplex_barge_in_aborts_bound_stage_requests_before_releasing_fence() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=False) - stage1 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0], [stage1]], - output_processors=[FakeOutputProcessor(), FakeOutputProcessor()], - stage_vllm_configs=[ - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - ], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - session = orchestrator.duplex_sessions.open_session( - DuplexFence("sid-stage-signal"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request(0, "req-stage0", fence=session.fence) - session.bind_stage_request(1, "req-stage1", fence=session.fence) - assert stage_pools[0].select_replica_id("req-stage0") == 0 - assert stage_pools[1].select_replica_id("req-stage1") == 0 - session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=session.fence, - ) - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="ctrl-signal", - fence=session.fence, - next_fence=DuplexFence("sid-stage-signal", epoch=1), - session_id="sid-stage-signal", - event="barge_in", - ), - ) - - assert stage0.abort_calls == [["req-stage0"]] - assert stage1.abort_calls == [["req-stage1"]] - assert stage0.collective_rpc_calls == [] - assert stage1.collective_rpc_calls == [] - assert session.stage_bindings == {} - assert session.fence == DuplexFence("sid-stage-signal", epoch=1) - result = rpc_q.get_nowait() - assert result.ok is True - - -@pytest.mark.asyncio -async def test_duplex_late_barge_in_releases_only_cancelled_fence_bindings() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=False) - stage1 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0], [stage1]], - output_processors=[FakeOutputProcessor(), FakeOutputProcessor()], - stage_vllm_configs=[ - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - ], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - enable_duplex_control=True, - ) - cancelled_fence = DuplexFence("sid-late-stage-signal") - session = orchestrator.duplex_sessions.open_session( - cancelled_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request(0, "req-stage0", fence=cancelled_fence) - session.bind_stage_request(1, "req-stage1", fence=cancelled_fence) - assert stage_pools[0].select_replica_id("req-stage0") == 0 - assert stage_pools[1].select_replica_id("req-stage1") == 0 - current_fence = DuplexFence("sid-late-stage-signal", epoch=1) - session.accept_fence(current_fence) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="ctrl-late-signal", - fence=cancelled_fence, - next_fence=current_fence, - session_id="sid-late-stage-signal", - event="barge_in", - ), - ) - - assert stage0.abort_calls == [["req-stage0"]] - assert stage1.abort_calls == [["req-stage1"]] - assert session.stage_bindings == {} - assert session.fence == current_fence - result = rpc_q.get_nowait() - assert result.ok is True - - -@pytest.mark.asyncio -async def test_duplex_cancel_without_next_fence_is_rejected_without_releasing_bindings() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - enable_duplex_control=True, - ) - fence = DuplexFence("sid-cancel-contract") - session = orchestrator.duplex_sessions.open_session(fence) - session.bind_stage_request(0, "req-live", fence=fence) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="cancel-without-next", - fence=fence, - session_id=fence.session_id, - event="input.cancel", - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert result.error_count == 1 - assert "next_fence" in result.stage_results[0]["result"]["error"] - assert session.fence == fence - assert session.stage_request_ids() == ["req-live"] - assert stage0.abort_calls == [] - - -@pytest.mark.asyncio -async def test_duplex_session_update_replaces_runtime_config() -> None: - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=[], - enable_duplex_control=True, - ) - fence = DuplexFence("sid-update-config") - session = orchestrator.duplex_sessions.open_session( - fence, - session_config={"temperature": 0.7}, - ) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="update-config", - fence=fence, - session_id=fence.session_id, - event="session.update", - session_config={"temperature": 0.0, "instructions": "updated"}, - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is True - assert session.session_config == {"temperature": 0.0, "instructions": "updated"} - - -@pytest.mark.asyncio -async def test_duplex_session_update_refreshes_next_append_sampling_params() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-update-policy") - session = orchestrator.duplex_sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - runtime_config={ - "duplex_stage_max_tokens": {"0": 2}, - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151645]}}, - }, - ) - req_state = _duplex_request_state(orchestrator, session, stage_id=0) - assert req_state is not None - assert req_state.sampling_params_list[0].max_tokens == 2 - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="update-policy", - fence=fence, - session_id=fence.session_id, - event="session.update", - runtime_config={ - "duplex_stage_max_tokens": {"0": 7}, - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151705]}}, - }, - ), - ) - assert rpc_q.get_nowait().ok is True - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-updated-policy", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - assert rpc_q.get_nowait().ok is True - assert req_state.sampling_params_list[0].max_tokens == 7 - assert req_state.sampling_params_list[0].stop_token_ids == [151705] - submitted_request = stage0.add_request_calls[0][0] - assert submitted_request.sampling_params.stop_token_ids == [151705] - assert submitted_request.sampling_params.max_tokens == 7 - - -@pytest.mark.asyncio -async def test_duplex_invalid_session_update_preserves_previous_config() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-invalid-update") - original_runtime_config = { - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151645]}}, - } - session = orchestrator.duplex_sessions.open_session( - fence, - runtime_config=original_runtime_config, - ) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="invalid-update", - fence=fence, - session_id=fence.session_id, - event="session.update", - runtime_config={ - "duplex_stage_sampling_params": {"0": {"stop_token_ids": 123}}, - }, - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert session.runtime_config == original_runtime_config - assert session.config_generation == 0 - - -@pytest.mark.asyncio -async def test_duplex_arbitrary_non_cancel_signal_is_rejected() -> None: - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=[], - enable_duplex_control=True, - ) - fence = DuplexFence("sid-unsupported-signal") - orchestrator.duplex_sessions.open_session(fence) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="unsupported-signal", - fence=fence, - session_id=fence.session_id, - event="turn.end", - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert "unsupported duplex runtime signal" in result.stage_results[0]["result"]["error"] - - -@pytest.mark.asyncio -async def test_duplex_cancel_rejects_late_old_append_and_accepts_next_epoch() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - cancelled_fence = DuplexFence("sid-cancel-late-append") - next_fence = DuplexFence("sid-cancel-late-append", epoch=1) - session = orchestrator.duplex_sessions.open_session( - cancelled_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="cancel-old-fence", - fence=cancelled_fence, - next_fence=next_fence, - session_id=session.session_id, - event="input.cancel", - ), - ) - - cancel_result = rpc_q.get_nowait() - assert cancel_result.ok is True - assert session.fence == next_fence - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="late-old-append", - fence=cancelled_fence, - session_id=session.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - stale_result = rpc_q.get_nowait() - assert stale_result.ok is False - assert stage0.add_request_calls == [] - assert session.stage_bindings == {} - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="next-epoch-append", - fence=next_fence, - session_id=session.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - next_result = rpc_q.get_nowait() - assert next_result.ok is True - assert len(stage0.add_request_calls) == 1 - assert session.stage_bindings[0].fence == next_fence - - -@pytest.mark.asyncio -async def test_duplex_append_updates_bridge_turn_id_on_long_lived_stage0_request() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=asyncio.Queue(), - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - session = orchestrator.duplex_sessions.open_session( - DuplexFence("sid-bridge-turn"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - session_config={ - "voice": "test", - }, - runtime_config={ - "duplex_stage_sampling_params": { - "0": {"stop_token_ids": [151645]}, - }, - }, - ) - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-1", - fence=DuplexFence("sid-bridge-turn"), - session_id="sid-bridge-turn", - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - req_state1 = _duplex_request_state(orchestrator, session, stage_id=0) - assert req_state1 is not None - duplex_state1 = req_state1.streaming.bridge_states["duplex"] - assert duplex_state1["session_id"] == "sid-bridge-turn" - assert duplex_state1["epoch"] == 0 - assert duplex_state1["turn_id"] == 0 - assert duplex_state1["session_config"] == session.session_config - assert req_state1.sampling_params_list[0].stop_token_ids == [151645] - submitted_request = stage0.add_request_calls[0][0] - assert submitted_request.sampling_params.stop_token_ids == [151645] - assert submitted_request.sampling_params.max_tokens == 1 - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-2", - fence=DuplexFence("sid-bridge-turn", turn_id=1, response_seq=1), - session_id="sid-bridge-turn", - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True, "new_user_turn": True}, - ), - ) - req_state2 = _duplex_request_state(orchestrator, session, stage_id=0) - assert req_state2 is req_state1 - duplex_state2 = req_state2.streaming.bridge_states["duplex"] - assert duplex_state2["turn_id"] == 1 - assert duplex_state2["epoch"] == 0 - expected_request_id = duplex_resource_request_id(DuplexFence("sid-bridge-turn"), "stage0") - assert [call[0].request_id for call in stage0.add_request_calls] == [ - expected_request_id, - expected_request_id, - ] - - # --------------------------------------------------------------------------- # Multi-replica tests # --------------------------------------------------------------------------- @@ -2641,61 +1854,6 @@ async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> N assert iteration_stats.num_generation_tokens == 3 -@pytest.mark.asyncio -async def test_duplex_reaper_loop_waits_between_ticks(): - class _Plane: - def __init__(self) -> None: - self.calls = 0 - - async def reap_expired(self) -> int: - self.calls += 1 - return 0 - - orchestrator = object.__new__(Orchestrator) - orchestrator.duplex_control_plane = _Plane() - orchestrator._duplex_reaper_interval_s = 0.01 - orchestrator._shutdown_event = asyncio.Event() - - task = asyncio.create_task(orchestrator._duplex_reaper_loop()) - await asyncio.sleep(0.035) - orchestrator._shutdown_event.set() - await task - - assert 2 <= orchestrator.duplex_control_plane.calls <= 5 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("first_cleanup_delay", [0.0, 0.05], ids=["immediate", "delayed"]) -async def test_duplex_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: float) -> None: - recovered = asyncio.Event() - - class _Plane: - def __init__(self) -> None: - self.calls = 0 - - async def reap_expired(self) -> int: - self.calls += 1 - if self.calls == 1: - await asyncio.sleep(first_cleanup_delay) - raise RuntimeError("transient cleanup failure") - recovered.set() - return 0 - - orchestrator = object.__new__(Orchestrator) - orchestrator.duplex_control_plane = _Plane() - orchestrator._duplex_reaper_interval_s = 0.01 - orchestrator._shutdown_event = asyncio.Event() - - task = asyncio.create_task(orchestrator._duplex_reaper_loop()) - try: - # Wait for recovery itself, including when cleanup exceeds the old 35 ms window. - await asyncio.wait_for(recovered.wait(), timeout=5.0) - assert orchestrator.duplex_control_plane.calls >= 2 - finally: - orchestrator._shutdown_event.set() - await asyncio.wait_for(task, timeout=5.0) - - @pytest.mark.asyncio async def test_abort_retry_does_not_repeat_successful_stage_abort(): class _Pool: @@ -2733,48 +1891,3 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort(): assert first.physical_abort_calls == 1 assert second.physical_abort_calls == 2 assert second.bound == set() - - -@pytest.mark.asyncio -async def test_request_cleanup_failure_is_deferred_to_control_plane(): - class _FailingPool: - async def abort_requests(self, request_ids: list[str]) -> None: - raise RuntimeError("stage abort failed") - - def release_bindings(self, request_ids: list[str]) -> None: - pass - - class _Plane: - def __init__(self) -> None: - self.deferred: list[str] = [] - self.finalized: list[str] = [] - - def close_sessions_for_request_ids(self, request_ids: list[str], **kwargs): - assert kwargs == {"abort": True, "cleanup_in_progress": True} - return {"sid-cleanup": ["req-a"]} - - def defer_request_cleanups(self, session_ids: list[str]) -> None: - self.deferred.extend(session_ids) - - def finalize_closed_sessions(self, session_ids: list[str]) -> None: - self.finalized.extend(session_ids) - - from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker - - orchestrator = object.__new__(Orchestrator) - orchestrator.stage_pools = [_FailingPool()] - orchestrator.duplex_control_plane = _Plane() - orchestrator._pd_kv_params = {} - orchestrator.request_states = {} - orchestrator._running_counter = None - orchestrator._cfg_tracker = CfgCompanionTracker() - - with pytest.raises(RuntimeError, match="stage abort failed"): - await orchestrator._cleanup_request_ids( - ["req-a"], - abort=True, - close_duplex_sessions=True, - ) - - assert orchestrator.duplex_control_plane.deferred == ["sid-cleanup"] - assert orchestrator.duplex_control_plane.finalized == [] diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index eaa246072..b8949786e 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -1054,7 +1054,7 @@ async def test_streaming_input_processor_client_error_does_not_forward_terminal_ final_stage_id=1, ) state.streaming.enabled = True - assert state.duplex_identity is None + assert state.session_owned is False orchestrator.request_states[state.request_id] = state try: await orchestrator._route_output(0, 0, _build_request_output(state.request_id), state, stage_metrics=None) @@ -1074,34 +1074,14 @@ async def test_streaming_input_processor_client_error_does_not_forward_terminal_ q.close() -@pytest.mark.asyncio -async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None: - class FailingStage(FakeStageClient): - def process_engine_inputs(self, *_args, **_kwargs): - raise ValueError("No latent or hidden_states found in thinker output") - - fixture = orchestrator_factory( - [FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)] - ) - state = OrchestratorRequestState( - request_id="bad", - prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]), - sampling_params_list=[_sampling_params(), _sampling_params()], - final_stage_id=1, - duplex_identity=SimpleNamespace(), - ) - - async def no_cleanup(*_args, **_kwargs): - pass - - monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup) - try: - await fixture.orchestrator._forward_to_next_stage_unguarded("bad", 0, _build_request_output("raw"), state) - error = await _wait_for_error_message(fixture, request_id="bad") - assert error.fatal is False - finally: - fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) - fixture.thread.join(timeout=5) +# ``test_duplex_input_processor_failure_is_request_scoped`` lived here upstream, +# where one Orchestrator served both turn-based and duplex requests. Duplex +# request ownership now belongs to DuplexOrchestrator, and the base +# ``_handle_forward_failure`` deliberately absorbs nothing, so the equivalent +# assertion is +# ``tests/engine/test_duplex_orchestrator.py::test_forward_failure_closes_the_owning_session`` +# -- which also checks the stage abort, the session teardown and the +# error-before-session.expired ordering. @pytest.mark.asyncio diff --git a/tests/engine/test_orchestrator_segment_stage_keying.py b/tests/engine/test_orchestrator_segment_stage_keying.py index 5d3952e4b..9c5290245 100644 --- a/tests/engine/test_orchestrator_segment_stage_keying.py +++ b/tests/engine/test_orchestrator_segment_stage_keying.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Streaming segment-boundary state must be keyed by stage. -``Orchestrator._orchestration_loop`` polls every stage into the same +``OrchestratorBase._orchestration_loop`` polls every stage into the same ``OrchestratorRequestState``. While the segment-boundary fields lived in a single flat slot per request, whichever stage polled most recently overwrote the others, and the per-stage consumers then read back a different stage's boundary. + +The duplex consumer of the per-stage segment (``DuplexOrchestrator``'s output +context) is covered in ``test_duplex_orchestrator.py``. """ from __future__ import annotations @@ -15,8 +18,6 @@ from dataclasses import dataclass, field import pytest from vllm.sampling_params import SamplingParams -from vllm_omni.engine.duplex.contracts import DuplexRequestIdentity -from vllm_omni.engine.duplex.messages import DuplexFence from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, @@ -54,22 +55,19 @@ class _RecordingPool: return _StageMetrics() -def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState: +def _streaming_req_state(request_id: str = "req-streaming") -> OrchestratorRequestState: req_state = OrchestratorRequestState( request_id=request_id, sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], final_stage_id=TALKER_STAGE, + session_owned=True, ) req_state.streaming.enabled = True - req_state.duplex_identity = DuplexRequestIdentity( - session_id="session-1", - fence=DuplexFence(session_id="session-1"), - ) return req_state -def test_duplex_output_context_reads_only_its_own_stage_segment() -> None: - req_state = _duplex_req_state() +def test_segment_accessor_reads_only_its_own_stage() -> None: + req_state = _streaming_req_state() req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState( finished=True, token_ids=[11, 22], @@ -78,26 +76,28 @@ def test_duplex_output_context_reads_only_its_own_stage_segment() -> None: # The talker polls after the thinker and is still mid-segment. req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False) - thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE) - talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE) - - assert thinker is not None - assert talker is not None + thinker = req_state.streaming.segment(THINKER_STAGE) + talker = req_state.streaming.segment(TALKER_STAGE) # These are the assertions that fail on the flat field: the talker is written # second, so the thinker loses its boundary to whichever stage polled last. - assert thinker.segment_finished is True - assert thinker.segment_token_ids == (11, 22) - assert thinker.segment_output_metadata == {"stage": "thinker"} + assert thinker.finished is True + assert thinker.token_ids == [11, 22] + assert thinker.output_metadata == {"stage": "thinker"} - assert talker.segment_finished is False - assert talker.segment_token_ids == () - assert talker.segment_output_metadata == {} + assert talker.finished is False + assert talker.token_ids == [] + assert talker.output_metadata == {} + + # A stage that never reported a boundary gets a fresh, unshared segment. + unreported = req_state.streaming.segment(2) + unreported.output_metadata["poisoned"] = True + assert req_state.streaming.segment(2).output_metadata == {} @pytest.mark.asyncio async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None: - req_state = _duplex_req_state() + req_state = _streaming_req_state() req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True) talker_pool = _RecordingPool() diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 69df71231..842824c2d 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -7,23 +7,17 @@ import asyncio import queue from types import SimpleNamespace from typing import Any -from unittest.mock import ANY, AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock import janus import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams -from vllm_omni.engine.duplex.contracts import ( - DuplexStageRequestContext, - DuplexStageSubmission, -) -from vllm_omni.engine.duplex.messages import DuplexFence from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, StreamingSegmentState, - _OrchestratorDuplexStagePort, ) from vllm_omni.engine.stage_pool import StagePool @@ -131,64 +125,6 @@ class FakePrewarmPool: return self.stage_client if self.get_bound_replica_id(request_id) is not None else None -def _duplex_stage_port_submission(): - stage_pools = [] - for stage_id in range(3): - pool = SimpleNamespace( - stage_client=SimpleNamespace(default_sampling_params=SamplingParams(max_tokens=1)), - stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - submit_initial=AsyncMock(return_value=stage_id + 10), - submit_update=AsyncMock(return_value=stage_id + 20), - ) - stage_pools.append(pool) - request_states: dict[str, OrchestratorRequestState] = {} - prewarm = AsyncMock() - port = _OrchestratorDuplexStagePort( - stage_pools=stage_pools, - request_states=request_states, - running_counter=None, - cleanup_request_ids=AsyncMock(), - async_chunk=True, - prewarm_async_chunk_stages=prewarm, - ) - context = DuplexStageRequestContext( - request_id="req-duplex", - session_id="session-duplex", - fence=DuplexFence("session-duplex"), - stage_id=0, - final_stage_id=2, - config_generation=0, - sampling_params=tuple(SamplingParams(max_tokens=1) for _ in range(3)), - ) - port.ensure_request(context) - submission = DuplexStageSubmission( - context=context, - prompt={"prompt_token_ids": [1, 2]}, - already_submitted=False, - ) - return port, stage_pools, request_states, prewarm, submission - - -def test_duplex_bridge_state_catches_up_after_serving_turn_boundary() -> None: - port, _, request_states, _, submission = _duplex_stage_port_submission() - request_state = request_states[submission.context.request_id] - bridge_state = request_state.streaming.bridge_states["duplex"] - bridge_state["model_turn_id"] = 0 - - next_context = DuplexStageRequestContext( - request_id=submission.context.request_id, - session_id=submission.context.session_id, - fence=DuplexFence(submission.context.session_id, turn_id=1), - stage_id=submission.context.stage_id, - final_stage_id=submission.context.final_stage_id, - config_generation=submission.context.config_generation, - sampling_params=submission.context.sampling_params, - ) - port.ensure_request(next_context) - - assert bridge_state["model_turn_id"] == 1 - - def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( index=0, @@ -279,13 +215,12 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None - orchestrator._record_duplex_stage_submission = MagicMock() + orchestrator._on_stage_submitted = MagicMock() req_state = OrchestratorRequestState( request_id="req-prewarm", prompt={"prompt_token_ids": [1, 2]}, sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, - duplex_identity=SimpleNamespace(), ) prewarmed = await orchestrator._prewarm_async_chunk_stages( @@ -306,7 +241,7 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N assert stage2.submitted[0].resumable is True assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts - orchestrator._record_duplex_stage_submission.assert_called_once_with( + orchestrator._on_stage_submitted.assert_called_once_with( 2, "req-prewarm", 0, @@ -314,42 +249,6 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N ) -@pytest.mark.asyncio -async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None: - port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() - prewarm.return_value = True - - result = await port.submit(submission) - - assert result.stage_id == 0 - stage_pools[0].submit_initial.assert_awaited_once() - prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"]) - - -@pytest.mark.asyncio -async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None: - """A failed prewarm has already aborted the request and popped its state. - - Returning a success result here would hand the control plane a replica for a - request that no longer exists, and the trailing bookkeeping would re-register - a running counter the cleanup just released -- a leak that never decrements. - """ - port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() - counter = MagicMock() - port._running_counter = counter - prewarm.return_value = False - request_state = request_states["req-duplex"] - - with pytest.raises(RuntimeError, match="prewarm failed"): - await port.submit(submission) - - prewarm.assert_awaited_once() - assert request_state.duplex_stage_fences == {} - assert request_state.stage_submit_ts == {} - assert request_state.running_counter_registered is False - counter.increment.assert_not_called() - - @pytest.mark.asyncio async def test_async_route_forwards_to_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) diff --git a/tests/engine/test_output_metadata_snapshots.py b/tests/engine/test_output_metadata_snapshots.py index 688f1252d..9670994c4 100644 --- a/tests/engine/test_output_metadata_snapshots.py +++ b/tests/engine/test_output_metadata_snapshots.py @@ -9,7 +9,7 @@ from vllm.sampling_params import RequestOutputKind from vllm.v1.engine import FinishReason from vllm_omni.outputs.mm_outputs import MultimodalPayload -from vllm_omni.outputs.output_modality import OutputModality, TensorAccumulationStrategy +from vllm_omni.outputs.output_modality import OutputModality from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -58,7 +58,7 @@ def test_rate_is_latest_snapshot_before_consolidation(key, representation): assert payload.get(key) is payload.to_dict()[key] payload = payload.merged_with(MultimodalPayload(tensors={"audio": torch.full((2,), 3)})) assert int(payload[key]) == 48000 # Missing metadata means retain, not clear. - payload.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + payload.consolidate_tensors(OutputModality.AUDIO) payload.consolidate_metadata() torch.testing.assert_close(payload["audio"], torch.arange(4).repeat_interleave(2)) @@ -86,7 +86,7 @@ def test_self_merge_retains_snapshot_and_content(key): assert payload is not None payload = payload.merged_with(payload) assert int(payload[key]) == 24000 - payload.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + payload.consolidate_tensors(OutputModality.AUDIO) torch.testing.assert_close(payload["audio"], torch.tensor([1.0, 2.0, 1.0, 2.0])) @@ -138,6 +138,6 @@ def test_non_snapshot_tensors_still_accumulate_and_empty_merge_keeps_identity(): assert first is not None and second is not None assert MultimodalPayload().merged_with(first) is first merged = first.merged_with(second) - merged.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + merged.consolidate_tensors(OutputModality.AUDIO) torch.testing.assert_close(merged["latent"], torch.tensor([1.0, 3.0])) torch.testing.assert_close(merged["audio"], torch.tensor([2.0, 4.0])) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 675a0650b..41d63af9e 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -389,7 +389,7 @@ class TestSingleStageModeDetection: return_value=None, ) mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), ) mocker.patch.object( @@ -446,11 +446,11 @@ class TestSingleStageModeDetection: def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture): duplex_session = DuplexSessionRuntimeConfig(max_sessions=2) get_pipeline_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", + "vllm_omni.engine.omni_engine_base.StageConfigFactory.get_pipeline_config", return_value=None, ) load_deploy_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=duplex_session), ) @@ -467,7 +467,8 @@ class TestSingleStageModeDetection: deploy_config_path="/fake/duplex.yaml", ) load_deploy_config.assert_called_once_with("/fake/duplex.yaml") - assert engine.duplex_session_config is duplex_session + # The turn-based engine only keeps the resolved deploy profile for introspection. + assert engine.deploy_config.duplex_session is duplex_session def test_auto_discovered_deploy_loads_duplex_runtime_config( self, @@ -476,11 +477,11 @@ class TestSingleStageModeDetection: deploy_path = "/resolved/qwen3_omni_moe.yaml" duplex_session = DuplexSessionRuntimeConfig(server_vad_model_path="/models/silero_vad.onnx") mocker.patch( - "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", + "vllm_omni.engine.omni_engine_base.StageConfigFactory.get_pipeline_config", return_value=None, ) load_deploy_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=duplex_session), ) @@ -557,7 +558,7 @@ class TestEndpointRestrictionsTrustRemoteCode: def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine: mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), ) mocker.patch.object( diff --git a/tests/entrypoints/duplex/test_duplex_serving.py b/tests/entrypoints/duplex/test_duplex_serving.py new file mode 100644 index 000000000..dde616d94 --- /dev/null +++ b/tests/entrypoints/duplex/test_duplex_serving.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The thin websocket handler: handshake, command translation, event pump, resume/takeover.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +from contextlib import suppress +from typing import Any + +import pytest +from fastapi import WebSocketDisconnect + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands +from vllm_omni.engine.duplex.config import DuplexCapabilities +from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, SessionClosed, SessionCreated +from vllm_omni.engine.duplex.messages import DuplexSessionError +from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request +from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler +from vllm_omni.entrypoints.duplex.websocket import MAX_EVENT_BYTES + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DISCONNECT = object() + + +class FakeWebSocket: + def __init__(self, query: dict[str, str] | None = None) -> None: + self.query_params = dict(query or {}) + self.sent: list[dict[str, Any]] = [] + self.accepted = False + self.closed: list[tuple[int, str]] = [] + self._send_failure: str | None = None + self._inbound: asyncio.Queue[Any] = asyncio.Queue() + + async def accept(self) -> None: + self.accepted = True + + async def send_json(self, payload: dict[str, Any]) -> None: + if self.closed or self._send_failure is not None: + raise RuntimeError( + self._send_failure or "Unexpected ASGI message 'websocket.send', after sending 'websocket.close'." + ) + self.sent.append(json.loads(json.dumps(payload))) + + async def receive_text(self) -> str: + item = await self._inbound.get() + if item is _DISCONNECT: + raise WebSocketDisconnect(code=1000) + return item + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.closed.append((code, reason)) + # A closed socket ends the reader exactly like a client disconnect. + self._inbound.put_nowait(_DISCONNECT) + + # ---- test helpers ---- + + def feed(self, payload: dict[str, Any] | str) -> None: + self._inbound.put_nowait(payload if isinstance(payload, str) else json.dumps(payload)) + + def disconnect(self) -> None: + self._inbound.put_nowait(_DISCONNECT) + + def break_sends(self) -> None: + """Kill the write half only: the reader stays parked, as it does in practice.""" + self._send_failure = "Unexpected ASGI message 'websocket.send', after sending 'websocket.close'." + + def types(self) -> list[str]: + return [payload["type"] for payload in self.sent] + + async def wait_for(self, wire_type: str, *, timeout_s: float = 2.0) -> dict[str, Any]: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + while loop.time() < deadline: + for payload in self.sent: + if payload["type"] == wire_type: + return payload + await asyncio.sleep(0.005) + raise AssertionError(f"no {wire_type!r} in {self.types()}") + + +class FakeHandle: + def __init__(self, session_id: str, capabilities: DuplexCapabilities, *, idle_timeout_s: float = 300) -> None: + self.session_id = session_id + self.capabilities = capabilities + self.public_session: dict[str, Any] = {"id": session_id, "idle_timeout_s": idle_timeout_s} + self.lease_generation = 0 + self.closed = False + self.close_reasons: list[str] = [] + self.commands: list[commands.DuplexCommand] = [] + self._outbox: asyncio.Queue[DuplexEvent] = asyncio.Queue() + + def deliver(self, event: DuplexEvent) -> None: + self._outbox.put_nowait(event) + + async def submit(self, command: commands.DuplexCommand) -> None: + if self.closed: + raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id) + self.commands.append(command) + + async def events(self): + while True: + event = await self._outbox.get() + yield event + if event.is_terminal: + return + + async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None: + self.close_reasons.append(reason) + if self.closed: + return + self.closed = True + self.deliver(SessionClosed(session_id=self.session_id, reason=reason)) + + +class FakeOmni: + def __init__( + self, *, resumable: bool = True, replay_max_bytes: int = 64 * 1024, idle_timeout_s: float = 300 + ) -> None: + self.duplex_session_config = DuplexSessionRuntimeConfig( + resume_replay_ttl_s=60.0, resume_replay_max_bytes_per_session=replay_max_bytes + ) + self.capabilities = DuplexCapabilities(supports_session_resume=resumable) + self.idle_timeout_s = idle_timeout_s + self.opened: list[dict[str, Any]] = [] + self.handles: dict[str, FakeHandle] = {} + self.resumed: list[tuple[str, int]] = [] + self.detached: list[str] = [] + self.open_error: DuplexSessionError | None = None + + async def open_session(self, config: Any) -> FakeHandle: + self.opened.append(dict(config)) + if self.open_error is not None: + raise self.open_error + session_id = f"duplex-{len(self.handles) + 1:032x}" + handle = FakeHandle(session_id, self.capabilities, idle_timeout_s=self.idle_timeout_s) + self.handles[session_id] = handle + handle.deliver(SessionCreated(session_id=session_id, session={"id": session_id, "model": config.get("model")})) + return handle + + def get_session(self, session_id: str) -> FakeHandle | None: + return self.handles.get(session_id) + + async def resume_session(self, session_id: str, *, expected_lease_generation: int) -> FakeHandle: + self.resumed.append((session_id, expected_lease_generation)) + handle = self.handles[session_id] + handle.lease_generation = expected_lease_generation + 1 + return handle + + async def detach_session(self, session_id: str) -> None: + self.detached.append(session_id) + + +def _handler(omni: FakeOmni, **kwargs: Any) -> OmniDuplexSessionHandler: + kwargs.setdefault("config_timeout_s", 1.0) + kwargs.setdefault("idle_timeout_s", 5.0) + return OmniDuplexSessionHandler(duplex_omni=omni, **kwargs) + + +def _session_update(**session: Any) -> dict[str, Any]: + return {"type": "session.update", "session": {"model": "test-model", "modalities": ["audio", "text"], **session}} + + +async def _open( + handler: OmniDuplexSessionHandler, omni: FakeOmni, **session: Any +) -> tuple[FakeWebSocket, FakeHandle, asyncio.Task]: + ws = FakeWebSocket({"duplex": "1", "autostart": "0"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed(_session_update(**session)) + created = await ws.wait_for("session.created") + handle = omni.handles[created["session"]["id"]] + return ws, handle, task + + +# --------------------------------------------------------------------------- # +# Handshake and commands # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_session_update_opens_a_session_and_announces_the_server_id() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni, session_id="mine", instructions="hi") + + assert ws.accepted + created = ws.sent[0] + assert created["type"] == "session.created" + assert created["session"]["id"] == handle.session_id != "mine" + assert created["attachment_generation"] == 1 + assert isinstance(created["resume_token"], str) and created["resume_token"] + assert "incarnation" not in created + assert "server_event_seq" not in created + # The whole session object went to open_session; the id inside it is ignored there. + assert omni.opened == [ + {"model": "test-model", "modalities": ["audio", "text"], "session_id": "mine", "instructions": "hi"} + ] + + ws.feed({"type": "input_audio_buffer.append", "audio": base64.b64encode(b"\x00\x01" * 8).decode("ascii")}) + ws.feed({"type": "input_audio_buffer.commit", "event_id": "evt-c"}) + await asyncio.sleep(0.05) + assert [type(command) for command in handle.commands] == [commands.AppendAudio, commands.Commit] + assert handle.commands[1].event_id == "evt-c" + + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + # A resumable session is detached (engine-owned grace), not closed. + assert omni.detached == [handle.session_id] + assert handle.close_reasons == [] + + +@pytest.mark.asyncio +async def test_session_events_are_journaled_and_sent_in_order() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + delta = await ws.wait_for("response.output_audio.delta") + assert delta["delta"] == "aGk=" and delta["server_event_seq"] == 1 + + ws.feed({"type": "session.event_ack", "server_event_seq": 1}) + ws.feed({"type": "session.event_ack", "server_event_seq": 9, "event_id": "evt-ack"}) + error = await ws.wait_for("error") + assert error["error"]["code"] == "invalid_event_ack" and error["error"]["event_id"] == "evt-ack" + + await handle.close(reason="client_close") + closed = await ws.wait_for("session.closed") + assert closed["reason"] == "client_close" + await asyncio.wait_for(task, timeout=2.0) + assert ws.closed[0][1] == "client_close" + + +@pytest.mark.asyncio +async def test_envelope_errors_are_answered_locally_without_touching_the_session() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + ws.feed("not json") + ws.feed("x" * (MAX_EVENT_BYTES + 1)) + ws.feed({"type": "totally.unknown", "event_id": "evt-u"}) + ws.feed({"type": "session.resume", "event_id": "evt-r"}) + ws.feed({"type": "playback.ack", "event_id": "evt-p"}) # played_ms missing + await asyncio.sleep(0.1) + + codes = [payload["error"]["code"] for payload in ws.sent if payload["type"] == "error"] + assert codes == ["invalid_json", "event_too_large", "unknown_event", "unsupported_session_resume", "bad_event"] + errors = [payload for payload in ws.sent if payload["type"] == "error"] + assert errors[2]["error"]["event_id"] == "evt-u" + assert errors[4]["error"]["event_id"] == "evt-p" + assert handle.commands == [] + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_open_rejection_is_reported_and_the_socket_ends() -> None: + omni = FakeOmni() + omni.open_error = DuplexSessionError("no room", code="resource_exhausted", retryable=True) + handler = _handler(omni) + ws = FakeWebSocket({"duplex": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed(_session_update()) + await asyncio.wait_for(task, timeout=2.0) + assert ws.types() == ["error"] + assert ws.sent[0]["error"]["code"] == "resource_exhausted" + + +@pytest.mark.asyncio +async def test_config_timeout_and_first_message_validation() -> None: + omni = FakeOmni() + handler = _handler(omni, config_timeout_s=0.05) + ws = FakeWebSocket({"duplex": "1"}) + await asyncio.wait_for(handler.handle_realtime_session(ws), timeout=2.0) + assert ws.sent[0]["error"]["code"] == "config_timeout" + + ws = FakeWebSocket({"duplex": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed("{not json") + await asyncio.wait_for(task, timeout=2.0) + assert ws.sent[0]["error"]["code"] == "invalid_json" + assert omni.opened == [] + + +@pytest.mark.asyncio +async def test_idle_timeout_detaches_a_resumable_session_like_a_disconnect() -> None: + omni = FakeOmni(idle_timeout_s=0.05) + handler = _handler(omni, idle_timeout_s=0.05) + ws, handle, task = await _open(handler, omni) + + await asyncio.wait_for(task, timeout=2.0) + # Serving makes no session-lifetime decision: the engine lease decides + # whether a silent, detached session expires. + assert handle.close_reasons == [] + assert omni.detached == [handle.session_id] + + +@pytest.mark.asyncio +async def test_idle_timeout_closes_a_non_resumable_session() -> None: + omni = FakeOmni(idle_timeout_s=0.05, resumable=False) + handler = _handler(omni, idle_timeout_s=0.05) + ws, handle, task = await _open(handler, omni) + + await asyncio.wait_for(task, timeout=2.0) + assert handle.close_reasons == ["disconnect"] + assert omni.detached == [] + + +@pytest.mark.asyncio +async def test_transport_send_failure_detaches_the_session_instead_of_closing_it() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + # Only the write half dies, so the pump's send is the one and only report + # of the broken socket (the reader is still parked on receive). + ws.break_sends() + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + await asyncio.sleep(0.05) + + # The engine session is alive and resumable, so this is a disconnect of + # the current attachment, not a session close; the pump keeps journaling. + assert handle.close_reasons == [] + assert omni.detached == [handle.session_id] + assert not handler._pumps[handle.session_id].done() + + # The reader sees the same broken socket a moment later. The attachment is + # already gone, so this must not detach a second time: another detach would + # restart the engine's disconnect grace window. + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + assert handle.close_reasons == [] + + handle.deliver(SessionClosed(session_id=handle.session_id, reason="client_close")) + await asyncio.sleep(0.05) + assert handle.session_id not in handler._pumps + + +@pytest.mark.asyncio +async def test_non_resumable_session_is_closed_on_disconnect() -> None: + omni = FakeOmni(resumable=False) + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + assert "resume_token" not in ws.sent[0] + + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert handle.close_reasons == ["disconnect"] + assert omni.detached == [] + + +@pytest.mark.asyncio +async def test_a_client_close_still_delivers_session_closed_before_the_socket_goes() -> None: + """The read loop must not outrun the pump's terminal event. + + ``DuplexSessionHandle._deliver`` queues ``SessionClosed`` and marks the + handle closed in one synchronous step, so ``handle.closed`` is already true + while the event is still sitting in the outbox. ``_read_loop`` loops on + exactly that flag: it returned, the endpoint returned, and the ASGI server + tore the socket down with ``session.closed`` unsent -- the client saw an + abrupt close (no close frame) instead of the terminal event it was waiting + for. + + The teardown is what makes this observable, so it is modelled here: a real + connection stops accepting writes the moment the endpoint returns, which is + why a pump that is merely *still scheduled* is already too late. + """ + omni = FakeOmni() + handler = _handler(omni) + ws = FakeWebSocket({"duplex": "1", "autostart": "0"}) + + async def serve() -> None: + try: + await handler.handle_realtime_session(ws) + finally: + # The ASGI server drops the connection when the endpoint returns; + # anything the pump writes after this never reaches the client. + ws.break_sends() + + task = asyncio.create_task(serve()) + ws.feed(_session_update()) + created = await ws.wait_for("session.created") + handle = omni.handles[created["session"]["id"]] + submit = handle.submit + + async def closing_submit(command: commands.DuplexCommand) -> None: + await submit(command) + if isinstance(command, commands.CloseSession): + # Same order as the real handle: queue the event, then flip the flag. + handle.deliver(SessionClosed(session_id=handle.session_id, reason="client_close")) + handle.closed = True + + handle.submit = closing_submit # type: ignore[method-assign] + + ws.feed({"type": "session.close"}) + await asyncio.wait_for(task, timeout=2.0) + + assert ws.types()[-1] == "session.closed", f"terminal event never reached the wire: {ws.types()}" + assert ws.closed and ws.closed[0][0] == 1000, "the socket must close with a normal close frame" + + +# --------------------------------------------------------------------------- # +# Resume and takeover # +# --------------------------------------------------------------------------- # + + +async def _resume( + handler: OmniDuplexSessionHandler, + session_id: str, + token: str, + *, + last_seq: int = 0, +) -> tuple[FakeWebSocket, asyncio.Task]: + ws = FakeWebSocket({"duplex": "1", "resume": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed( + { + "type": "session.resume", + "session_id": session_id, + "resume_token": token, + "last_received_server_event_seq": last_seq, + } + ) + return ws, task + + +@pytest.mark.asyncio +async def test_resume_after_disconnect_replays_missed_events_and_rotates_the_token() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="one")) + await ws.wait_for("response.output_audio.delta") + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + # Emitted while detached: journaled for replay only. + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="two")) + await asyncio.sleep(0.05) + + ws2, task2 = await _resume(handler, handle.session_id, token, last_seq=1) + resumed = await ws2.wait_for("session.resumed") + assert resumed["session_id"] == handle.session_id + assert resumed["attachment_generation"] == 2 + assert resumed["resume_token"] != token + assert "incarnation" not in resumed + replayed = await ws2.wait_for("response.output_audio.delta") + assert replayed["delta"] == "two" and replayed["server_event_seq"] == 2 + assert omni.resumed == [(handle.session_id, 0)] + + # The old token is revoked by the rotation. + ws3, task3 = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(task3, timeout=2.0) + assert ws3.sent[0]["error"]["code"] == "invalid_resume_token" + + ws2.feed({"type": "session.close"}) + await asyncio.sleep(0.05) + ws2.disconnect() + await asyncio.wait_for(task2, timeout=2.0) + + +@pytest.mark.asyncio +async def test_resume_takes_over_a_live_attachment() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + + ws2, task2 = await _resume(handler, handle.session_id, token) + await ws2.wait_for("session.resumed") + replaced = await ws.wait_for("session.replaced") + assert replaced["attachment_generation"] == 1 + assert ws.closed and ws.closed[0][1] == "session_replaced" + # The replaced socket's later input is ignored; the new one drives the session. + ws.feed({"type": "input_audio_buffer.clear"}) + ws2.feed({"type": "input_audio_buffer.clear", "event_id": "evt-new"}) + await asyncio.sleep(0.05) + assert [command.event_id for command in handle.commands] == ["evt-new"] + await asyncio.wait_for(task, timeout=2.0) + + ws2.disconnect() + await asyncio.wait_for(task2, timeout=2.0) + assert omni.detached == [handle.session_id] + + +@pytest.mark.asyncio +async def test_resume_validation_errors() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws = FakeWebSocket({"duplex": "1", "resume": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed({"type": "session.resume", "session_id": "duplex-x"}) + await asyncio.wait_for(task, timeout=2.0) + assert ws.sent[0]["error"]["code"] == "invalid_session_resume" + + ws, task = await _resume(handler, "duplex-unknown", "token") + await asyncio.wait_for(task, timeout=2.0) + assert ws.sent[0]["error"]["code"] == "session_resume_expired" + + +@pytest.mark.asyncio +async def test_journal_overflow_degrades_to_live_delivery_with_resync_required() -> None: + omni = FakeOmni(replay_max_bytes=256) + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="x" * 400)) + resync = await ws.wait_for("session.resync_required") + assert resync["reason"] == "journal_overflow" + delta = await ws.wait_for("response.output_audio.delta") + assert "server_event_seq" not in delta + assert ws.types().index("session.resync_required") < ws.types().index("response.output_audio.delta") + + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + + +# --------------------------------------------------------------------------- # +# Envelope helpers # +# --------------------------------------------------------------------------- # + + +def test_realtime_envelope_query_rules() -> None: + envelope = RealtimeEnvelope.from_query_params({"model": "m"}) + assert envelope.initial_open_payload() == {"model": "m"} + assert envelope.initial_open_payload() is None # autostart happens once + + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + assert envelope.resume_only is True + assert envelope.initial_open_payload() is None + assert RealtimeEnvelope.from_query_params({"resume": "1"}).resume_only is True + assert RealtimeEnvelope.from_query_params({"session_id": "mine"}).default_session_payload() == {"model": None} + + +def test_realtime_envelope_first_message_classification() -> None: + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + resume = envelope.first_message({"type": "session.resume", "session_id": "s"}) + assert resume.kind == "resume" and resume.resume_payload["session_id"] == "s" + + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + opened = envelope.first_message({"type": "session.update", "session": {"model": "x", "instructions": "hi"}}) + assert opened.kind == "open" and opened.session_payload == {"model": "x", "instructions": "hi"} + assert opened.pending_command_payload is None + + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + autostarted = envelope.first_message({"type": "input_audio_buffer.commit"}) + assert autostarted.kind == "open" and autostarted.session_payload == {"model": "m"} + assert autostarted.pending_command_payload == {"type": "input_audio_buffer.commit"} + + +def test_parse_resume_request_requires_the_three_fields_only() -> None: + request = parse_resume_request({"session_id": "s", "resume_token": "t", "last_received_server_event_seq": 3}) + assert request is not None + assert (request.session_id, request.resume_token, request.last_received_server_event_seq) == ("s", "t", 3) + assert parse_resume_request({"session_id": "s", "resume_token": "t"}).last_received_server_event_seq == 0 + assert parse_resume_request({"session_id": "s", "resume_token": "t", "incarnation": 1}) is not None + assert parse_resume_request({"session_id": "s"}) is None + assert parse_resume_request({"session_id": "s", "resume_token": "t", "last_received_server_event_seq": -1}) is None + + +@pytest.mark.asyncio +async def test_a_resume_that_fails_to_activate_does_not_detach_the_live_attachment() -> None: + """A rejected resume must roll back only what it owns. + + Two reconnects can both authenticate and both complete the engine resume; + only one activates. The loser's activation raises, and the rollback used to + detach the *session*, starting the engine's disconnect grace for the winner + -- a grace ordinary heartbeats do not clear. + + The activation failure is injected rather than raced for: what is under test + is the rollback's precondition, not the window that produces it. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + try: + assert await registry.has_attachment(handle.session_id), "the opener is attached" + detached_before = list(omni.detached) + + async def failing_resume(*args, **kwargs): + raise RuntimeError("transport activation lost the race") + + original_resume = registry.resume + registry.resume = failing_resume # type: ignore[method-assign] + try: + ws_lose, task_lose = await _resume(handler, handle.session_id, token) + error = await ws_lose.wait_for("error") + finally: + registry.resume = original_resume # type: ignore[method-assign] + + assert error["error"]["code"] == "session_resume_conflict" + assert omni.detached == detached_before, "the loser must not detach the winner" + assert await registry.has_attachment(handle.session_id), "the winner is still attached" + finally: + for pending in (task, locals().get("task_lose")): + if pending is not None: + pending.cancel() + with suppress(asyncio.CancelledError, Exception): + await pending diff --git a/tests/entrypoints/duplex/test_runtime_adapter_boundary.py b/tests/entrypoints/duplex/test_runtime_adapter_boundary.py deleted file mode 100644 index cb90e4bd5..000000000 --- a/tests/entrypoints/duplex/test_runtime_adapter_boundary.py +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import subprocess -import sys -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - validate_serving_runtime_adapter, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -@pytest.mark.parametrize( - "module_name", - [ - "vllm_omni.entrypoints.duplex.session_runner", - "vllm_omni.entrypoints.duplex.runtime_bridge", - "vllm_omni.entrypoints.duplex.serving", - ], -) -def test_generic_openai_runtime_import_does_not_load_model_adapters(module_name: str) -> None: - repo_root = Path(__file__).resolve().parents[3] - script = f""" -import sys - -import {module_name} - -model_duplex_prefixes = ( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex", - "vllm_omni.model_executor.models.nemotron_voicechat.duplex", - "vllm_omni.model_executor.models.personaplex.duplex", -) -loaded = sorted( - name - for name in sys.modules - if any( - name == prefix or name.startswith(prefix + ".") - for prefix in model_duplex_prefixes - ) -) -if loaded: - raise SystemExit("generic OpenAI runtime loaded model duplex modules: " + ", ".join(loaded)) -""" - - result = subprocess.run( - [sys.executable, "-c", script], - cwd=repo_root, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - - assert result.returncode == 0, result.stderr - - -def test_generic_handler_supports_adapterless_chat_fallback() -> None: - repo_root = Path(__file__).resolve().parents[3] - script = """ -from types import SimpleNamespace - -from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler - -chat_service = SimpleNamespace( - engine_client=SimpleNamespace(), - model_config=SimpleNamespace(model="test-model"), -) -handler = OmniDuplexSessionHandler(chat_service=chat_service) -if handler._serving_runtime_adapter is not None: - raise SystemExit("generic handler silently selected a model serving adapter") -""" - - result = subprocess.run( - [sys.executable, "-c", script], - cwd=repo_root, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - - assert result.returncode == 0, result.stderr - - -def _valid_runtime_adapter() -> SimpleNamespace: - data_plane = SimpleNamespace( - begin_request=lambda request_id: None, - is_terminal=lambda request_id: False, - mark_terminal=lambda request_id: None, - close_stream=lambda request_id: None, - close_session=lambda session_id, **kwargs: None, - project=lambda result, **kwargs: (), - ) - return SimpleNamespace( - adapter_id="test-adapter", - session_states={}, - data_plane=data_plane, - clean_response_done_prefix="clean", - interrupted_tts_prefix="interrupted", - private_runtime_config_keys=frozenset({"private_key"}), - create_session_state=lambda: object(), - session_state=lambda session_id: object(), - remove_session_state=lambda session_id: None, - is_enabled=lambda config: True, - capabilities=lambda **kwargs: object(), - validate_client_extra_body=lambda extra_body: None, - prepare_runtime_config=lambda config, **kwargs: {}, - runtime_config_for_update=lambda config, current: {}, - data_plane_context=lambda **kwargs: object(), - ) - - -@pytest.mark.parametrize( - ("attribute", "invalid_value"), - [ - ("session_states", None), - ("clean_response_done_prefix", None), - ("interrupted_tts_prefix", None), - ("private_runtime_config_keys", {"private_key"}), - ], -) -def test_runtime_adapter_validator_rejects_invalid_protocol_attributes( - attribute: str, - invalid_value: object, -) -> None: - adapter = _valid_runtime_adapter() - setattr(adapter, attribute, invalid_value) - - with pytest.raises(TypeError, match=attribute): - validate_serving_runtime_adapter(adapter) - - -def test_runtime_adapter_validator_rejects_incomplete_data_plane() -> None: - adapter = _valid_runtime_adapter() - adapter.data_plane.close_session = None - - with pytest.raises(TypeError, match="data_plane.*close_session"): - validate_serving_runtime_adapter(adapter) diff --git a/tests/entrypoints/duplex/test_runtime_control_redaction.py b/tests/entrypoints/duplex/test_runtime_control_redaction.py deleted file mode 100644 index cfc5b68c5..000000000 --- a/tests/entrypoints/duplex/test_runtime_control_redaction.py +++ /dev/null @@ -1,115 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Regression tests for duplex ``runtime_control`` payload redaction. - -The redacted result is emitted on the WebSocket (for example -``session.created.runtime_control``), so it must always survive -``json.dumps``. Control results carry typed values straight from the engine -- -``fence`` and ``accepted_fence`` are ``DuplexFence`` instances added by -``DuplexControlClient.execute`` -- which previously reached -``WebSocket.send_json`` unconverted and killed the session at handshake. -""" - -from __future__ import annotations - -import json -from dataclasses import dataclass -from enum import Enum - -import pytest - -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.duplex.runtime_bridge import NativeRuntimeBridgeMixin - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - -_redact = NativeRuntimeBridgeMixin._redact_runtime_control_result - - -def _control_result() -> dict[str, object]: - """The shape ``DuplexControlClient.execute`` returns for an open.""" - return { - "ok": True, - "operation": "open", - "session_id": "sess-1", - "fence": DuplexFence("sess-1", epoch=0, turn_id=0, response_seq=0, incarnation=0), - "accepted_fence": DuplexFence("sess-1", epoch=1, turn_id=2, response_seq=3, incarnation=4), - "stage_results": [{"stage_id": -1, "replica_id": -1, "result": {"supported": True}}], - } - - -def test_redacted_control_result_is_json_serializable() -> None: - """The whole payload must serialize the way send_json does.""" - redacted = _redact(_control_result()) - # Must not raise. - json.dumps({"type": "session.created", "runtime_control": redacted}) - - -def test_fences_are_converted_to_plain_dicts() -> None: - redacted = _redact(_control_result()) - assert redacted["fence"] == { - "session_id": "sess-1", - "epoch": 0, - "turn_id": 0, - "response_seq": 0, - "incarnation": 0, - } - # Field values are preserved, not just the type. - assert redacted["accepted_fence"]["epoch"] == 1 - assert redacted["accepted_fence"]["turn_id"] == 2 - assert redacted["accepted_fence"]["response_seq"] == 3 - assert redacted["accepted_fence"]["incarnation"] == 4 - - -def test_nested_dataclasses_are_converted() -> None: - """Fences can appear inside stage results, not only at the top level.""" - redacted = _redact( - { - "stage_results": [ - {"stage_id": 0, "result": {"fence": DuplexFence("sess-2", epoch=7)}}, - ] - } - ) - json.dumps(redacted) - assert redacted["stage_results"][0]["result"]["fence"]["epoch"] == 7 - - -def test_enums_are_reduced_to_values() -> None: - class _Mode(str, Enum): - APPEND = "append_audio_chunk" - - redacted = _redact({"mode": _Mode.APPEND}) - json.dumps(redacted) - assert redacted["mode"] == "append_audio_chunk" - - -def test_redaction_still_drops_heavy_payload_keys() -> None: - """The pre-existing contract must be preserved.""" - redacted = _redact( - { - "ok": True, - "stage_handoff": object(), - "tts_handoff": object(), - "omni_payload": object(), - "tts_hidden_states": object(), - "tts_token_ids": object(), - "traceback": "long trace", - } - ) - assert redacted == {"ok": True} - - -def test_plain_values_pass_through_unchanged() -> None: - payload = {"ok": True, "count": 3, "name": "x", "missing": None, "items": [1, "two"]} - assert _redact(payload) == payload - - -def test_dataclass_types_are_not_treated_as_instances() -> None: - """A class object is not a dataclass instance and must not be asdict'd.""" - - @dataclass - class _Shape: - x: int = 1 - - assert _redact(_Shape) is _Shape diff --git a/tests/entrypoints/duplex/test_websocket_actor.py b/tests/entrypoints/duplex/test_websocket_actor.py deleted file mode 100644 index 4e5d151b9..000000000 --- a/tests/entrypoints/duplex/test_websocket_actor.py +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import asyncio - -import pytest - -from vllm_omni.entrypoints.duplex.websocket import ( - DuplexWebSocketActor, - normalize_duplex_input_event, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class FakeWebSocket: - def __init__(self) -> None: - self.sent: list[dict[str, object]] = [] - - async def send_json(self, payload: dict[str, object]) -> None: - self.sent.append(dict(payload)) - - -@pytest.mark.parametrize( - ("event", "expected"), - [ - ({"type": "signal_turn", "event": "barge_in"}, {"type": "turn.signal", "event": "barge_in"}), - ({"type": "close_session"}, {"type": "session.close"}), - ({"type": "audio.playback_ack", "played_ms": 1}, {"type": "playback.ack", "played_ms": 1}), - ({"type": "input_text.append", "text": "a"}, {"type": "input.text.append", "text": "a"}), - ({"type": "push_text", "text": "b"}, {"type": "input.text.append", "text": "b"}), - ( - {"type": "input.audio.append", "audio": "wav"}, - {"type": "input_audio_buffer.append", "audio": "wav", "format": "wav"}, - ), - ( - {"type": "push_chunk", "audio": "wav", "format": "pcm_f32le"}, - {"type": "input_audio_buffer.append", "audio": "wav", "format": "pcm_f32le"}, - ), - ], -) -def test_input_aliases_normalize_once_at_mailbox_boundary(event, expected): - assert normalize_duplex_input_event(event) == expected - - -def test_actor_uses_one_fifo_mailbox_for_inbound_events(): - actor = DuplexWebSocketActor(FakeWebSocket()) - - assert isinstance(actor.mailbox, asyncio.Queue) - assert not isinstance(actor.mailbox, asyncio.PriorityQueue) - assert not hasattr(actor, "control_queue") - assert not hasattr(actor, "input_queue") - assert not hasattr(actor, "event_queue") - assert not hasattr(actor, "lifecycle_state") - assert not hasattr(actor, "last_response_id") - assert not hasattr(actor, "overlap_speech_ms") - assert not hasattr(actor, "transition") - assert not hasattr(actor, "output_generation_in_flight") - assert not hasattr(actor, "session") - assert not hasattr(actor, "runtime_opened") - assert not hasattr(actor, "runtime_closed") - assert not hasattr(actor, "drain_input_queue") - assert not hasattr(actor, "control_events_seen") - assert not hasattr(actor, "input_events_seen") - assert not hasattr(actor, "cancel_count") - - -@pytest.mark.asyncio -async def test_control_event_does_not_overtake_earlier_audio_input(): - actor = DuplexWebSocketActor(FakeWebSocket()) - await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"}) - await actor.enqueue_event({"type": "response.cancel"}) - - assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"} - assert await actor.next_event() == {"type": "response.cancel"} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("event_type", ["input.cancel", "session.close", "close_session"]) -async def test_terminal_control_preserves_wire_order(event_type: str): - actor = DuplexWebSocketActor(FakeWebSocket()) - await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"}) - await actor.enqueue_event({"type": event_type}) - - assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"} - assert await actor.next_event() == {"type": event_type} - - -@pytest.mark.asyncio -async def test_mailbox_delivers_every_enqueued_event_exactly_once(): - actor = DuplexWebSocketActor(FakeWebSocket()) - events = [ - {"type": "input_audio_buffer.append", "audio": str(index)} - if index % 2 - else {"type": "response.created", "response_id": f"resp-{index}"} - for index in range(20) - ] - - await asyncio.gather(*(actor.enqueue_event(dict(event)) for event in events)) - received = [await actor.next_event() for _ in events] - - assert sorted(received, key=repr) == sorted(events, key=repr) - - -@pytest.mark.asyncio -async def test_writer_is_single_owner_of_websocket_send(): - websocket = FakeWebSocket() - actor = DuplexWebSocketActor(websocket) - writer = asyncio.create_task(actor.writer_loop()) - - await actor.send_json({"type": "one"}) - await actor.send_json({"type": "two"}) - await actor.close_writer() - await writer - - assert websocket.sent == [{"type": "one"}, {"type": "two"}] - - -@pytest.mark.asyncio -async def test_stale_fence_payload_is_dropped_before_websocket_send(): - websocket = FakeWebSocket() - actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2) - writer = asyncio.create_task(actor.writer_loop()) - - await actor.send_json({"type": "response.output_audio.delta", "epoch": 1}) - await actor.send_json({"type": "response.output_audio.delta", "epoch": 2}) - await actor.close_writer() - await writer - - assert websocket.sent == [{"type": "response.output_audio.delta", "epoch": 2}] - assert actor.stale_output_dropped == 1 - - -@pytest.mark.asyncio -async def test_writer_does_not_revoke_accepted_terminal_after_close_starts(): - websocket = FakeWebSocket() - actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 1) - await actor.send_json({"type": "response.done", "epoch": 1, "response_id": "resp-1"}) - actor.closing = True - await actor.close_writer() - - await actor.writer_loop() - - assert websocket.sent == [{"type": "response.done", "epoch": 1, "response_id": "resp-1"}] - assert actor.stale_output_dropped == 0 diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py deleted file mode 100644 index 77c8f2cb0..000000000 --- a/tests/entrypoints/openai/test_duplex_protocol.py +++ /dev/null @@ -1,601 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import asyncio -import json - -import pytest - -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexOverlapPolicy, - DuplexSession, - DuplexSessionConfig, - DuplexSessionRegistry, - DuplexTurnController, - DuplexTurnEventType, - DuplexTurnState, - ResponseCreateOptions, -) -from vllm_omni.entrypoints.duplex.realtime_session import ( - NativeRealtimeSessionProtocol, -) -from vllm_omni.entrypoints.duplex.realtime_state import RealtimeStateOwner -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( - minicpmo45_native_capabilities, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_realtime_state_owner_uses_explicit_bindings_not_dynamic_attribute_proxy(): - assert "__getattr__" not in RealtimeStateOwner.__dict__ - assert "__setattr__" not in RealtimeStateOwner.__dict__ - - -class _ProtocolWebSocket: - def __init__(self, *events: dict[str, object]) -> None: - self._events = asyncio.Queue() - for event in events: - self._events.put_nowait(json.dumps(event)) - - async def receive_text(self) -> str: - return await self._events.get() - - -def test_duplex_session_commits_text_and_audio_as_one_turn(): - registry = DuplexSessionRegistry() - session = registry.create(DuplexSessionConfig(model="test-model")) - - session.append_text("hello") - session.append_audio("YWJj", fmt="wav", sample_rate_hz=16000) - committed = session.commit_user_input() - - assert committed is not None - assert committed.turn_id == 0 - assert committed.input_commit_seq == 1 - assert committed.epoch == 0 - assert len(session.history) == 1 - content = session.history[0]["content"] - assert isinstance(content, list) - assert content[0] == {"type": "text", "text": "hello"} - assert content[1]["type"] == "audio_url" - assert content[1]["audio_url"]["url"] == "data:audio/wav;base64,YWJj" - assert session.turn_state == DuplexTurnState.USER_COMMITTED - - -def test_response_options_apply_to_one_response_without_mutating_session_defaults(): - session = DuplexSession( - session_id="sid-response-options", - config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64), - ) - session.reserve_response_options( - ResponseCreateOptions( - instructions="one response", - voice="override-voice", - max_tokens=8, - ) - ) - - assert session.config.instructions == "base" - assert session.config.voice == "base-voice" - session.begin_response() - assert session.config.instructions == "base" - assert session.config.voice == "base-voice" - assert session.response_config.instructions == "one response" - assert session.response_config.voice == "override-voice" - assert session.response_config.max_tokens == 8 - - session.end_response() - assert session.config.instructions == "base" - assert session.config.voice == "base-voice" - assert session.config.max_tokens == 64 - - -def test_response_options_cannot_overwrite_an_unconsumed_reservation(): - session = DuplexSession( - session_id="sid-response-options-pending", - config=DuplexSessionConfig(instructions="base"), - ) - session.reserve_response_options(ResponseCreateOptions(instructions="first")) - - with pytest.raises(RuntimeError, match="already reserved"): - session.reserve_response_options(ResponseCreateOptions(instructions="second")) - - session.begin_response() - assert session.response_config.instructions == "first" - - -def test_response_options_cannot_be_reserved_while_response_is_active(): - session = DuplexSession( - session_id="sid-response-options-active", - config=DuplexSessionConfig(instructions="base"), - ) - session.begin_response() - - with pytest.raises(RuntimeError, match="active"): - session.reserve_response_options(ResponseCreateOptions(instructions="too late")) - - -def test_realtime_model_name_does_not_implicitly_enable_native_duplex(): - protocol = NativeRealtimeSessionProtocol({}) - - event = protocol._session_create_from_realtime({"model": "openbmb/MiniCPM-o-4_5"}) - - assert "native_duplex" not in event["session"]["extra_body"] - assert "minicpmo45_native_duplex" not in event["session"]["extra_body"] - - -def test_realtime_explicit_native_duplex_flag_is_preserved(): - protocol = NativeRealtimeSessionProtocol({}) - - event = protocol._session_create_from_realtime( - { - "model": "openbmb/MiniCPM-o-4_5", - "extra_body": {"native_duplex": True}, - } - ) - - assert event["session"]["extra_body"]["native_duplex"] is True - - -def test_realtime_legacy_native_duplex_alias_is_normalized(): - # The deprecated model-prefixed spelling still opts in, but everything - # downstream (and every session echo) sees only the canonical key. - protocol = NativeRealtimeSessionProtocol({}) - - event = protocol._session_create_from_realtime( - { - "model": "openbmb/MiniCPM-o-4_5", - "extra_body": {"minicpmo45_native_duplex": True}, - } - ) - - assert event["session"]["extra_body"]["native_duplex"] is True - assert "minicpmo45_native_duplex" not in event["session"]["extra_body"] - - -def test_realtime_explicit_query_native_duplex_flag_is_available_before_autostart(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - } - ) - - event = json.loads(asyncio.run(protocol.receive_internal_event_text(None))) - - assert event["type"] == "session.create" - assert event["session"]["extra_body"]["native_duplex"] is True - - -def test_realtime_legacy_query_native_duplex_alias_seeds_canonical_key(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "minicpmo45_native_duplex": "1", - } - ) - - event = json.loads(asyncio.run(protocol.receive_internal_event_text(None))) - - assert event["type"] == "session.create" - assert event["session"]["extra_body"]["native_duplex"] is True - assert "minicpmo45_native_duplex" not in event["session"]["extra_body"] - - -@pytest.mark.asyncio -async def test_realtime_resume_heartbeat_and_event_ack_translate_without_session_autostart(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - } - ) - - resume = await protocol._to_duplex_event( - { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - ) - heartbeat = await protocol._to_duplex_event( - { - "type": "session.heartbeat", - "event_id": "event-heartbeat", - } - ) - event_ack = await protocol._to_duplex_event( - { - "type": "session.event_ack", - "server_event_seq": 19, - } - ) - - assert resume == { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - assert heartbeat == { - "type": "session.heartbeat", - "event_id": "event-heartbeat", - } - assert event_ack == { - "type": "session.event_ack", - "server_event_seq": 19, - } - assert protocol._opened is False - - wire_protocol = NativeRealtimeSessionProtocol({}) - first_wire_event = json.loads( - await wire_protocol.receive_internal_event_text( - _ProtocolWebSocket( - { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - ) - ) - ) - - assert first_wire_event["type"] == "session.resume" - - -@pytest.mark.asyncio -async def test_realtime_resume_query_suppresses_model_autostart_on_same_url(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - "resume": "1", - } - ) - resume = { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - - first_wire_event = json.loads(await protocol.receive_internal_event_text(_ProtocolWebSocket(resume))) - - assert first_wire_event == resume - assert protocol._autostarted_default_session is False - - -def test_realtime_projects_resume_lifecycle_events_without_duplex_prefix(): - protocol = NativeRealtimeSessionProtocol({}) - - created = protocol.encode_outbound_event( - { - "type": "session.created", - "session": {"id": "sid-resume"}, - "incarnation": 3, - "attachment_generation": 1, - "resume_token": "first-token", - } - )[0] - resumed = protocol.encode_outbound_event( - { - "type": "session.resumed", - "session_id": "sid-resume", - "incarnation": 3, - "attachment_generation": 2, - "resume_token": "second-token", - } - )[0] - heartbeat = protocol.encode_outbound_event( - { - "type": "session.heartbeat_ack", - "session_id": "sid-resume", - } - )[0] - replaced = protocol.encode_outbound_event( - { - "type": "session.replaced", - "session_id": "sid-resume", - "attachment_generation": 1, - } - )[0] - resync = protocol.encode_outbound_event( - { - "type": "session.resync_required", - "session_id": "sid-resume", - "reason": "journal_gap", - } - )[0] - - assert created["type"] == "session.created" - assert created["incarnation"] == 3 - assert created["attachment_generation"] == 1 - assert created["resume_token"] == "first-token" - assert resumed["type"] == "session.resumed" - assert resumed["resume_token"] == "second-token" - assert heartbeat["type"] == "session.heartbeat_ack" - assert replaced["type"] == "session.replaced" - assert resync == { - "type": "session.resync_required", - "session_id": "sid-resume", - "reason": "journal_gap", - "event_id": resync["event_id"], - } - - -def test_realtime_projects_listen_response_identity(): - protocol = NativeRealtimeSessionProtocol({}) - protocol.encode_outbound_event({"type": "session.created", "session": {"id": "sid-listen"}}) - - bound = protocol.encode_outbound_event( - { - "type": "response.listen", - "session_id": "sid-listen", - "epoch": 0, - "reason": "model_listen", - "model_listen": True, - "response_id": "resp_1", - } - )[0] - # A listen that terminates a precreated response keeps the response - # identity on the Realtime dialect; clients demultiplex on response_id. - assert bound["type"] == "response.listen" - assert bound["response_id"] == "resp_1" - assert bound["response"]["id"] == "resp_1" - assert bound["response"]["status"] == "listening" - - anonymous = protocol.encode_outbound_event( - { - "type": "response.listen", - "session_id": "sid-listen", - "epoch": 0, - "reason": "silence_or_noise", - } - )[0] - assert "response_id" not in anonymous - assert "id" not in anonymous["response"] - - -def test_duplex_session_registry_advances_incarnation_when_id_is_reused(): - registry = DuplexSessionRegistry() - - first = registry.create(session_id="reused-session") - registry.close(first.session_id) - second = registry.create(session_id="reused-session") - - assert first.incarnation == 0 - assert second.incarnation == 1 - - -def test_native_input_commit_does_not_advance_model_turn_identity(): - session = DuplexSession( - session_id="sid-model-owned-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - - first = session.commit_native_audio_input(transcript="first chunk") - second = session.commit_native_audio_input(transcript="second chunk") - - assert session.input_commit_seq == 2 - assert first.input_commit_seq == 1 - assert second.input_commit_seq == 2 - assert first.turn_id == second.turn_id == session.turn_id == 0 - - -def test_duplex_session_owns_response_and_overlap_identity(): - session = DuplexSession(session_id="sid-state-owner", config=DuplexSessionConfig()) - - response_id = session.begin_response() - session.accumulate_overlap_speech(320) - session.accumulate_overlap_speech(180) - session.end_response() - - assert session.active_response_id is None - assert session.last_response_id == response_id - assert session.overlap_speech_ms == 500 - assert session.reset_overlap_speech() == 500 - assert session.overlap_speech_ms == 0 - - -def test_duplex_session_composes_single_owner_ledgers(): - session = DuplexSession(session_id="sid-ledgers", config=DuplexSessionConfig()) - - session.append_text("hello") - session.bind_request("req-1") - response_id = session.begin_response(turn_id=3) - session.mark_audio_sent(duration_ms=240) - - assert not hasattr(session, "input") - assert not hasattr(session, "response") - assert not hasattr(session, "playback_ledger") - assert not hasattr(session, "conversation") - assert session.pending_text == ("hello",) - assert session.active_request_id == "req-1" - assert session.active_response_id == response_id - assert session.active_response_turn_id == 3 - assert session.playback.sent_ms == 240 - - try: - session.playback.sent_ms = 480 - except (AttributeError, TypeError): - pass - else: - raise AssertionError("playback view must be immutable") - - -def test_duplex_barge_in_advances_epoch_and_drops_uncommitted_assistant_text(): - registry = DuplexSessionRegistry() - session = registry.create() - response_id = session.begin_response() - session.bind_request("chatcmpl-duplex-test") - session.append_assistant_text("unplayed answer") - - new_epoch = session.barge_in() - - assert response_id is not None - assert new_epoch == 1 - assert session.epoch == 1 - assert session.active_request_id is None - assert session.active_response_id is None - assert session.assistant_text_buffer == () - assert session.history == () - assert session.turn_state == DuplexTurnState.BARGE_IN - - -def test_duplex_playback_ack_tracks_committed_cursor_separately(): - registry = DuplexSessionRegistry() - session = registry.create() - session.mark_audio_sent(duration_ms=10_000) - - session.acknowledge_playback(played_ms=2_000) - - assert session.playback.generated_ms == 10_000 - assert session.playback.sent_ms == 10_000 - assert session.playback.played_ms == 2_000 - assert session.playback.committed_ms == 2_000 - - -@pytest.mark.parametrize("sample_rate_hz,frame_count", [(22_050, 23), (24_000, 24)]) -def test_duplex_generated_audio_rounds_only_cumulative_frames(sample_rate_hz: int, frame_count: int): - session = DuplexSessionRegistry().create() - response_id = session.begin_response() - - durations = [ - session.record_generated_audio(response_id, frame_count=1, sample_rate_hz=sample_rate_hz) - for _ in range(frame_count) - ] - - assert durations[:-1] == [0] * (frame_count - 1) - assert durations[-1] == 1 - assert session.playback.generated_ms == 1 - assert session.playback.sent_ms == 0 - assert session.playback.played_ms == 0 - - -def test_duplex_generated_audio_rejects_rate_change_without_advancing(): - session = DuplexSessionRegistry().create() - response_id = session.begin_response() - session.record_generated_audio(response_id, frame_count=2205, sample_rate_hz=22_050) - - with pytest.raises(ValueError, match="sample rate changed"): - session.record_generated_audio(response_id, frame_count=2400, sample_rate_hz=24_000) - - assert session.playback.generated_ms == 100 - assert session.playback.sent_ms == 0 - - -def test_duplex_late_audio_accounting_cannot_change_new_response(): - session = DuplexSessionRegistry().create() - old_response = session.begin_response() - session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) - session.barge_in() - new_response = session.begin_response() - session.record_generated_audio(new_response, frame_count=2400, sample_rate_hz=48_000) - current_state = session.turn_state - - assert session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) is None - session.mark_audio_sent(100, response_id=old_response, text_chars=99) - - assert session.playback_for_response(old_response).sent_ms == 100 - assert session.playback.generated_ms == 50 - assert session.playback.sent_ms == 0 - assert session.assistant_audio_text_marks == () - assert session.turn_state == current_state - - -def test_duplex_history_commit_uses_audio_text_alignment_marks(): - registry = DuplexSessionRegistry() - session = registry.create() - session.begin_response() - - session.append_assistant_text("hello ") - session.mark_audio_sent(duration_ms=1_000, text_chars=6) - session.append_assistant_text("world") - session.mark_audio_sent(duration_ms=2_000, text_chars=11) - session.acknowledge_playback(played_ms=1_200, committed_ms=1_200) - - committed = session.end_response(commit_text=True) - - assert committed == {"role": "assistant", "content": "hello w"} - - -def test_turn_controller_accepts_external_signal_source(): - registry = DuplexSessionRegistry() - session = registry.create() - controller = DuplexTurnController() - - event = controller.signal(session, DuplexTurnEventType.USER_STARTED.value) - - assert event["type"] == "turn.event" - assert event["event"] == "user_started" - assert event["turn_state"] == "user_speaking" - assert session.turn_state == DuplexTurnState.USER_SPEAKING - - -def test_duplex_capabilities_do_not_claim_core_kv_or_input_append(): - registry = DuplexSessionRegistry() - session = registry.create() - caps = session.capabilities.as_dict() - - assert caps["implementation_level"] == "serving_session_adapter" - assert caps["supports_kv_lease"] is False - assert caps["supports_input_append"] is False - assert caps["supports_reencode_context"] is True - assert caps["adapter_patterns"] == ["chunk_group_append"] - assert "turn_commit_only" in caps["input_modes"] - assert "client_event" in caps["signal_sources"] - - -def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): - caps = minicpmo45_native_capabilities(max_sessions=2).as_dict() - - assert caps["implementation_level"] == "model_native_duplex" - assert caps["supports_input_append"] is True - assert caps["input_modes"] == ["append_audio_chunk"] - assert caps["adapter_patterns"] == ["scheduler_data_plane"] - assert caps["supports_model_internal_state"] is True - assert caps["requires_model_runner_kv"] is True - assert caps["requires_native_stage_role"] is True - assert caps["supports_kv_lease"] is False - assert caps["supports_core_kv_lease"] is False - assert caps["supports_stage_resumption"] is True - assert caps["supports_scheduler_native_append"] is False - assert caps["supports_core_resumable_request"] is True - assert caps["supports_stage_connector_handoff"] is True - assert caps["supports_audio_truncate"] is True - assert caps["supports_barge_in"] is True - assert caps["target_barge_in_latency_ms"] is None - assert caps["supports_multi_session"] is True - assert caps["supports_multi_session_same_replica"] is True - assert caps["supports_session_lease"] is True - assert caps["supports_session_resume"] is True - assert caps["session_admission_mode"] == "engine_managed" - assert caps["stage_handoff_transport"] == "scheduler_data_plane" - - -def test_duplex_overlap_policy_auto_falls_back_to_listen_only(): - from vllm_omni.entrypoints.duplex.protocol import ( - DuplexOverlapPolicy, - DuplexSessionConfig, - ) - - assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value - - -def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment(): - caps = minicpmo45_native_capabilities(max_sessions=1).as_dict() - - assert caps["supports_multi_session"] is False - assert caps["supports_multi_session_same_replica"] is False - - -def test_duplex_overlap_policy_defaults_and_invalid_values_to_listen_only(): - assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value - assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index 13bd32761..9d7a208b5 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -133,7 +133,7 @@ async def test_registry_failed_send_is_accepted_only_when_journaled(mocker, jour assert accepted.call_args_list == expected_calls raise RuntimeError("transport lost") - await registry.create("sid-send-failure", incarnation=0, send=failing_send, close=mocker.AsyncMock()) + await registry.create("sid-send-failure", send=failing_send, close=mocker.AsyncMock()) with pytest.raises(RuntimeError, match="transport lost"): await registry.send_event( @@ -152,7 +152,7 @@ async def test_registry_detached_event_is_accepted_only_when_journaled(mocker, j registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) send = mocker.AsyncMock() accepted = mocker.Mock() - created = await registry.create("sid-detached-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + created = await registry.create("sid-detached-acceptance", send=send, close=mocker.AsyncMock()) await registry.detach("sid-detached-acceptance", attachment_generation=created.attachment_generation) entry = await registry.send_event( @@ -175,7 +175,7 @@ async def test_registry_unjournaled_event_is_accepted_after_successful_send(mock async def send(payload): accepted.assert_not_called() - await registry.create("sid-live-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + await registry.create("sid-live-acceptance", send=send, close=mocker.AsyncMock()) entry = await registry.send_event( "sid-live-acceptance", @@ -193,7 +193,7 @@ async def test_registry_overflow_does_not_accept_or_send_event(mocker) -> None: registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=100) accepted = mocker.Mock() send = mocker.AsyncMock() - await registry.create("sid-overflow-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + await registry.create("sid-overflow-acceptance", send=send, close=mocker.AsyncMock()) with pytest.raises(DuplexJournalOverflowError, match="byte limit"): await registry.send_event( @@ -231,13 +231,12 @@ async def test_registry_resume_rotates_token_replays_and_atomically_replaces_att async def close_b(reason): closes_b.append(reason) - created = await registry.create("sid", incarnation=2, send=send_a, close=close_a) + created = await registry.create("sid", send=send_a, close=close_a) await registry.send_event("sid", {"type": "event-1"}) await registry.send_event("sid", {"type": "event-2"}) resumed = await registry.resume( "sid", - incarnation=2, resume_token=created.resume_token.plaintext, last_received_server_event_seq=1, send=send_b, @@ -260,6 +259,34 @@ async def test_registry_resume_rotates_token_replays_and_atomically_replaces_att assert closes_a == [] and sends_b == [] and closes_b == [] +@pytest.mark.asyncio +async def test_registry_detach_without_a_generation_targets_the_current_attachment_once() -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + sends: list[dict] = [] + + async def send(payload): + sends.append(payload) + + async def close(reason): + return None + + await registry.create("sid-current", send=send, close=close) + + # The outbound pump outlives any one connection and only knows the session: + # ``None`` means "whichever socket is attached right now". + assert await registry.detach("sid-current", attachment_generation=None) is True + # A second report of the same dead socket is not a new detach, so it cannot + # restart the engine's disconnect grace window. + assert await registry.detach("sid-current", attachment_generation=None) is False + assert await registry.detach("sid-current", attachment_generation=1) is False + assert await registry.is_current_attachment("sid-current", 1) is False + + # A detached session still journals for a later resume. + await registry.send_event("sid-current", {"type": "event-1"}) + assert sends == [] + assert await registry.detach("missing-session", attachment_generation=None) is False + + @pytest.mark.asyncio async def test_registry_resume_sends_activation_then_replay_before_new_live_events() -> None: registry = DuplexSessionAttachmentRegistry( @@ -282,13 +309,12 @@ async def test_registry_resume_sends_activation_then_replay_before_new_live_even activation_started.set() await release_activation.wait() - created = await registry.create("sid-order", incarnation=0, send=old_send, close=close) + created = await registry.create("sid-order", send=old_send, close=close) await registry.detach("sid-order", attachment_generation=1) await registry.send_event("sid-order", {"type": "replayed"}) resume_task = asyncio.create_task( registry.resume( "sid-order", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=new_send, @@ -322,13 +348,12 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover async def close(reason): del reason - created = await registry.create("sid-recovery", incarnation=0, send=send, close=close) + created = await registry.create("sid-recovery", send=send, close=close) await registry.detach("sid-recovery", attachment_generation=1) with pytest.raises(RuntimeError, match="transport lost"): await registry.resume( "sid-recovery", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=failing_send, @@ -342,7 +367,6 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover recovered = await registry.resume( "sid-recovery", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -353,7 +377,6 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover with pytest.raises(InvalidResumeTokenError): await registry.resume( "sid-recovery", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -374,13 +397,12 @@ async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winne async def close(reason): del reason - created = await registry.create("sid-race", incarnation=0, send=send, close=close) + created = await registry.create("sid-race", send=send, close=close) async def attempt(): try: return await registry.resume( "sid-race", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -396,7 +418,7 @@ async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winne @pytest.mark.asyncio -async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> None: +async def test_registry_keeps_sessions_journals_and_tokens_isolated() -> None: registry = DuplexSessionAttachmentRegistry( replay_ttl_s=60.0, replay_max_bytes_per_session=4096, @@ -408,8 +430,8 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N async def close(reason): del reason - created_a = await registry.create("sid-a", incarnation=1, send=send, close=close) - created_b = await registry.create("sid-b", incarnation=4, send=send, close=close) + created_a = await registry.create("sid-a", send=send, close=close) + created_b = await registry.create("sid-b", send=send, close=close) await registry.detach("sid-a", attachment_generation=1) await registry.detach("sid-b", attachment_generation=1) await registry.send_event("sid-a", {"type": "a-only"}) @@ -417,7 +439,6 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N resumed_a = await registry.resume( "sid-a", - incarnation=1, resume_token=created_a.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -425,7 +446,6 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N ) resumed_b = await registry.resume( "sid-b", - incarnation=4, resume_token=created_b.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -434,11 +454,10 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N assert [entry.payload["type"] for entry in resumed_a.replay_entries] == ["a-only"] assert [entry.payload["type"] for entry in resumed_b.replay_entries] == ["b-only"] - with pytest.raises(ValueError, match="incarnation mismatch"): + with pytest.raises(InvalidResumeTokenError): await registry.resume( "sid-a", - incarnation=2, - resume_token=resumed_a.resume_token.plaintext, + resume_token=resumed_b.resume_token.plaintext, last_received_server_event_seq=1, send=send, close=close, @@ -458,7 +477,7 @@ async def test_registry_repr_never_contains_plaintext_tokens() -> None: async def close(reason): del reason - created = await registry.create("sid-repr", incarnation=0, send=send, close=close) + created = await registry.create("sid-repr", send=send, close=close) await registry.detach("sid-repr", attachment_generation=1) entry = await registry.send_event( "sid-repr", @@ -469,3 +488,67 @@ async def test_registry_repr_never_contains_plaintext_tokens() -> None: assert created.resume_token.plaintext not in repr(created) assert created.resume_token.plaintext not in repr(entry) assert created.resume_token.plaintext not in repr(registry) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_on", ["activation", "replay"]) +async def test_registry_resume_cancelled_mid_delivery_rolls_back_like_a_failure(cancel_on: str) -> None: + """Cancellation must take the same rollback as an ordinary delivery failure. + + ``CancelledError`` derives from ``BaseException``, so an ``except + Exception`` around the activation and replay sends let a cancelled resume + keep its attachment as the current generation while the token it rotated + past was already gone — the session left attached to an unusable transport + with no way back. Both send points are covered because either can be the + one that gets cancelled. + """ + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + + async def send(payload): + del payload + + async def close(reason): + del reason + + created = await registry.create("sid-cancel", send=send, close=close) + # One journaled event, so the replay loop has something to send. + await registry.send_event("sid-cancel", {"type": "response.output_text.delta", "delta": "a"}) + await registry.detach("sid-cancel", attachment_generation=1) + + sends = 0 + cancel_at = 1 if cancel_on == "activation" else 2 + + async def cancelling_send(payload): + del payload + nonlocal sends + sends += 1 + if sends == cancel_at: + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await registry.resume( + "sid-cancel", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=cancelling_send, + close=close, + activation_payload_factory=lambda token, generation: { + "type": "session.resumed", + "resume_token": token.plaintext, + "attachment_generation": generation, + }, + ) + + assert sends == cancel_at, "the send that was cancelled is the one under test" + # The cancelled attachment must not remain current... + assert not await registry.has_attachment("sid-cancel") + # ...and the token the caller presented has to work again, or the session + # is unreachable: resume() had already rotated past it. + recovered = await registry.resume( + "sid-cancel", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=close, + ) + assert recovered.attachment_generation == 3 diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 21725183a..37e2dc430 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -920,7 +920,6 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video")) monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor) monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor) - monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False) state = State() await api_server.omni_init_app_state(engine, state, _minimal_args()) diff --git a/tests/entrypoints/openai_api/test_duplex_api_server.py b/tests/entrypoints/openai_api/test_duplex_api_server.py new file mode 100644 index 000000000..b528cbc2a --- /dev/null +++ b/tests/entrypoints/openai_api/test_duplex_api_server.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Guards for the duplex-only API server (unified full-duplex framework). + +A model whose pipeline declares ``duplex_plugin`` is served through +``DuplexOmni`` and nothing else: sessions run over ``/v1/realtime?duplex=1`` +(alias ``/v1/duplex``) and every turn-based route reports "not available". +These tests lock that shape, which ``test_api_server_guards.py`` covers for the +diffusion and multi-stage servers: + +* the pipeline probe that decides duplex vs turn-based at startup, +* the ``app.state`` snapshot of a duplex server (what is wired, what stays + ``None``), +* both duplex websocket routes reaching the session handler, and + ``/v1/realtime`` without the flag still reporting the turn-based Realtime API + as unavailable, +* the startup-warmup gate that holds real clients while the warmup session runs. + +CPU-only: no engine is started; ``DuplexOmni`` is subclassed with just the +accessors the app-state wiring reads. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.datastructures import State +from starlette.websockets import WebSocketDisconnect + +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.config import DuplexCapabilities +from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler +from vllm_omni.entrypoints.duplex_omni import DuplexOmni +from vllm_omni.entrypoints.openai import api_server + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# ``app.state`` of a duplex-only server. Presence is not enough: a route whose +# handler is missing answers "not available", so "not wired" is written as None. +_DUPLEX_APP_STATE_KEYS = { + "engine_client", + "log_stats", + "args", + "sleeping_stages", + "stage_configs", + "vllm_config", + "diffusion_engine", + "openai_serving_models", + "serving_tokenization", + "serving_tokens", + "online_renderer", + "openai_serving_chat", + "openai_serving_chat_batch", + "openai_serving_completion", + "openai_serving_responses", + "openai_serving_embedding", + "openai_serving_pooling", + "openai_serving_classification", + "openai_serving_scores", + "openai_serving_transcription", + "openai_serving_translation", + "openai_serving_speech", + "openai_serving_audio_generate", + "openai_serving_video", + "openai_streaming_speech", + "openai_streaming_video", + "openai_streaming_video_output", + "openai_serving_realtime", + "openai_serving_realtime_robot", + "anthropic_serving_messages", + "openai_serving_duplex", + "enable_server_load_tracking", + "server_load_metrics", +} +#: Every turn-based service, plus the Realtime route that is not the duplex one. +#: ``openai_serving_chat`` is absent: a duplex engine also serves chat, through +#: the ordinary turn-based service, because DuplexOmni extends AsyncOmni. +#: ``online_renderer`` is absent for the same reason -- the chat service needs it. +_DUPLEX_MUST_BE_NONE = _DUPLEX_APP_STATE_KEYS - { + "engine_client", + "log_stats", + "args", + "sleeping_stages", + "stage_configs", + "vllm_config", + "openai_serving_models", + "openai_serving_duplex", + "openai_serving_chat", + "online_renderer", + "enable_server_load_tracking", + "server_load_metrics", +} +_DUPLEX_MUST_BE_WIRED = { + "engine_client", + "args", + "stage_configs", + "vllm_config", + "openai_serving_models", + "openai_serving_duplex", + "openai_serving_chat", +} + + +class _FakeDuplexOmni(DuplexOmni): + """A ``DuplexOmni`` without an engine: only what the app-state wiring reads.""" + + def __init__(self) -> None: + self.model = "demo-duplex-model" + self._stage_configs = [object(), object(), object()] + self._vllm_config = SimpleNamespace( + lora_config=None, + model_config=SimpleNamespace(), + parallel_config=SimpleNamespace(_api_process_rank=0), + ) + self._duplex_session_config = DuplexSessionRuntimeConfig() + self._capabilities = DuplexCapabilities(supports_chat_completions=True) + + async def get_vllm_config(self): + return self._vllm_config + + async def get_supported_tasks(self) -> tuple[str, ...]: + return ("generate",) + + async def get_tokenizer(self): + return SimpleNamespace(chat_template="{{ messages }}") + + @property + def stage_configs(self) -> list[object]: + return self._stage_configs + + @property + def duplex_session_config(self) -> DuplexSessionRuntimeConfig: + return self._duplex_session_config + + @property + def duplex_capabilities(self) -> DuplexCapabilities: + return self._capabilities + + # Properties on AsyncOmni, so the fake overrides rather than assigns them. + @property + def model_config(self) -> SimpleNamespace: + return SimpleNamespace() + + @property + def renderer(self) -> SimpleNamespace: + return SimpleNamespace() + + +class _FakeModels: + def __init__(self, *args, **kwargs) -> None: + self.base_model_paths = kwargs.get("base_model_paths") or [] + + +def _minimal_args(**overrides) -> SimpleNamespace: + args = SimpleNamespace( + model="demo-duplex-model", + served_model_name=None, + disable_log_stats=True, + enable_log_requests=False, + max_log_len=None, + enable_server_load_tracking=False, + trust_remote_code=True, + deploy_config=None, + # What the chat service reads when a duplex model also serves chat. + chat_template=None, + chat_template_content_format="auto", + trust_request_chat_template=False, + default_chat_template_kwargs=None, + response_role="assistant", + return_tokens_as_token_ids=False, + enable_auto_tool_choice=False, + exclude_tools_when_tool_choice_none=False, + tool_call_parser=None, + structured_outputs_config=SimpleNamespace(reasoning_parser=None), + enable_prompt_tokens_details=False, + enable_force_include_usage=False, + enable_log_outputs=False, + enable_log_deltas=False, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +# --------------------------------------------------------------------------- # +# Startup: which entrypoint class serves the model # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("pipeline_config", "expected"), + [ + (SimpleNamespace(duplex_plugin="pkg.mod.Plugin"), True), + (SimpleNamespace(duplex_plugin=None), False), + (SimpleNamespace(duplex_plugin=""), False), + (SimpleNamespace(), False), + (None, False), + ], +) +def test_duplex_model_probe_follows_the_pipeline_plugin(monkeypatch, pipeline_config, expected: bool) -> None: + """``duplex_plugin`` alone decides whether ``serve`` builds ``DuplexOmni``.""" + seen: dict[str, object] = {} + + def fake_get_pipeline_config(*, model, trust_remote_code, deploy_config_path=None, **_kwargs): + seen.update(model=model, trust_remote_code=trust_remote_code, deploy_config_path=deploy_config_path) + return pipeline_config + + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config) + + kwargs = {"trust_remote_code": True, "deploy_config": "deploy.yaml"} + assert api_server._is_duplex_model("demo-duplex-model", kwargs) is expected + assert seen == { + "model": "demo-duplex-model", + "trust_remote_code": True, + "deploy_config_path": "deploy.yaml", + } + + +def test_duplex_model_probe_propagates_resolution_errors(monkeypatch) -> None: + """A pipeline that cannot be resolved must fail startup, not fall back to turn-based.""" + + def fake_get_pipeline_config(**_kwargs): + raise ValueError("needs trust_remote_code=True") + + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config) + + with pytest.raises(ValueError, match="trust_remote_code"): + api_server._is_duplex_model("demo-duplex-model", {}) + + +# --------------------------------------------------------------------------- # +# app.state of a duplex-only server # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) -> None: + """Lock the duplex ``app.state``: the two session-backed surfaces are live, every turn-based service is None. + + Fails if a turn-based service is wired into a duplex server (its route + would answer instead of reporting "not available"), or if a key the routes + read disappears. ``openai_serving_chat`` is the one shared name: the route + is served, but by the adapter that runs it on a duplex session, never by + the turn-based chat service. + """ + monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) + monkeypatch.setattr(api_server, "OnlineRenderer", lambda **kwargs: SimpleNamespace(**kwargs)) + monkeypatch.setattr(api_server, "OmniOpenAIServingChat", lambda **kwargs: SimpleNamespace(**kwargs)) + engine = _FakeDuplexOmni() + state = State() + + await api_server.omni_init_app_state(engine, state, _minimal_args()) + + present = {key for key in _DUPLEX_APP_STATE_KEYS if hasattr(state, key)} + assert present == _DUPLEX_APP_STATE_KEYS + not_wired = sorted(key for key in _DUPLEX_MUST_BE_WIRED if getattr(state, key) is None) + assert not not_wired, f"duplex app.state keys registered but not wired: {not_wired}" + unexpectedly_set = sorted(key for key in _DUPLEX_MUST_BE_NONE if getattr(state, key) is not None) + assert not unexpectedly_set, f"turn-based services wired into a duplex server: {unexpectedly_set}" + assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) + assert state.openai_serving_chat is not None + assert state.engine_client is engine + assert state.vllm_config is engine._vllm_config + + +# --------------------------------------------------------------------------- # +# Routes # +# --------------------------------------------------------------------------- # + + +class _RecordingHandler: + """Stands in for ``OmniDuplexSessionHandler`` on the app state.""" + + def __init__(self) -> None: + self.queries: list[dict[str, str]] = [] + + async def handle_realtime_session(self, websocket) -> None: + await websocket.accept() + self.queries.append(dict(websocket.query_params)) + await websocket.send_json({"type": "session.created", "session": {"id": "duplex-test"}}) + await websocket.close() + + +def _duplex_app(handler: object) -> FastAPI: + app = FastAPI() + app.include_router(api_server.router) + app.state.openai_serving_duplex = handler + # A duplex server serves no turn-based Realtime route. + app.state.openai_serving_realtime = None + return app + + +@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"]) +def test_both_duplex_routes_reach_the_session_handler(path: str) -> None: + """``/v1/duplex`` is an alias of ``/v1/realtime?duplex=1``: same handler, same protocol.""" + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(path) as websocket: + assert websocket.receive_json()["type"] == "session.created" + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert len(handler.queries) == 1 + + +@pytest.mark.parametrize("flag", ["1", "true", "on"]) +def test_realtime_route_switches_to_duplex_only_for_the_documented_flag_values(flag: str) -> None: + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(f"/v1/realtime?duplex={flag}") as websocket: + assert websocket.receive_json()["type"] == "session.created" + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert handler.queries == [{"duplex": flag}] + + +@pytest.mark.parametrize("query", ["", "?model=openbmb%2FMiniCPM-o-4_5"]) +def test_realtime_without_a_duplex_flag_reaches_the_session_handler(query: str) -> None: + """A bare connection to a duplex deployment is a duplex session. + + ``?duplex=1`` is the spelling the server advertises, but it is not a + requirement: a deployment that declares a duplex plugin answers plain + ``/v1/realtime`` too, so a stock Realtime client needs no vendor query + parameter. + """ + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(f"/v1/realtime{query}") as websocket: + assert websocket.receive_json()["type"] == "session.created" + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert len(handler.queries) == 1 + + +@pytest.mark.parametrize("query", ["?duplex=0", "?duplex=false", "?duplex=off"]) +def test_realtime_opted_out_of_duplex_falls_through_to_the_turn_based_route(query: str) -> None: + """An explicit opt-out selects the legacy handler, which a duplex server does not mount.""" + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(f"/v1/realtime{query}") as websocket: + assert websocket.receive_json() == { + "type": "error", + "error": "Realtime API is not available", + "code": "unsupported", + } + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert handler.queries == [] + + +# --------------------------------------------------------------------------- # +# Startup warmup gate # +# --------------------------------------------------------------------------- # + + +def _warmup_websocket(warmup_done: asyncio.Event | None, query: dict[str, str]) -> SimpleNamespace: + return SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(duplex_warmup_done=warmup_done)), + query_params=query, + ) + + +@pytest.mark.asyncio +async def test_warmup_gate_holds_clients_until_the_warmup_session_finishes() -> None: + warmup_done = asyncio.Event() + websocket = _warmup_websocket(warmup_done, {}) + + waiting = asyncio.create_task(api_server._wait_for_duplex_warmup(websocket)) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(waiting), timeout=0.05) + + warmup_done.set() + await asyncio.wait_for(waiting, timeout=1.0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("warmup_done", "query"), + [ + (None, {}), + (asyncio.Event(), {"vllm_omni_warmup": "1"}), + ], + ids=["no-warmup-configured", "the-warmup-connection-itself"], +) +async def test_warmup_gate_lets_the_warmup_connection_and_plain_servers_through(warmup_done, query) -> None: + websocket = _warmup_websocket(warmup_done, query) + + await asyncio.wait_for(api_server._wait_for_duplex_warmup(websocket), timeout=1.0) + + +@pytest.mark.asyncio +async def test_a_model_that_does_not_declare_chat_completions_does_not_get_the_route(monkeypatch) -> None: + """The decision stays the model's: no capability, no chat service, and the route says so.""" + monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) + engine = _FakeDuplexOmni() + engine._capabilities = DuplexCapabilities(supports_chat_completions=False) + state = State() + + await api_server.omni_init_app_state(engine, state, _minimal_args()) + + assert state.openai_serving_chat is None + assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py deleted file mode 100644 index 0b64e7f14..000000000 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ /dev/null @@ -1,9491 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import base64 -import io -import json -import struct -import wave -from collections.abc import Callable -from contextlib import nullcontext -from types import SimpleNamespace -from typing import Any - -import numpy as np -import pytest -from starlette.websockets import WebSocketDisconnect - -# Import the OpenAI entrypoints package at collection time, before pytest -# attaches its caplog handlers to the root logger. vllm's api_server -# transitively imports model_hosting_container_standards, whose import-time -# configure_root_logger() raises the level of every handler then attached to -# the root logger to ERROR. The serving runtime under test reaches this chain -# lazily (runtime_bridge -> entrypoints.openai.protocol.audio), so without -# this eager import the side effect fires mid-test and silences caplog for -# the remainder of the session. Before the capability predicate moved to -# vllm_omni.entrypoints.duplex.capability, this file imported the package at -# module scope implicitly. -import vllm_omni.entrypoints.openai # noqa: F401 -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig -from vllm_omni.engine.duplex.control_client import DuplexControlRequestError -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence, DuplexSessionLifecycleMessage -from vllm_omni.engine.duplex.runtime import duplex_resource_request_id -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexCapabilities, - DuplexOverlapPolicy, - DuplexPlaybackCommitPolicy, - DuplexSession, - DuplexSessionConfig, - ResponseCreateOptions, -) -from vllm_omni.entrypoints.duplex.realtime_session import NativeRealtimeSessionProtocol -from vllm_omni.entrypoints.duplex.runtime_adapter import ServingRuntimeConfigError -from vllm_omni.entrypoints.duplex.serving import ( - OmniDuplexSessionHandler, - should_enable_duplex_endpoint, -) -from vllm_omni.entrypoints.duplex.websocket import DuplexWebSocketActor -from vllm_omni.model_executor.models.minicpmo_4_5.duplex import ( - MiniCPMO45NativeDuplexServingAdapter, - MiniCPMO45PcmAppendBuffer, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( - minicpmo45_native_capabilities, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( - MiniCPMO45DataPlaneContext, - MiniCPMO45DataPlaneSession, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import ( - MiniCPMO45DuplexRuntimeExtension, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( - MiniCPMO45ServingSessionState, -) -from vllm_omni.model_executor.models.personaplex.duplex.serving_adapter import ( - PersonaPlexServingRuntimeAdapter, -) -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import attach_duplex_output_decision - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -INVALID_VALIDATED_SAMPLE_RATES = ( - pytest.param(16_000.5, id="fractional"), - pytest.param(float("inf"), id="infinite"), - pytest.param(float("nan"), id="nan"), -) - - -def test_native_input_append_supports_explicit_session_opt_out(): - session = DuplexSession( - "sid-explicit-tts", - DuplexSessionConfig(model="test-model"), - capabilities=minicpmo45_native_capabilities(), - ) - - assert OmniDuplexSessionHandler._uses_native_input_append(session) is True - - session.config.extra_body["native_duplex"] = False - assert OmniDuplexSessionHandler._uses_native_input_append(session) is False - - -class _ModelConfig: - model = "test-model" - - -class FakeEngineClient: - output_modalities = ["text", "audio"] - - def __init__( - self, - *, - fail_open: bool = False, - fail_signal: bool = False, - fail_touch: bool = False, - fail_signal_events: set[str] | None = None, - fail_close: bool = False, - fail_abort: bool = False, - control_result: dict[str, object] | None = None, - open_result: dict[str, object] | None = None, - append_result: dict[str, object] | None = None, - append_results: list[dict[str, object]] | None = None, - collect_outputs: list[list[object]] | None = None, - collect_delay_s: float = 0.0, - signal_result: dict[str, object] | None = None, - close_result: dict[str, object] | None = None, - ) -> None: - self.fail_open = fail_open - self.fail_signal = fail_signal - self.fail_touch = fail_touch - self.fail_signal_events = set(fail_signal_events or ()) - self.fail_close = fail_close - self.fail_abort = fail_abort - self.control_result = control_result - self.open_result = open_result - self.append_result = append_result - self.append_results = list(append_results or []) - self.collect_outputs = list(collect_outputs or []) - self.collect_delay_s = collect_delay_s - self.signal_result = signal_result - self.close_result = close_result - self.opened: list[str] = [] - self.opened_fences: list[DuplexFence | None] = [] - self.appended: list[tuple[str, str, Any, bool]] = [] - self.append_operation_ids: list[str | None] = [] - self.appended_fences: list[DuplexFence | None] = [] - self.opened_configs: list[dict[str, Any]] = [] - self.opened_runtime_configs: list[dict[str, Any]] = [] - self.signals: list[tuple[str, str]] = [] - self.signal_fences: list[DuplexFence | None] = [] - self.signal_next_fences: list[DuplexFence | None] = [] - self.signal_session_configs: list[dict[str, Any] | None] = [] - self.signal_runtime_configs: list[dict[str, Any] | None] = [] - self.closed: list[tuple[str, str]] = [] - self.touched: list[tuple[str, DuplexLeaseActivity]] = [] - self.touch_fences: list[DuplexFence] = [] - self.resumed: list[tuple[str, int]] = [] - self.resume_fences: list[DuplexFence] = [] - self.aborted: list[str] = [] - self.abort_batches: list[list[str]] = [] - self.internal_abort_batches: list[list[str]] = [] - self.collected: list[tuple[str, int | None]] = [] - self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue() - - async def open_duplex_session_async( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = None, - fence: DuplexFence | None = None, - ) -> Any: - del timeout - if self.fail_open: - raise RuntimeError("open failed") - self.opened.append(session_id) - self.opened_fences.append(fence) - self.opened_configs.append(dict(session_config or {})) - self.opened_runtime_configs.append(dict(runtime_config or {})) - return self.open_result if self.open_result is not None else self.control_result - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - timeout: float | None = None, - collect_outputs: bool = True, - fence: DuplexFence | None = None, - ) -> Any: - del timeout, collect_outputs - self.appended.append((session_id, mode, payload, final)) - self.append_operation_ids.append(operation_id) - self.appended_fences.append(fence) - if self.append_results: - return self.append_results.pop(0) - return self.append_result if self.append_result is not None else self.control_result - - async def collect_duplex_data_plane_outputs_async( - self, - request_id: str, - *, - response_stage_id: int | None = None, - timeout: float | None = None, - ) -> list[object]: - del timeout - self.collected.append((request_id, response_stage_id)) - if self.collect_delay_s > 0: - await asyncio.sleep(self.collect_delay_s) - if not self.collect_outputs: - return [] - return self.collect_outputs.pop(0) - - async def signal_duplex_turn_async( - self, - session_id: str, - *, - event: str, - timeout: float | None = None, - fence: DuplexFence | None = None, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - ) -> Any: - del timeout - if self.fail_signal or event in self.fail_signal_events: - raise RuntimeError("signal failed") - self.signals.append((session_id, event)) - self.signal_fences.append(fence) - self.signal_next_fences.append(next_fence) - self.signal_session_configs.append(session_config) - self.signal_runtime_configs.append(runtime_config) - return self.signal_result if self.signal_result is not None else self.control_result - - async def close_duplex_session_async( - self, - session_id: str, - *, - reason: str = "client_close", - timeout: float | None = None, - fence: DuplexFence | None = None, - ) -> Any: - del timeout - if self.fail_close: - raise RuntimeError("close failed") - self.closed.append((session_id, reason)) - return self.close_result if self.close_result is not None else self.control_result - - async def touch_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = None, - ) -> dict[str, object]: - del timeout - if self.fail_touch: - raise RuntimeError("touch failed") - self.touched.append((session_id, activity)) - self.touch_fences.append(fence) - return { - "operation": "touch", - "session_id": session_id, - "ok": True, - "stage_results": [ - { - "stage_id": -1, - "replica_id": -1, - "result": {"supported": True, "lease_generation": 0}, - } - ], - } - - async def resume_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = None, - ) -> dict[str, object]: - del timeout - self.resumed.append((session_id, expected_lease_generation)) - self.resume_fences.append(fence) - return { - "operation": "resume", - "session_id": session_id, - "ok": True, - "stage_results": [ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "lease_generation": expected_lease_generation + 1, - }, - } - ], - } - - async def abort(self, request_ids: list[str]) -> None: - if self.fail_abort: - raise RuntimeError("abort failed") - assert isinstance(request_ids, list) - self.abort_batches.append(list(request_ids)) - self.aborted.extend(request_ids) - - async def _abort_internal_requests(self, request_ids: list[str]) -> None: - if self.fail_abort: - raise RuntimeError("abort failed") - assert isinstance(request_ids, list) - self.internal_abort_batches.append(list(request_ids)) - self.aborted.extend(request_ids) - - -class FakeChatService: - duplex_serving_adapter_path = ( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter" - ) - - def __init__(self, engine_client: FakeEngineClient) -> None: - self.engine_client = engine_client - self.model_config = _ModelConfig() - self.seen_request_ids: list[str] = [] - - async def create_chat_completion(self, request, raw_request=None): - self.seen_request_ids.append(request.request_id) - - async def _gen(): - await asyncio.sleep(999) - yield "data: [DONE]\n\n" - - return _gen() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - -class TurnBasedFakeChatService(FakeChatService): - """Chat-fallback service without a model-native duplex adapter.""" - - duplex_serving_adapter_path = None - - def __init__(self, engine_client: FakeEngineClient) -> None: - super().__init__(engine_client) - self.requests: list[Any] = [] - - async def create_chat_completion(self, request, raw_request=None): - del raw_request - self.requests.append(request) - self.seen_request_ids.append(request.request_id) - - async def _gen(): - yield "data: [DONE]\n\n" - - return _gen() - - -class GateChatService(FakeChatService): - def __init__(self, engine_client: FakeEngineClient) -> None: - super().__init__(engine_client) - self.first_response_gate = asyncio.Event() - self.response_count = 0 - - async def create_chat_completion(self, request, raw_request=None): - del raw_request - self.seen_request_ids.append(request.request_id) - self.response_count += 1 - response_number = self.response_count - - async def _gen(): - if response_number == 1: - await self.first_response_gate.wait() - yield "data: [DONE]\n\n" - - return _gen() - - -def _test_data_plane() -> MiniCPMO45DataPlaneSession: - def encode_audio( - audio: object, - sample_rate_hz: int, - response_format: str, - speed: float | None, - ) -> str | None: - del sample_rate_hz, response_format, speed - if audio is None: - return None - if hasattr(audio, "numel"): - return f"wav-{int(audio.numel())}" - return f"wav-{int(np.asarray(audio).size)}" - - return MiniCPMO45DataPlaneSession(encode_audio) - - -def _data_plane_context(session: DuplexSession | None = None) -> MiniCPMO45DataPlaneContext: - if session is None: - return MiniCPMO45DataPlaneContext() - return MiniCPMO45DataPlaneContext( - epoch=session.epoch, - turn_id=session.turn_id, - active_response_turn_id=session.active_response_turn_id, - active_response_id=session.active_response_id, - auto_responds=bool(session.config.extra_body.get("auto_response", False)), - response_format=session.config.response_format, - speed=session.config.speed, - modalities=tuple(session.config.modalities), - ) - - -def _project_data_plane( - data_plane: MiniCPMO45DataPlaneSession, - result: object, - *, - session: DuplexSession | None = None, -) -> list[dict[str, Any]]: - return list(data_plane.project(result, context=_data_plane_context(session))) - - -class TimedWebSocket: - def __init__(self, *, on_send=None, receive_timeout_s: float = 1.0): - self._q: asyncio.Queue[str | BaseException] = asyncio.Queue() - self.sent: list[dict[str, Any]] = [] - self.accepted = False - self._on_send = on_send - self._receive_timeout_s = receive_timeout_s - self.query_params: dict[str, str] = {} - - async def accept(self): - self.accepted = True - - async def receive_text(self) -> str: - try: - item = await asyncio.wait_for(self._q.get(), timeout=self._receive_timeout_s) - except asyncio.TimeoutError as exc: - raise WebSocketDisconnect(code=1000) from exc - if isinstance(item, BaseException): - raise item - return item - - async def send_json(self, data: dict[str, Any]): - self.sent.append(data) - if self._on_send is not None: - self._on_send(self, data) - - def put(self, payload: dict[str, Any]) -> None: - self._q.put_nowait(json.dumps(payload)) - - async def close(self, code: int = 1000, reason: str | None = None) -> None: - del reason - self._q.put_nowait(WebSocketDisconnect(code=code)) - - def sent_types(self) -> list[str]: - return [m.get("type", "") for m in self.sent] - - -class FakeServerVADBackend: - frame_samples = 160 - - def __init__(self, probabilities: list[float]) -> None: - self.probabilities = probabilities - self.calls = 0 - - def new_state(self) -> int: - return 0 - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - del frame - self.calls += 1 - index = int(state) - probability = self.probabilities[index] if index < len(self.probabilities) else 0.0 - return probability, index + 1 - - -class FakeServerVADProvider: - def __init__(self, backend: FakeServerVADBackend) -> None: - self.backend = backend - - def get(self) -> FakeServerVADBackend: - return self.backend - - -class FailingServerVADProvider: - def __init__(self) -> None: - self.calls = 0 - - def get(self) -> FakeServerVADBackend: - self.calls += 1 - raise RuntimeError("server VAD backend unavailable") - - -class FailOnceServerVADBackend(FakeServerVADBackend): - def __init__(self, probabilities: list[float]) -> None: - super().__init__(probabilities) - self._failed = False - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - if not self._failed and int(state) == 2: - self._failed = True - raise RuntimeError("server VAD inference failed") - return super().infer(frame, state) - - -def _server_vad_turn_detection(*, create_response: bool = True) -> dict[str, Any]: - return { - "type": "server_vad", - "prefix_padding_ms": 10, - "silence_duration_ms": 20, - "create_response": create_response, - "interrupt_response": False, - } - - -def _server_vad_session_update( - turn_detection: dict[str, Any] | None, - *, - event_id: str | None = None, - sample_rate_hz: int = 16_000, -) -> dict[str, Any]: - event: dict[str, Any] = { - "type": "session.update", - "session": { - "model": "test-model", - "audio": { - "input": { - "format": {"type": "audio/pcm", "rate": sample_rate_hz}, - "turn_detection": turn_detection, - } - }, - }, - } - if event_id is not None: - event["event_id"] = event_id - return event - - -def _server_vad_audio_append( - frame_count: int = 5, - *, - event_id: str | None = None, - is_speech: bool | None = None, - sample_rate_hz: int = 16_000, -) -> dict[str, Any]: - sample_count = frame_count * 160 * sample_rate_hz // 16_000 - event: dict[str, Any] = { - "type": "input_audio_buffer.append", - "audio": base64.b64encode(np.zeros(sample_count, dtype=" None: - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - if not at_creation: - ws.put(_server_vad_session_update(_server_vad_turn_detection())) - update = _server_vad_session_update(_server_vad_turn_detection(), event_id="bad-vad-audio-format") - update["session"]["audio"]["input"].update(input_update) - ws.put(update) - if not at_creation: - ws.put(_server_vad_audio_append()) - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - (error,) = [event["error"] for event in ws.sent if event.get("type") == "error"] - assert error["code"] == "unsupported_turn_detection" - assert error["event_id"] == "bad-vad-audio-format" - assert "mono PCM16" in error["message"] - assert ("session.created" in ws.sent_types()) is not at_creation - assert ("response.done" in ws.sent_types()) is not at_creation - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sample_rate_hz", [16_000, 24_000]) -@pytest.mark.parametrize("create_response", [False, True], ids=["commit-only", "minimal-config"]) -@pytest.mark.parametrize("at_creation", [True, False], ids=["create", "update"]) -async def test_realtime_server_vad_commits_without_client_commit( - sample_rate_hz: int, create_response: bool, at_creation: bool -): - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - terminal_event = "response.done" if create_response else "input_audio_buffer.committed" - if data.get("type") == terminal_event: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - if not at_creation: - ws.put(_server_vad_session_update(None, sample_rate_hz=sample_rate_hz)) - ws.put( - _server_vad_session_update( - {"type": "server_vad"} if create_response else _server_vad_turn_detection(create_response=False), - sample_rate_hz=sample_rate_hz, - ) - ) - # Cover the default 500 ms silence timeout, plus streaming FIR right context. - ws.put(_server_vad_audio_append(54, sample_rate_hz=sample_rate_hz)) - engine = FakeEngineClient() - chat_service = TurnBasedFakeChatService(engine) - handler = OmniDuplexSessionHandler( - chat_service=chat_service, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - event_types = ws.sent_types() - assert "error" not in event_types - effective_session = [ - event["session"] for event in ws.sent if event["type"] in {"session.created", "session.updated"} - ][-1] - assert effective_session["turn_detection"]["interrupt_response"] is False - assert effective_session["audio"]["input"]["turn_detection"] == effective_session["turn_detection"] - assert effective_session["overlap_policy"] == "listen_only" - assert event_types.index("input_audio_buffer.speech_started") < event_types.index( - "input_audio_buffer.speech_stopped" - ) - assert event_types.index("input_audio_buffer.speech_stopped") < event_types.index("input_audio_buffer.committed") - assert event_types.index("input_audio_buffer.committed") < event_types.index("conversation.item.added") - assert event_types.index("conversation.item.added") < event_types.index("conversation.item.done") - assert "conversation.item.created" not in event_types - committed = next(event for event in ws.sent if event.get("type") == "input_audio_buffer.committed") - added = next(event for event in ws.sent if event.get("type") == "conversation.item.added") - done = next(event for event in ws.sent if event.get("type") == "conversation.item.done") - assert committed["item_id"] == added["item"]["id"] == done["item"]["id"] - assert event_types.count("response.created") == int(create_response) - assert len(chat_service.seen_request_ids) == int(create_response) - if create_response: - assert event_types.index("input_audio_buffer.committed") < event_types.index("response.created") - assert event_types.index("response.created") < event_types.index("response.done") - assert engine.opened == [] - assert engine.signals == [] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_server_vad_create_response_false_accepts_response_create(): - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "input_audio_buffer.committed": - ws.put({"type": "response.create"}) - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection(create_response=False))) - ws.put(_server_vad_audio_append()) - engine = FakeEngineClient() - chat_service = TurnBasedFakeChatService(engine) - handler = OmniDuplexSessionHandler( - chat_service=chat_service, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - event_types = ws.sent_types() - assert event_types.index("input_audio_buffer.committed") < event_types.index("response.created") - assert event_types.index("response.created") < event_types.index("response.done") - assert event_types.count("response.created") == 1 - assert len(chat_service.seen_request_ids) == 1 - assert engine.opened == engine.signals == engine.closed == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sample_rate_hz", [16_000, 24_000]) -async def test_realtime_server_vad_complete_audio_item_bypasses_vad(sample_rate_hz: int): - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - item = data.get("item") - if ( - data.get("type") == "conversation.item.done" - and isinstance(item, dict) - and item.get("id") == "item-complete-audio" - ): - ws.put({"type": "response.create"}) - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection(), sample_rate_hz=sample_rate_hz)) - samples = np.arange(sample_rate_hz // 10, dtype=np.float32) - pcm16 = np.rint(0.5 * np.sin(2 * np.pi * 1_000 * samples / sample_rate_hz) * 32_767).astype(" None: - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection())) - first_append = _server_vad_audio_append(frame_count=2) - first_append["audio"] = base64.b64encode(np.full(320, 12_000, dtype=" None: - nonlocal committed_count, response_count - if data.get("type") == "input_audio_buffer.committed": - committed_count += 1 - if data.get("type") == "error" and data.get("error", {}).get("code") == "input_backpressure": - chat_service.first_response_gate.set() - if data.get("type") == "response.created": - response_count += 1 - if response_count == 2: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection())) - for _ in range(3): - ws.put(_server_vad_audio_append()) - handler = OmniDuplexSessionHandler( - chat_service=chat_service, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0] * 3)), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - assert committed_count == 2 - assert response_count == 2 - event_types = ws.sent_types() - committed_indices = [index for index, event in enumerate(event_types) if event == "input_audio_buffer.committed"] - response_created_indices = [index for index, event in enumerate(event_types) if event == "response.created"] - assert committed_indices[1] < event_types.index("response.done") < response_created_indices[1] - backpressure = [ - event - for event in ws.sent - if event.get("type") == "error" and event.get("error", {}).get("code") == "input_backpressure" - ] - assert len(backpressure) == 1 - - -@pytest.mark.asyncio -async def test_native_session_update_rejects_client_runtime_config(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-runtime-update-reject")) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": { - "extra_body": {"duplex_stage_max_tokens": {"0": 999}}, - }, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "invalid_duplex_runtime_config" - assert ("sid-runtime-update-reject", "session.update") not in engine.signals - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_requires_ref_audio_before_enabling_audio(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-runtime-update-ref-required", modalities=["text"])) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"modalities": ["text", "audio"]}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "ref_audio_required" - assert "session.updated" not in ws.sent_types() - assert ("sid-runtime-update-ref-required", "session.update") not in engine.signals - - -def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check(): - handler = OmniDuplexSessionHandler( - chat_service=SimpleNamespace(engine_client=SimpleNamespace()), - served_model_name="test-model", - serving_runtime_adapter=PersonaPlexServingRuntimeAdapter(lambda *_: None), - ) - session = DuplexSession( - session_id="sid-personaplex-voice-change", - config=DuplexSessionConfig(modalities=["audio"], voice="NATF2.pt"), - capabilities=PersonaPlexServingRuntimeAdapter.capabilities(max_sessions=1), - ) - candidate_config = DuplexSessionConfig(modalities=["audio"], voice="OtherVoice.pt") - - error = handler._runtime_session_candidate_update_error(session, candidate_config) - - assert error is None - - -@pytest.mark.parametrize( - ("native", "turn_detection", "expected_code"), - [ - (False, {"type": "server_vad", "interrupt_response": False}, None), - (False, {"type": "server_vad"}, None), - (False, {"type": "server_vad", "interrupt_response": True}, "server_vad_requires_native_duplex"), - (True, {"type": "server_vad", "interrupt_response": False}, "unsupported_turn_detection"), - (True, {"type": "server_vad"}, None), - ], - ids=[ - "turn-based-endpointing", - "turn-based-default-endpointing", - "turn-based-explicit-interruption", - "native-non-interrupting", - "native-default-interruption", - ], -) -def test_server_vad_runtime_capability_matrix( - native: bool, - turn_detection: dict[str, object], - expected_code: str | None, -): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - ) - capabilities = minicpmo45_native_capabilities() if native else DuplexCapabilities() - session = DuplexSession( - "sid-server-vad-capability", - DuplexSessionConfig(), - capabilities=capabilities, - ) - candidate = DuplexSessionConfig.from_event({"session": {"turn_detection": turn_detection}}) - handler._resolve_server_vad_defaults(session, candidate, {"turn_detection": turn_detection}) - - error = handler._runtime_session_candidate_update_error(session, candidate) - - assert (error or {}).get("code") == expected_code - if expected_code is None: - assert candidate.server_vad is not None - assert candidate.server_vad.interrupt_response is native - assert candidate.overlap_policy == ("barge_in_on_speech" if native else "listen_only") - elif expected_code == "server_vad_requires_native_duplex": - assert "set interrupt_response=false" in error["error"] - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_rejects_instructions_after_native_context_locked(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-instructions-locked", modalities=["text"])) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 1000, - "is_speech": True, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "You are now a pirate."}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "instructions_update_unsupported" - assert "session.updated" not in ws.sent_types() - assert ("sid-instructions-locked", "session.update") not in engine.signals - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_allows_unchanged_instructions_after_lock(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-instructions-unchanged", modalities=["text"])) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 1000, - "is_speech": True, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "You are a concise assistant.", "temperature": 0.3}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "session.updated" in ws.sent_types() - assert not any( - message.get("type") == "error" and message.get("code") == "instructions_update_unsupported" - for message in ws.sent - ) - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_rejects_instructions_before_first_real_append(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-buffered-before-lock", modalities=["text"])) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3200, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 200, - "is_speech": True, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": { - "instructions": ( - "You are now a swashbuckling pirate captain who speaks entirely in " - "nautical slang, references buried treasure constantly, and never " - "breaks character no matter what the user asks." - ) - }, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True, "response_create": True}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "instructions_update_unsupported" - assert "session.updated" not in ws.sent_types() - assert engine.appended, "the commit must still flush the buffered chunk despite the rejected update" - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-duplex-flag-flip", modalities=["text"])) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"extra_body": {"native_duplex": False}}, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "You are now a pirate."}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - errors = [message for message in ws.sent if message.get("type") == "error"] - assert [error["code"] for error in errors] == [ - "native_duplex_mode_update_unsupported", - "instructions_update_unsupported", - ] - assert "session.updated" not in ws.sent_types() - - -@pytest.mark.parametrize("sample_rate_hz", [22_050, 24_000]) -def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(sample_rate_hz: int): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - payloads = protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-a", - "audio": "AAAA", - "format": "pcm", - "sample_rate_hz": sample_rate_hz, - "audio_duration_ms": 135, - } - ) - - audio_events = [ - payload for payload in payloads if payload["type"] in {"response.output_audio.delta", "response.audio.delta"} - ] - assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} - assert {payload["format"] for payload in audio_events} == {"pcm16"} - assert {payload["sample_rate_hz"] for payload in audio_events} == {sample_rate_hz} - assert audio_events[0]["metadata"]["audio_duration_ms"] == 135 - - -def test_native_realtime_protocol_ignores_removed_legacy_event_switches(): - ws = TimedWebSocket() - ws.query_params = { - "legacy_audio_events": "true", - "vllm_omni_legacy_events": "true", - "output_audio_events": "true", - "vllm_omni_output_audio_events": "true", - } - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - payloads = protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-canonical", - "audio": "AAAA", - "text": "hello", - "format": "pcm", - "sample_rate_hz": 24000, - "end_of_turn": True, - } - ) - - event_types = {payload["type"] for payload in payloads} - assert "response.output_audio.delta" in event_types - assert "response.output_audio.done" in event_types - assert "response.output_audio_transcript.delta" in event_types - assert "response.output_audio_transcript.done" in event_types - assert "response.audio.delta" not in event_types - assert "response.audio.done" not in event_types - assert "response.audio_transcript.delta" not in event_types - assert "response.audio_transcript.done" not in event_types - assert "response.output_item.created" not in event_types - assert "response.text.delta" not in event_types - - -def test_native_realtime_protocol_emits_terminal_audio_transcript_events(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - payloads = protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-a", - "audio": "AAAA", - "text": "hello", - "format": "pcm", - "sample_rate_hz": 24000, - "end_of_turn": True, - } - ) - - by_type = {payload["type"]: payload for payload in payloads} - assert "response.output_audio.done" in by_type - assert "response.output_audio_transcript.delta" in by_type - assert "response.output_audio_transcript.done" in by_type - assert by_type["response.content_part.done"]["part"]["transcript"] == "hello" - assert by_type["response.output_item.done"]["item"]["object"] == "realtime.item" - assert by_type["response.done"]["response"]["output"][0]["object"] == "realtime.item" - - -def test_native_realtime_protocol_terminal_transcript_equals_joined_deltas(): - protocol = NativeRealtimeSessionProtocol(TimedWebSocket()) # type: ignore[arg-type] - response_id = "resp-transcript-contract" - events: list[dict[str, object]] = [] - - for text, end_of_turn in (("It's Canberra.", False), (" Next question.", True)): - events.extend( - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": response_id, - "audio": "AAAA", - "text": text, - "format": "pcm", - "sample_rate_hz": 24000, - "end_of_turn": end_of_turn, - } - ) - ) - - transcript = "".join( - str(event["delta"]) for event in events if event["type"] == "response.output_audio_transcript.delta" - ) - done = [str(event["transcript"]) for event in events if event["type"] == "response.output_audio_transcript.done"] - - assert transcript == "It's Canberra. Next question." - assert done == [transcript] - - -def test_native_realtime_protocol_updates_in_progress_item_for_audio_truncate(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-truncate"}) - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-truncate", - "audio": "AAAA", - "text": "hello", - "format": "pcm", - "sample_rate_hz": 24000, - "audio_duration_ms": 100, - "audio_text_marks": [{"text_chars": 5, "audio_end_ms": 100}], - } - ) - - item = protocol._conversation_items["item_resp-truncate"] - assert item["content"][0]["type"] == "output_audio" - assert item["content"][0]["transcript"] == "hello" - assert item["content"][0]["audio_duration_ms"] == 100 - - translated = asyncio.run( - protocol._to_duplex_event( - { - "type": "conversation.item.truncate", - "item_id": "item_resp-truncate", - "content_index": 0, - "audio_end_ms": 50, - } - ) - ) - - assert translated is None - assert protocol._conversation_items["item_resp-truncate"]["content"][0]["transcript"] == "he" - - -@pytest.mark.asyncio -async def test_native_realtime_protocol_rejects_truncate_for_user_item(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol.bind_sender(ws.send_json) - protocol._conversation_items["item-user"] = { - "id": "item-user", - "object": "realtime.item", - "type": "message", - "role": "user", - "status": "completed", - "content": [{"type": "input_audio", "transcript": "hello"}], - } - - translated = await protocol._to_duplex_event( - { - "type": "conversation.item.truncate", - "item_id": "item-user", - "content_index": 0, - "audio_end_ms": 10, - } - ) - - assert translated is None - error = next(event for event in ws.sent if event.get("type") == "error") - assert error["error"]["code"] == "bad_event" - - -def test_native_realtime_protocol_response_done_emits_audio_done_before_terminal_events(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-done"}) - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-done", - "audio": "AAAA", - "format": "pcm", - "sample_rate_hz": 24000, - } - ) - - payloads = protocol._from_duplex_event({"type": "response.done", "response_id": "resp-done"}) - event_types = [payload["type"] for payload in payloads] - - assert event_types.index("response.output_audio.done") < event_types.index("response.content_part.done") - assert event_types.index("response.content_part.done") < event_types.index("response.output_item.done") - assert event_types.index("conversation.item.done") < event_types.index("response.done") - assert event_types.index("response.output_item.done") < event_types.index("response.done") - - -def test_native_realtime_protocol_audio_cancelled_uses_active_response_id(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-cancel"}) - - payloads = protocol._from_duplex_event( - { - "type": "audio.cancelled", - "reason": "barge_in", - "committed_ms": 0, - } - ) - - done = [payload for payload in payloads if payload.get("type") == "response.done"] - assert len(done) == 1 - assert done[0]["response_id"] == "resp-cancel" - assert done[0]["response"]["id"] == "resp-cancel" - - -def test_native_realtime_protocol_audio_cancelled_does_not_reopen_completed_response(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-complete"}) - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-complete", - "audio": "AAAA", - "format": "pcm", - "sample_rate_hz": 24000, - } - ) - protocol._from_duplex_event({"type": "response.done", "response_id": "resp-complete"}) - - payloads = protocol._from_duplex_event( - { - "type": "audio.cancelled", - "reason": "session_close", - "committed_ms": 999999, - } - ) - - assert [payload.get("type") for payload in payloads] == [] - - -def test_native_audio_text_marks_are_normalized_to_session_cumulative_offsets(): - marks = OmniDuplexSessionHandler._normalize_native_audio_text_marks( - [{"text_chars": 3, "audio_end_ms": 200}], - audio_offset_ms=800, - text_offset_chars=5, - ) - - assert marks == [{"text_chars": 8, "audio_end_ms": 1000}] - - -def test_duplex_session_playback_commit_uses_multi_delta_audio_text_marks(): - session = DuplexSession( - session_id="sid-marks", - config=DuplexSessionConfig(playback_commit_policy=DuplexPlaybackCommitPolicy.ACK_ONLY.value), - ) - session.begin_response() - session.append_assistant_text("hello ") - session.mark_audio_sent(1000, text_chars=6) - session.append_assistant_text("world") - session.mark_audio_sent(2000, text_chars=11) - session.acknowledge_playback(played_ms=1500, committed_ms=1500) - - committed = session.end_response(commit_text=True) - - assert committed == {"role": "assistant", "content": "hello wo"} - assert session.history == (committed,) - - -def test_duplex_session_preserves_response_history_order_when_user_item_arrives_during_generation(): - session = DuplexSession(session_id="sid-history-order", config=DuplexSessionConfig()) - first_user = {"role": "user", "content": "first turn"} - next_user = { - "role": "user", - "content": [{"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,AAAA"}}], - } - session.append_history_message(first_user) - session.begin_response() - session.append_assistant_text("first answer") - session.append_history_message(next_user) - - assistant = session.end_response(commit_text=True) - - assert session.history == (first_user, assistant, next_user) - - -@pytest.mark.asyncio -async def test_duplex_session_actor_preserves_wire_order_before_control(): - ws = TimedWebSocket() - actor = DuplexWebSocketActor(ws) - - await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "AAAA"}) - await actor.enqueue_event({"type": "response.cancel"}) - - first = await actor.next_event() - second = await actor.next_event() - - assert first["type"] == "input_audio_buffer.append" - assert second["type"] == "response.cancel" - - -def _pcm_f32_b64(samples: int, *, value: float = 0.05) -> str: - return base64.b64encode(struct.pack(f"<{samples}f", *([value] * samples))).decode("ascii") - - -def _pcm16_b64(samples: int = 1, *, value: int = 1000) -> str: - return base64.b64encode(struct.pack(f"<{samples}h", *([value] * samples))).decode("ascii") - - -def _session_create(session_id: str = "duplex-test") -> dict[str, Any]: - return { - "type": "session.create", - "session_id": session_id, - "session": { - "model": "test-model", - "modalities": ["text", "audio"], - "idle_timeout_s": 1, - }, - } - - -def _native_session_create( - session_id: str = "duplex-native", - *, - modalities: list[str] | None = None, -) -> dict[str, Any]: - event = _session_create(session_id) - event["session"]["model"] = "openbmb/MiniCPM-o-4_5" - event["session"]["modalities"] = list(modalities or ["text"]) - event["session"]["instructions"] = "You are a concise assistant." - event["session"]["extra_body"] = {"native_duplex": True} - return event - - -def _native_realtime_session_update( - session_id: str, - *, - modalities: list[str] | None = None, -) -> dict[str, Any]: - return { - "type": "session.update", - "session": { - "session_id": session_id, - "model": "openbmb/MiniCPM-o-4_5", - "modalities": list(modalities or ["text"]), - "instructions": "You are a concise assistant.", - "idle_timeout_s": 1, - "extra_body": {"native_duplex": True}, - }, - } - - -def _native_server_vad_update(event_id: str, **session: object) -> dict[str, object]: - return { - "type": "session.update", - "event_id": event_id, - "session": {"turn_detection": {"type": "server_vad"}, **session}, - } - - -def test_native_stage0_request_id_matches_engine_after_same_session_id_reopen(): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - first = handler._registry.create(session_id="sid-reused") - handler._registry.close(first.session_id) - reopened = handler._registry.create(session_id="sid-reused") - expected = duplex_resource_request_id( - DuplexFence( - reopened.session_id, - epoch=reopened.epoch, - incarnation=reopened.incarnation, - ), - "stage0", - ) - - assert reopened.incarnation == 1 - assert handler._native_stage0_request_id(reopened, reopened.epoch) == expected - - -def _native_audio_payload( - *, - samples: int = 16000, - value: float = 0.05, - is_speech: bool | None = True, - force_listen: bool | None = None, -) -> dict[str, Any]: - payload: dict[str, Any] = { - "type": "audio", - "audio": _pcm_f32_b64(samples, value=value), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - if is_speech is not None: - payload["is_speech"] = is_speech - if force_listen is not None: - payload["force_listen"] = force_listen - return payload - - -def _auto_response_context( - session_id: str, - *, - playback_active: bool = False, -) -> tuple[OmniDuplexSessionHandler, DuplexSession]: - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id=session_id, - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - if playback_active: - session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value - session.begin_response() - session.mark_audio_sent(duration_ms=1000) - return handler, session - - -def _install_direct_silence_scheduler( - handler: OmniDuplexSessionHandler, - session: DuplexSession, -) -> None: - native = handler._runtime_session_state(session) - - async def _schedule(payload: object, **kwargs: Any) -> bool: - if handler._native_silence_continuation_is_stale( - session, - request_id=kwargs["request_id"], - response_id=kwargs["response_id"], - response_owned=kwargs["response_owned"], - expected_epoch=kwargs["expected_epoch"], - expected_incarnation=kwargs["expected_incarnation"], - expected_model_turn_id=kwargs["expected_model_turn_id"], - ): - return False - append_ok, _ = await handler._append_runtime_input( - session, - payload, - final=False, - send_json=kwargs["send_json"], - mode="append_audio_chunk", - expected_epoch=kwargs["expected_epoch"], - ) - return append_ok - - native.silence_continuation_scheduler = _schedule - - -def test_minicpmo_pcm_append_buffer_flush_preserves_accumulated_speech_marker(): - buffer = MiniCPMO45PcmAppendBuffer() - payload = _native_audio_payload(samples=8000) - - assert buffer.append(payload, chunk_period_ms=1000) is None - flushed = buffer.flush(chunk_period_ms=1000) - - assert flushed is not None - assert flushed["is_speech"] is True - - -def test_minicpmo_pcm_append_buffer_drops_serving_new_user_turn_marker(): - buffer = MiniCPMO45PcmAppendBuffer() - first = _native_audio_payload(samples=8000) - first.update( - new_user_turn=True, - new_user_turn_prefix_variant="clean_response_done", - force_speak=True, - ) - second = _native_audio_payload(samples=8000) - - assert buffer.append(first, chunk_period_ms=1000) is None - emitted = buffer.append(second, chunk_period_ms=1000) - - assert emitted is not None - assert emitted["is_speech"] is True - assert "new_user_turn" not in emitted - assert "new_user_turn_prefix_variant" not in emitted - assert "force_speak" not in emitted - - -def test_minicpmo_merge_native_audio_payloads_preserves_speech_marker(): - first = _native_audio_payload(samples=8000) - second = _native_audio_payload(samples=8000, value=0.0, is_speech=False) - - merged = OmniDuplexSessionHandler._merge_native_audio_payloads(first, second) - - assert merged["is_speech"] is True - - -def test_minicpmo_merge_drops_serving_new_user_turn_marker(): - first = _native_audio_payload(samples=8000) - first.update( - new_user_turn=True, - new_user_turn_prefix_variant="clean_response_done", - force_speak=True, - ) - second = _native_audio_payload(samples=8000, is_speech=False) - - merged = OmniDuplexSessionHandler._merge_native_audio_payloads(first, second) - - assert merged["is_speech"] is True - assert "new_user_turn" not in merged - assert "new_user_turn_prefix_variant" not in merged - assert "force_speak" not in merged - - -@pytest.mark.asyncio -async def test_minicpmo_clear_continuation_does_not_cancel_pending_silence_task(): - async def _pending() -> bool: - await asyncio.sleep(3600) - return True - - native = MiniCPMO45ServingSessionState() - task = asyncio.create_task(_pending()) - native.continuation_owner_id = "owner" - native.continuation_units = 1 - native.pending_silence_task = task - native.pending_silence_owner_id = "owner" - - native.clear_continuation() - - assert native.continuation_owner_id is None - assert native.continuation_units == 0 - assert native.pending_silence_task is None - assert native.pending_silence_owner_id is None - assert not task.cancelled() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - -def test_auto_response_playback_overlap_keeps_model_owned_listen_speak_decision(): - handler, session = _auto_response_context("sid-auto-silent-overlap", playback_active=True) - session.mark_audio_sent(duration_ms=1000) - silent_payload = _native_audio_payload(value=0.0, is_speech=None, force_listen=False) - speech_payload = _native_audio_payload(is_speech=None, force_listen=False) - - assert handler._should_force_listen_for_auto_response_overlap(session, {}, silent_payload) is False - assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False - assert ( - handler._should_force_listen_for_auto_response_overlap( - session, - {"force_listen": True}, - speech_payload, - ) - is True - ) - assert ( - handler._should_force_listen_for_auto_response_overlap( - session, - {"force_barge_in": True}, - speech_payload, - ) - is False - ) - - session.acknowledge_playback(played_ms=1000, committed_ms=1000) - assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False - - -def test_legacy_auto_overlap_policy_falls_back_to_listen_only(): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-auto-policy-fallback", - config=DuplexSessionConfig( - overlap_policy=DuplexSessionConfig._normalize_overlap_policy("auto"), - ), - ) - session.capabilities = minicpmo45_native_capabilities() - - decision = handler._overlap_decision( - session, - {"duration_ms": session.config.overlap_barge_in_ms, "is_speech": True}, - _native_audio_payload(is_speech=True), - ) - - assert session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value - assert decision["action"] == "listen" - assert decision["reason"] == "policy_listen_only" - - -def test_auto_response_overlap_silence_advances_model_unit_and_preserves_realtime_input(): - handler, session = _auto_response_context( - "sid-auto-silent-wire-buffer", - playback_active=True, - ) - decision = handler._overlap_decision( - session, - {"is_speech": False}, - _native_audio_payload(samples=3200, is_speech=False), - ) - - assert decision["action"] == "listen" - assert decision["buffer_audio"] is True - assert decision["defer_runtime_append"] is False - assert decision["force_listen"] is False - assert decision["preserve_realtime_input"] is True - - -def test_auto_response_playback_overlap_admits_model_units_and_tracks_speech(): - handler, session = _auto_response_context("sid-auto-deferred-overlap", playback_active=True) - payload = _native_audio_payload(samples=8640) - - decision = handler._overlap_decision( - session, - {"duration_ms": 540, "is_speech": True}, - payload, - ) - - assert decision["action"] == "listen" - assert decision["defer_runtime_append"] is False - assert decision["force_listen"] is False - assert decision["buffer_audio"] is True - assert session.overlap_speech_ms == 540 - - -@pytest.mark.asyncio -async def test_auto_response_overlap_exact_unit_commit_does_not_block_or_replay_next_unit(): - class ModelUnitEngine(FakeEngineClient): - def __init__(self) -> None: - super().__init__() - self.append_started = asyncio.Event() - self.active_response_ids_at_append: list[str | None] = [] - self.response_done_seen_at_append = False - self.session: DuplexSession | None = None - self.websocket: TimedWebSocket | None = None - - async def append_duplex_input_async(self, session_id: str, **kwargs): - assert self.session is not None - assert self.websocket is not None - self.active_response_ids_at_append.append(self.session.active_response_id) - self.response_done_seen_at_append = "response.done" in self.websocket.sent_types() - kwargs.pop("expected_epoch", None) - result = await super().append_duplex_input_async(session_id, **kwargs) - self.append_started.set() - return result - - engine = ModelUnitEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - engine.websocket = ws - create = _native_session_create("sid-auto-overlap-model-unit") - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - try: - for _ in range(100): - session = handler._registry.get("sid-auto-overlap-model-unit") - if session is not None and "session.created" in ws.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("native duplex session did not open") - - engine.session = session - session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value - response_id = session.begin_response() - session.mark_audio_sent(duration_ms=2000) - request_id = handler._native_stage0_request_id(session, session.epoch) - session.bind_request(request_id) - expected_fence = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - - chunk = { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3200, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 200, - "is_speech": True, - } - for _ in range(4): - ws.put(chunk) - await asyncio.sleep(0.05) - assert engine.appended == [] - - ws.put(chunk) - await asyncio.wait_for(engine.append_started.wait(), timeout=0.5) - - ws.put({"type": "input_audio_buffer.commit", "final": True}) - for _ in range(100): - if "input.committed" in ws.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("exact model unit input did not commit") - - engine.append_started.clear() - for _ in range(5): - ws.put(chunk) - await asyncio.wait_for(engine.append_started.wait(), timeout=0.5) - await asyncio.sleep(0.05) - assert "error" not in ws.sent_types() - finally: - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert len(engine.appended) == 2 - for appended_session_id, mode, payload, final in engine.appended: - assert appended_session_id == session.session_id - assert mode == "append_audio_chunk" - assert final is False - assert isinstance(payload, dict) - assert np.frombuffer(base64.b64decode(payload["audio"]), dtype=" None: - nonlocal overlap_sent - if data.get("type") != "response.output_audio.delta" or overlap_sent: - return - overlap_sent = True - chunk = { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3200, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 200, - "is_speech": True, - } - for _ in range(4): - ws.put(chunk) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - - ws = TimedWebSocket(on_send=commit_overlap_after_first_response, receive_timeout_s=3) - create = _native_realtime_session_update(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 1000, - "is_speech": True, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - - handler_task = asyncio.create_task(handler.handle_realtime_session(ws)) - try: - for _ in range(300): - if len(engine.appended) >= 2: - break - await asyncio.sleep(0.01) - else: - pytest.fail( - f"committed overlap was not promoted after response.done: " - f"appends={len(engine.appended)}, events={ws.sent_types()}" - ) - - assert overlap_sent is True - assert len(engine.appended) == 2 - assert engine.appended[1][3] is True - assert ws.sent_types().count("response.created") == 1 - finally: - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=3) - - assert ws.sent_types().count("response.done") >= 1 - - -def test_auto_response_barge_in_does_not_mark_serving_new_user_turn_payload(): - handler, session = _auto_response_context("sid-auto-barge-new-turn", playback_active=True) - payload = _native_audio_payload() - - assert handler._should_force_listen_for_auto_response_overlap(session, {"force_barge_in": True}, payload) is False - assert "new_user_turn" not in payload - assert "new_user_turn_prefix_variant" not in payload - assert session.turn_id == 0 - assert payload.get("force_speak") is not True - - -def test_native_realtime_barge_in_payload_does_not_create_serving_turn_boundary(): - payload: dict[str, object] = { - "type": "audio", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "is_speech": True, - } - session = DuplexSession(session_id="sid-realtime-barge", config=DuplexSessionConfig()) - session.turn_id = 2 - - assert "new_user_turn" not in payload - assert "new_user_turn_prefix_variant" not in payload - assert session.turn_id == 2 - - -def test_auto_response_force_barge_in_has_no_waiting_turn_variant(): - handler, session = _auto_response_context("sid-auto-barge-keeps-variant") - payload = _native_audio_payload() - - assert ( - handler._should_force_listen_for_auto_response_overlap( - session, - {"force_barge_in": True}, - payload, - ) - is False - ) - - -def test_duplex_endpoint_requires_explicit_realtime_configuration(): - assert should_enable_duplex_endpoint(None) is False - assert should_enable_duplex_endpoint([]) is False - assert should_enable_duplex_endpoint([SimpleNamespace(session_mode="turn")]) is False - assert should_enable_duplex_endpoint([{"session_mode": "turn"}]) is False - assert ( - should_enable_duplex_endpoint( - [ - SimpleNamespace(session_mode="turn"), - SimpleNamespace(session_mode="duplex"), - ] - ) - is True - ) - assert should_enable_duplex_endpoint([{"session_mode": "duplex"}]) is True - - -def test_duplex_endpoint_supports_top_level_session_mode(tmp_path): - config_path = tmp_path / "stage_config.yaml" - config_path.write_text( - """ -session_mode: duplex -stage_args: - - stage_id: 0 - engine_args: {} -""", - encoding="utf-8", - ) - - assert should_enable_duplex_endpoint([], config_path=str(config_path)) is True - - -def test_duplex_endpoint_supports_inherited_duplex_session(tmp_path): - base_path = tmp_path / "base.yaml" - base_path.write_text( - """ -session_mode: turn -duplex_session: {} -stages: [] -""", - encoding="utf-8", - ) - config_path = tmp_path / "deploy.yaml" - config_path.write_text("base_config: base.yaml\n", encoding="utf-8") - - assert should_enable_duplex_endpoint([], config_path=str(config_path)) is True - - -def test_duplex_handler_splits_data_plane_audio_list_into_deltas(): - import torch - - data_plane = _test_data_plane() - output = SimpleNamespace( - finished=True, - outputs=[ - SimpleNamespace( - text="hello", - multimodal_output={ - "audio": [ - torch.zeros(10, dtype=torch.float32), - torch.zeros(20, dtype=torch.float32), - ], - "sr": 24000, - }, - ) - ], - ) - - native_results = list(data_plane.project_output(output)) - - assert [result["audio_data"] for result in native_results] == ["wav-10", "wav-20"] - assert [result["text"] for result in native_results] == ["hello", ""] - assert [result["end_of_turn"] for result in native_results] == [False, True] - - -def test_duplex_listen_latent_does_not_poison_cumulative_audio_offset(): - import torch - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - - # A model-listen decision wraps the segment with a latent tensor that is - # NOT reply audio; it must not advance the cumulative audio offset. - listen_output = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[], - multimodal_output={ - "duplex_native_decision": "listen", - "model_listen": True, - "latent": torch.zeros(331776, dtype=torch.float32), - "meta": {"sr": 24000}, - }, - ) - listen_results = list(data_plane.project_output(listen_output)) - assert [result.get("is_listen") for result in listen_results] == [True] - - # The first speak unit carries cumulative stage-1 audio far smaller than - # the listen latent; it must still be delivered from sample 0. - speak_output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[ - SimpleNamespace( - text=" It was a very", - multimodal_output={}, - ) - ], - multimodal_output={ - "audio": torch.zeros(32768, dtype=torch.float32), - "sr": 24000, - }, - ) - speak_results = list(data_plane.project_output(speak_output)) - assert [result.get("audio_data") for result in speak_results] == ["wav-32768"] - assert speak_results[0]["text"] == " It was a very" - - -def test_direct_listen_decision_survives_inner_completion_metadata(): - inner_output = SimpleNamespace( - outputs=[ - SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - ) - ] - ) - decision = MiniCPMO45DuplexRuntimeExtension().decide_output( - stage_id=0, - final_stage_id=1, - segment_finished=True, - segment_token_ids=(151705,), - segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, - output=inner_output, - ) - assert decision is not None - output = attach_duplex_output_decision( - OmniRequestOutput.from_stage_output( - inner_output, - request_id="duplex-direct-listen", - finished=True, - stage_id=0, - final_output_type=decision.final_output_type, - metrics={ - "stage_metrics": { - "0": { - "vllm_ttft_ms": 125.0, - "vllm_tpot_ms": 18.0, - "vllm_itl_ms": 17.5, - "vllm_itls_ms": [17.0, 18.0], - } - } - }, - ), - decision, - ) - - results = list(_test_data_plane().project_output(output)) - - assert len(results) == 1 - assert results[0]["is_listen"] is True - assert results[0]["model_listen"] is True - assert results[0]["listen_source"] == "model_listen" - assert results[0]["stage_metrics"] == output.metrics["stage_metrics"] - - -def test_duplex_segment_text_is_attached_once_across_streaming_batches(): - import torch - - class _ChatService: - model_config = _ModelConfig() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - - def _speak(total_samples: int, text: str, *, finished: bool): - output = SimpleNamespace( - request_id=request_id, - finished=finished, - outputs=[SimpleNamespace(text=text, multimodal_output={})], - multimodal_output={ - "audio": torch.zeros(total_samples, dtype=torch.float32), - "sr": 24000, - }, - ) - return _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - def _audio(results): - return [r for r in results if r.get("is_listen") is False] - - def _listen(results): - return [r for r in results if r.get("is_listen") is True] - - # One talker segment streams several cumulative-audio batches, each - # carrying the SAME segment text; only the first may attach it. - assert [r["text"] for r in _audio(_speak(100, " movie called", finished=False))] == [" movie called"] - assert [r["text"] for r in _audio(_speak(220, " movie called", finished=False))] == [""] - finished_results = _speak(300, " movie called", finished=True) - assert [r["text"] for r in _audio(finished_results)] == [""] - assert _listen(finished_results) == [] - - # Continuation units re-run the talker with the SAME handed text past a - # finished boundary (every engine segment ends finished=True); the text - # must stay suppressed or each continuation duplicates the transcript. - assert [r["text"] for r in _audio(_speak(400, " movie called", finished=False))] == [""] - assert [r["text"] for r in _audio(_speak(450, " movie called", finished=True))] == [""] - - # A genuinely new segment text is attached in full. - assert [r["text"] for r in _audio(_speak(500, " Titanic", finished=True))] == [" Titanic"] - - # Text growing within a segment is delivered as suffix deltas. - assert [r["text"] for r in _audio(_speak(600, " by", finished=False))] == [" by"] - assert [r["text"] for r in _audio(_speak(700, " by James", finished=True))] == [" James"] - - # A segment whose finished batch slices to an EMPTY audio delta (all - # samples already delivered) must not block the next segment's text. - assert [r["text"] for r in _audio(_speak(800, " James Cameron. The", finished=False))] == [" James Cameron. The"] - assert _speak(800, " James Cameron. The", finished=True) == [] - assert [r["text"] for r in _audio(_speak(900, " 997.", finished=False))] == [" 997."] - - -def test_duplex_data_plane_output_prefers_audio_segment_text_metadata(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-segment-text-e0-stage0" - segment_text = "你好，有什么想聊的吗？" - cumulative_text = segment_text * 2 - output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=cumulative_text, multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - "meta.llm_output_text_utf8": np.frombuffer(segment_text.encode("utf-8"), dtype=np.uint8), - }, - ) - - results = list(data_plane.project_output(output)) - - assert len(results) == 1 - assert results[0]["audio_data"] == "wav-24000" - assert results[0]["text"] == segment_text - - next_output_without_segment_metadata = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=segment_text * 3, multimodal_output={})], - multimodal_output={ - "audio": np.zeros(30720, dtype=np.float32), - "sr": 24000, - }, - ) - next_results = list(data_plane.project_output(next_output_without_segment_metadata)) - - assert len(next_results) == 1 - assert next_results[0]["audio_data"] == "wav-6720" - assert next_results[0]["text"] == "" - - -@pytest.mark.parametrize( - ("request_id", "second_samples", "expected_second_audio"), - [ - ("duplex-sid-model-turn-audio-e0-stage0", 36000, "wav-12000"), - ("duplex-sid-model-turn-audio-restart-e0-stage0", 12000, "wav-12000"), - ], - ids=["cumulative-across-model-turns", "upstream-cumulative-audio-restarts"], -) -def test_duplex_audio_cursor_across_model_turns( - request_id: str, - second_samples: int, - expected_second_audio: str, -): - data_plane = _test_data_plane() - - def output(turn_id: int, samples: int): - return SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([turn_id], dtype=np.int32), - }, - ) - - first = list(data_plane.project_output(output(0, 24000))) - second = list(data_plane.project_output(output(1, second_samples))) - - assert first[0]["audio_data"] == "wav-24000" - assert second[0]["audio_data"] == expected_second_audio - - -def test_duplex_data_plane_accepts_model_turn_output_across_client_commits(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-auto-stale-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.commit_native_audio_input() - session.commit_native_audio_input() - session.turn_id = 1 - session.epoch = 0 - request_id = "duplex-sid-auto-stale-turn-stage0" - output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text="old tail", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([1], dtype=np.int32), - "meta.llm_output_text_utf8": np.frombuffer("旧尾巴".encode(), dtype=np.uint8), - }, - ) - - results = list(data_plane.project_output(output, context=_data_plane_context(session))) - - assert len(results) == 1 - assert results[0]["audio_data"] == "wav-24000" - assert results[0]["model_turn_id"] == 1 - - -def test_duplex_data_plane_accepts_active_response_turn_identity(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-active-response-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.turn_id = 1 - session.epoch = 0 - session.begin_response(turn_id=0) - output = SimpleNamespace( - request_id="duplex-sid-active-response-turn-stage0", - finished=True, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(10, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([0], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - - -def test_duplex_data_plane_listen_preserves_model_turn_identity(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-listen-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - output = SimpleNamespace( - request_id="duplex-sid-listen-turn-stage0", - finished=True, - outputs=[SimpleNamespace(text="", token_ids=[], multimodal_output={})], - multimodal_output={ - "duplex_native_decision": "listen", - "meta.duplex_turn_id": np.array([2], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["is_listen"] is True - assert results[0]["model_turn_id"] == 2 - - -def test_duplex_data_plane_does_not_drop_future_model_turn_while_response_active(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-future-response-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.turn_id = 0 - session.epoch = 0 - session.begin_response(turn_id=0) - output = SimpleNamespace( - request_id="duplex-sid-future-response-turn-stage0", - finished=False, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(10, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([1], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["model_turn_id"] == 1 - assert results[0]["audio_data"] == "wav-10" - - -def _duplex_tts_output( - *, - request_id: str, - samples: int, - finished: bool, - text: str = "hello", - tts_is_last_chunk: bool = False, - turn_end: bool = False, - turn_id: int = 0, - token_ids: list[int] | None = None, -): - return SimpleNamespace( - request_id=request_id, - finished=finished, - outputs=[SimpleNamespace(text=text, token_ids=list(token_ids or []), multimodal_output={})], - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.tts_is_last_chunk": np.array([int(tts_is_last_chunk)], dtype=np.int32), - "meta.turn_end": np.array([int(turn_end)], dtype=np.int32), - "meta.duplex_turn_id": np.array([turn_id], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - -@pytest.mark.parametrize( - ( - "finished", - "tts_is_last_chunk", - "turn_end", - "consume_audio_first", - "expected_end_of_turn", - "expected_audio", - ), - [ - (False, True, False, False, False, "wav-24000"), - (True, True, False, True, False, ""), - (True, True, True, True, True, ""), - ], - ids=[ - "last-audio-batch-does-not-end-response", - "scheduler-eos-only-ends-tts-segment", - "explicit-turn-end-metadata-ends-response", - ], -) -def test_duplex_auto_response_tts_boundary( - finished: bool, - tts_is_last_chunk: bool, - turn_end: bool, - consume_audio_first: bool, - expected_end_of_turn: bool, - expected_audio: str, -): - data_plane = _test_data_plane() - request_id = "duplex-sid-tts-boundary-e0-stage0" - session = DuplexSession( - session_id="sid-tts-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_response_turn(0) - if consume_audio_first: - data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) - output = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=finished, - tts_is_last_chunk=tts_is_last_chunk, - turn_end=turn_end, - token_ids=[151645] if finished else None, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - boundary = results[0] - assert boundary["stage_role"] == "tts" - assert boundary["end_of_turn"] is expected_end_of_turn - assert boundary["abort_data_plane_request"] is True - assert boundary["audio_data"] == expected_audio - - -def test_duplex_auto_response_tts_scheduler_eos_fallback(): - data_plane = _test_data_plane() - request_id = "duplex-sid-tts-eos-no-profile-e0-stage0" - session = DuplexSession( - session_id="sid-tts-eos-no-profile", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_response_turn(0) - data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) - output = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=True, - tts_is_last_chunk=False, - token_ids=[151645], - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["stage_role"] == "tts" - assert results[0]["abort_data_plane_request"] is True - - -def test_duplex_auto_response_rearms_tts_boundary_for_next_segment_in_same_turn(): - data_plane = _test_data_plane() - request_id = "duplex-sid-next-segment-same-turn-e0-stage0" - session = DuplexSession( - session_id="sid-next-segment-same-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_response_turn(0) - data_plane.begin_request(request_id) - - first_segment = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=False, - tts_is_last_chunk=True, - turn_id=0, - ) - first_results = _project_data_plane( - data_plane, - {"data_plane_outputs": [first_segment]}, - session=session, - ) - - assert len(first_results) == 1 - assert first_results[0]["abort_data_plane_request"] is True - - data_plane.begin_request(request_id) - second_segment = _duplex_tts_output( - request_id=request_id, - samples=48000, - finished=False, - text="hello again", - tts_is_last_chunk=True, - turn_id=0, - ) - second_results = _project_data_plane( - data_plane, - {"data_plane_outputs": [second_segment]}, - session=session, - ) - - assert len(second_results) == 1 - assert second_results[0]["audio_data"] == "wav-24000" - assert second_results[0]["abort_data_plane_request"] is True - - -def test_duplex_turn_end_is_not_swallowed_by_finished_segment_fallback(): - data_plane = _test_data_plane() - request_id = "duplex-sid-turn-end-after-segment-e0-stage0" - session = DuplexSession( - session_id="sid-turn-end-after-segment", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response(turn_id=0) - session.bind_request(request_id) - _project_data_plane( - data_plane, - { - "data_plane_outputs": [ - _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=False, - tts_is_last_chunk=True, - ) - ] - }, - session=session, - ) - output = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=True, - text="", - tts_is_last_chunk=False, - turn_end=True, - token_ids=[151645], - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["stage_role"] == "tts" - assert results[0]["end_of_turn"] is True - - duplicate = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert not any(result.get("end_of_turn") is True for result in duplicate) - - -def test_duplex_auto_response_discards_terminal_only_audio_before_response_creation(): - data_plane = _test_data_plane() - request_id = "duplex-sid-terminal-only-audio-e0-stage0" - session = DuplexSession( - session_id="sid-terminal-only-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_request(request_id) - first_audio = _duplex_tts_output( - request_id=request_id, - samples=3840, - finished=False, - text="", - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [first_audio]}, session=session) == [] - - terminal = _duplex_tts_output( - request_id=request_id, - samples=3840, - finished=True, - text="", - tts_is_last_chunk=True, - turn_end=True, - token_ids=[151645], - ) - results = _project_data_plane(data_plane, {"data_plane_outputs": [terminal]}, session=session) - - assert len(results) == 1 - assert results[0]["audio_data"] == "" - assert results[0]["end_of_turn"] is True - assert session.active_request_id == request_id - assert not data_plane.is_terminal(request_id) - - -def test_duplex_auto_response_releases_buffered_audio_when_transcript_arrives(): - data_plane = _test_data_plane() - request_id = "duplex-sid-delayed-transcript-e0-stage0" - session = DuplexSession( - session_id="sid-delayed-transcript", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - audio_before_text = _duplex_tts_output( - request_id=request_id, - samples=100, - finished=False, - text="", - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [audio_before_text]}, session=session) == [] - - audio_with_text = _duplex_tts_output( - request_id=request_id, - samples=200, - finished=False, - text="hello", - ) - results = _project_data_plane(data_plane, {"data_plane_outputs": [audio_with_text]}, session=session) - - assert [result["audio_data"] for result in results] == ["wav-100", "wav-100"] - assert [result["text"] for result in results] == ["", "hello"] - - -def test_duplex_auto_response_releases_buffered_audio_on_text_only_terminal(): - data_plane = _test_data_plane() - request_id = "duplex-sid-delayed-terminal-text-e0-stage0" - session = DuplexSession( - session_id="sid-delayed-terminal-text", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - audio_before_text = _duplex_tts_output( - request_id=request_id, - samples=100, - finished=False, - text="", - turn_id=1, - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [audio_before_text]}, session=session) == [] - - textless_turn_terminal = _duplex_tts_output( - request_id=request_id, - samples=100, - finished=True, - text="", - tts_is_last_chunk=True, - turn_end=True, - turn_id=1, - token_ids=[151645], - ) - terminal_results = list( - data_plane.project( - {"data_plane_outputs": [textless_turn_terminal]}, - context=_data_plane_context(session), - ) - ) - - assert len(terminal_results) == 1 - assert terminal_results[0]["audio_data"] == "" - assert terminal_results[0]["end_of_turn"] is True - - text_only_terminal = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[SimpleNamespace(text="hello", token_ids=[151645], multimodal_output={})], - multimodal_output={ - "sr": 24000, - "meta.tts_is_last_chunk": np.array([1], dtype=np.int32), - "meta.turn_end": np.array([1], dtype=np.int32), - "meta.duplex_turn_id": np.array([2], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - results = _project_data_plane(data_plane, {"data_plane_outputs": [text_only_terminal]}, session=session) - - assert len(results) == 1 - assert results[0]["audio_data"] == "wav-100" - assert results[0]["text"] == "hello" - assert results[0]["end_of_turn"] is True - assert results[0]["abort_data_plane_request"] is True - assert not data_plane.has_pending_audio(request_id) - - -@pytest.mark.asyncio -async def test_native_append_propagates_current_turn_fence_to_engine(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-fenced-append", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.turn_id = 2 - session.begin_response(turn_id=2) - ws = TimedWebSocket() - - await handler._append_runtime_input( - session, - {"duplex_turn_id": 2, "audio": "", "format": "pcm_f32le"}, - final=True, - send_json=ws.send_json, - mode="append_audio_chunk", - expected_epoch=0, - ) - - assert engine.appended_fences == [DuplexFence("sid-fenced-append", epoch=0, turn_id=2)] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_tts_segment_boundary_appends_silence_unit(): - request_id = "duplex-sid-segment-boundary-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-segment-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - session.mark_audio_sent(100) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "", - "audio_data": "", - "audio_format": "pcm16", - "end_of_turn": False, - "abort_data_plane_request": True, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert payload["format"] == "pcm_f32le" - assert payload["duplex_turn_id"] == 0 - assert "force_speak" not in payload - assert final is False - assert session.active_response_id is not None - assert "response.done" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_pre_response_tts_boundary_continues_model_turn(): - request_id = "duplex-sid-pre-response-boundary-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-pre-response-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "model_turn_id": session.turn_id, - "text": "", - "audio_data": "", - "audio_format": "pcm16", - "end_of_turn": False, - "abort_data_plane_request": True, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert close_reason is None - assert emitted is False - assert session.active_response_id is None - assert ws.sent == [] - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert payload["duplex_turn_id"] == session.turn_id - assert final is False - - -@pytest.mark.asyncio -async def test_minicpmo_pre_response_continuation_drops_after_model_turn_ends(): - request_id = "duplex-sid-stale-pre-response-boundary-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-stale-pre-response-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - model_turn_id = session.turn_id - - native = handler._runtime_session_state(session) - - async def _stale_before_append(payload: object, **kwargs: Any) -> bool: - session.complete_model_turn(model_turn_id) - if handler._native_silence_continuation_is_stale( - session, - request_id=kwargs["request_id"], - response_id=kwargs["response_id"], - response_owned=kwargs["response_owned"], - expected_epoch=kwargs["expected_epoch"], - expected_incarnation=kwargs["expected_incarnation"], - expected_model_turn_id=kwargs["expected_model_turn_id"], - ): - return False - append_ok, _ = await handler._append_runtime_input( - session, - payload, - final=False, - send_json=kwargs["send_json"], - mode="append_audio_chunk", - expected_epoch=kwargs["expected_epoch"], - ) - return append_ok - - native.silence_continuation_scheduler = _stale_before_append - - await handler._maybe_continue_native_response( - TimedWebSocket().send_json, - session=session, - expected_epoch=session.epoch, - expected_model_turn_id=model_turn_id, - ) - await asyncio.sleep(0.01) - - assert engine.appended == [] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_continuation_stops_at_large_safety_boundary(): - request_id = "duplex-sid-long-response-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-long-response", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - session.mark_audio_sent(100) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - for _ in range(handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS): - await handler._maybe_continue_native_response( - ws.send_json, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.02) - - response_id = session.active_response_id - assert response_id is not None - assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - assert all(payload["duplex_turn_id"] == 0 for _, _, payload, _ in engine.appended) - assert all(not {"force_speak", "force_listen"} & payload.keys() for _, _, payload, _ in engine.appended[:-1]) - assert engine.appended[-1][2]["force_listen"] is True - assert handler._native_response_continuations_remaining(session, response_id) is False - - await handler._maybe_continue_native_response( - ws.send_json, - session=session, - expected_epoch=session.epoch, - ) - - assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - assert ws.sent_types()[-2:] == ["response.listen", "response.done"] - assert session.active_response_id is None - assert session.active_request_id == request_id - assert session.turn_id == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_boundary_listen_closes_response(): - request_id = "duplex-sid-boundary-listen-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-boundary-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response(turn_id=0) - session.bind_request(request_id) - native = handler._runtime_session_state(session) - native.continuation_owner_id = f"response:{response_id}" - native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert engine.appended == [] - assert ws.sent_types() == ["response.listen", "response.done"] - # The listen terminates the precreated response: clients demultiplex the - # decision by its id, so the payload must carry the response identity. - listen_payload, done_payload = ws.sent[-2], ws.sent[-1] - assert listen_payload["response_id"] == response_id - assert done_payload["response_id"] == response_id - assert session.active_response_id is None - assert session.active_request_id == request_id - assert session.turn_id == 1 - assert native.continuation_owner_id is None - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_boundary_allows_next_model_turn_response(): - request_id = "duplex-sid-boundary-next-turn-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-boundary-next-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.bind_request(request_id) - native = handler._runtime_session_state(session) - native.continuation_owner_id = f"response:{session.active_response_id}" - native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - ws = TimedWebSocket() - - await handler._maybe_continue_native_response( - ws.send_json, - session=session, - expected_epoch=session.epoch, - ) - - assert session.active_response_id is None - assert session.turn_id == 1 - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "next response", - "audio_data": "next-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 1, - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert ws.sent_types() == [ - "response.listen", - "response.done", - "response.created", - "response.speak", - "response.output_audio.delta", - ] - assert session.active_response_turn_id == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_boundary_does_not_close_new_epoch_response(): - request_id = "duplex-sid-boundary-race-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-boundary-race", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - old_response_id = session.begin_response(turn_id=0) - session.bind_request(request_id) - native = handler._runtime_session_state(session) - native.continuation_owner_id = f"response:{old_response_id}" - native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - sent: list[dict[str, Any]] = [] - new_response_ids: list[str] = [] - - async def _barge_in_on_listen(payload: dict[str, Any]) -> None: - sent.append(payload) - if payload["type"] == "response.listen": - session.barge_in() - session.bind_request("duplex-sid-boundary-race-e1-stage0") - new_response_ids.append(session.begin_response(turn_id=1)) - - await handler._maybe_continue_native_response( - _barge_in_on_listen, - session=session, - expected_epoch=session.epoch, - ) - - assert [payload["type"] for payload in sent] == ["response.listen"] - assert len(new_response_ids) == 1 - assert session.active_response_id == new_response_ids[0] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_pre_speak_listen_continues_same_response(): - request_id = "duplex-sid-pre-speak-listen-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-pre-speak-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert len(engine.appended) == 1 - assert session.active_response_id == response_id - assert session.active_request_id == request_id - assert "response.listen" not in ws.sent_types() - assert "response.done" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_listen_before_response_keeps_resumable_request(): - request_id = "duplex-sid-listen-before-response-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-listen-before-response", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert session.active_request_id == request_id - assert not handler._serving_runtime_adapter.data_plane.is_terminal(request_id) - assert ws.sent_types().count("response.listen") == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_listen_without_response_does_not_defer_commit(): - session_id = "sid-listen-before-response-commit" - request_id = duplex_resource_request_id(DuplexFence(session_id), "stage0") - inner_output = SimpleNamespace( - outputs=[ - SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - ) - ] - ) - decision = MiniCPMO45DuplexRuntimeExtension().decide_output( - stage_id=0, - final_stage_id=1, - segment_finished=True, - segment_token_ids=(151705,), - segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, - output=inner_output, - ) - assert decision is not None - listen_output = attach_duplex_output_decision( - OmniRequestOutput.from_stage_output( - inner_output, - request_id=request_id, - finished=True, - stage_id=0, - final_output_type=decision.final_output_type, - ), - decision, - ) - append_result = { - "operation": "append", - "session_id": session_id, - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [listen_output], - } - engine = FakeEngineClient(append_result=append_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - listen_count = 0 - - def append_residual_after_listen(ws: TimedWebSocket, data: dict[str, Any]) -> None: - nonlocal listen_count - if data.get("type") == "response.listen": - listen_count += 1 - if listen_count == 1: - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3_200), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - "duration_ms": 200, - "is_speech": True, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - elif data.get("type") == "input_audio_buffer.committed": - asyncio.get_running_loop().call_later(0.05, ws.put, {"type": "session.close"}) - - ws = TimedWebSocket(on_send=append_residual_after_listen, receive_timeout_s=2) - create = _native_realtime_session_update(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16_000), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - "duration_ms": 1_000, - "is_speech": True, - } - ) - - await handler.handle_realtime_session(ws) - - committed = [event for event in ws.sent if event.get("type") == "input_audio_buffer.committed"] - assert len(engine.appended) == 2 - assert engine.appended[1][3] is True - assert len(committed) == 1 - assert committed[0].get("event", {}).get("overlap_deferred") is not True - assert committed[0].get("event", {}).get("response_create_deferred") is not True - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_pcm_f32_residual_commit_releases_pending_input_bytes(): - session_id = "sid-pcm-residual-byte-accounting" - control_result = { - "operation": "append", - "session_id": session_id, - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [], - } - engine = FakeEngineClient(control_result=control_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=2) - create = _native_realtime_session_update(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - handler_task = asyncio.create_task(handler.handle_realtime_session(ws)) - - try: - for _ in range(100): - session = handler._registry.get(session_id) - if session is not None and "session.created" in ws.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("native duplex session did not open") - - pending_bytes_after_turn: list[int] = [] - for turn_index in range(3): - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(19_200), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - "duration_ms": 1_200, - "is_speech": True, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - expected_appends = (turn_index + 1) * 2 - for _ in range(100): - if len(engine.appended) >= expected_appends: - await asyncio.sleep(0) - break - await asyncio.sleep(0.01) - else: - pytest.fail(f"turn {turn_index} residual append did not complete") - pending_bytes_after_turn.append(session.pending_input_bytes) - finally: - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert pending_bytes_after_turn == [0, 0, 0] - assert not any( - event.get("type") == "error" - and isinstance(event.get("error"), dict) - and event["error"].get("code") == "input_backpressure" - for event in ws.sent - ) - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_post_speak_listen_continues_same_response(): - request_id = "duplex-sid-post-speak-listen-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-post-speak-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - session.mark_audio_sent(100) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert len(engine.appended) == 1 - assert session.active_response_id == response_id - assert session.active_request_id == request_id - assert "response.listen" not in ws.sent_types() - assert "response.done" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_turn_end_preserves_resumable_request(): - request_id = "duplex-sid-turn-end-preserve-request-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-turn-end-preserve-request", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.bind_request(request_id) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "done", - "audio_data": "audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": True, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert session.active_response_id is None - assert session.active_request_id == request_id - assert "response.done" in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_empty_turn_end_emits_model_listen(): - request_id = "duplex-sid-empty-turn-end-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-empty-turn-end", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "", - "audio_data": "", - "audio_format": "pcm16", - "end_of_turn": True, - "abort_data_plane_request": True, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert session.turn_id == 1 - assert session.active_request_id == request_id - assert session.active_response_id is None - assert ws.sent_types() == ["response.listen"] - assert ws.sent[0]["model_listen"] is True - assert ws.sent[0]["reason"] == "model_turn_completed_without_output" - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_drops_late_audio_from_completed_model_turn(): - request_id = "duplex-sid-late-completed-turn-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-late-completed-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - session.complete_model_turn(0) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "late duplicate", - "audio_data": "late-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": True, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is False - assert session.turn_id == 1 - assert session.active_response_id is None - assert "response.created" not in ws.sent_types() - assert "response.done" not in ws.sent_types() - - -def test_duplex_data_plane_drops_stale_epoch_audio_after_barge_in(): - class _ChatService: - model_config = _ModelConfig() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-auto-stale-epoch", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.turn_id = 1 - session.epoch = 2 - output = SimpleNamespace( - request_id="duplex-sid-auto-stale-epoch-stage0", - finished=False, - outputs=[SimpleNamespace(text="old epoch", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - "meta.duplex_epoch": np.array([1], dtype=np.int32), - "meta.duplex_turn_id": np.array([1], dtype=np.int32), - "meta.llm_output_text_utf8": np.frombuffer("旧轮".encode(), dtype=np.uint8), - }, - ) - - assert list(data_plane.project_output(output, context=_data_plane_context(session))) == [] - - -def test_duplex_data_plane_text_delta_preserves_repeated_suffix_growth(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-repeated-word-e0-stage0" - - assert data_plane.segment_text_delta(request_id, "好的") == "好的" - assert data_plane.segment_text_delta(request_id, "好的好的") == "好的" - - -def test_duplex_data_plane_text_delta_appends_distinct_non_prefix_segments(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-rewrite-e0-stage0" - - deltas = [ - data_plane.segment_text_delta(request_id, "It's Canberra."), - data_plane.segment_text_delta(request_id, " Next question."), - ] - - assert deltas == ["It's Canberra.", " Next question."] - assert "".join(deltas) == "It's Canberra. Next question." - - -def test_duplex_data_plane_close_session_removes_encoded_request_state(): - data_plane = _test_data_plane() - target_request_id = duplex_resource_request_id( - DuplexFence("sid-to-close", incarnation=1, epoch=2), - "stage0", - ) - other_request_id = duplex_resource_request_id( - DuplexFence("sid-to-keep", incarnation=1, epoch=2), - "stage0", - ) - data_plane.begin_request(target_request_id) - data_plane.begin_request(other_request_id) - - data_plane.close_session("sid-to-close") - - assert data_plane.has_request(target_request_id) is False - assert data_plane.has_request(other_request_id) is True - - -def test_duplex_auto_response_segment_complete_keeps_data_plane_request_open(): - import numpy as np - - class _ChatService: - model_config = _ModelConfig() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response() - session.mark_audio_sent(1280) - data_plane.slice_cumulative_audio(request_id, np.zeros(30720, dtype=np.float32)) - - output = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[SimpleNamespace(text="hello", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(30720, dtype=np.float32), - "sr": 24000, - }, - ) - - native_results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(native_results) == 1 - assert native_results[0]["is_listen"] is True - assert native_results[0]["reason"] == "auto_response_segment_complete" - assert native_results[0]["data_plane_request_id"] == request_id - assert "abort_data_plane_request" not in native_results[0] - - -def test_duplex_auto_response_prior_playback_does_not_abort_before_current_response(): - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond-next-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_audio_sent(1280) - session.acknowledge_playback(played_ms=1280, committed_ms=1280) - - output = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={}, - ) - - native_results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert native_results == [] - - -def test_duplex_auto_response_text_only_flush_does_not_consume_audio_transcript(): - import numpy as np - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond-delayed-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - text = "你好呀，有什莫问的吗？" - - text_only_output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=text, multimodal_output={})], - multimodal_output={"sr": 24000}, - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [text_only_output]}, session=session) == [] - - audio_output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=text, multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - }, - ) - - native_results = _project_data_plane(data_plane, {"data_plane_outputs": [audio_output]}, session=session) - - assert len(native_results) == 1 - assert native_results[0]["audio_data"].startswith("wav-") - assert native_results[0]["text"] == text - - -@pytest.mark.asyncio -async def test_duplex_auto_response_empty_terminal_does_not_create_empty_response(): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - request_id = "duplex-sid-native-empty-terminal-e0-stage0" - session = DuplexSession( - session_id="sid-native-empty-terminal", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_audio_sent(1280) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text="hello", multimodal_output={})], - multimodal_output={ - "sr": 24000, - "meta": {"turn_end": True, "duplex_turn_id": 0}, - }, - ) - - native_results = list( - handler._serving_runtime_adapter.data_plane.project( - {"data_plane_outputs": [output]}, - context=_data_plane_context(session), - ) - ) - assert len(native_results) == 1 - assert native_results[0]["end_of_turn"] is True - assert native_results[0]["model_turn_id"] == 0 - - for native_result in native_results: - await handler._send_one_native_duplex_event(send_json, native_result, session=session) - - assert "response.created" not in [m.get("type") for m in sent] - assert "response.done" not in [m.get("type") for m in sent] - assert session.turn_id == 1 - assert not handler._serving_runtime_adapter.data_plane.is_terminal(request_id) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_session_close_cleans_auto_response_data_plane_state(): - session_id = "sid-native-cleanup" - request_id = f"duplex-{session_id}-e0-stage0" - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - handler._serving_runtime_adapter.data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) - handler._serving_runtime_adapter.data_plane.segment_text_delta(request_id, "hello") - handler._serving_runtime_adapter.data_plane.mark_terminal(request_id) - native = MiniCPMO45ServingSessionState() - handler._serving_runtime_adapter.session_states[session_id] = native - - event = _native_session_create(session_id) - event["session"]["extra_body"]["auto_response"] = True - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert not handler._serving_runtime_adapter.data_plane.has_request(request_id) - assert session_id not in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - session = DuplexSession(session_id="sid-chat-audio", config=DuplexSessionConfig()) - response_id = session.begin_response() - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._emit_chat_payload( - session, - { - "modality": "audio", - "choices": [ - { - "delta": { - "content": "AAAA", - }, - "audio_metadata": { - "format": "wav", - "sample_rate_hz": 22_050, - "frame_count": 2205, - "channels": 1, - }, - } - ], - }, - session.epoch, - response_id, - send_json, - ) - - assert sent == [ - { - "type": "response.output_audio.delta", - "session_id": "sid-chat-audio", - "response_id": response_id, - "epoch": 0, - "audio": "AAAA", - "format": "wav", - "sample_rate_hz": 22_050, - "channels": 1, - "audio_duration_ms": 100, - } - ] - assert session.playback.generated_ms == 100 - assert session.playback.sent_ms == 100 - assert session.playback.played_ms == 0 - - -@pytest.mark.parametrize("response_format", ["pcm", "wav"]) -def test_duplex_chat_request_uses_response_audio_format(response_format: str): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - extra_audio = {"format": "flac", "voice": "alloy"} - session = DuplexSession( - session_id="sid-chat-audio-format", - config=DuplexSessionConfig(response_format=response_format, extra_body={"audio": extra_audio}), - ) - session.begin_response() - - request = handler._build_chat_request(session, "format-request") - - assert request.audio == {"format": response_format, "voice": "alloy"} - assert extra_audio == {"format": "flac", "voice": "alloy"} - - -def _chat_audio_completion_payload(*, frame_count: int = 2205, sample_rate_hz: int = 22_050) -> dict[str, object]: - return { - "modality": "audio", - "choices": [ - { - "delta": {"content": "AAAA"}, - "audio_metadata": { - "format": "pcm", - "sample_rate_hz": sample_rate_hz, - "frame_count": frame_count, - "channels": 1, - }, - } - ], - } - - -@pytest.mark.asyncio -async def test_duplex_chat_full_audio_message_preserves_actual_format_and_duration(): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-full-audio", config=DuplexSessionConfig(response_format="wav")) - response_id = session.begin_response() - sent: list[dict[str, object]] = [] - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._emit_chat_payload( - session, - { - "choices": [ - { - "message": {"role": "assistant", "content": "hello", "audio": {"data": "AAAA"}}, - "audio_metadata": { - "format": "pcm", - "sample_rate_hz": 22_050, - "frame_count": 2205, - "channels": 1, - }, - } - ], - }, - session.epoch, - response_id, - send_json, - ) - - assert [event["type"] for event in sent] == ["response.text.delta", "response.output_audio.delta"] - assert sent[1]["audio"] == "AAAA" - assert sent[1]["format"] == "pcm" - assert sent[1]["sample_rate_hz"] == 22_050 - assert sent[1]["audio_duration_ms"] == 100 - assert session.assistant_text_buffer == ("hello",) - assert session.playback.sent_ms == 100 - - -@pytest.mark.asyncio -async def test_duplex_chat_audio_does_not_round_each_fragment(): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-sub-ms", config=DuplexSessionConfig()) - response_id = session.begin_response() - sent: list[dict[str, object]] = [] - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - for _ in range(23): - await handler._emit_chat_payload( - session, - _chat_audio_completion_payload(frame_count=1), - session.epoch, - response_id, - send_json, - ) - - assert [event["audio_duration_ms"] for event in sent] == [0] * 22 + [1] - assert session.playback.generated_ms == 1 - assert session.playback.sent_ms == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("send_error", [None, RuntimeError("output unavailable")], ids=["rejected", "failed"]) -async def test_duplex_chat_unaccepted_audio_is_generated_but_not_sent(send_error: RuntimeError | None): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-unaccepted-audio", config=DuplexSessionConfig()) - response_id = session.begin_response() - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - if send_error is not None: - raise send_error - - with pytest.raises(RuntimeError, match="output unavailable") if send_error is not None else nullcontext(): - await handler._emit_chat_payload( - session, _chat_audio_completion_payload(), session.epoch, response_id, send_json - ) - - assert session.playback.generated_ms == 100 - assert session.playback.sent_ms == 0 - assert session.playback.played_ms == 0 - - -@pytest.mark.asyncio -async def test_duplex_chat_late_send_completion_only_updates_old_response(): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-late-audio", config=DuplexSessionConfig()) - old_response = session.begin_response() - send_started = asyncio.Event() - finish_send = asyncio.Event() - - async def delayed_send(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - send_started.set() - await finish_send.wait() - assert on_accepted is not None - on_accepted() - - emit_task = asyncio.create_task( - handler._emit_chat_payload(session, _chat_audio_completion_payload(), session.epoch, old_response, delayed_send) - ) - try: - await asyncio.wait_for(send_started.wait(), timeout=1) - session.barge_in() - new_response = session.begin_response() - current_state = session.turn_state - finish_send.set() - await asyncio.wait_for(emit_task, timeout=1) - - assert session.active_response_id == new_response - assert session.playback_for_response(old_response).sent_ms == 100 - assert session.playback.generated_ms == 0 - assert session.playback.sent_ms == 0 - assert session.turn_state == current_state - finally: - emit_task.cancel() - await asyncio.gather(emit_task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_duplex_chat_missing_audio_metadata_reports_error(): - class MissingMetadataChatService(TurnBasedFakeChatService): - async def create_chat_completion(self, request, raw_request=None): - async def generate(): - payload = {"modality": "audio", "choices": [{"delta": {"content": "AAAA"}}]} - yield f"data: {json.dumps(payload)}\n\n" - - return generate() - - handler = OmniDuplexSessionHandler( - chat_service=MissingMetadataChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-missing-metadata", config=DuplexSessionConfig()) - sent: list[dict[str, object]] = [] - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._run_response(session, send_json) - - assert [event["type"] for event in sent] == ["response.created", "error"] - assert sent[-1]["code"] == "response_error" - assert "AudioChunkMetadata" in sent[-1]["error"] - assert session.playback.generated_ms == 0 - assert session.playback.sent_ms == 0 - - -@pytest.mark.asyncio -async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): - class SnapshotChatService(FakeChatService): - async def create_chat_completion(self, request, raw_request=None): - del raw_request - self.seen_request_ids.append(request.request_id) - snapshots = [ - {"num_tokens_out": 1, "vllm_tpot_ms": 0.0, "vllm_itls_ms": []}, - {"num_tokens_out": 2, "vllm_tpot_ms": 10.0, "vllm_itls_ms": [10.0]}, - {"num_tokens_out": 3, "vllm_tpot_ms": 11.0, "vllm_itls_ms": [10.0, 12.0]}, - ] - - async def generate(): - for index, snapshot in enumerate(snapshots): - payload = { - "modality": "audio", - "choices": [ - { - "delta": {"content": f"audio-{index}"}, - "audio_metadata": { - "format": "pcm", - "sample_rate_hz": 24_000, - "frame_count": 2400, - "channels": 1, - }, - } - ], - "metrics": {"stage_metrics": {"0": snapshot}}, - } - yield f"data: {json.dumps(payload)}\n\n" - - return generate() - - handler = OmniDuplexSessionHandler( - chat_service=SnapshotChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-chat-snapshots", - config=DuplexSessionConfig(instructions="speak"), - ) - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._run_response(session, send_json) - - audio_deltas = [event for event in sent if event.get("type") == "response.output_audio.delta"] - done = next(event for event in sent if event.get("type") == "response.done") - assert [event["audio_duration_ms"] for event in audio_deltas] == [100, 200, 300] - assert done["playback"] == { - "generated_ms": 300, - "sent_ms": 300, - "played_ms": 0, - "committed_ms": 0, - } - assert all("vllm_omni" not in event for event in audio_deltas) - assert done["vllm_omni"]["stage_metrics"]["0"] == { - "num_tokens_out": 3, - "vllm_tpot_ms": 11.0, - "vllm_itls_ms": [10.0, 12.0], - } - - -@pytest.mark.asyncio -async def test_duplex_chat_rejects_unknown_output_modality(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-unknown", config=DuplexSessionConfig()) - response_id = session.begin_response() - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any]) -> None: - sent.append(data) - - await handler._emit_chat_payload( - session, - { - "modality": "video", - "choices": [{"delta": {"content": "not-text"}}], - }, - session.epoch, - response_id, - send_json, - ) - - assert sent == [ - { - "type": "error", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "code": "unsupported_response_modality", - "error": "Unsupported chat response modality: video", - } - ] - assert session.assistant_text_buffer == () - - -@pytest.mark.asyncio -async def test_minicpmo_model_name_does_not_auto_enable_experimental_native_duplex(): - event = _session_create("sid-minicpmo-default-protocol") - event["session"]["model"] = "openbmb/MiniCPM-o-4_5" - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - created = next(m for m in ws.sent if m.get("type") == "session.created") - capabilities = created["session"]["capabilities"] - assert capabilities["implementation_level"] == "serving_session_adapter" - assert capabilities["supports_input_append"] is False - - -@pytest.mark.asyncio -async def test_duplex_handler_aborts_current_chat_request_id_on_barge_in(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.created": - ws.put({"type": "input.cancel", "reason": "test_barge_in"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_session_create("sid-a")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - assert ws.accepted - assert "response.created" in ws.sent_types() - assert "audio.cancelled" in ws.sent_types() - assert chat_service.seen_request_ids == ["duplex-sid-a-0-1"] - assert engine.aborted == ["chatcmpl-duplex-sid-a-0-1"] - assert engine.internal_abort_batches == [["chatcmpl-duplex-sid-a-0-1"]] - - -@pytest.mark.asyncio -async def test_duplex_handler_runtime_open_failure_is_reported_to_client(): - engine = FakeEngineClient(fail_open=True) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-open-fail")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "runtime_open_failed" - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_duplex_handler_preserves_typed_admission_error(): - class AdmissionRejectingEngine(FakeEngineClient): - async def open_duplex_session_async(self, session_id: str, **kwargs): - del kwargs - raise DuplexControlRequestError( - { - "operation": "open", - "session_id": session_id, - "error": { - "code": "resource_exhausted", - "message": "duplex session capacity exhausted", - "retryable": True, - }, - "accepted_fence": None, - "lease_generation": None, - } - ) - - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(AdmissionRejectingEngine()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-admission-rejected")) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "resource_exhausted" - assert ws.sent[0]["retryable"] is True - - -@pytest.mark.asyncio -async def test_duplex_cancel_reports_playback_committed_cursor(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.created": - ws.put({"type": "playback.ack", "played_ms": 1200, "committed_ms": 1000}) - ws.put({"type": "input.cancel", "reason": "test_barge_in"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_session_create("sid-playback")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - ack = next(m for m in ws.sent if m.get("type") == "playback.acknowledged") - cancelled = next(m for m in ws.sent if m.get("type") == "audio.cancelled") - assert ack["playback"]["played_ms"] == 1200 - assert ack["playback"]["committed_ms"] == 1000 - assert cancelled["committed_ms"] == 1000 - assert cancelled["epoch"] == 1 - - -@pytest.mark.asyncio -async def test_playback_ack_response_id_selects_matching_pending_history_item(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-playback-response-id", - config=DuplexSessionConfig(), - ) - first_item_id = "item_resp-first" - second_item_id = "item_resp-second" - session.stage_pending_history_item( - first_item_id, - {"role": "assistant", "content": "first response"}, - ) - session.stage_pending_history_item( - second_item_id, - {"role": "assistant", "content": "second response"}, - ) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": "resp-second", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["item_id"] == second_item_id - assert ack["history_committed"] is True - assert first_item_id in session.pending_history_item_ids - assert second_item_id not in session.pending_history_item_ids - - -@pytest.mark.asyncio -async def test_playback_ack_recovers_missing_response_history_registration(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-playback-missing-registration", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - response_id = session.begin_response() - session.append_assistant_text("response text") - session.mark_audio_sent(1000, text_chars=len("response text")) - session.end_response(commit_text=True, preserve_request=True) - assert not session.history_item_ids - assert f"item_{response_id}" in session.pending_history_item_ids - - ws = TimedWebSocket() - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": f"item_{response_id}", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["history_committed"] is True - assert f"item_{response_id}" in session.history_item_ids - assert session.history[-1] == {"role": "assistant", "content": "response text"} - - -@pytest.mark.asyncio -async def test_playback_ack_rejects_response_after_followup_user_commit(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-playback-after-followup-user", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.append_text("user A") - session.commit_user_input() - response_id = session.begin_response() - session.append_assistant_text("assistant A") - session.mark_audio_sent(1000, text_chars=len("assistant A")) - session.end_response(commit_text=False, preserve_request=True) - session.append_text("user B") - session.commit_user_input() - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": f"item_{response_id}", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "playback_ack_too_late" - assert session.history == ( - {"role": "user", "content": "user A"}, - {"role": "user", "content": "user B"}, - ) - - -@pytest.mark.asyncio -async def test_precommit_zero_audio_ack_reserves_history_for_response_that_finishes_after_commit(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-zero-audio-checkpoint", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.mark_user_input_activity() - response_id = session.begin_response() - item_id = f"item_{response_id}" - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 0, - "committed_ms": 0, - }, - ws.send_json, - ) - session.commit_native_audio_input(transcript="later user input") - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - committed_message = session.end_response(commit_text=True, preserve_request=True) - session.register_history_item(item_id, committed_message) - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert not any(message.get("type") == "error" for message in ws.sent) - assert session.history == ( - {"role": "assistant", "content": "hello world"}, - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "native-duplex:input-audio"}, - "transcript": "later user input", - } - ], - "transcript": "later user input", - }, - ) - - -@pytest.mark.asyncio -async def test_precommit_partial_ack_extends_history_after_response_finishes(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-partial-checkpoint", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.mark_user_input_activity() - response_id = session.begin_response() - item_id = f"item_{response_id}" - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - session.commit_native_audio_input(transcript="later user input") - committed_message = session.end_response(commit_text=True, preserve_request=True) - session.register_history_item(item_id, committed_message) - - assert session.history[0] == {"role": "assistant", "content": "hello"} - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert not any(message.get("type") == "error" for message in ws.sent) - assert session.history[0] == {"role": "assistant", "content": "hello world"} - assert session.history[1]["role"] == "user" - - -@pytest.mark.asyncio -async def test_explicit_truncate_seals_snapshot_against_later_playback_ack(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-hard-playback-truncate", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - response_id = session.begin_response() - item_id = f"item_{response_id}" - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - session.end_response(commit_text=False, preserve_request=True) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 500, - "committed_ms": 500, - "truncate": True, - }, - ws.send_json, - ) - assert session.history == ({"role": "assistant", "content": "hello"},) - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert not any(message.get("type") == "error" for message in ws.sent) - assert session.history == ({"role": "assistant", "content": "hello"},) - - -@pytest.mark.asyncio -async def test_cancel_does_not_append_old_assistant_after_followup_user(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-cancel-after-followup-user", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.append_text("user A") - session.commit_user_input() - session.begin_response() - session.append_assistant_text("assistant A") - session.mark_audio_sent(1000, text_chars=len("assistant A")) - session.acknowledge_playback(500, 500) - session.append_text("user B") - session.commit_user_input() - ws = TimedWebSocket() - - cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="new_response") - - assert cancelled is True - assert session.history == ( - {"role": "user", "content": "user A"}, - {"role": "user", "content": "user B"}, - ) - - -@pytest.mark.asyncio -async def test_playback_ack_rejects_item_from_another_response(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-playback-mismatch", config=DuplexSessionConfig()) - first_response_id = session.begin_response() - session.append_assistant_text("first response") - session.mark_audio_sent(1000) - session.end_response(commit_text=False, preserve_request=True) - session.register_history_item(f"item_{first_response_id}", None) - second_response_id = session.begin_response() - session.mark_audio_sent(1000) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": second_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "playback_item_mismatch" - assert session.active_response_id == second_response_id - assert session.playback.as_dict() == { - "generated_ms": 1000, - "sent_ms": 1000, - "played_ms": 0, - "committed_ms": 0, - } - - -@pytest.mark.asyncio -async def test_playback_ack_rejects_user_history_item(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-playback-user-item", config=DuplexSessionConfig()) - user_message = {"role": "user", "content": "user input"} - session.append_history_message(user_message) - session.register_history_item("item_user", user_message) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": "user", - "item_id": "item_user", - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "playback_item_not_found" - assert session.history == ({"role": "user", "content": "user input"},) - - -@pytest.mark.asyncio -async def test_late_playback_ack_recovers_response_local_history_after_next_response_begins(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-playback-response-local", config=DuplexSessionConfig()) - first_response_id = session.begin_response() - session.append_assistant_text("first response") - session.mark_audio_sent(1000, text_chars=len("first response")) - session.end_response(commit_text=False, preserve_request=True) - second_response_id = session.begin_response() - session.mark_audio_sent(2000, text_chars=len("second response")) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": first_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["history_committed"] is True - assert session.history[-1] == {"role": "assistant", "content": "first response"} - assert session.active_response_id == second_response_id - assert session.playback.as_dict() == { - "generated_ms": 2000, - "sent_ms": 2000, - "played_ms": 0, - "committed_ms": 0, - } - - -@pytest.mark.asyncio -async def test_late_playback_ack_does_not_advance_new_response_cursor(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-late-playback-ack", - config=DuplexSessionConfig(), - ) - first_response_id = session.begin_response() - session.append_assistant_text("first response") - session.mark_audio_sent(11480, text_chars=len("first response")) - session.end_response(commit_text=False, preserve_request=True) - session.register_history_item(f"item_{first_response_id}", None) - - second_response_id = session.begin_response() - session.mark_audio_sent(2200, text_chars=len("second response")) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": first_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 11480, - "committed_ms": 11480, - }, - ws.send_json, - ) - - assert session.active_response_id == second_response_id - assert session.playback.as_dict() == { - "generated_ms": 2200, - "sent_ms": 2200, - "played_ms": 0, - "committed_ms": 0, - } - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["playback"] == { - "generated_ms": 11480, - "sent_ms": 11480, - "played_ms": 11480, - "committed_ms": 11480, - } - assert first_response_id not in session.response_playbacks - assert session.response_playbacks[second_response_id] == session.playback - - -@pytest.mark.asyncio -async def test_late_playback_ack_truncates_history_with_matching_response_cursor(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-late-playback-history", - config=DuplexSessionConfig(), - ) - first_response_id = session.begin_response() - session.append_assistant_text("abcdefghij") - session.mark_audio_sent(1000) - session.end_response(commit_text=False, preserve_request=True) - session.register_history_item(f"item_{first_response_id}", None) - - session.begin_response() - session.mark_audio_sent(10000) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": first_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - - assert session.history[-1] == {"role": "assistant", "content": "abcde"} - assert first_response_id in session.response_playbacks - - -@pytest.mark.asyncio -async def test_cancel_active_response_commits_only_played_assistant_history(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - session = DuplexSession( - session_id="sid-playback-partial", - config=DuplexSessionConfig(playback_commit_policy=DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value), - ) - response_id = session.begin_response() - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - session.acknowledge_playback(played_ms=500, committed_ms=500) - - cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="barge_in") - - assert cancelled is True - assert session.epoch == 1 - assert session.history == ({"role": "assistant", "content": "hello"},) - assert session.history_item_ids[f"item_{response_id}"] == session.history[0] - event = next(m for m in ws.sent if m.get("type") == "audio.cancelled") - assert event["committed_ms"] == 500 - - -@pytest.mark.asyncio -async def test_cancel_active_response_keeps_truncated_history_in_serving(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - session = DuplexSession( - session_id="sid-playback-runtime-rebuild", - config=DuplexSessionConfig(playback_commit_policy=DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value), - ) - response_id = session.begin_response() - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - session.acknowledge_playback(played_ms=500, committed_ms=500) - - cancelled = await handler._cancel_active_response( - session, - None, - ws.send_json, - reason="barge_in", - ) - - assert cancelled is True - assert session.history == ({"role": "assistant", "content": "hello"},) - assert session.history_item_ids[f"item_{response_id}"] == session.history[0] - assert engine.signals == [] - assert "error" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_cancel_active_native_data_plane_request_aborts_stage_request_id(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - session = DuplexSession(session_id="sid-native-abort", config=DuplexSessionConfig()) - session.bind_request("duplex-sid-native-abort-e0-stage0-s1") - - cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="barge_in") - - assert cancelled is True - assert engine.aborted == ["duplex-sid-native-abort-e0-stage0-s1"] - assert engine.internal_abort_batches == [["duplex-sid-native-abort-e0-stage0-s1"]] - assert "audio.cancelled" in ws.sent_types() - - -@pytest.mark.asyncio -async def test_cancel_chat_fallback_response_aborts_request_and_task(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=TurnBasedFakeChatService(engine)) - session = DuplexSession("sid-chat-cancel", DuplexSessionConfig()) - response_id = session.begin_response() - session.bind_request("chatcmpl-cancel") - ws = TimedWebSocket() - task = asyncio.create_task(asyncio.sleep(60)) - try: - cancelled = await handler._cancel_active_response(session, task, ws.send_json, reason="response.cancel") - - assert cancelled is True - assert task.cancelled() - assert engine.internal_abort_batches == [["chatcmpl-cancel"]] - assert "error" not in ws.sent_types() - (event,) = [event for event in ws.sent if event["type"] == "audio.cancelled"] - assert event["response_id"] == response_id - assert event["reason"] == "response.cancel" - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_duplex_handler_explicit_close_closes_runtime_once_with_client_reason(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-close")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "session.closed"] - assert engine.closed == [("sid-close", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_disconnect_detaches_resumable_session_without_runtime_close(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-resumable-disconnect")) - - await handler.handle_realtime_session(ws) - - created = next(message for message in ws.sent if message.get("type") == "session.created") - assert created["session"]["id"] == "sid-resumable-disconnect" - assert isinstance(created["resume_token"], str) - assert created["resume_token"] - assert created["attachment_generation"] == 1 - assert engine.closed == [] - assert engine.touched == [ - ("sid-resumable-disconnect", DuplexLeaseActivity.DETACH), - ] - assert handler._registry.get("sid-resumable-disconnect") is not None - assert "sid-resumable-disconnect" in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_initial_resume_token_delivery_failure_closes_unrecoverable_session(): - class FailingCreatedWebSocket(TimedWebSocket): - async def send_json(self, data: dict[str, Any]): - if data.get("type") == "session.created": - raise WebSocketDisconnect(code=1006) - await super().send_json(data) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = FailingCreatedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-created-delivery-failure")) - - await handler.handle_realtime_session(ws) - - assert engine.closed == [("sid-created-delivery-failure", "disconnect")] - assert handler._registry.get("sid-created-delivery-failure") is None - assert "sid-created-delivery-failure" not in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_send_after_created_disconnect_preserves_resumable_session(): - class ClosedAfterCreatedWebSocket(TimedWebSocket): - def __init__(self) -> None: - super().__init__(receive_timeout_s=0.01) - self.transport_closed = False - - async def send_json(self, data: dict[str, Any]): - if self.transport_closed: - raise RuntimeError( - "Unexpected ASGI message 'websocket.send', " - "after sending 'websocket.close' or response already completed." - ) - await super().send_json(data) - if data.get("type") == "session.created": - self.transport_closed = True - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = ClosedAfterCreatedWebSocket() - ws.put(_native_realtime_session_update("sid-created-then-disconnected")) - - await handler.handle_realtime_session(ws) - - assert engine.closed == [] - assert engine.touched == [ - ("sid-created-then-disconnected", DuplexLeaseActivity.DETACH), - ] - assert handler._registry.get("sid-created-then-disconnected") is not None - assert "sid-created-then-disconnected" in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_resume_rotates_token_replays_and_preserves_runtime_identity(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-realtime-resume")) - first.put({"type": "session.heartbeat"}) - - await handler.handle_realtime_session(first) - - created = next(message for message in first.sent if message.get("type") == "session.created") - heartbeat = next(message for message in first.sent if message.get("type") == "session.heartbeat_ack") - token = created["resume_token"] - incarnation = created["incarnation"] - second = TimedWebSocket(receive_timeout_s=0.1) - second.query_params = { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - "resume": "1", - } - second.put( - { - "type": "session.resume", - "session_id": "sid-realtime-resume", - "incarnation": incarnation, - "resume_token": token, - "last_received_server_event_seq": heartbeat["server_event_seq"] - 1, - } - ) - second.put({"type": "session.close"}) - - await handler.handle_realtime_session(second) - - resumed = next(message for message in second.sent if message.get("type") == "session.resumed") - replayed = [ - message - for message in second.sent - if message.get("type") == "session.heartbeat_ack" - and message.get("server_event_seq") == heartbeat["server_event_seq"] - ] - assert resumed["session_id"] == "sid-realtime-resume" - assert resumed["incarnation"] == incarnation - assert resumed["attachment_generation"] == 2 - assert resumed["resume_token"] != token - assert replayed == [heartbeat] - assert engine.opened == ["sid-realtime-resume"] - assert engine.resumed == [("sid-realtime-resume", 0)] - assert engine.resume_fences[0] == engine.opened_fences[0] - assert engine.closed == [("sid-realtime-resume", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_resume_preserves_append_tail_order_across_connections(): - class SlowFirstAppendEngine(FakeEngineClient): - def __init__(self) -> None: - super().__init__() - self.first_append_started = asyncio.Event() - self.second_append_started = asyncio.Event() - self.release_first_append = asyncio.Event() - self.append_start_order: list[int] = [] - - async def append_duplex_input_async(self, session_id: str, **kwargs): - call_index = len(self.append_start_order) + 1 - self.append_start_order.append(call_index) - if call_index == 1: - self.first_append_started.set() - await self.release_first_append.wait() - else: - self.second_append_started.set() - kwargs.pop("expected_epoch", None) - return await super().append_duplex_input_async(session_id, **kwargs) - - engine = SlowFirstAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=1.0) - session_update = _native_realtime_session_update("sid-resume-append-order") - session_update["session"]["extra_body"]["auto_response"] = True - first.put(session_update) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.01) - else: - pytest.fail("resumable session did not open") - assert created is not None - session = handler._registry.get("sid-resume-append-order") - assert session is not None - native_state = handler._runtime_session_state(session) - first_scheduler = native_state.silence_continuation_scheduler - assert first_scheduler is not None - first.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - await asyncio.wait_for(engine.first_append_started.wait(), timeout=1) - await first.close(code=1006) - await asyncio.wait_for(first_task, timeout=1) - - second = TimedWebSocket(receive_timeout_s=0.5) - second.query_params = { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - "resume": "1", - } - second.put( - { - "type": "session.resume", - "session_id": "sid-resume-append-order", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": created.get("server_event_seq", 0), - } - ) - second_task = asyncio.create_task(handler.handle_realtime_session(second)) - for _ in range(100): - if "session.resumed" in second.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("session did not resume") - assert native_state.silence_continuation_scheduler is not first_scheduler - second.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - try: - await asyncio.sleep(0.05) - assert not engine.second_append_started.is_set() - finally: - engine.release_first_append.set() - - await asyncio.wait_for(engine.second_append_started.wait(), timeout=1) - second.put({"type": "session.close"}) - await asyncio.wait_for(second_task, timeout=2) - assert engine.append_start_order == [1, 2] - - -@pytest.mark.asyncio -async def test_realtime_takeover_ignores_late_detach_from_replaced_attachment(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=1.0) - first.put(_native_realtime_session_update("sid-takeover")) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.001) - else: - raise AssertionError("first attachment did not create the session") - - second = TimedWebSocket(receive_timeout_s=0.03) - second.put( - { - "type": "session.resume", - "session_id": "sid-takeover", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - await handler.handle_realtime_session(second) - await first_task - - detach_touches = [item for item in engine.touched if item[1] is DuplexLeaseActivity.DETACH] - assert detach_touches == [("sid-takeover", DuplexLeaseActivity.DETACH)] - assert "session.replaced" in first.sent_types() - assert engine.resumed == [("sid-takeover", 0)] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_takeover_does_not_clear_new_attachment_pending_turns(): - class DelayedReplacementCloseWebSocket(TimedWebSocket): - def __init__(self) -> None: - super().__init__(receive_timeout_s=1.0) - self.replacement_close_started = asyncio.Event() - self.release_replacement_close = asyncio.Event() - - async def close(self, code: int = 1000, reason: str | None = None) -> None: - if not self.replacement_close_started.is_set(): - self.replacement_close_started.set() - await self.release_replacement_close.wait() - await super().close(code=code, reason=reason) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = DelayedReplacementCloseWebSocket() - first.put(_native_realtime_session_update("sid-takeover-pending-turn")) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.001) - else: - pytest.fail("first attachment did not create the session") - assert created is not None - - second = TimedWebSocket(receive_timeout_s=1.0) - second.put( - { - "type": "session.resume", - "session_id": "sid-takeover-pending-turn", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - second_task = asyncio.create_task(handler.handle_realtime_session(second)) - await asyncio.wait_for(first.replacement_close_started.wait(), timeout=1) - - session = handler._registry.get("sid-takeover-pending-turn") - assert session is not None - assert session.reserve_pending_turn(limit=2) - first.release_replacement_close.set() - await asyncio.wait_for(first_task, timeout=1) - - assert session.pending_input_turns == 1 - session.release_pending_turn() - second.put({"type": "session.close"}) - await asyncio.wait_for(second_task, timeout=2) - - -@pytest.mark.asyncio -async def test_realtime_takeover_drops_input_queued_by_replaced_attachment(): - class DelayedReplacementCloseWebSocket(TimedWebSocket): - def __init__(self) -> None: - super().__init__(receive_timeout_s=1.0) - self.replacement_close_started = asyncio.Event() - self.release_replacement_close = asyncio.Event() - self.stale_heartbeat_read = asyncio.Event() - - async def receive_text(self) -> str: - payload = await super().receive_text() - if json.loads(payload).get("type") == "session.heartbeat": - self.stale_heartbeat_read.set() - return payload - - async def close(self, code: int = 1000, reason: str | None = None) -> None: - if not self.replacement_close_started.is_set(): - self.replacement_close_started.set() - await self.release_replacement_close.wait() - await super().close(code=code, reason=reason) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = DelayedReplacementCloseWebSocket() - first.put(_native_realtime_session_update("sid-takeover-stale-input")) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.001) - else: - pytest.fail("first attachment did not create the session") - assert created is not None - - second = TimedWebSocket(receive_timeout_s=1.0) - second.put( - { - "type": "session.resume", - "session_id": "sid-takeover-stale-input", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - second_task = asyncio.create_task(handler.handle_realtime_session(second)) - await asyncio.wait_for(first.replacement_close_started.wait(), timeout=1) - - first.put({"type": "session.heartbeat"}) - await asyncio.wait_for(first.stale_heartbeat_read.wait(), timeout=1) - first.release_replacement_close.set() - await asyncio.wait_for(first_task, timeout=1) - - second.put({"type": "session.close"}) - await asyncio.wait_for(second_task, timeout=2) - assert all(activity is not DuplexLeaseActivity.HEARTBEAT for _, activity in engine.touched) - assert "session.heartbeat_ack" not in second.sent_types() - - -@pytest.mark.asyncio -async def test_realtime_resume_rejects_invalid_token_before_engine_control(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-invalid-resume")) - await handler.handle_realtime_session(first) - - created = next(message for message in first.sent if message.get("type") == "session.created") - second = TimedWebSocket(receive_timeout_s=0.01) - second.put( - { - "type": "session.resume", - "session_id": "sid-invalid-resume", - "incarnation": created["incarnation"], - "resume_token": "invalid-token", - "last_received_server_event_seq": 0, - } - ) - - await handler.handle_realtime_session(second) - - error = next(message for message in second.sent if message.get("type") == "error") - assert error["error"]["code"] == "invalid_resume_token" - assert engine.resumed == [] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_resume_recovers_when_rotated_token_delivery_is_lost(): - class FailingResumeWebSocket(TimedWebSocket): - async def send_json(self, data: dict[str, Any]): - if data.get("type") == "session.resumed": - raise RuntimeError("resume transport disappeared") - await super().send_json(data) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-token-delivery")) - await handler.handle_realtime_session(first) - created = next(message for message in first.sent if message.get("type") == "session.created") - resume_event = { - "type": "session.resume", - "session_id": "sid-token-delivery", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - - failed_transport = FailingResumeWebSocket(receive_timeout_s=0.01) - failed_transport.put(resume_event) - await handler.handle_realtime_session(failed_transport) - - recovered_transport = TimedWebSocket(receive_timeout_s=0.1) - recovered_transport.put(resume_event) - recovered_transport.put({"type": "session.close"}) - await handler.handle_realtime_session(recovered_transport) - - assert "session.resumed" in recovered_transport.sent_types() - assert engine.resumed == [ - ("sid-token-delivery", 0), - ("sid-token-delivery", 1), - ] - assert engine.closed == [("sid-token-delivery", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_disconnect_grace_cancels_only_orphan_response(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig(disconnect_grace_s=0.02), - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-grace-response")) - - await handler.handle_realtime_session(ws) - session = handler._registry.get("sid-grace-response") - assert session is not None - session.begin_response() - session.bind_request("orphan-request") - orphan_task = asyncio.create_task(asyncio.sleep(10)) - handler._session_tasks[session.session_id].active_response_task = orphan_task - assert engine.closed == [] - - await asyncio.sleep(0.05) - - assert engine.closed == [] - assert engine.aborted == ["orphan-request"] - assert orphan_task.cancelled() - assert session.active_response_id is None - assert "sid-grace-response" in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_heartbeat_control_failure_is_non_terminal(): - engine = FakeEngineClient(fail_touch=True) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=0.1) - ws.put(_native_realtime_session_update("sid-heartbeat-failure")) - ws.put({"type": "session.heartbeat"}) - ws.put({"type": "session.close"}) - - await handler.handle_realtime_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["error"]["code"] == "runtime_touch_failed" - assert engine.closed == [("sid-heartbeat-failure", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_idle_ttl_lifecycle_removes_detached_serving_projection(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-idle-expired")) - - await handler.handle_realtime_session(ws) - fence = engine.opened_fences[0] - assert fence is not None - await engine.duplex_lifecycle_events.put( - DuplexSessionLifecycleMessage( - fence=fence, - session_id="sid-idle-expired", - event="expired", - reason="idle_ttl_expired", - lease_generation=1, - submitted_request_ids=[], - reserved_request_ids=[], - ) - ) - await asyncio.sleep(0.05) - - assert handler._registry.get("sid-idle-expired") is None - assert "sid-idle-expired" not in handler._serving_runtime_adapter.session_states - assert "sid-idle-expired" not in handler._session_tasks - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_resume_reports_resync_required_after_journal_ttl_gap(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig(resume_replay_ttl_s=0.01), - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-journal-gap")) - first.put({"type": "session.heartbeat"}) - await handler.handle_realtime_session(first) - created = next(message for message in first.sent if message.get("type") == "session.created") - await asyncio.sleep(0.02) - - second = TimedWebSocket(receive_timeout_s=0.01) - second.put( - { - "type": "session.resume", - "session_id": "sid-journal-gap", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - await handler.handle_realtime_session(second) - - assert "session.resync_required" in second.sent_types() - assert engine.resumed == [] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_journal_overflow_degrades_to_live_resync_without_closing_runtime(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig(resume_replay_max_bytes_per_session=128), - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-journal-overflow")) - - await handler.handle_realtime_session(ws) - - assert "session.created" in ws.sent_types() - assert "session.resync_required" in ws.sent_types() - assert engine.closed == [] - assert handler._registry.get("sid-journal-overflow") is not None - - -@pytest.mark.asyncio -async def test_realtime_input_backpressure_is_per_session_and_rejects_before_mutation(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig( - max_pending_input_bytes_per_session=100, - ), - ) - first = TimedWebSocket(receive_timeout_s=0.1) - first.put(_native_realtime_session_update("sid-backpressure-a")) - for _ in range(2): - first.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(20), - "format": "pcm16", - "sample_rate_hz": 16000, - } - ) - first.put({"type": "session.close"}) - - await handler.handle_realtime_session(first) - - backpressure = [ - event - for event in first.sent - if event.get("type") == "error" - and isinstance(event.get("error"), dict) - and event["error"].get("code") == "input_backpressure" - ] - assert len(backpressure) == 1 - - second = TimedWebSocket(receive_timeout_s=0.1) - second.put(_native_realtime_session_update("sid-backpressure-b")) - second.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(20), - "format": "pcm16", - "sample_rate_hz": 16000, - } - ) - second.put({"type": "session.close"}) - - await handler.handle_realtime_session(second) - - assert not any( - event.get("type") == "error" - and isinstance(event.get("error"), dict) - and event["error"].get("code") == "input_backpressure" - for event in second.sent - ) - - -@pytest.mark.asyncio -async def test_duplex_handler_idle_timeout_close_does_not_emit_runtime_control(): - control_result = { - "operation": "close", - "session_id": "sid-disconnect", - "ok": True, - "unsupported_count": 1, - "error_count": 0, - "stage_results": [{"stage_id": 0, "replica_id": 0, "result": {"supported": False}}], - } - engine = FakeEngineClient(append_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=0.1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-disconnect")) - - await handler.handle_session(ws) - - assert engine.closed == [("sid-disconnect", "timeout")] - - -@pytest.mark.asyncio -async def test_idle_output_audio_clear_does_not_advance_runtime_fence(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_session_create("sid-idle-output-clear")) - ws.put({"type": "output_audio_buffer.clear"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - cancelled = next(event for event in ws.sent if event.get("type") == "audio.cancelled") - assert cancelled["cancelled_epoch"] == cancelled["epoch"] == 0 - assert engine.signals == [] - - -@pytest.mark.asyncio -async def test_duplex_handler_runtime_close_failure_is_reported_without_closed_ack(): - engine = FakeEngineClient(fail_close=True) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-close-fail")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "error"] - assert ws.sent[-1]["code"] == "runtime_close_failed" - - -@pytest.mark.asyncio -async def test_duplex_handler_control_close_failure_is_reported_without_closed_ack(): - control_result = { - "operation": "close", - "session_id": "sid-control-close-fail", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "error": "stage close failed"}, - } - ], - } - engine = FakeEngineClient(close_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-control-close-fail")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "error"] - assert ws.sent[-1]["code"] == "runtime_close_failed" - assert ws.sent[-1]["runtime_control"]["error_count"] == 1 - - -@pytest.mark.asyncio -async def test_duplex_cancel_without_active_response_clears_pending_input_and_acks(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_session_create("sid-pending-cancel")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.cancel", "reason": "user_cancel"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - cancelled = next(m for m in ws.sent if m.get("type") == "input.cancelled") - assert cancelled["cancelled"] == {"text_chunks": 1, "audio_chunks": 0} - assert cancelled["epoch"] == 1 - assert chat_service.seen_request_ids == [] - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_turn_signal_input_cancel_uses_epoch_fence_transition(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-turn-signal-cancel")) - ws.put({"type": "input.text.append", "text": "discard me"}) - ws.put({"type": "turn.signal", "event": "input.cancel"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert engine.signal_fences[-1] == DuplexFence("sid-turn-signal-cancel") - assert engine.signal_next_fences[-1] == DuplexFence("sid-turn-signal-cancel", epoch=1) - assert any(message.get("type") == "input.cancelled" for message in ws.sent) - - -@pytest.mark.asyncio -async def test_duplex_handler_local_turn_signal_does_not_round_trip_runtime(): - engine = FakeEngineClient(fail_signal=True) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_session_create("sid-signal-fail")) - ws.put({"type": "turn.signal", "event": "user_started"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - turn_event = next(m for m in ws.sent if m.get("type") == "turn.event") - assert turn_event["event"] == "user_started" - assert not engine.signals - assert "runtime_signal_failed" not in {m.get("code") for m in ws.sent} - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_append_without_format_defaults_to_pcm16(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-default-pcm16")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert final is False - assert isinstance(payload, dict) - assert payload["format"] == "pcm_f32le" - assert payload["sample_rate_hz"] == 16000 - samples = np.frombuffer(base64.b64decode(payload["audio"]), dtype=" None: - super().__init__() - self.append_started = asyncio.Event() - self.append_cancelled = asyncio.Event() - self.release_append = asyncio.Event() - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - timeout: float | None = None, - collect_outputs: bool = True, - fence: DuplexFence | None = None, - ) -> Any: - self.append_started.set() - try: - await self.release_append.wait() - except asyncio.CancelledError: - self.append_cancelled.set() - raise - return await super().append_duplex_input_async( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - timeout=timeout, - collect_outputs=collect_outputs, - fence=fence, - ) - - rollback = mocker.spy(MiniCPMO45PcmAppendBuffer, "_rollback_reservation") - engine = SlowAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-cancel-slow-append")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - ws.put({"type": "input.cancel"}) - ws.put({"type": "session.close"}) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - await asyncio.wait_for(engine.append_started.wait(), timeout=1) - try: - await asyncio.wait_for(engine.append_cancelled.wait(), timeout=0.2) - except TimeoutError: - pass - finally: - engine.release_append.set() - await asyncio.wait_for(handler_task, timeout=1) - - assert engine.append_cancelled.is_set() - rollback.assert_called_once() - assert engine.signal_fences[-1] == DuplexFence("sid-cancel-slow-append") - assert engine.signal_next_fences[-1] == DuplexFence("sid-cancel-slow-append", epoch=1) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_append_control_error_does_not_emit_model_delta(): - control_result = { - "operation": "append", - "session_id": "sid-native-append-error", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": False, - "error": "duplex_stage_handoff_target_missing:tts", - }, - }, - ], - } - engine = FakeEngineClient(append_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-append-error")) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "response.output_audio.delta" not in ws.sent_types() - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "runtime_append_failed" - assert error["runtime_control"]["error_count"] == 1 - - -@pytest.mark.asyncio -async def test_duplex_handler_session_update_control_failure_rolls_back_config(): - control_result = { - "operation": "signal", - "session_id": "sid-control-signal-fail", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "error": "stage signal failed"}, - } - ], - } - engine = FakeEngineClient(signal_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-control-signal-fail")) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"temperature": 0.2}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "session.updated" not in ws.sent_types() - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "runtime_signal_failed" - assert error["runtime_control"]["error_count"] == 1 - signaled_config = engine.signal_session_configs[-1] - assert signaled_config is not None - assert signaled_config["temperature"] == 0.2 - - -@pytest.mark.asyncio -async def test_native_session_update_commits_config_only_after_runtime_ack(): - class BlockingUpdateEngine(FakeEngineClient): - def __init__(self): - super().__init__() - self.update_started = asyncio.Event() - self.release_update = asyncio.Event() - - async def signal_duplex_turn_async(self, session_id: str, **kwargs): - if kwargs.get("event") == "session.update": - self.update_started.set() - await self.release_update.wait() - return await super().signal_duplex_turn_async(session_id, **kwargs) - - engine = BlockingUpdateEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([])), - ) - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - ws.put(_native_realtime_session_update("sid-update-two-phase")) - ws.put(_native_server_vad_update("update-two-phase")) - - handler_task = asyncio.create_task(handler.handle_session(ws, realtime_protocol=protocol)) - await asyncio.wait_for(engine.update_started.wait(), timeout=10) - live_session = handler._registry.get("sid-update-two-phase") - assert live_session is not None - assert live_session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value - assert protocol._pending_turn_detection_update is not None - assert protocol._turn_detection is None - - protocol.encode_outbound_event({"type": "error", "code": "unrelated", "error": "unrelated"}) - assert protocol._pending_turn_detection_update is not None - - engine.release_update.set() - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - updated = next(message for message in reversed(ws.sent) if message.get("type") == "session.updated") - assert updated["session"]["overlap_policy"] == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - assert protocol._turn_detection is not None - - -@pytest.mark.parametrize( - ("append_fails", "expected_code"), - [(True, "session_update_aborted"), (False, "instructions_update_unsupported")], -) -@pytest.mark.asyncio -async def test_native_server_vad_update_rejection_does_not_stall(append_fails, expected_code): - append_result = {"operation": "append", "ok": False, "error_count": 1} if append_fails else None - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient(append_result=append_result)), - config_timeout_s=0.1, - idle_timeout_s=1, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([])), - ) - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - create = _native_realtime_session_update("sid-rejected-vad-update") - create["session"]["extra_body"]["auto_response"] = append_fails - ws.put(create) - if append_fails: - ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(16_000), "format": "pcm_f32le"}) - update = _native_server_vad_update("rejected-vad-update") - if not append_fails: - update["session"]["instructions"] = "You are now a pirate." - ws.put(update) - ws.put({"type": "session.close"}) - - await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=10) - - errors = [message["error"] for message in ws.sent if isinstance(message.get("error"), dict)] - error = next(error for error in errors if error.get("code") == expected_code) - assert error["event_id"] == "rejected-vad-update" - assert protocol._pending_turn_detection_update is None - - -@pytest.mark.asyncio -async def test_native_session_update_waits_for_prior_append_effect(): - class OrderedUpdateEngine(FakeEngineClient): - def __init__(self): - super().__init__() - self.append_started = asyncio.Event() - self.release_append = asyncio.Event() - self.update_started = asyncio.Event() - self.effects: list[str] = [] - - async def append_duplex_input_async(self, session_id: str, **kwargs): - self.effects.append("append.started") - self.append_started.set() - await self.release_append.wait() - kwargs.pop("expected_epoch", None) - result = await super().append_duplex_input_async(session_id, **kwargs) - self.effects.append("append.done") - return result - - async def signal_duplex_turn_async(self, session_id: str, **kwargs): - if kwargs.get("event") == "session.update": - self.effects.append("session.update") - self.update_started.set() - return await super().signal_duplex_turn_async(session_id, **kwargs) - - engine = OrderedUpdateEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-ordered-update")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"temperature": 0.2}, - } - ) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - await asyncio.wait_for(engine.append_started.wait(), timeout=1) - try: - await asyncio.sleep(0.05) - assert not engine.update_started.is_set() - finally: - engine.release_append.set() - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.effects.index("append.done") < engine.effects.index("session.update") - - -@pytest.mark.asyncio -async def test_failed_native_append_prevents_queued_append_from_running(): - failed_result = { - "operation": "append", - "session_id": "sid-append-chain-failure", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "error": "append failed"}, - } - ], - } - - class BlockingFailureEngine(FakeEngineClient): - def __init__(self): - super().__init__(append_result=failed_result) - self.first_started = asyncio.Event() - self.release_first = asyncio.Event() - - async def append_duplex_input_async(self, session_id: str, **kwargs): - if not self.appended: - self.first_started.set() - await self.release_first.wait() - kwargs.pop("expected_epoch", None) - return await super().append_duplex_input_async(session_id, **kwargs) - - engine = BlockingFailureEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - create = _native_session_create("sid-append-chain-failure") - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - for _ in range(2): - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16_000, value=0.05), - "format": "pcm_f32le", - } - ) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - await asyncio.wait_for(engine.first_started.wait(), timeout=1) - engine.release_first.set() - await asyncio.sleep(0.1) - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert len(engine.appended) == 1 - - -@pytest.mark.asyncio -async def test_failed_final_append_discards_retained_audio_without_retry(): - failed_result = { - "operation": "append", - "session_id": "sid-final-append-retry", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [], - } - engine = FakeEngineClient(append_result=failed_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def retry_after_runtime_error(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "error" and data.get("code") == "runtime_append_failed": - ws.put({"type": "response.create"}) - - ws = TimedWebSocket(on_send=retry_after_runtime_error) - ws.put(_native_session_create("sid-final-append-retry")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(8_000, value=0.05), - "format": "pcm_f32le", - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "response.create"}) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - for _ in range(100): - state = handler._serving_runtime_adapter.session_states.get("sid-final-append-retry") - if len(engine.appended) >= 1 and state is not None and state.committed_audio_payload is None: - break - await asyncio.sleep(0.01) - state = handler._serving_runtime_adapter.session_states["sid-final-append-retry"] - assert state.committed_audio_payload is None - assert state.committed_audio_reserved_bytes == 0 - assert handler._registry.get("sid-final-append-retry").pending_input_bytes == 0 - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert len(engine.appended) == 1 - - -@pytest.mark.asyncio -async def test_duplex_handler_signal_unsupported_workers_with_data_plane_ack_is_not_fatal(): - control_result = { - "operation": "signal", - "session_id": "sid-control-signal-data-plane", - "ok": False, - "unsupported_count": 2, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "reason": "worker_duplex_signal_not_implemented"}, - }, - { - "stage_id": 1, - "replica_id": 0, - "result": {"supported": False, "reason": "worker_duplex_signal_not_implemented"}, - }, - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_signal": True, - "event": "barge_in", - "epoch": 1, - }, - }, - ], - } - engine = FakeEngineClient(signal_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-control-signal-data-plane")) - ws.put({"type": "turn.signal", "event": "barge_in"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "runtime_signal_failed" not in {m.get("code") for m in ws.sent} - assert "input.cancelled" in ws.sent_types() - runtime_control = next(m for m in ws.sent if m.get("type") == "runtime.control") - assert runtime_control["result"]["unsupported_count"] == 2 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_open_unsupported_fails_session_create(): - control_result = { - "operation": "open", - "session_id": "sid-native-busy", - "ok": True, - "unsupported_count": 1, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "reason": "native_duplex_session_busy"}, - } - ], - } - engine = FakeEngineClient(control_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-busy")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "runtime_open_unsupported" - assert ws.sent[0]["runtime_control"]["unsupported_count"] == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract(): - class LegacyEngineClient(FakeEngineClient): - async def open_duplex_session_async( # type: ignore[override] - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - timeout: float | None = None, - ) -> Any: - return await super().open_duplex_session_async( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - timeout=timeout, - ) - - engine = LegacyEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-no-fence-contract")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "runtime_contract_invalid" - assert "open_duplex_session_async" in ws.sent[0]["error"] - assert "fence" in ws.sent[0]["error"] - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_audio_append_does_not_retain_pending_audio(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-no-pending")) - short_chunk = base64.b64encode(b"\x00" * (800 * 4)).decode("ascii") - ws.put({"type": "input_audio_buffer.append", "audio": short_chunk, "format": "pcm_f32le"}) - ws.put({"type": "input.cancel", "reason": "barge_in"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - cancelled = next(m for m in ws.sent if m.get("type") == "input.cancelled") - assert cancelled["cancelled"] == {"text_chunks": 0, "audio_chunks": 0} - assert cancelled["epoch"] == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_text_append_before_runtime_call(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-text")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "native_text_append_unsupported" - assert engine.appended == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_separates_public_and_runtime_config(monkeypatch): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - async def fake_resolve_ref_audio(_, *, model_config): - del model_config - return [0.25] * 1600, 16000 - - monkeypatch.setattr( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.MiniCPMO45NativeDuplexServingAdapter.resolve_ref_audio", - fake_resolve_ref_audio, - ) - event = _native_session_create("sid-native-ref-audio", modalities=["text", "audio"]) - event["session"]["ref_audio"] = "data:audio/wav;base64,AAAA" - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types()[0] == "session.created" - opened_config = engine.opened_configs[0] - opened_runtime_config = engine.opened_runtime_configs[0] - extra_body = opened_config["extra_body"] - assert "ref_audio_data" not in extra_body - assert "duplex_stage_max_tokens" not in extra_body - assert opened_runtime_config["ref_audio_format"] == "pcm_f32le" - assert opened_runtime_config["ref_audio_sample_rate_hz"] == 16000 - assert opened_runtime_config["duplex_stage_max_tokens"] == {"0": 20, "1": 8192} - assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20 - assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8 - # The Talker YAML carries upstream's min_new_token=50 for chat TTS; a duplex - # chunk is 26 codec samples, so check_stop would never release the request. - assert opened_runtime_config["duplex_stage_sampling_params"]["1"] == {"min_tokens": 0} - assert base64.b64decode(opened_runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.25] * 1600)) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_client_runtime_config(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - event = _native_session_create("sid-native-private-config") - event["session"]["extra_body"]["duplex_stage_sampling_params"] = {"0": {"temperature": 0.0}} - ws = TimedWebSocket() - ws.put(event) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "invalid_duplex_runtime_config" - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_audio_output_requires_ref_audio(monkeypatch): - monkeypatch.setattr( - MiniCPMO45NativeDuplexServingAdapter, - "_load_native_tokenizer", - staticmethod(lambda model_config: None), - ) - - event = _native_session_create( - "sid-native-ref-required", - modalities=["text", "audio"], - ) - config = DuplexSessionConfig.from_event(event) - - with pytest.raises(ServingRuntimeConfigError) as exc_info: - await MiniCPMO45NativeDuplexServingAdapter.prepare_runtime_config( - config, - model_config=SimpleNamespace(model="openbmb/MiniCPM-o-4_5"), - ) - - assert exc_info.value.code == "ref_audio_required" - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_open_session_maps_missing_ref_audio_to_typed_error(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put( - _native_session_create( - "sid-native-ref-required-open", - modalities=["text", "audio"], - ) - ) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "ref_audio_required" - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_text_only_omits_ref_audio_when_client_does_not_provide_it(monkeypatch): - monkeypatch.setattr( - MiniCPMO45NativeDuplexServingAdapter, - "_load_native_tokenizer", - staticmethod(lambda model_config: None), - ) - - event = _native_session_create("sid-native-default-ref", modalities=["text"]) - config = DuplexSessionConfig.from_event(event) - - runtime_config = await MiniCPMO45NativeDuplexServingAdapter.prepare_runtime_config( - config, - model_config=SimpleNamespace(model="openbmb/MiniCPM-o-4_5"), - ) - - assert "ref_audio_data" not in runtime_config - assert "ref_audio_format" not in runtime_config - assert "ref_audio_sample_rate_hz" not in runtime_config - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_preserves_ref_audio_channels_until_normalize(monkeypatch): - class FakeMediaConnector: - def __init__(self, **kwargs): - self.kwargs = kwargs - - async def fetch_audio_async(self, ref_audio): - assert ref_audio == "data:audio/wav;base64,AAAA" - return np.tile(np.array([[0.25, -0.25], [0.5, -0.5]], dtype=np.float32), (800, 1)), 16000 - - monkeypatch.setattr( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.adapter.MediaConnector", - FakeMediaConnector, - ) - - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - event = _native_session_create("sid-native-ref-audio-stereo", modalities=["text", "audio"]) - event["session"]["ref_audio"] = "data:audio/wav;base64,AAAA" - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - runtime_config = engine.opened_runtime_configs[0] - assert base64.b64decode(runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.0] * 1600)) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_ref_audio_path(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - event = _native_session_create("sid-native-ref-path") - event["session"]["extra_body"] = { - "native_duplex": True, - "ref_audio_path": "/tmp/ref.wav", - } - ws = TimedWebSocket() - ws.put(event) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "unsupported_ref_audio_path" - assert engine.opened == [] - - -def test_minicpmo_native_duplex_explicit_barge_in_request_interrupts(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-native-explicit-barge-disabled", - config=DuplexSessionConfig(overlap_policy=DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value), - capabilities=minicpmo45_native_capabilities(), - ) - payload = { - "type": "audio", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "is_speech": True, - } - - decision = handler._overlap_decision( - session, - {"force_barge_in": True}, - payload, - ) - - assert decision == { - "action": "barge_in", - "reason": "client_force_barge_in", - "duration_ms": 1000, - "buffer_audio": True, - } - - -@pytest.mark.asyncio -async def test_native_server_vad_complete_audio_item_reaches_runtime(): - engine = FakeEngineClient() - backend = FakeServerVADBackend([0.9] * 10) - ws = TimedWebSocket() - create = _native_realtime_session_update("sid-native-complete-audio") - create["session"]["turn_detection"] = {"type": "server_vad"} - create["session"]["input_audio_format"] = "pcm16" - ws.put(create) - ws.put( - { - "type": "conversation.item.create", - "item": { - "id": "item-native-audio", - "type": "message", - "role": "user", - "content": [{"type": "input_audio", "audio": _pcm16_b64(2400)}], - }, - } - ) - ws.put({"type": "response.create"}) - ws.put({"type": "session.close"}) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - server_vad_backend_provider=FakeServerVADProvider(backend), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - assert any(final for _, _, _, final in engine.appended) - assert "error" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_native_server_vad_uses_shared_pipeline_and_emits_boundaries(): - backend = FakeServerVADBackend([0.9] * 10 + [0.0] * 50) - engine = FakeEngineClient() - - def commit_after_speech_stops(ws: TimedWebSocket, event: dict[str, Any]) -> None: - if event.get("type") == "input_audio_buffer.speech_stopped": - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "response.create"}) - ws.put({"type": "session.close"}) - - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - server_vad_backend_provider=FakeServerVADProvider(backend), - ) - ws = TimedWebSocket(on_send=commit_after_speech_stops) - create = _native_realtime_session_update("sid-native-shared-server-vad") - create["session"]["turn_detection"] = {"type": "server_vad"} - ws.put(create) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(60 * backend.frame_samples, value=0.0), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - } - ) - await asyncio.wait_for( - handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})), - timeout=10, - ) - - boundary_events = [ - event - for event in ws.sent - if event.get("type") in {"input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"} - ] - assert [event["type"] for event in boundary_events] == [ - "input_audio_buffer.speech_started", - "input_audio_buffer.speech_stopped", - ] - assert boundary_events[0]["item_id"] == boundary_events[1]["item_id"] - created_session = next(event["session"] for event in ws.sent if event["type"] == "session.created") - assert created_session["audio"]["input"]["turn_detection"]["interrupt_response"] is True - assert created_session["overlap_policy"] == "barge_in_on_speech" - assert backend.calls == 60 - assert any(final for _, _, _, final in engine.appended) - assert not any(event.get("error", {}).get("code") == "server_vad_manual_commit_unsupported" for event in ws.sent) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "barge_in_event", - [ - {"type": "barge_in"}, - {"type": "turn.signal", "event": "barge_in"}, - ], -) -async def test_minicpmo_native_duplex_accepts_explicit_barge_in_control( - barge_in_event: dict[str, object], -): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-barge-control-disabled")) - ws.put(barge_in_event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert not any(message.get("code") == "barge_in_unsupported" for message in ws.sent) - assert ("sid-native-barge-control-disabled", "barge_in") in engine.signals - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_explicit_non_speech_stays_listening_without_append(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-silence")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": "AAAA", - "format": "pcm_f32le", - "is_speech": False, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert engine.appended == [] - assert "response.listen" in ws.sent_types() - assert "response.created" not in ws.sent_types() - assert "response.output_audio.delta" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_drains_data_plane_stream_until_done(): - def _stage_output(samples: int, *, finished: bool): - return SimpleNamespace( - request_id="duplex-sid-native-stream-e0-stage0-s1", - finished=finished, - outputs=[ - SimpleNamespace( - text="", - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - }, - ) - ], - ) - - control_result = { - "operation": "append", - "session_id": "sid-native-stream", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": "duplex-sid-native-stream-e0-stage0-s1", - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [_stage_output(10, finished=False)], - } - engine = FakeEngineClient( - append_result=control_result, - collect_outputs=[ - [_stage_output(20, finished=False)], - [_stage_output(30, finished=True)], - ], - ) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def close_on_second_created(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if ws.sent_types().count("response.created") >= 2: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=close_on_second_created) - ws.put(_native_session_create("sid-native-stream")) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - - await handler.handle_session(ws) - - deltas = [m for m in ws.sent if m.get("type") == "response.output_audio.delta"] - assert [m["audio"] for m in deltas] == ["wav-10", "wav-10", "wav-10"] - assert [m["end_of_turn"] for m in deltas] == [False, False, True] - assert len([m for m in ws.sent if m.get("type") == "response.done"]) == 1 - assert engine.collected == [ - ("duplex-sid-native-stream-e0-stage0-s1", 1), - ("duplex-sid-native-stream-e0-stage0-s1", 1), - ] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_restarts_drain_when_append_races_idle_exit( - monkeypatch: pytest.MonkeyPatch, -): - request_id = "duplex-sid-native-late-output-e0-stage0" - late_output = object() - engine = FakeEngineClient( - collect_outputs=[[], [], [], [late_output]], - collect_delay_s=0.05, - ) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-native-late-output", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - native = handler._runtime_session_state(session) - projected_batches: list[object] = [] - projected = asyncio.Event() - - async def project_late_output(_send_json, result, **_kwargs): - projected_batches.append(result) - projected.set() - session.close() - return None, True - - monkeypatch.setattr(handler, "_send_native_duplex_events", project_late_output) - result = { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ] - } - - assert ( - await handler._start_native_data_plane_stream_task( - None, - result, - session=session, - expected_epoch=session.epoch, - ) - is True - ) - first_drain = native.data_plane_task - assert first_drain is not None - await asyncio.sleep(0.01) - - assert ( - await handler._start_native_data_plane_stream_task( - None, - result, - session=session, - expected_epoch=session.epoch, - ) - is False - ) - await asyncio.wait_for(projected.wait(), timeout=1) - await asyncio.wait_for(first_drain, timeout=1) - - assert len(engine.collected) == 4 - assert projected_batches == [{"data_plane_outputs": [late_output]}] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_ignores_outputs_after_data_plane_turn_done(): - request_id = "duplex-sid-native-post-done-e0-stage0-s1" - - def _stage_output(samples: int, *, turn_end: bool = False): - return SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[ - SimpleNamespace( - text="hello", - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta": { - "turn_end": turn_end, - "duplex_turn_id": 0, - "duplex_epoch": 0, - }, - }, - ) - ], - ) - - control_result = { - "operation": "append", - "session_id": "sid-native-post-done", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [_stage_output(10)], - } - engine = FakeEngineClient( - append_result=control_result, - collect_outputs=[ - [_stage_output(20, turn_end=True)], - [_stage_output(30)], - ], - ) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def close_on_done(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - event = _native_session_create("sid-native-post-done") - event["session"]["extra_body"]["auto_response"] = True - ws = TimedWebSocket(on_send=close_on_done) - ws.put(event) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - - await handler.handle_session(ws) - - assert ws.sent_types().count("response.created") == 1 - assert ws.sent_types().count("response.done") == 1 - deltas = [m for m in ws.sent if m.get("type") == "response.output_audio.delta"] - assert [m["audio"] for m in deltas] == ["wav-10", "wav-10"] - assert engine.collected - assert all(collected == (request_id, 1) for collected in engine.collected) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_accepts_next_append_after_response_done(): - def _stage_output(*, finished: bool): - return SimpleNamespace( - request_id="duplex-sid-native-next-turn-e0-stage0-s1", - finished=finished, - outputs=[ - SimpleNamespace( - text="", - multimodal_output={ - "audio": np.zeros(10, dtype=np.float32), - "sr": 24000, - }, - ) - ], - ) - - control_result = { - "operation": "append", - "session_id": "sid-native-next-turn", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": "duplex-sid-native-next-turn-e0-stage0-s1", - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [_stage_output(finished=True)], - } - engine = FakeEngineClient(append_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - done_count = 0 - - def send_next_turn_or_close(ws: TimedWebSocket, data: dict[str, Any]) -> None: - nonlocal done_count - if data.get("type") != "response.done": - return - done_count += 1 - if done_count == 1: - ws.put({"type": "input_audio_buffer.append", "audio": "BBBB", "format": "pcm_f32le"}) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "response.create"}) - else: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=send_next_turn_or_close, receive_timeout_s=5) - create = _native_session_create("sid-native-next-turn") - create["session"]["idle_timeout_s"] = 5 - ws.put(create) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - - await handler.handle_session(ws) - - assert len(engine.appended) >= 2, [ - (message.get("type"), message.get("code"), message.get("reason")) for message in ws.sent - ] - assert len([m for m in ws.sent if m.get("type") == "response.done"]) == 2 - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_commit_finalizes_after_streamed_full_chunk(): - control_result = { - "operation": "append", - "session_id": "sid-native-auto-finalize", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [], - } - engine = FakeEngineClient(control_result=control_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - event = _native_session_create("sid-native-auto-finalize") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - assert engine.appended[0][3] is False - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_commit_finalizes_with_active_data_plane_stream(): - request_id = "duplex-sid-native-auto-active-stream-stage0" - control_result = { - "operation": "append", - "session_id": "sid-native-auto-active-stream", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - engine = FakeEngineClient( - control_result=control_result, - collect_outputs=[[], [], []], - collect_delay_s=0.05, - ) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - event = _native_session_create("sid-native-auto-active-stream") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - assert engine.appended[0][3] is False - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_accepts_realtime_commit_after_stream_consumed(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - event = _native_session_create("sid-native-auto-consumed-stream") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put( - { - "type": "input_audio_buffer.commit", - "final": True, - "item_id": "item-consumed-stream", - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert not any(message.get("code") == "input_audio_buffer_empty" for message in ws.sent) - committed = next(message for message in ws.sent if message.get("type") == "input.committed") - assert committed["item_id"] == "item-consumed-stream" - assert committed["native_audio"] is True - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_post_response_commit_ignores_stale_progress_state(): - session_id = "sid-native-auto-post-response" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get(session_id) - assert session is not None - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(8000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - create = _native_session_create(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - assert engine.appended[0][3] is True - assert "input.committed" in ws.sent_types() - assert not any(message.get("code") == "response_already_active" for message in ws.sent) - - -def test_realtime_overlap_commit_uses_current_generation_identity(): - handler, session = _auto_response_context("sid-realtime-reserved-input-turn") - session.turn_id = 1 - session.bind_response_turn(1) - session.turn_id = 3 - committed = handler._commit_native_audio_input(session) - - assert committed.turn_id == 3 - assert session.turn_id == 3 - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_continuation(): - request_id = "duplex-sid-native-auto-continuation-e0-stage0" - control_result = { - "operation": "append", - "session_id": "sid-native-auto-continuation", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - terminal_segment = _duplex_tts_output( - request_id=request_id, - samples=0, - finished=True, - tts_is_last_chunk=True, - token_ids=[151645], - ) - engine = FakeEngineClient( - control_result=control_result, - collect_outputs=[[terminal_segment], [], []], - collect_delay_s=0.05, - ) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get("sid-native-auto-continuation") - assert session is not None - session.capabilities.chunk_period_ms = 50 - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - - ws = TimedWebSocket(on_send=on_send) - event = _native_session_create("sid-native-auto-continuation") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - - await handler.handle_session(ws) - - assert len(engine.appended) >= 3 - _, mode, payload, final = engine.appended[2] - assert mode == "append_audio_chunk" - assert payload["duplex_turn_id"] == 0 - assert final is False - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_real_input_preempts_unsent_silence(): - request_id = "duplex-sid-native-real-before-silence-e0-stage0" - - class SequencedAppendEngine(FakeEngineClient): - def __init__(self) -> None: - control_result = { - "operation": "append", - "session_id": "sid-native-real-before-silence", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - terminal_segment = _duplex_tts_output( - request_id=request_id, - samples=0, - finished=True, - tts_is_last_chunk=True, - token_ids=[151645], - ) - super().__init__( - control_result=control_result, - collect_outputs=[[terminal_segment], [], []], - collect_delay_s=0.01, - ) - self.append_sequence: list[str] = [] - self.first_real_started = asyncio.Event() - self.second_real_started = asyncio.Event() - self.real_count = 0 - - async def append_duplex_input_async(self, session_id: str, **kwargs): - kwargs.pop("expected_epoch", None) - result = await super().append_duplex_input_async(session_id, **kwargs) - payload = kwargs.get("payload") - is_silence = ( - isinstance(payload, dict) - and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - ) - if is_silence: - self.append_sequence.append("silence") - else: - self.real_count += 1 - self.append_sequence.append(f"real-{self.real_count}") - if self.real_count == 1: - self.first_real_started.set() - elif self.real_count == 2: - self.second_real_started.set() - return result - - session_created = asyncio.Event() - engine = SequencedAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get("sid-native-real-before-silence") - assert session is not None - session.capabilities.chunk_period_ms = 50 - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - session_created.set() - - ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) - event = _native_session_create("sid-native-real-before-silence") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - await asyncio.wait_for(session_created.wait(), timeout=1) - await asyncio.wait_for(engine.first_real_started.wait(), timeout=1) - await asyncio.sleep(0.02) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800, value=0.07), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - await asyncio.wait_for(engine.second_real_started.wait(), timeout=1) - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.append_sequence[:2] == ["real-1", "real-2"] - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_silence_tail(): - request_id = "duplex-sid-native-silence-tail-e0-stage0" - - class SubmittedSilenceEngine(FakeEngineClient): - def __init__(self) -> None: - control_result = { - "operation": "append", - "session_id": "sid-native-silence-tail", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - terminal_segment = _duplex_tts_output( - request_id=request_id, - samples=0, - finished=True, - tts_is_last_chunk=True, - token_ids=[151645], - ) - super().__init__( - control_result=control_result, - collect_outputs=[[terminal_segment], [], []], - collect_delay_s=0.01, - ) - self.silence_append_started = asyncio.Event() - self.release_silence_append = asyncio.Event() - self.append_sequence: list[str] = [] - self.real_append_count = 0 - - async def append_duplex_input_async(self, session_id: str, **kwargs): - kwargs.pop("expected_epoch", None) - payload = kwargs.get("payload") - is_silence_continuation = ( - isinstance(payload, dict) - and payload.get("format") == "pcm_f32le" - and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - ) - if is_silence_continuation: - result = await super().append_duplex_input_async(session_id, **kwargs) - self.append_sequence.append("silence-started") - self.silence_append_started.set() - await self.release_silence_append.wait() - self.append_sequence.append("silence-returned") - return result - result = await super().append_duplex_input_async(session_id, **kwargs) - if ( - isinstance(payload, dict) - and payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - ): - self.real_append_count += 1 - self.append_sequence.append(f"real-{self.real_append_count}") - return result - - engine = SubmittedSilenceEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get("sid-native-silence-tail") - assert session is not None - session.capabilities.chunk_period_ms = 50 - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - - ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) - event = _native_session_create("sid-native-silence-tail") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - await asyncio.wait_for(engine.silence_append_started.wait(), timeout=1) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.07), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - await asyncio.sleep(0.05) - - assert engine.real_append_count == 1 - assert engine.append_sequence == ["real-1", "silence-started"] - - engine.release_silence_append.set() - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.real_append_count == 2 - assert engine.append_sequence == ["real-1", "silence-started", "silence-returned", "real-2"] - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_preserves_silence_continuations_across_append_tails(): - request_id = "duplex-sid-native-continuation-tails-e0-stage0" - - class SequencedAppendEngine(FakeEngineClient): - def __init__(self) -> None: - control_result = { - "operation": "append", - "session_id": "sid-native-continuation-tails", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - super().__init__(control_result=control_result) - self.real_append_started = asyncio.Event() - self.release_real_append = asyncio.Event() - self.first_silence_started = asyncio.Event() - self.release_first_silence = asyncio.Event() - self.second_silence_started = asyncio.Event() - self.release_collect = asyncio.Event() - self.silence_count = 0 - - async def append_duplex_input_async(self, session_id: str, **kwargs): - kwargs.pop("expected_epoch", None) - payload = kwargs.get("payload") - assert isinstance(payload, dict) - result = await super().append_duplex_input_async(session_id, **kwargs) - if payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO: - self.real_append_started.set() - await self.release_real_append.wait() - return result - - self.silence_count += 1 - if self.silence_count == 1: - self.first_silence_started.set() - await self.release_first_silence.wait() - else: - self.second_silence_started.set() - return result - - async def collect_duplex_data_plane_outputs_async( - self, - request_id: str, - *, - response_stage_id: int | None = None, - timeout: float | None = None, - ) -> list[object]: - del timeout - self.collected.append((request_id, response_stage_id)) - await self.release_collect.wait() - return [] - - session_created = asyncio.Event() - - def on_send(_ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "session.created": - session_created.set() - - engine = SequencedAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) - event = _native_session_create("sid-native-continuation-tails") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - await asyncio.wait_for(session_created.wait(), timeout=1) - await asyncio.wait_for(engine.real_append_started.wait(), timeout=1) - session = handler._registry.get("sid-native-continuation-tails") - assert session is not None - session.capabilities.chunk_period_ms = 50 - session.bind_request(request_id) - response_id = session.begin_response(turn_id=session.turn_id) - native = handler._runtime_session_state(session) - scheduler = native.silence_continuation_scheduler - assert scheduler is not None - - payload = handler._native_silence_unit_payload() - payload["duplex_turn_id"] = session.turn_id - kwargs = { - "request_id": request_id, - "owner_id": f"response:{response_id}", - "response_id": response_id, - "response_owned": True, - "expected_epoch": session.epoch, - "expected_incarnation": session.incarnation, - "expected_model_turn_id": None, - "send_json": ws.send_json, - } - assert await scheduler(payload, **kwargs) is True - - next_schedule = asyncio.create_task(scheduler(payload, **kwargs)) - await asyncio.sleep(0.05) - assert not next_schedule.done() - - engine.release_real_append.set() - await asyncio.wait_for(engine.first_silence_started.wait(), timeout=1) - assert not next_schedule.done() - - engine.release_first_silence.set() - assert await asyncio.wait_for(next_schedule, timeout=1) is True - await asyncio.wait_for(engine.second_silence_started.wait(), timeout=1) - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.silence_count == 2 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_uses_segment_text_metadata_for_transcript_cursor(): - request_id = "duplex-sid-native-repeat-text-e0-stage0" - - def _stage_output( - text: str, - samples: int, - *, - segment_text: str, - turn_id: int = 0, - turn_end: bool = False, - token_ids: list[int] | None = None, - ): - return SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[ - SimpleNamespace( - text=text, - token_ids=list(token_ids or []), - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.llm_output_text_utf8": np.frombuffer(segment_text.encode("utf-8"), dtype=np.uint8), - "meta": {"duplex_turn_id": turn_id, "turn_end": turn_end}, - }, - ) - ], - ) - - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-repeat-text", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - async def emit(output: object) -> None: - result = {"data_plane_outputs": [output]} - for native_result in handler._serving_runtime_adapter.data_plane.project( - result, - context=_data_plane_context(session), - ): - await handler._send_one_native_duplex_event(send_json, native_result, session=session) - - await emit(_stage_output("same reply", 24000, segment_text="same reply")) - await emit( - _stage_output("same replysame reply", 24000, segment_text="same reply", turn_end=True, token_ids=[151645]) - ) - - await emit(_stage_output("same replysame replysame reply", 48000, segment_text="same reply", turn_id=1)) - await emit( - _stage_output( - "same replysame replysame replysame reply", - 48000, - segment_text="same reply", - turn_id=1, - turn_end=True, - token_ids=[151645], - ) - ) - - created = [m for m in sent if m.get("type") == "response.created"] - done = [m for m in sent if m.get("type") == "response.done"] - deltas = [m for m in sent if m.get("type") == "response.output_audio.delta" and m.get("audio")] - protocol = NativeRealtimeSessionProtocol(TimedWebSocket()) # type: ignore[arg-type] - realtime_events = [event for payload in sent for event in protocol._from_duplex_event(payload)] - transcript_deltas = [ - m for m in realtime_events if m.get("type") == "response.output_audio_transcript.delta" and m.get("delta") - ] - assert len(created) == 2 - assert len(done) == 2 - assert len(deltas) == 2 - assert [m["text"] for m in deltas] == ["same reply", "same reply"] - assert [m["audio"] for m in deltas] == ["wav-24000", "wav-24000"] - assert [m["delta"] for m in transcript_deltas] == ["same reply", "same reply"] - assert [m["response_id"] for m in deltas] == [m["response_id"] for m in created] - assert [m["response_id"] for m in done] == [m["response_id"] for m in created] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_continuous_speak_reuses_active_response_until_turn_end(): - request_id = "duplex-sid-native-continuous-speak-e0-stage0" - - def _stage_output( - text: str, - samples: int, - *, - turn_end: bool = False, - token_ids: list[int] | None = None, - ): - return SimpleNamespace( - request_id=request_id, - finished=turn_end, - outputs=[ - SimpleNamespace( - text=text, - token_ids=list(token_ids or []), - multimodal_output={}, - ) - ], - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([0], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - "meta.turn_end": np.array([int(turn_end)], dtype=np.int32), - "meta.tts_is_last_chunk": np.array([int(turn_end)], dtype=np.int32), - }, - ) - - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-continuous-speak", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - async def emit(output: object) -> None: - result = {"data_plane_outputs": [output]} - for native_result in handler._serving_runtime_adapter.data_plane.project( - result, - context=_data_plane_context(session), - ): - await handler._send_one_native_duplex_event(send_json, native_result, session=session) - - await emit(_stage_output("hello", 24000)) - await emit(_stage_output("hello again", 48000)) - - created = [m for m in sent if m.get("type") == "response.created"] - deltas = [m for m in sent if m.get("type") == "response.output_audio.delta" and m.get("audio")] - assert len(created) == 1 - assert len(deltas) == 2 - assert {m["response_id"] for m in deltas} == {created[0]["response_id"]} - assert "response.done" not in [m.get("type") for m in sent] - - await emit(_stage_output("hello again", 48000, turn_end=True, token_ids=[151645])) - - done = [m for m in sent if m.get("type") == "response.done"] - assert len(done) == 1 - assert done[0]["response_id"] == created[0]["response_id"] - assert session.active_response_id is None - assert session.turn_id == 1 - - -def test_response_timing_binds_latest_request_start_for_model_turn(): - session = DuplexSession( - session_id="sid-response-request-timing", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_model_turn_request_started(0, 10.0) - session.mark_model_turn_request_started(0, 11.0) - session.begin_response(turn_id=0) - - assert session.mark_response_first_outputs( - observed_at_s=11.2, - has_text=True, - has_audio=False, - )["ttft_ms"] == pytest.approx(200.0) - - session.mark_model_turn_request_started(0, 12.0) - assert session.mark_response_first_outputs( - observed_at_s=12.3, - has_text=False, - has_audio=True, - )["ttfp_ms"] == pytest.approx(1300.0) - - -def test_response_tpot_fallback_ignores_single_token_segment(): - session = DuplexSession( - session_id="sid-response-tpot", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response(turn_id=0) - - first_metrics = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 1, "vllm_tpot_ms": 900.0}}) - assert "vllm_tpot_ms" not in first_metrics["0"] - - second_metrics = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 2, "vllm_tpot_ms": 15.0}}) - assert second_metrics["0"]["num_tokens_out"] == 3 - assert second_metrics["0"]["vllm_tpot_ms"] == 15.0 - - -def test_response_tpot_keeps_token_weighted_value_when_itls_exist(): - session = DuplexSession( - session_id="sid-response-tpot-multi-token", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response(turn_id=0) - metrics = session.accumulate_response_stage_metrics( - { - "0": { - "num_tokens_out": 4, - "vllm_tpot_ms": 10.0, - "vllm_itls_ms": [30.0], - } - } - ) - assert metrics["0"]["vllm_itls_ms"] == [30.0] - assert metrics["0"]["vllm_itl_ms"] == 30.0 - assert metrics["0"]["vllm_tpot_ms"] == 10.0 - - -@pytest.mark.asyncio -async def test_continuous_response_metrics_accumulate_only_owned_model_units(monkeypatch): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-response-metrics", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_model_turn_request_started(0, 10.0) - monkeypatch.setattr( - "vllm_omni.entrypoints.duplex.runtime_bridge.time.monotonic", - lambda: 10.25, - ) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - async def emit(*, text: str, audio: str, tokens: int, ttft_ms: float, itls: list[float], tpot_ms: float) -> None: - await handler._send_one_native_duplex_event( - send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": "duplex-sid-response-metrics-e0-stage0", - "text": text, - "audio_data": audio, - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 0, - "stage_metrics": { - "0": { - "num_tokens_out": tokens, - "stage_gen_time_ms": 80.0, - "vllm_ttft_ms": ttft_ms, - "vllm_tpot_ms": tpot_ms, - "vllm_itls_ms": itls, - } - }, - }, - session=session, - ) - - await emit(text="first", audio="audio-a", tokens=3, ttft_ms=50.0, itls=[10.0, 11.0], tpot_ms=10.5) - await emit(text="second", audio="audio-b", tokens=4, ttft_ms=70.0, itls=[12.0, 13.0, 14.0], tpot_ms=13.0) - - deltas = [payload for payload in sent if payload.get("type") == "response.output_audio.delta"] - assert len(deltas) == 2 - first_metrics = deltas[0]["vllm_omni"]["stage_metrics"]["0"] - second_metrics = deltas[1]["vllm_omni"]["stage_metrics"]["0"] - assert first_metrics["num_tokens_out"] == 3 - assert first_metrics["vllm_tpot_ms"] == pytest.approx(10.5) - assert deltas[0]["vllm_omni"]["response_request_metrics"] == { - "source": "server_monotonic_request_start", - "measurement_origin": { - "ttft": "native model-turn request execution start to first non-empty text output", - "ttfp": "native model-turn request execution start to first audio output", - }, - "ttft_ms": 250.0, - "ttfp_ms": 250.0, - } - assert second_metrics["num_tokens_out"] == 7 - assert second_metrics["vllm_ttft_ms"] == 50.0 - assert second_metrics["vllm_itls_ms"] == [10.0, 11.0, 12.0, 13.0, 14.0] - assert second_metrics["vllm_tpot_ms"] == pytest.approx(12.0) - - first_response_id = deltas[0]["response_id"] - session.end_response() - await emit(text="third", audio="audio-c", tokens=3, ttft_ms=90.0, itls=[15.0, 16.0], tpot_ms=15.5) - - third_delta = [payload for payload in sent if payload.get("type") == "response.output_audio.delta"][-1] - third_metrics = third_delta["vllm_omni"]["stage_metrics"]["0"] - assert third_delta["response_id"] != first_response_id - assert third_metrics["num_tokens_out"] == 3 - assert third_metrics["vllm_ttft_ms"] == 90.0 - assert third_metrics["vllm_itls_ms"] == [15.0, 16.0] - assert third_metrics["vllm_tpot_ms"] == pytest.approx(15.5) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_drops_old_turn_audio_while_new_response_is_active(): - request_id = "duplex-sid-native-old-turn-audio-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-old-turn-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - response_id = session.begin_response(turn_id=1) - session.turn_id = 2 - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "stale", - "audio_data": "old-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is False - assert session.active_response_id == response_id - assert ws.sent == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_future_turn_audio_starts_next_response(): - request_id = "duplex-sid-native-future-turn-audio-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-future-turn-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - first_response_id = session.begin_response(turn_id=0) - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "next", - "audio_data": "next-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 1, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - sent_types = [event["type"] for event in ws.sent] - assert sent_types == [ - "response.done", - "response.created", - "response.speak", - "response.output_audio.delta", - ] - assert ws.sent[0]["response_id"] == first_response_id - assert ws.sent[1]["response_id"] != first_response_id - assert ws.sent[3]["response_id"] == ws.sent[1]["response_id"] - assert ws.sent[3]["vllm_omni"]["model_turn_id"] == 1 - assert session.active_response_turn_id == 1 - assert session.turn_id == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_drops_old_turn_listen_while_new_response_is_active(): - request_id = "duplex-sid-native-old-turn-listen-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-old-turn-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - response_id = session.begin_response(turn_id=1) - session.turn_id = 2 - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is False - assert session.active_response_id == response_id - assert ws.sent == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_silence_continuation_requires_active_request(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - session = DuplexSession( - session_id="sid-native-stale-continuation", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response() - session.bind_request("duplex-sid-native-stale-continuation-stage0") - session.mark_audio_sent(100) - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any]) -> None: - sent.append(data) - - _install_direct_silence_scheduler(handler, session) - - await handler._maybe_continue_native_response(send_json, session=session, expected_epoch=session.epoch) - session.clear_request() - await asyncio.sleep(0.01) - - assert engine.appended == [] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_segment_complete_continues_until_model_listen(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - session = DuplexSession( - session_id="sid-native-segment-complete", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response() - request_id = "duplex-sid-native-segment-complete-stage0" - session.bind_request(request_id) - session.mark_audio_sent(1000) - _install_direct_silence_scheduler(handler, session) - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any]) -> None: - sent.append(data) - - close_reason, emitted = await handler._send_one_native_duplex_event( - send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": False, - "listen_source": "auto_response_segment_complete", - "reason": "auto_response_segment_complete", - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - "owned_runtime": False, - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert close_reason is None - assert emitted is False - assert session.active_response_id == response_id - assert session.active_request_id == request_id - assert not handler._serving_runtime_adapter.data_plane.is_terminal(request_id) - assert sent == [] - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert payload["type"] == "audio" - assert payload.get("force_listen") is not True - assert final is False - - -@pytest.mark.parametrize( - ("steps", "expected_deltas"), - [ - ( - [("你好呀", 1), ("你好呀", 1), ("你好呀", 2)], - ["你好呀", "", "你好呀"], - ), - ( - [("你好呀", 1), ("你好呀", 2)], - ["你好呀", "你好呀"], - ), - ], - ids=["turn-scoped", "exact-repeated-turn-text"], -) -def test_minicpmo_data_plane_text_cursor( - steps: list[tuple[str, int]], - expected_deltas: list[str], -): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-e0-stage0" - - deltas = [data_plane.segment_text_delta(request_id, text, turn_id=turn_id) for text, turn_id in steps] - - assert deltas == expected_deltas - - -def test_minicpmo_data_plane_text_cursor_does_not_clip_prior_turn_overlap(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-e0-stage0" - - assert data_plane.segment_text_delta(request_id, "你好呀。你有什么想聊的吗？", turn_id=1) - assert ( - data_plane.segment_text_delta(request_id, "你有什么想聊的吗？你好呀。", turn_id=2) - == "你有什么想聊的吗？你好呀。" - ) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_cancel_interrupts_background_data_plane_stream(): - control_result = { - "operation": "append", - "session_id": "sid-native-cancel-stream", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": "duplex-sid-native-cancel-stream-e0-stage0-s1", - "response_stage_id": 1, - }, - } - ], - } - engine = FakeEngineClient( - append_result=control_result, - collect_delay_s=0.2, - collect_outputs=[ - [ - SimpleNamespace( - request_id="duplex-sid-native-cancel-stream-e0-stage0-s1", - finished=True, - outputs=[ - SimpleNamespace( - text="", - multimodal_output={"audio": np.zeros(10, dtype=np.float32), "sr": 24000}, - ) - ], - ) - ] - ], - ) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-cancel-stream")) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - ws.put({"type": "input.cancel"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "audio.cancelled" in ws.sent_types() - assert "response.output_audio.delta" not in ws.sent_types() - assert engine.signals == [("sid-native-cancel-stream", "barge_in")] - assert engine.signal_fences == [DuplexFence("sid-native-cancel-stream")] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_session_close_discards_partial_pcm_tail(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-short-chunks")) - chunk = _pcm_f32_b64(3200) - for _ in range(6): - ws.put( - { - "type": "input_audio_buffer.append", - "audio": chunk, - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert final is False - assert isinstance(payload, dict) - assert len(base64.b64decode(payload["audio"])) == 16000 * 4 - assert "duplex_num_input_tokens" not in payload - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_idle_timeout_closes_runtime_with_timeout_reason(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=0.1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-timeout")) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "session.closed"] - assert ws.sent[-1]["reason"] == "timeout" - assert engine.closed == [("sid-native-timeout", "timeout")] diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index be83ea9db..d0476b56c 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -32,7 +32,7 @@ from tests.helpers.media import ( from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.config.stage_config import load_deploy_config -from vllm_omni.entrypoints.duplex.server_vad import ( +from vllm_omni.engine.duplex.vad import ( SILERO_VAD_FILENAME, SILERO_VAD_REPO_ID, SILERO_VAD_REVISION, diff --git a/tests/entrypoints/openai_api/test_server_vad.py b/tests/entrypoints/openai_api/test_server_vad.py deleted file mode 100644 index 67476be73..000000000 --- a/tests/entrypoints/openai_api/test_server_vad.py +++ /dev/null @@ -1,429 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import base64 -import io -import sys -import threading -import wave -import weakref -from concurrent.futures import ThreadPoolExecutor -from types import SimpleNamespace - -import numpy as np -import pytest - -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexSession, - DuplexSessionConfig, -) -from vllm_omni.entrypoints.duplex.server_vad import ( - ServerVADConfig, - ServerVADFrame, - ServerVADPipeline, - SileroVADBackend, - SileroVADBackendProvider, - SpeechEndpointDecision, - ThresholdEndpointPolicy, - parse_session_turn_detection, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class SequenceDetector: - frame_samples = 160 - - def __init__(self, probabilities: list[float]) -> None: - self.probabilities = probabilities - - def new_state(self) -> int: - return 0 - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - index = int(state) - probability = self.probabilities[index] if index < len(self.probabilities) else 0.0 - return probability, index + 1 - - -def test_server_vad_config_defaults_and_validation(): - config = ServerVADConfig.from_value({"type": "server_vad"}) - - assert config.type == "server_vad" - assert config.threshold == 0.5 - assert config.prefix_padding_ms == 300 - assert config.silence_duration_ms == 500 - assert config.create_response is True - assert config.interrupt_response is None - assert "interrupt_response" not in config.as_dict() - assert config.min_speech_duration_ms is None - assert config.as_dict()["type"] == config.type - - interrupting = ServerVADConfig.from_value( - { - "type": "server_vad", - "interrupt_response": True, - "min_speech_duration_ms": 96, - } - ) - assert interrupting.interrupt_response is True - assert interrupting.min_speech_duration_ms == 96 - for value in (None, 0, "true"): - with pytest.raises(ValueError, match="interrupt_response must be a boolean"): - ServerVADConfig.from_value({"type": "server_vad", "interrupt_response": value}) - for enabled in (False, True): - _, aliased = parse_session_turn_detection( - { - "turn_detection": {"type": "server_vad"}, - "audio": {"input": {"turn_detection": {"type": "server_vad", "interrupt_response": enabled}}}, - } - ) - assert aliased is not None and aliased.interrupt_response is enabled - with pytest.raises(ValueError, match="Unknown"): - ServerVADConfig.from_value({"type": "server_vad", "semantic_eagerness": "high"}) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prefix_padding_ms", [0, 10, 30]) -async def test_server_vad_pipeline_handles_arbitrary_chunk_boundaries(prefix_padding_ms: int): - detector = SequenceDetector([0.0, 0.9, 0.8, 0.0, 0.0]) - pipeline = ServerVADPipeline( - detector, - ServerVADConfig(prefix_padding_ms=prefix_padding_ms, silence_duration_ms=20), - ) - samples = np.zeros(detector.frame_samples * 5, dtype=np.float32) - - batches = [ - await pipeline.push(samples[:73]), - await pipeline.push(samples[73:431]), - await pipeline.push(samples[431:]), - ] - frames = [frame for batch in batches for frame in batch.frames] - - assert len(frames) == 5 - assert [index for index, frame in enumerate(frames) if frame.decision.speech_started] == [1] - assert [index for index, frame in enumerate(frames) if frame.decision.speech_stopped] == [4] - assert frames[1].decision.audio_start_ms == max(0, 10 - prefix_padding_ms) - assert frames[4].decision.audio_end_ms == 50 - - pipeline.reset() - reset_batch = await pipeline.push(samples) - assert [index for index, frame in enumerate(reset_batch.frames) if frame.decision.speech_started] == [1] - assert [index for index, frame in enumerate(reset_batch.frames) if frame.decision.speech_stopped] == [4] - # Keep session-relative timestamps, but never reference prefix audio - # discarded by reset at 50 ms. - assert reset_batch.frames[1].decision.audio_start_ms == max(50, 60 - prefix_padding_ms) - assert reset_batch.frames[4].decision.audio_end_ms == 100 - - -@pytest.mark.asyncio -async def test_server_vad_retained_audio_does_not_keep_input_chunk(): - pipeline = ServerVADPipeline(SequenceDetector([]), ServerVADConfig()) - session = DuplexSession("server-vad-memory", DuplexSessionConfig()) - samples = np.zeros(160 * 100 + 1, dtype=np.float32) - source_ref = weakref.ref(samples) - batch = await pipeline.push(samples) - session.append_server_vad_frame( - batch.frames[-1].samples, - speech_started=False, - speech_stopped=False, - prefix_samples=80, - ) - session.append_server_vad_frame( - batch.frames[-1].samples, - speech_started=True, - speech_stopped=False, - prefix_samples=80, - ) - del batch, samples - # Allow the completed to_thread callback to release its result, too. - await asyncio.sleep(0) - - assert source_ref() is None - assert pipeline.scratch_bytes == 4 - assert session.server_vad_utterance_bytes == (80 + 160) * 4 - - -@pytest.mark.asyncio -async def test_server_vad_pipeline_pcm16_resampling_is_split_invariant(): - sample_index = np.arange(2_400, dtype=np.int32) - source = ((sample_index * 7919) % 65_536 - 32_768).astype(" tuple[np.ndarray, list[SpeechEndpointDecision]]: - pipeline = ServerVADPipeline( - SequenceDetector([0.0] * 10), - ServerVADConfig(), - ) - frames: list[ServerVADFrame] = [] - offset = 0 - for chunk_size in chunk_sizes: - chunk = continuous_source[offset : offset + chunk_size] - batch = await pipeline.push_pcm16(chunk.tobytes(), source_sample_rate_hz=24_000) - frames.extend(batch.frames) - offset += chunk_size - assert offset == continuous_source.size - return np.concatenate([frame.samples for frame in frames]), [frame.decision for frame in frames] - - whole_samples, whole_decisions = await process([continuous_source.size]) - split_samples, split_decisions = await process([7, 13, 511, 512, 513, continuous_source.size - 1_556]) - - np.testing.assert_array_equal(split_samples, whole_samples) - assert whole_samples.size == 1_600 - assert split_decisions == whole_decisions - - -@pytest.mark.asyncio -async def test_server_vad_pipeline_reset_clears_resampler_and_source_rate_lock(): - pipeline = ServerVADPipeline( - SequenceDetector([0.0]), - ServerVADConfig(), - ) - - batch = await pipeline.push_pcm16( - np.asarray([30_000, -30_000], dtype=" None: - self.providers = providers - self.sess_options = sess_options - self.calls: list[dict[str, np.ndarray]] = [] - self.run_barrier: threading.Barrier | None = None - sessions.append(self) - - def get_inputs(self) -> list[SimpleNamespace]: - return [SimpleNamespace(name=name) for name in ("input", "state", "sr")] - - def run(self, _outputs: object, inputs: dict[str, np.ndarray]) -> list[np.ndarray]: - self.calls.append({name: np.array(value, copy=True) for name, value in inputs.items()}) - if self.run_barrier is not None: - self.run_barrier.wait(timeout=5) - return [ - np.asarray([[0.75]], dtype=np.float32), - np.full((2, 1, 128), len(self.calls), dtype=np.float32), - ] - - sessions: list[FakeInferenceSession] = [] - monkeypatch.setitem( - sys.modules, - "onnxruntime", - SimpleNamespace( - SessionOptions=FakeSessionOptions, - InferenceSession=FakeInferenceSession, - ), - ) - - model_path = tmp_path / "silero_vad.onnx" - model_path.write_bytes(b"fake-model") - backend = SileroVADBackend(model_path) - session = sessions[0] - - assert session.providers == ["CPUExecutionProvider"] - assert session.sess_options.inter_op_num_threads == 1 - assert session.sess_options.intra_op_num_threads == 1 - assert session.calls[0]["input"].shape == (1, 576) - assert session.calls[0]["state"].shape == (2, 1, 128) - np.testing.assert_array_equal(session.calls[0]["state"], np.zeros((2, 1, 128), dtype=np.float32)) - assert session.calls[0]["sr"].shape == () - assert session.calls[0]["sr"].item() == 16_000 - - state = backend.new_state() - first_frame = np.arange(backend.frame_samples, dtype=np.float32) - probability, state = backend.infer(first_frame, state) - first_call = session.calls[1] - - assert probability == pytest.approx(0.75) - np.testing.assert_array_equal( - first_call["input"][:, : backend.context_samples], - np.zeros((1, backend.context_samples), dtype=np.float32), - ) - np.testing.assert_array_equal( - first_call["input"][:, backend.context_samples :], - first_frame[None, :], - ) - np.testing.assert_array_equal(first_call["state"], np.zeros((2, 1, 128), dtype=np.float32)) - - second_frame = -first_frame - backend.infer(second_frame, state) - second_call = session.calls[2] - np.testing.assert_array_equal( - second_call["input"][:, : backend.context_samples], - first_frame[None, -backend.context_samples :], - ) - np.testing.assert_array_equal( - second_call["input"][:, backend.context_samples :], - second_frame[None, :], - ) - np.testing.assert_array_equal( - second_call["state"], - np.full((2, 1, 128), 2, dtype=np.float32), - ) - - # Independent pipeline states may enter the shared ORT session concurrently. - session.run_barrier = threading.Barrier(2) - with ThreadPoolExecutor(max_workers=2) as executor: - futures = [ - executor.submit(backend.infer, first_frame, backend.new_state()), - executor.submit(backend.infer, second_frame, backend.new_state()), - ] - assert [future.result()[0] for future in futures] == pytest.approx([0.75, 0.75]) - - -def test_duplex_session_owns_server_vad_prefix_and_utterance_audio(): - session = DuplexSession("server-vad-session", DuplexSessionConfig()) - frame = np.zeros(160, dtype=np.float32) - - session.reserve_input_bytes(frame.nbytes * 3, limit=frame.nbytes * 4) - released = session.append_server_vad_frame( - frame, - speech_started=False, - speech_stopped=False, - prefix_samples=160, - ) - assert released == 0 - released = session.append_server_vad_frame( - frame, - speech_started=False, - speech_stopped=False, - prefix_samples=160, - ) - assert released == frame.nbytes - session.release_input_bytes(released) - session.append_server_vad_frame( - frame, - speech_started=True, - speech_stopped=False, - prefix_samples=160, - ) - - assert session.server_vad_utterance_bytes == frame.nbytes * 2 - assert session.stage_server_vad_audio_for_commit() is True - committed = session.commit_user_input() - assert committed is not None - assert committed.message["role"] == "user" - audio_url = committed.message["content"][0]["audio_url"]["url"] - assert audio_url.startswith("data:audio/wav;base64,") - wav_bytes = base64.b64decode(audio_url.partition(",")[2]) - with wave.open(io.BytesIO(wav_bytes), "rb") as wav_file: - assert wav_file.getnchannels() == 1 - assert wav_file.getsampwidth() == 2 - assert wav_file.getframerate() == 16_000 - assert wav_file.getnframes() == frame.size * 2 - - -def test_silero_provider_rejects_missing_or_invalid_local_artifact(tmp_path): - missing = SileroVADBackendProvider(model_path=str(tmp_path / "missing.onnx")) - with pytest.raises(RuntimeError, match="does not exist"): - missing.get() - - invalid_path = tmp_path / "silero_vad.onnx" - invalid_path.write_bytes(b"not-the-pinned-model") - invalid = SileroVADBackendProvider(model_path=str(invalid_path)) - with pytest.raises(RuntimeError, match="checksum mismatch"): - invalid.get() diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index dda9ab38c..3b378930a 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -674,7 +674,7 @@ def test_resolve_stage_configs_delegates_overrides_to_resolver(mocker): engine_args=SimpleNamespace(additional_config=additional_config), ) resolve_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_config", + "vllm_omni.engine.omni_engine_base.resolve_omni_config", return_value=OmniConfigResolution( config_path="dummy.yaml", stage_configs=(fake_diffusion_stage,), diff --git a/tests/entrypoints/test_async_omni_duplex.py b/tests/entrypoints/test_async_omni_duplex.py deleted file mode 100644 index 81e0d75cb..000000000 --- a/tests/entrypoints/test_async_omni_duplex.py +++ /dev/null @@ -1,601 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import pytest - -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexOutputAction, - DuplexOutputDecision, - duplex_resource_request_id, -) -from vllm_omni.engine.messages import OutputMessage -from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.duplex_request_client import DuplexRequestClient -from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import attach_duplex_output_decision - - -def test_async_omni_uses_extracted_duplex_request_client(): - app = object.__new__(AsyncOmni) - app.engine = SimpleNamespace(num_stages=2) - app.request_states = {} - app._duplex_request_client = None - - client = app._get_duplex_request_client() - - assert isinstance(client, DuplexRequestClient) - assert not hasattr(client, "_async_omni") - - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -@pytest.mark.asyncio -async def test_async_omni_open_duplex_forwards_session_config_and_timeout(): - calls = [] - output_handler_starts = 0 - - async def open_duplex_session_async(session_id, **kwargs): - calls.append((session_id, kwargs)) - return {"ok": True} - - app = object.__new__(AsyncOmni) - app.engine = SimpleNamespace(open_duplex_session_async=open_duplex_session_async) - - def start_output_handler(): - nonlocal output_handler_starts - output_handler_starts += 1 - - app._final_output_handler = start_output_handler - - result = await app.open_duplex_session_async( - "sid", - fence=DuplexFence("sid"), - capabilities={"implementation_level": "model_native_duplex"}, - session_config={"instructions": "Be brief.", "idle_timeout_s": 30}, - timeout=7.5, - ) - - assert result == {"ok": True} - assert output_handler_starts == 1 - assert calls == [ - ( - "sid", - { - "session_mode": "duplex", - "capabilities": {"implementation_level": "model_native_duplex"}, - "session_config": {"instructions": "Be brief.", "idle_timeout_s": 30}, - "fence": DuplexFence("sid"), - "timeout": 7.5, - }, - ) - ] - - -@pytest.mark.asyncio -async def test_async_omni_duplex_runtime_controls_forward_timeout(): - calls = [] - cancelled_fence = DuplexFence("sid", epoch=2, turn_id=3) - next_fence = DuplexFence("sid", epoch=3, turn_id=3) - - async def append_duplex_input_async(session_id, **kwargs): - calls.append(("append", session_id, kwargs)) - return {"ok": True} - - async def signal_duplex_turn_async(session_id, **kwargs): - calls.append(("signal", session_id, kwargs)) - return {"ok": True} - - async def close_duplex_session_async(session_id, **kwargs): - calls.append(("close", session_id, kwargs)) - return {"ok": True} - - async def touch_duplex_session_async(session_id, **kwargs): - calls.append(("touch", session_id, kwargs)) - return {"ok": True} - - async def resume_duplex_session_async(session_id, **kwargs): - calls.append(("resume", session_id, kwargs)) - return {"ok": True} - - app = object.__new__(AsyncOmni) - app.request_states = {} - app._final_output_handler = lambda: None - app.engine = SimpleNamespace( - append_duplex_input_async=append_duplex_input_async, - signal_duplex_turn_async=signal_duplex_turn_async, - close_duplex_session_async=close_duplex_session_async, - touch_duplex_session_async=touch_duplex_session_async, - resume_duplex_session_async=resume_duplex_session_async, - ) - - await app.append_duplex_input_async( - "sid", - mode="append_audio_chunk", - payload={}, - fence=cancelled_fence, - timeout=12.5, - ) - await app.signal_duplex_turn_async( - "sid", - event="barge_in", - fence=cancelled_fence, - next_fence=next_fence, - timeout=13.5, - ) - await app.close_duplex_session_async("sid", reason="done", fence=next_fence, timeout=14.5) - await app.touch_duplex_session_async( - "sid", - fence=next_fence, - activity=DuplexLeaseActivity.HEARTBEAT, - timeout=15.5, - ) - await app.resume_duplex_session_async( - "sid", - fence=next_fence, - expected_lease_generation=3, - timeout=16.5, - ) - - assert calls == [ - ( - "append", - "sid", - { - "mode": "append_audio_chunk", - "payload": {}, - "final": False, - "expected_epoch": None, - "fence": cancelled_fence, - "timeout": 12.5, - }, - ), - ( - "signal", - "sid", - { - "event": "barge_in", - "fence": cancelled_fence, - "next_fence": next_fence, - "timeout": 13.5, - }, - ), - ( - "close", - "sid", - { - "reason": "done", - "fence": next_fence, - "timeout": 14.5, - }, - ), - ( - "touch", - "sid", - { - "fence": next_fence, - "activity": DuplexLeaseActivity.HEARTBEAT, - "timeout": 15.5, - }, - ), - ( - "resume", - "sid", - { - "fence": next_fence, - "expected_lease_generation": 3, - "timeout": 16.5, - }, - ), - ] - - -@pytest.mark.asyncio -async def test_async_omni_duplex_append_can_defer_data_plane_output_collection(): - app = object.__new__(AsyncOmni) - request_id = duplex_resource_request_id(DuplexFence("sid"), "stage0") - - async def append_duplex_input_async(session_id, **kwargs): - del session_id, kwargs - request_state = app.request_states.get(request_id) - assert request_state is not None - assert request_state.metrics is not None - return { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ] - } - - app.engine = SimpleNamespace(append_duplex_input_async=append_duplex_input_async) - app.request_states = {} - app._final_output_handler = lambda: None - - result = await app.append_duplex_input_async( - "sid", - mode="append_audio_chunk", - payload={}, - fence=DuplexFence("sid"), - collect_outputs=False, - ) - - assert "data_plane_outputs" not in result - - -@pytest.mark.asyncio -async def test_async_omni_duplex_append_forwards_fence(): - calls = [] - - async def append_duplex_input_async(session_id, **kwargs): - calls.append((session_id, kwargs)) - return {"ok": True} - - app = object.__new__(AsyncOmni) - app.request_states = {} - app._final_output_handler = lambda: None - app.engine = SimpleNamespace(append_duplex_input_async=append_duplex_input_async) - fence = DuplexFence("sid", epoch=1, turn_id=2) - - await app.append_duplex_input_async( - "sid", - mode="append_audio_chunk", - payload={}, - fence=fence, - collect_outputs=False, - ) - - assert calls[0][1]["fence"] is fence - - -@pytest.mark.asyncio -async def test_async_omni_duplex_collect_waits_for_response_stage(): - app = object.__new__(AsyncOmni) - req_state = ClientRequestState("duplex-sid-e0-stage0") - stage0_output = OmniRequestOutput( - request_id="duplex-sid-e0-stage0", - stage_id=0, - finished=False, - ) - stage1_output = OmniRequestOutput( - request_id="duplex-sid-e0-stage0", - stage_id=1, - final_output_type="audio", - finished=False, - ) - await req_state.queue.put( - OutputMessage( - request_id="duplex-sid-e0-stage0", - stage_id=0, - engine_outputs=stage0_output, - finished=False, - ) - ) - await req_state.queue.put( - OutputMessage( - request_id="duplex-sid-e0-stage0", - stage_id=1, - engine_outputs=stage1_output, - finished=False, - ) - ) - - outputs = await app._collect_duplex_data_plane_outputs( - "duplex-sid-e0-stage0", - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert outputs == [stage1_output] - - -@pytest.mark.asyncio -async def test_duplex_request_client_retains_output_route_at_segment_terminal(): - request_id = "duplex-sid-e0-stage0" - request_state = ClientRequestState(request_id) - request_states = {request_id: request_state} - output = OmniRequestOutput( - request_id=request_id, - stage_id=1, - final_output_type="audio", - finished=True, - ) - await request_state.queue.put( - OutputMessage( - request_id=request_id, - stage_id=1, - engine_outputs=output, - finished=True, - ) - ) - client = DuplexRequestClient( - SimpleNamespace(), - SimpleNamespace( - request_states=request_states, - num_stages=2, - log_stats=False, - ), - ) - - outputs = await client.collect_outputs( - request_id, - request_state, - response_stage_id=1, - timeout=1.0, - ) - - assert outputs == [output] - assert request_states[request_id] is request_state - - -@pytest.mark.asyncio -async def test_async_omni_duplex_direct_output_keeps_stage_metrics(): - app = object.__new__(AsyncOmni) - app._duplex_request_client = None - app._final_output_handler = lambda: None - request_id = "duplex-sid-e0-stage0" - req_state = ClientRequestState(request_id) - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={ - "duplex_direct_response": True, - "duplex_native_decision": "speak", - }, - ) - direct_output = attach_duplex_output_decision( - OmniRequestOutput( - request_id=request_id, - stage_id=0, - final_output_type="text", - finished=True, - ), - decision, - ) - await req_state.queue.put( - OutputMessage( - request_id=request_id, - stage_id=0, - engine_outputs=direct_output, - metrics=StageRequestStats( - batch_id=1, - batch_size=1, - num_tokens_in=160, - num_tokens_out=8, - stage_gen_time_ms=80.0, - rx_transfer_bytes=0, - rx_decode_time_ms=0.0, - rx_in_flight_time_ms=0.0, - stage_stats=StageStats(total_token=8, total_gen_time_ms=80.0), - vllm_ttft_ms=12.5, - vllm_tpot_ms=9.0, - vllm_itl_ms=8.75, - vllm_itls_ms=[8.0, 9.5], - ), - finished=True, - ) - ) - - outputs = await app._collect_duplex_data_plane_outputs( - request_id, - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert outputs == [direct_output] - assert direct_output.metrics["stage_metrics"]["0"]["vllm_ttft_ms"] == 12.5 - assert direct_output.metrics["stage_metrics"]["0"]["vllm_tpot_ms"] == 9.0 - assert direct_output.metrics["stage_metrics"]["0"]["vllm_itls_ms"] == [8.0, 9.5] - - -@pytest.mark.asyncio -async def test_duplex_metrics_cursor_isolates_resumable_collect_windows(): - app = object.__new__(AsyncOmni) - app._duplex_request_client = None - app._final_output_handler = lambda: None - request_id = "duplex-sid-e0-stage0" - req_state = ClientRequestState(request_id) - req_state.metrics = OrchestratorAggregator( - num_stages=2, - log_stats=False, - wall_start_ts=0.0, - final_stage_id_for_e2e=1, - ) - - def stage_message(*, batch_id: int, tokens: int, ttft_ms: float, itls: list[float]) -> OutputMessage: - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={"duplex_direct_response": True, "duplex_native_decision": "speak"}, - ) - output = attach_duplex_output_decision( - OmniRequestOutput( - request_id=request_id, - stage_id=0, - final_output_type="text", - finished=True, - ), - decision, - ) - return OutputMessage( - request_id=request_id, - stage_id=0, - engine_outputs=output, - metrics=StageRequestStats( - batch_id=batch_id, - batch_size=1, - num_tokens_in=160, - num_tokens_out=tokens, - stage_gen_time_ms=80.0, - rx_transfer_bytes=0, - rx_decode_time_ms=0.0, - rx_in_flight_time_ms=0.0, - stage_stats=StageStats(total_token=tokens, total_gen_time_ms=80.0), - vllm_ttft_ms=ttft_ms, - vllm_tpot_ms=9.0, - vllm_itl_ms=sum(itls) / len(itls), - vllm_itls_ms=itls, - ), - finished=True, - ) - - await req_state.queue.put(stage_message(batch_id=1, tokens=71, ttft_ms=62.5, itls=[10.0, 11.0])) - first = await app._collect_duplex_data_plane_outputs( - request_id, - req_state, - response_stage_id=1, - timeout=1.0, - ) - await req_state.queue.put(stage_message(batch_id=2, tokens=155, ttft_ms=81.0, itls=[12.0, 13.0, 14.0])) - second = await app._collect_duplex_data_plane_outputs( - request_id, - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert first[0].metrics["stage_metrics"]["0"]["num_tokens_out"] == 71 - second_stage0 = second[0].metrics["stage_metrics"]["0"] - assert second_stage0["num_tokens_out"] == 155 - assert second_stage0["vllm_ttft_ms"] == 81.0 - assert second_stage0["vllm_itls_ms"] == [12.0, 13.0, 14.0] - - -def test_async_omni_duplex_direct_output_prefers_outer_control_metadata(): - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={ - "duplex_direct_response": True, - "duplex_native_decision": "listen", - }, - ) - output = attach_duplex_output_decision( - OmniRequestOutput( - request_id="duplex-sid-e0-stage0", - stage_id=0, - outputs=[], - _multimodal_output={"kind": "processed-payload"}, - _custom_output={"inner": "completion"}, - ), - decision, - ) - - assert AsyncOmni._is_direct_duplex_data_plane_response(output) - - -def test_async_omni_drops_unregistered_duplex_prefixed_output(): - app = object.__new__(AsyncOmni) - app.request_states = {} - output = OmniRequestOutput(request_id="duplex-unregistered-e0-stage0", stage_id=0) - - keep_processing, request_id, stage_id, request_state = app._handle_output_message( - OutputMessage( - request_id=output.request_id, - stage_id=0, - engine_outputs=output, - finished=False, - ) - ) - - assert keep_processing is True - assert request_id is stage_id is request_state is None - assert app.request_states == {} - - -@pytest.mark.asyncio -async def test_async_omni_duplex_collect_wraps_raw_response_stage_output(): - app = object.__new__(AsyncOmni) - app.log_stats = False - app._enable_ar_profiler = False - app.request_states = {} - app.prom_metrics = SimpleNamespace( - set_running=lambda value: None, - set_waiting=lambda value: None, - ) - app.engine = SimpleNamespace( - num_stages=2, - get_stage_metadata=lambda stage_id: SimpleNamespace( - final_output=stage_id == 1, - final_output_type="audio", - ), - ) - req_state = ClientRequestState("duplex-sid-e0-stage0") - raw_stage1_output = SimpleNamespace( - finished=False, - outputs=[], - stage_durations={}, - peak_memory_mb=0.0, - final_output_type="audio", - ) - await req_state.queue.put( - OutputMessage( - request_id="duplex-sid-e0-stage0", - stage_id=1, - engine_outputs=raw_stage1_output, - finished=False, - ) - ) - - outputs = await app._collect_duplex_data_plane_outputs( - "duplex-sid-e0-stage0", - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert len(outputs) == 1 - assert outputs[0].request_id == "duplex-sid-e0-stage0" - assert outputs[0].finished is False - assert outputs[0].stage_id == 1 - assert outputs[0].final_output_type == "audio" - assert outputs[0].peak_memory_mb == raw_stage1_output.peak_memory_mb - - -def test_async_omni_duplex_request_info_includes_response_stage(): - request_id, response_stage_id = AsyncOmni._duplex_data_plane_request_info( - { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": "duplex-sid-e0-stage0", - "response_stage_id": 1, - } - } - ] - } - ) - - assert request_id == "duplex-sid-e0-stage0" - assert response_stage_id == 1 - - -@pytest.mark.asyncio -async def test_duplex_request_client_retains_output_route_when_close_fails(): - fence = DuplexFence("sid-close-route") - request_id = duplex_resource_request_id(fence, "stage0") - request_states = {request_id: ClientRequestState(request_id)} - - async def close_duplex_session_async(session_id, **kwargs): - del session_id, kwargs - raise RuntimeError("close failed") - - client = DuplexRequestClient( - SimpleNamespace(close_duplex_session_async=close_duplex_session_async), - SimpleNamespace(request_states=request_states), - ) - - with pytest.raises(RuntimeError, match="close failed"): - await client.close(fence.session_id, reason="test", fence=fence, timeout=1.0) - - assert request_id in request_states diff --git a/tests/entrypoints/test_disable_log_stats_routing.py b/tests/entrypoints/test_disable_log_stats_routing.py index f9dd2da4a..623cc9af4 100644 --- a/tests/entrypoints/test_disable_log_stats_routing.py +++ b/tests/entrypoints/test_disable_log_stats_routing.py @@ -53,7 +53,7 @@ def test_async_engine_consumes_disable_log_stats_before_resolution(monkeypatch: ) monkeypatch.setattr( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + "vllm_omni.engine.omni_engine_base.load_and_resolve_stage_configs", fake_resolve, ) engine = SimpleNamespace(_apply_strategy_lb_policy=lambda *args, **kwargs: None) diff --git a/tests/entrypoints/test_duplex_fence_propagation.py b/tests/entrypoints/test_duplex_fence_propagation.py deleted file mode 100644 index 550723ec8..000000000 --- a/tests/entrypoints/test_duplex_fence_propagation.py +++ /dev/null @@ -1,126 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import pytest - -from tests.entrypoints.openai_api.test_duplex_handler import FakeChatService, TimedWebSocket -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.duplex.protocol import DuplexSession, DuplexSessionConfig -from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class FenceRecordingEngine: - output_modalities = ["text", "audio"] - - def __init__(self) -> None: - self.calls: list[tuple[str, DuplexFence | None]] = [] - self.signal_next_fences: list[DuplexFence | None] = [] - - async def open_duplex_session_async(self, session_id, *, fence=None, **kwargs): - del session_id, kwargs - self.calls.append(("open", fence)) - return {"ok": True} - - async def append_duplex_input_async(self, session_id, *, fence=None, **kwargs): - del session_id, kwargs - self.calls.append(("append", fence)) - return {"ok": True} - - async def signal_duplex_turn_async( - self, - session_id, - *, - event, - fence=None, - next_fence=None, - timeout=None, - ): - del session_id, event, timeout - self.calls.append(("signal", fence)) - self.signal_next_fences.append(next_fence) - return {"ok": True} - - async def close_duplex_session_async(self, session_id, *, fence=None, **kwargs): - del session_id, kwargs - self.calls.append(("close", fence)) - return {"ok": True} - - -@pytest.mark.asyncio -async def test_async_omni_forwards_fence_for_all_runtime_operations(): - engine = FenceRecordingEngine() - app = object.__new__(AsyncOmni) - app.engine = engine - app.request_states = {} - app._final_output_handler = lambda: None - fence = DuplexFence("sid", epoch=1, turn_id=2) - - await app.open_duplex_session_async("sid", fence=fence) - await app.append_duplex_input_async( - "sid", mode="append_audio_chunk", payload={}, fence=fence, collect_outputs=False - ) - await app.signal_duplex_turn_async("sid", event="turn.end", fence=fence) - await app.close_duplex_session_async("sid", reason="done", fence=fence) - - assert engine.calls == [(operation, fence) for operation in ("open", "append", "signal", "close")] - - -@pytest.mark.asyncio -async def test_openai_handler_passes_current_session_fence_for_runtime_controls(): - engine = FenceRecordingEngine() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid", - config=DuplexSessionConfig(extra_body={"native_duplex": True}), - epoch=1, - turn_id=2, - ) - ws = TimedWebSocket() - fence = DuplexFence("sid", epoch=1, turn_id=2) - - await handler._open_runtime_session(session, ws.send_json) - await handler._signal_runtime_session(session, "turn.end", send_json=ws.send_json) - await handler._close_runtime_session(session, reason="done", send_json=ws.send_json) - - assert engine.calls == [(operation, fence) for operation in ("open", "signal", "close")] - - -@pytest.mark.asyncio -async def test_runtime_signal_uses_fence_captured_before_session_turn_advances(): - engine = FenceRecordingEngine() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession(session_id="sid", config=DuplexSessionConfig(extra_body={"native_duplex": True})) - captured_fence = DuplexFence("sid", epoch=0, turn_id=0) - session.turn_id = 1 - - await handler._signal_runtime_session( - session, - "input.commit", - fence=captured_fence, - ) - - assert engine.calls == [("signal", captured_fence)] - - -@pytest.mark.asyncio -async def test_openai_handler_forwards_cancel_fence_through_async_omni_facade(): - engine = FenceRecordingEngine() - app = object.__new__(AsyncOmni) - app.engine = engine - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(app)) - session = DuplexSession(session_id="sid-cancel", config=DuplexSessionConfig(extra_body={"native_duplex": True})) - cancelled_fence = DuplexFence("sid-cancel", epoch=2, turn_id=3) - next_fence = DuplexFence("sid-cancel", epoch=3, turn_id=3) - - assert await handler._signal_runtime_session( - session, - "input.cancel", - fence=cancelled_fence, - next_fence=next_fence, - ) - - assert engine.calls == [("signal", cancelled_fence)] - assert engine.signal_next_fences == [next_fence] diff --git a/tests/entrypoints/test_duplex_omni.py b/tests/entrypoints/test_duplex_omni.py new file mode 100644 index 000000000..f33b3ea4d --- /dev/null +++ b/tests/entrypoints/test_duplex_omni.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexOmni`` / ``DuplexSessionHandle``: the thin Python API over a fake engine.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from typing import Any + +import pytest + +from tests.entrypoints.test_omni_entrypoints import FakeAsyncOmniEngine +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.events import ( + AudioDelta, + DuplexEvent, + SessionClosed, + SessionCreated, + SessionExpired, +) +from vllm_omni.engine.duplex.messages import ( + DuplexControlResultMessage, + DuplexSessionError, + DuplexSessionEventMessage, +) +from vllm_omni.engine.messages import ErrorMessage +from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeDuplexEngine(FakeAsyncOmniEngine): + """The engine surface ``DuplexOmni`` uses; every call is recorded.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.duplex_session_config = DuplexSessionRuntimeConfig(max_sessions=2) + self.duplex_capabilities = DuplexCapabilities(supports_session_resume=True) + self.opened: list[tuple[str, DuplexSessionConfig]] = [] + self.closed: list[tuple[str, str]] = [] + self.resumed: list[tuple[str, int]] = [] + self.touched: list[tuple[str, str]] = [] + self.commands: list[tuple[str, commands.DuplexCommand]] = [] + self.open_error: DuplexSessionError | None = None + self.close_error: DuplexSessionError | None = None + + def _result(self, operation: str, session_id: str, **fields: Any) -> DuplexControlResultMessage: + return DuplexControlResultMessage(control_id="c", operation=operation, session_id=session_id, ok=True, **fields) + + async def open_session_async(self, session_id, session_config, *, timeout=None): + self.opened.append((session_id, session_config)) + if self.open_error is not None: + raise self.open_error + return self._result( + "open", + session_id, + capabilities=self.duplex_capabilities, + public_session={"id": session_id, "idle_timeout_s": 7}, + lease_generation=0, + ) + + async def close_session_async(self, session_id, *, reason="client_close", timeout=None): + self.closed.append((session_id, reason)) + if self.close_error is not None: + raise self.close_error + # The runner emits session.closed before answering the RPC. + self.emit(session_id, SessionClosed(reason=reason)) + return self._result("close", session_id) + + async def resume_session_async(self, session_id, *, expected_lease_generation, timeout=None): + self.resumed.append((session_id, expected_lease_generation)) + return self._result("resume", session_id, lease_generation=expected_lease_generation + 1) + + async def touch_session_async(self, session_id, *, activity, timeout=None): + self.touched.append((session_id, activity)) + return self._result("touch", session_id) + + async def submit_command_async(self, session_id, command): + self.commands.append((session_id, command)) + + def emit(self, session_id: str, event: DuplexEvent) -> None: + # The session manager binds the session identity to every event it emits. + event = replace(event, session_id=session_id) + self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event)) + + +def _make_omni(monkeypatch: pytest.MonkeyPatch) -> tuple[DuplexOmni, FakeDuplexEngine]: + engine = FakeDuplexEngine() + monkeypatch.setattr("vllm_omni.entrypoints.duplex_omni.DuplexOmniEngine", lambda *args, **kwargs: engine) + monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) + return DuplexOmni("dummy-model"), engine + + +async def _collect(handle: DuplexSessionHandle, count: int, *, timeout_s: float = 2.0) -> list[DuplexEvent]: + events: list[DuplexEvent] = [] + + async def consume() -> None: + async for event in handle.events(): + events.append(event) + if len(events) >= count: + break + + await asyncio.wait_for(consume(), timeout=timeout_s) + return events + + +@pytest.mark.asyncio +async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session({"session_id": "mine", "id": "mine-too", "instructions": "hi"}) + + assert handle.session_id.startswith("duplex-") and len(handle.session_id) == len("duplex-") + 32 + assert handle.session_id != "mine" + ((session_id, config),) = engine.opened + assert session_id == handle.session_id + assert isinstance(config, DuplexSessionConfig) + assert config.model == "dummy-model" and config.instructions == "hi" + assert handle.capabilities is engine.duplex_capabilities + assert handle.public_session["id"] == handle.session_id + assert omni.sessions == {handle.session_id: handle} + assert omni.get_session(handle.session_id) is handle + assert omni.active_session_count() == 1 + + second = await omni.open_session(DuplexSessionConfig(instructions="typed")) + assert second.session_id != handle.session_id + assert engine.opened[1][1].model == "dummy-model" + assert omni.duplex_session_config.max_sessions == 2 + assert omni.duplex_capabilities.supports_session_resume is True + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_open_failure_drops_the_pending_handle_and_raises_the_engine_code(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + engine.open_error = DuplexSessionError("full", code="resource_exhausted", retryable=True) + try: + with pytest.raises(DuplexSessionError) as excinfo: + await omni.open_session() + assert excinfo.value.code == "resource_exhausted" + assert omni.sessions == {} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_events_are_routed_to_their_handle_and_end_on_session_closed(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.emit("duplex-unknown", AudioDelta(delta="zzz")) # dropped + engine.emit(handle.session_id, SessionCreated(session={"id": handle.session_id})) + engine.emit(handle.session_id, AudioDelta(response_id="r1", delta="abc")) + engine.emit(handle.session_id, SessionClosed(reason="client_close")) + + events = await _collect(handle, 10) + + assert [event.type for event in events] == ["session.created", "response.output_audio.delta", "session.closed"] + assert events[0].session_id == handle.session_id + assert handle.closed and handle.close_reason == "client_close" + assert await handle.wait_closed() == "client_close" + assert handle.session_id not in omni.sessions + with pytest.raises(DuplexSessionError) as excinfo: + await handle.submit(commands.Heartbeat()) + assert excinfo.value.code == "session_closed" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_events_is_single_consumer_but_may_be_reentered(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.emit(handle.session_id, AudioDelta(delta="one")) + first = handle.events() + assert (await first.__anext__()).type == "response.output_audio.delta" + with pytest.raises(RuntimeError, match="active events"): + await handle.events().__anext__() + await first.aclose() + + engine.emit(handle.session_id, SessionExpired(reason="idle_ttl_expired")) + events = await _collect(handle, 1) + assert events[0].type == "session.expired" and events[0].is_terminal + assert handle.closed and handle.close_reason == "idle_ttl_expired" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_close_session_waits_for_the_typed_session_closed_event(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + consumer = asyncio.create_task(_collect(handle, 1)) + await handle.close(reason="done") + + assert engine.closed == [(handle.session_id, "done")] + assert handle.closed and handle.close_reason == "done" + assert (await consumer)[0].type == "session.closed" + # Idempotent: a second close is a no-op. + await omni.close_session(handle.session_id) + assert len(engine.closed) == 1 + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_close_failure_marks_the_handle_closed_and_propagates(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + engine.close_error = DuplexSessionError("boom", code="engine_error") + try: + handle = await omni.open_session() + with pytest.raises(DuplexSessionError, match="boom"): + await omni.close_session(handle.session_id) + assert handle.closed + # A failed close must not leave a dead handle shadowing the id. + assert omni.sessions == {} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_engine_death_ends_every_session_handle(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.output_q.put(ErrorMessage(error="orchestrator died", fatal=True)) + + await asyncio.wait_for(handle.wait_closed(), timeout=2.0) + assert handle.closed + assert "engine_dead" in (handle.close_reason or "") + assert omni.sessions == {} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_resume_detach_and_touch_forward_to_the_engine(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + resumed = await omni.resume_session(handle.session_id, expected_lease_generation=0) + assert resumed is handle and handle.lease_generation == 1 + assert engine.resumed == [(handle.session_id, 0)] + + await omni.detach_session(handle.session_id) + await omni.touch_session(handle.session_id, activity="heartbeat") + assert engine.touched == [(handle.session_id, "detach"), (handle.session_id, "heartbeat")] + + for call in ( + omni.resume_session("duplex-nope", expected_lease_generation=0), + omni.detach_session("duplex-nope"), + ): + with pytest.raises(DuplexSessionError) as excinfo: + await call + assert excinfo.value.code == "unknown_session" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_handle_wrappers_build_typed_commands_in_caller_order(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + await handle.append_audio(b"\x00\x01", sample_rate_hz=16000, is_speech=True) + await handle.commit(create_response=True) + await handle.create_response({"instructions": "x"}) + await handle.cancel_response("resp-1") + await handle.barge_in() + await handle.ack_playback(120.7, response_id="resp-1", committed_ms=100) + await handle.update({"temperature": 0.2}) + await handle.create_item({"id": "item_1", "type": "message"}, previous_item_id="root") + await handle.delete_item("item_1") + await handle.truncate_item("item_1", audio_end_ms=50) + await handle.signal_turn("user_started", {"x": 1}) + await handle.heartbeat() + + submitted = [command for session_id, command in engine.commands if session_id == handle.session_id] + assert [type(command) for command in submitted] == [ + commands.AppendAudio, + commands.Commit, + commands.CreateResponse, + commands.CancelResponse, + commands.BargeIn, + commands.AckPlayback, + commands.UpdateSession, + commands.CreateItem, + commands.DeleteItem, + commands.TruncateItem, + commands.SignalTurn, + commands.Heartbeat, + ] + append, commit, create, cancel, _, ack, update, item, delete, truncate, signal, _ = submitted + assert append.audio == b"\x00\x01" and append.sample_rate_hz == 16000 and append.is_speech is True + assert commit.create_response is True and commit.final is True + assert create.options == {"instructions": "x"} + assert cancel.response_id == "resp-1" + assert ack.played_ms == 120 and ack.committed_ms == 100 and ack.response_id == "resp-1" + assert update.patch == {"temperature": 0.2} + assert item.previous_item_id == "root" and item.item["id"] == "item_1" + assert delete.item_id == "item_1" + assert truncate.audio_end_ms == 50 + assert signal.event == "user_started" and signal.signal_payload == {"x": 1} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_closes_every_handle(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + first = await omni.open_session() + second = await omni.open_session() + + omni.shutdown() + + assert first.closed and second.closed + assert first.close_reason == "shutdown" + assert omni.sessions == {} + assert engine.shutdown_called diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py index 827ab67b6..d343e2f03 100644 --- a/tests/entrypoints/test_omni_base_profiler.py +++ b/tests/entrypoints/test_omni_base_profiler.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for OmniBase and AsyncOmni profiler methods.""" from types import SimpleNamespace @@ -28,10 +31,14 @@ class TestOmniBaseProfiler: @pytest.fixture def omni_base_instance(self, mock_engine, mocker: MockerFixture): """Create an OmniBase instance with mocked dependencies.""" - mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + from vllm_omni.entrypoints.omni_base import OmniBase + + # ``OmniBase`` builds its engine through the ``_create_engine`` seam that + # ``AsyncOmni`` / ``DuplexOmni`` implement; there is no module-level + # engine class to patch any more. + mocker.patch.object(OmniBase, "_create_engine", return_value=mock_engine) mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") - from vllm_omni.entrypoints.omni_base import OmniBase instance = OmniBase(model="test-model") return instance @@ -48,10 +55,14 @@ class TestOmniBaseProfiler: def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture): """Legacy diffusion_batch_size callers receive an actionable error.""" - mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + from vllm_omni.entrypoints.omni_base import OmniBase + + # ``OmniBase`` builds its engine through the ``_create_engine`` seam that + # ``AsyncOmni`` / ``DuplexOmni`` implement; there is no module-level + # engine class to patch any more. + mocker.patch.object(OmniBase, "_create_engine", return_value=mock_engine) mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") - from vllm_omni.entrypoints.omni_base import OmniBase with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"): OmniBase(model="test-model", diffusion_batch_size=8) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 01264eaea..32d7d9988 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -118,8 +118,6 @@ class FakeAsyncOmniEngine: self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)] self.input_processor = None self.endpoint_restrictions = () - self.duplex_session_config = None - self.duplex_serving_adapter_path = None self.output_q: queue.Queue[Any] = queue.Queue() self.submitted: list[dict[str, Any]] = [] @@ -183,7 +181,8 @@ class FakeAsyncOmniEngine: def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) -> None: - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", lambda *args, **kwargs: engine) + monkeypatch.setattr("vllm_omni.entrypoints.omni.AsyncOmniEngine", lambda *args, **kwargs: engine) + monkeypatch.setattr("vllm_omni.entrypoints.async_omni.AsyncOmniEngine", lambda *args, **kwargs: engine) monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) # Don't add random UUIDs to requests calling .generate since we usually # just want to check for present requests anyway, and would need to just diff --git a/tests/examples/test_minicpmo_realtime_web_static.py b/tests/examples/test_minicpmo_realtime_web_static.py index 926ae36ef..2027d5aa3 100644 --- a/tests/examples/test_minicpmo_realtime_web_static.py +++ b/tests/examples/test_minicpmo_realtime_web_static.py @@ -36,7 +36,7 @@ def test_client_uses_proxy_relative_realtime_url_and_model_policy_session(): assert "new URL(config.realtimePath, window.location.href)" in source assert "url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'" in source assert "url.searchParams.set('autostart', '0')" in source - assert "url.searchParams.set('native_duplex', '1')" in source + assert "native_duplex" not in source assert "auto_response: true" in source assert "input_audio_buffer.append" in source assert "input_audio_buffer.commit" not in source diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py index b933f7173..d7f3db633 100644 --- a/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest -from vllm_omni.entrypoints.duplex.commit_policy import ( +from vllm_omni.engine.duplex.session.commit_policy import ( CommitAction, CommitSnapshot, decide_commit_action, @@ -24,7 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=0, - native_response_in_progress=False, + response_in_progress=False, playback_active=False, ), CommitAction.START_AUTO_RESPONSE, @@ -36,7 +36,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=800, - native_response_in_progress=True, + response_in_progress=True, playback_active=True, ), CommitAction.START_AUTO_RESPONSE, @@ -48,7 +48,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=False, active_response_id=None, overlap_speech_ms=0, - native_response_in_progress=True, + response_in_progress=True, playback_active=False, ), CommitAction.COMMIT_ONLY, @@ -60,7 +60,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=800, - native_response_in_progress=True, + response_in_progress=True, playback_active=False, ), CommitAction.DEFER_ACTIVE_RESPONSE, @@ -72,7 +72,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=800, - native_response_in_progress=False, + response_in_progress=False, playback_active=True, ), CommitAction.DEFER_ACTIVE_RESPONSE, @@ -84,7 +84,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=0, - native_response_in_progress=False, + response_in_progress=False, playback_active=False, ), CommitAction.COMMIT_ONLY, diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py new file mode 100644 index 000000000..b2be91366 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The cumulative Stage1 audio cursor. + +Stage1 hands the data plane a snapshot of every sample it has generated for the +request so far, and only the new tail may reach the client. The cursor that +makes that possible is the one piece of per-request state whose loss is silent +and expensive: rewinding it re-sends the whole session as a single delta. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( + MiniCPMO45DataPlaneSession, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REQUEST_ID = "duplex-s.c2Vzc2lvbg==.e.0.r.stage0" + + +def _session() -> MiniCPMO45DataPlaneSession: + return MiniCPMO45DataPlaneSession(lambda *args, **kwargs: "") + + +def _snapshot(num_samples: int) -> np.ndarray: + return np.arange(num_samples, dtype=np.float32) + + +def _sliced_len(session: MiniCPMO45DataPlaneSession, num_samples: int) -> int: + sliced = session.slice_cumulative_audio(REQUEST_ID, _snapshot(num_samples)) + if sliced is None: + return 0 + return int(np.asarray(sliced).reshape(-1).size) + + +def test_a_growing_snapshot_yields_only_the_new_tail() -> None: + session = _session() + assert _sliced_len(session, 24000) == 24000 + assert _sliced_len(session, 48000) == 24000 + assert _sliced_len(session, 48000) == 0 # unchanged snapshot emits nothing + + +def test_a_placeholder_unit_does_not_rewind_the_cursor() -> None: + """The failure behind a 300s benchmark case producing 1373s of audio. + + A continuation unit arrives as a one-sample snapshot. Taking that as "the + producer restarted" drops the cursor to 1, so the next real snapshot is + sliced from the start and the client is handed the entire session again. + """ + session = _session() + assert _sliced_len(session, 637440) == 637440 + assert _sliced_len(session, 1) == 1 + assert _sliced_len(session, 1) == 1 + assert _sliced_len(session, 1) == 1 + + # 649920 - 637440, not 649920 - 1. + assert _sliced_len(session, 649920) == 12480 + + +def test_a_genuinely_restarted_producer_is_still_picked_up() -> None: + """Two consecutive below-cursor snapshots that grow are a new buffer.""" + session = _session() + assert _sliced_len(session, 48000) == 48000 + assert _sliced_len(session, 1000) == 1000 # head of the new buffer + assert _sliced_len(session, 3000) == 2000 # its tail, not 3000 - 48000 + assert _sliced_len(session, 5000) == 2000 + + +def test_close_stream_releases_the_cursor_for_the_next_stream() -> None: + session = _session() + assert _sliced_len(session, 48000) == 48000 + session.close_stream(REQUEST_ID) + assert _sliced_len(session, 24000) == 24000 diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index e12010ce4..6c22b818b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -55,20 +55,18 @@ class TestRegistryDeclaration: assert pipeline.model_type == _PIPELINE_KEY assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration" - def test_native_duplex_control_is_explicit_without_a_fixed_session_cap(self) -> None: + def test_duplex_plugin_is_declared_without_a_fixed_session_cap(self) -> None: pipeline = OMNI_PIPELINES[_PIPELINE_KEY] - assert pipeline.duplex_control_enabled is True - assert pipeline.duplex_serving_adapter == ( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter" + assert pipeline.duplex_plugin == ( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin.MiniCPMO45DuplexPlugin" ) assert not hasattr(pipeline, "max_native_duplex_sessions") - def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None: + def test_ordinary_pipeline_declares_no_duplex_plugin(self) -> None: pipeline = PipelineConfig( model_type="ordinary", stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),) ) - assert pipeline.duplex_control_enabled is False - assert pipeline.duplex_serving_adapter is None + assert pipeline.duplex_plugin is None assert not hasattr(pipeline, "max_native_duplex_sessions") diff --git a/tests/model_executor/models/nemotron_voicechat/duplex/conftest.py b/tests/model_executor/models/nemotron_voicechat/duplex/conftest.py new file mode 100644 index 000000000..134731bf6 --- /dev/null +++ b/tests/model_executor/models/nemotron_voicechat/duplex/conftest.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Nemotron VoiceChat duplex is not on the duplex plugin framework yet. + +RFC vllm-omni#7181 splits the port: this PR lands the unified framework with +MiniCPM-o 4.5, and the follow-up PR moves Nemotron VoiceChat onto +``DuplexModelPlugin``. Until then the modules under test still import the +pre-framework duplex runtime / serving adapter that this PR removed, so +collecting them would fail at import time and abort the whole pytest session. +The tests themselves are unchanged and come back with the port. + +Listed per file rather than globbed, so the coverage loss is visible and shrinks +file by file as the port lands. ``test_input.py`` imports nothing that was +removed and keeps running. +""" + +#: Both import ``DuplexFence`` from ``engine.duplex.messages``, which no longer +#: exports it; ``test_data_plane`` also imports ``engine.duplex.runtime``. +collect_ignore = ["test_data_plane.py", "test_runtime_contract.py"] diff --git a/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py b/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py index 4bdc835d9..04e57c397 100644 --- a/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py +++ b/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py @@ -7,9 +7,9 @@ from types import SimpleNamespace import numpy as np import pytest +from vllm_omni.engine.duplex.runtime import duplex_resource_request_id from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import duplex_resource_request_id from vllm_omni.model_executor.models.nemotron_voicechat.duplex.data_plane import ( NemotronVoiceChatDataPlaneContext, NemotronVoiceChatDataPlaneSession, diff --git a/tests/model_executor/models/personaplex/duplex/conftest.py b/tests/model_executor/models/personaplex/duplex/conftest.py new file mode 100644 index 000000000..d49e8c617 --- /dev/null +++ b/tests/model_executor/models/personaplex/duplex/conftest.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""PersonaPlex duplex is not on the duplex plugin framework yet. + +RFC vllm-omni#7181 splits the port: this PR lands the unified framework with +MiniCPM-o 4.5, and the follow-up PR moves PersonaPlex onto ``DuplexModelPlugin``. +Until then the modules under test still import the pre-framework duplex +runtime / serving adapter that this PR removed, so collecting them would fail +at import time and abort the whole pytest session. The tests themselves are +unchanged and come back with the port. + +Listed per file rather than globbed, so the coverage loss is visible and shrinks +file by file as the port lands. +""" + +#: ``test_unified_runtime`` imports ``engine.duplex.runtime`` and +#: ``entrypoints.duplex.runtime_adapter`` directly. ``test_stage0_runtime`` fails +#: transitively: ``personaplex.duplex.runtime_extension`` imports ``DuplexFence`` +#: from ``engine.duplex.messages``, which no longer exports it. +collect_ignore = ["test_stage0_runtime.py", "test_unified_runtime.py"] diff --git a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py index 1f03e894b..e4f9ad9b1 100644 --- a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py @@ -8,6 +8,12 @@ from types import SimpleNamespace import numpy as np import pytest from vllm.sampling_params import SamplingParams +from vllm_omni.engine.duplex.runtime import ( + DuplexInputMode, +) +from vllm_omni.entrypoints.duplex.runtime_adapter import ( + ServingRuntimeConfigError, +) from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver from vllm_omni.config.stage_config import ( @@ -15,12 +21,6 @@ from vllm_omni.config.stage_config import ( merge_pipeline_deploy, ) from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexInputMode, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - ServingRuntimeConfigError, -) from vllm_omni.model_executor.models.personaplex.duplex.config import DEFAULT_PERSONA from vllm_omni.model_executor.models.personaplex.duplex.data_plane import ( PersonaPlexDataPlaneContext, diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index c47802fbc..388e334cc 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -23,7 +23,7 @@ def _minicpmo_duplex_policy_case( MiniCPMO45OmniForConditionalGeneration, ) - session_key = ("sid-policy", 1) + session_key = "sid-policy" model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model_stage = "llm" model._minicpmo45_native_duplex_token_ids_cache = { @@ -35,8 +35,7 @@ def _minicpmo_duplex_policy_case( row = DuplexSamplingRow( row_idx=0, request_id="req-policy", - session_id=session_key[0], - incarnation=session_key[1], + session_id=session_key, seq=3, payload=payload, max_tokens=20, @@ -63,13 +62,12 @@ def test_minicpmo_model_hook_owns_duplex_sampling_rows_and_force_listen(): "tts_bos_token_id": 8, "turn_eos_token_id": 9, } - model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-hook", 2): state}) + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={"sid-hook": state}) logits = torch.zeros((1, 16), dtype=torch.float32) row = DuplexSamplingRow( row_idx=0, request_id="req-hook", session_id="sid-hook", - incarnation=2, seq=3, payload={"force_listen": True, "is_speech": True}, max_tokens=20, @@ -78,7 +76,7 @@ def test_minicpmo_model_hook_owns_duplex_sampling_rows_and_force_listen(): model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) assert model._minicpmo45_active_duplex_rows == [0] - assert model._minicpmo45_duplex_row_sessions == {0: ("sid-hook", 2)} + assert model._minicpmo45_duplex_row_sessions == {0: "sid-hook"} assert model._minicpmo45_duplex_row_payloads == {0: row.payload} assert model._minicpmo45_duplex_row_max_tokens == {0: 20} assert logits[0, listen_id].item() == 0.0 @@ -152,7 +150,7 @@ def test_minicpmo_model_hook_old_response_output_does_not_clear_pending_speech_c pending_speech_context=True, ) model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True}) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( @@ -173,7 +171,7 @@ def test_minicpmo_model_hook_new_response_output_clears_pending_speech_context() pending_speech_context=True, ) model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( @@ -195,7 +193,7 @@ def test_minicpmo_model_hook_empty_speak_envelope_preserves_pending_speech_conte pending_speech_response_open=False, ) model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} token_ids = { "listen_token_id": 7, @@ -236,7 +234,7 @@ def test_minicpmo_model_hook_second_new_response_step_is_not_forced_to_listen(): original_logits = logits.clone() model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( 0, @@ -269,7 +267,7 @@ def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_ logits[0, 10] = 20.0 model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( 0, @@ -339,7 +337,6 @@ def test_generic_ar_runner_builds_typed_duplex_sampling_rows(): "duplex": { "data_plane": True, "session_id": "sid-runner-hook", - "incarnation": 4, "seq": 4, "payload": {"is_speech": True}, } @@ -359,7 +356,6 @@ def test_generic_ar_runner_builds_typed_duplex_sampling_rows(): assert rows[0].row_idx == 0 assert rows[0].request_id == "req-duplex" assert rows[0].session_id == "sid-runner-hook" - assert rows[0].incarnation == 4 assert rows[0].seq == 4 assert rows[0].payload == {"is_speech": True} assert rows[0].max_tokens == 32 @@ -458,13 +454,13 @@ def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe(): assert 'if "duplex_rows" not in str(exc)' not in source -def test_minicpmo_model_cleans_incarnation_state_when_request_finishes(): +def test_minicpmo_model_cleans_session_state_when_request_finishes(): from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) request_id = "duplex-sid-cleanup-i3-e0-stage0" - session_key = ("sid-cleanup", 3) + session_key = "sid-cleanup" model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model = SimpleNamespace() model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: object()}) @@ -1731,8 +1727,8 @@ def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk(): model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model_stage = "llm" model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) - session_key = ("sid-cross-chunk-repetition", 0) - state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + session_key = "sid-cross-chunk-repetition" + state = _MiniCPMO45Stage0SessionState(session_id=session_key) state.generated_tokens = [198] * 8 model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) model._minicpmo45_duplex_row_sessions = {0: session_key} @@ -1791,8 +1787,8 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty( model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model_stage = "llm" model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) - session_key = ("sid-negative-logit-repetition", 0) - state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + session_key = "sid-negative-logit-repetition" + state = _MiniCPMO45Stage0SessionState(session_id=session_key) repeated = 198 alternative = 1234 state.generated_tokens = [repeated] @@ -1827,8 +1823,8 @@ def test_minicpmo_stage0_records_bounded_model_policy_history(): MiniCPMO45OmniForConditionalGeneration, ) - session_key = ("sid-text-history", 0) - state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + session_key = "sid-text-history" + state = _MiniCPMO45Stage0SessionState(session_id=session_key) model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) model._minicpmo45_duplex_row_sessions = {0: session_key} diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index c5a96148e..959007c30 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -79,16 +79,15 @@ CHECK_IMPORTS = { "vllm_omni/entrypoints/openai/video_stream_base.py", "vllm_omni/entrypoints/openai/video_stream_session.py", "vllm_omni/experimental/fullduplex/client.py", + "vllm_omni/engine/duplex/audio.py", + "vllm_omni/engine/duplex/commands.py", "vllm_omni/engine/duplex/contracts.py", + "vllm_omni/engine/duplex/events.py", + "vllm_omni/engine/duplex/realtime_commands.py", "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py", - "vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py", "vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py", - "vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py", "vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py", - "vllm_omni/entrypoints/duplex/audio.py", - "vllm_omni/entrypoints/duplex/realtime_input.py", - "vllm_omni/entrypoints/duplex/runtime_bridge.py", - "vllm_omni/entrypoints/duplex/serving.py", "vllm_omni/model_executor/models/personaplex/duplex/input.py", "vllm_omni/model_executor/models/personaplex/duplex/runtime_extension.py", "vllm_omni/model_executor/models/personaplex/duplex/stage0.py", diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py index 83e224f34..e566ac12b 100644 --- a/vllm_omni/benchmarks/omniinteract.py +++ b/vllm_omni/benchmarks/omniinteract.py @@ -417,6 +417,13 @@ def _has_post_commit_decision( # refusal is harmless — record it instead of failing the case. _TOLERATED_ERROR_CODES = frozenset({"playback_ack_too_late"}) +# ``session.close`` carries no reason of its own, so the server stamps the +# terminal event with the reason it inferred for it. Anything else on a +# ``session.closed`` after our own close request means the session ended for a +# reason we did not ask for (a disconnect, a lease timeout, a shutdown), which +# is what this guard is looking for. +_CLIENT_CLOSE_REASONS = frozenset({"client_close"}) + def _error_code(event: dict[str, object]) -> str | None: error = event.get("error") @@ -452,7 +459,7 @@ def _raise_if_session_terminated( event_type == "session.closed" and explicit_close_from is not None and index >= explicit_close_from - and reason is None + and (reason is None or reason in _CLIENT_CLOSE_REASONS) ) if expected: continue @@ -826,7 +833,7 @@ def write_batch_artifacts( ) -def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str: +def _websocket_url(config: OmniInteractBenchmarkConfig) -> str: endpoint = ( config.endpoint if urlsplit(config.endpoint).scheme @@ -834,13 +841,7 @@ def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str: ) from vllm_omni.clients.duplex import build_realtime_url - return build_realtime_url( - endpoint, - config.model, - autostart=False, - native_duplex=True, - session_id=session_id, - ) + return build_realtime_url(endpoint, config.model, autostart=False) class _RealtimeSession: @@ -855,7 +856,7 @@ class _RealtimeSession: _MAX_FRAME_BYTES = 64 * 1024 * 1024 - def __init__(self, config: OmniInteractBenchmarkConfig, session_id: str, reference_audio: str) -> None: + def __init__(self, config: OmniInteractBenchmarkConfig, reference_audio: str) -> None: from vllm_omni.clients.duplex import DuplexClient, EventCollector, SessionConfig self.session_config = SessionConfig( @@ -864,12 +865,11 @@ class _RealtimeSession: playback_commit_policy="ack_only", idle_timeout_s=float(config.timeout_s), extra_body={ - "native_duplex": True, "force_listen_count": 0, **(config.extra_body or {}), }, ) - self.url = _websocket_url(config, session_id) + self.url = _websocket_url(config) headers = dict(config.extra_headers or {}) async def connect(url: str) -> WebSocketTransport: @@ -886,7 +886,6 @@ class _RealtimeSession: self.url, model=config.model, config=self.session_config, - session_id=session_id, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=min(config.timeout_s, 20.0), @@ -1013,7 +1012,7 @@ async def run_omniinteract_case( frames = prepared_input.video_frames if not any(frames): raise ValueError(f"No video frames were decoded from {case.video_path}") - async with _RealtimeSession(config, session_id, reference_audio) as client: + async with _RealtimeSession(config, reference_audio) as client: session_from = 0 # the collector holds only this session's events pcm = _ensure_final_commit_tail(pcm, client.events.events) playback = _Playback() diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index a230b9c2e..bee910c90 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -14,7 +14,7 @@ import time import traceback import uuid import wave -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, replace from datetime import datetime from pathlib import Path @@ -2153,10 +2153,68 @@ async def async_request_openai_audio_speech( return output +#: Silence budget per Seed-TTS turn: a model-native duplex session generates +#: per audio unit, and the target text rides the session context, so the +#: silence only advances the clock. It stops at the turn's response.done -- +#: a native model that keeps hearing silence after its turn may decide to +#: speak again, and the benchmark measures one response per utterance -- so +#: the budget is only spent on a turn the model is slow to take: it may +#: choose to listen on a few units first, and a turn that has not settled +#: when the budget runs out is reported with what the model did. +_SEED_TTS_SILENCE_SECONDS = 30.0 +#: A native model normally answers the seeded text within this much silence; +#: a turn that needs more is logged so a slow-to-speak model shows in the run. +_SEED_TTS_PROMPT_RESPONSE_S = 12.0 +#: MiniCPM-o emits 24 kHz mono; used to report audio_frames after the session closed. +_SEED_TTS_OUTPUT_SAMPLE_RATE_HZ = 24_000 + + +def _seed_tts_turn_stall_report(events: object, response_offset: int, request_index: int, silence_s: float) -> str: + """Explain a Seed-TTS turn that never settled: what the model did with the silence.""" + response_ids = list(getattr(events, "response_ids")[response_offset:]) + if not response_ids: + return ( + f"Seed-TTS Realtime TTS turn {request_index} never started a response: the model listened " + f"through {silence_s:.1f}s of silence and the wait that followed" + ) + audio_bytes = getattr(events, "audio_bytes") + response_text = getattr(events, "response_text") + started = ", ".join( + f"{response_id} ({len(audio_bytes(response_id))} audio bytes, text {response_text(response_id)!r})" + for response_id in response_ids + ) + return ( + f"Seed-TTS Realtime TTS turn {request_index} started {len(response_ids)} response(s) after " + f"{silence_s:.1f}s of silence but none reached response.done: {started}" + ) + + +def _seed_tts_turn_response_id(events: object, response_offset: int, request_index: int) -> str: + """The response id the Seed-TTS turn is measured on: the first one with audio. + + A model-native session answers the seeded text once, but nothing in the + protocol stops it from speaking again on silence it hears afterwards, so + a later audio response is the model's own and not a failed turn. + """ + response_ids = getattr(events, "response_ids") + audio_bytes = getattr(events, "audio_bytes") + audio_response_ids = [response_id for response_id in response_ids[response_offset:] if audio_bytes(response_id)] + if not audio_response_ids: + raise RuntimeError(f"Seed-TTS Realtime TTS turn {request_index} produced no audio response") + if len(audio_response_ids) > 1: + logger.warning( + "Seed-TTS Realtime TTS turn %d: model spoke again after its response (%d audio responses); " + "measuring the first", + request_index, + len(audio_response_ids), + ) + return audio_response_ids[0] + + def _realtime_websocket_url(api_url: str) -> str: from vllm_omni.clients.duplex import build_realtime_url - return build_realtime_url(api_url, None, native_duplex=None) + return build_realtime_url(api_url, None) def _nonnegative_number(value: object) -> bool: @@ -2321,6 +2379,13 @@ class _RealtimeTTSProbe: request time. """ + #: How long ``configure`` waits for a free duplex session. Every Seed-TTS + #: utterance is one session and the deploy config admits ``max_sessions`` + #: of them, so a benchmark run above that concurrency queues for a slot + #: rather than counting the server's (retryable) refusal as a failed + #: request. + _SESSION_SLOT_WAIT_S = 120.0 + def __init__(self, url: str) -> None: from vllm_omni.clients.duplex import EventCollector @@ -2347,16 +2412,13 @@ class _RealtimeTTSProbe: *, output_audio_format: str = "pcm16", instructions: str | None = None, - native_duplex: bool = False, auto_response: bool = False, extra_body: dict[str, object] | None = None, - session_id: str | None = None, timeout_s: float = 120.0, ) -> None: - from vllm_omni.clients.duplex import AudioFormat, DuplexClient, SessionConfig + from vllm_omni.clients.duplex import AudioFormat, DuplexClient, DuplexProtocolError, SessionConfig session_extra_body: dict[str, object] = dict(extra_body or {}) - session_extra_body["native_duplex"] = bool(native_duplex) config = SessionConfig( output_audio=AudioFormat(output_audio_format, 24_000), instructions=instructions, @@ -2365,16 +2427,33 @@ class _RealtimeTTSProbe: playback_commit_policy="ack_only", extra_body=session_extra_body, ) - self._client = DuplexClient( - self._url, - model=model, - config=config, - session_id=session_id, - reconnect=None, - heartbeat_interval_s=None, - handshake_timeout_s=timeout_s, - ) - await self._client.__aenter__() + deadline = time.monotonic() + self._SESSION_SLOT_WAIT_S + delay_s = 0.25 + waited = False + while True: + # A failed handshake closes the socket on the client's side, so + # a refused attempt leaves nothing behind to clean up. + client = DuplexClient( + self._url, + model=model, + config=config, + reconnect=None, + heartbeat_interval_s=None, + handshake_timeout_s=timeout_s, + ) + try: + await client.__aenter__() + except DuplexProtocolError as exc: + if exc.code != "resource_exhausted" or time.monotonic() >= deadline: + raise + if not waited: + logger.info("Seed-TTS Realtime TTS: no free duplex session (%s); waiting for a slot", exc) + waited = True + await asyncio.sleep(delay_s) + delay_s = min(delay_s * 2.0, 2.0) + continue + self._client = client + break self._consume_task = asyncio.create_task(self.events.consume(self._client)) async def send(self, event: dict[str, object]) -> None: @@ -2387,6 +2466,32 @@ class _RealtimeTTSProbe: assert self._client is not None await acknowledge_collected_playback(self._client, self.events) + async def stream_silence( + self, + *, + seconds: float, + chunk_ms: int = 200, + until: Callable[[], bool] | None = None, + ) -> float: + """Append silent PCM16 units so a model-native session has units to speak on. + + A duplex model generates per audio unit. The target text rides the + session context (``duplex_initial_user_text``), so the audio only has + to advance the clock; silence keeps it from adding content of its own. + Streams in real time for at most ``seconds``, stopping as soon as + ``until`` holds, and returns the seconds actually appended. + """ + assert self._client is not None + input_format = self._client.config.input_audio + chunk = bytes(max(input_format.byte_count(chunk_ms), input_format.bytes_per_sample)) + chunk_s = input_format.duration_ms(len(chunk)) / 1000.0 + streamed_s = 0.0 + while streamed_s < seconds and not (until is not None and until()): + await self._client.append_audio(chunk, is_speech=False) + streamed_s += chunk_s + await asyncio.sleep(chunk_s) + return streamed_s + async def close_session(self, *, timeout_s: float = 20.0) -> None: assert self._client is not None await self._client.close(timeout_s=timeout_s) @@ -2423,69 +2528,70 @@ async def async_request_openai_realtime_duplex( if not turn_prompts: turn_prompts = [("", request_func_input.prompt)] session_id = f"seed-tts-{request_func_input.request_id or uuid.uuid4().hex}" + silence_seconds = float(getattr(request_func_input, "seed_tts_silence_seconds", 0.0) or _SEED_TTS_SILENCE_SECONDS) + turn_metrics: list[dict[str, object]] = [] + turn_timings: list[dict[str, object]] = [] + turn_pcm_bytes: list[bytes] = [] + turn_transcripts: list[str] = [] + measurement_origin = { + "ttft": "first silence append client send to first non-empty text delta", + "ttfp": "first silence append client send to first audio packet", + "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", + } try: - async with _RealtimeTTSProbe(_realtime_websocket_url(request_func_input.api_url)) as client: - await client.configure( - request_func_input.model_name or request_func_input.model, - output_audio_format="pcm16", - instructions=getattr( - request_func_input, - "seed_tts_system_prompt", - SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, - ), - native_duplex=False, - auto_response=False, - extra_body=speech_extra, - session_id=session_id, - timeout_s=120.0, - ) - turn_metrics: list[dict[str, object]] = [] - turn_timings: list[dict[str, object]] = [] - turn_pcm_bytes: list[bytes] = [] - turn_transcripts: list[str] = [] - measurement_origin = { - "ttft": "conversation.item.create client send to first non-empty text delta", - "ttfp": "conversation.item.create client send to first audio packet", - "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", - } - for request_index, (utterance_id, target_text) in enumerate(turn_prompts): + # One session per utterance. A model-native duplex session takes its + # text once, in the session context (``duplex_initial_user_text``), so + # a session cannot be re-seeded for a second target text. + for request_index, (utterance_id, target_text) in enumerate(turn_prompts): + async with _RealtimeTTSProbe(_realtime_websocket_url(request_func_input.api_url)) as client: + await client.configure( + request_func_input.model_name or request_func_input.model, + output_audio_format="pcm16", + instructions=getattr( + request_func_input, + "seed_tts_system_prompt", + SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, + ), + auto_response=True, + extra_body={ + **speech_extra, + "duplex_initial_user_text": target_text, + "force_listen_count": 0, + }, + timeout_s=120.0, + ) response_offset = len(client.events.response_ids) done_before = client.events.count("response.done") errors_before = len(client.events.errors()) - turn_started_at_s = time.monotonic() - await client.send( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": target_text}], - }, - } - ) - await client.send({"type": "response.create"}) - await wait_for_condition( - lambda: ( + + def turn_settled() -> bool: + return ( client.events.count("response.done") > done_before or len(client.events.errors()) > errors_before - ), - timeout_s=180.0, - label=f"Seed-TTS Realtime TTS turn {request_index} response.done", - ) + ) + + turn_started_at_s = time.monotonic() + silence_s = await client.stream_silence(seconds=silence_seconds, until=turn_settled) + if silence_s > _SEED_TTS_PROMPT_RESPONSE_S: + logger.warning( + "Seed-TTS Realtime TTS turn %d: the model took %.1fs of silence to settle its response", + request_index, + silence_s, + ) + try: + await wait_for_condition( + turn_settled, + timeout_s=180.0, + label=f"Seed-TTS Realtime TTS turn {request_index} response.done", + ) + except TimeoutError as exc: + raise RuntimeError( + _seed_tts_turn_stall_report(client.events, response_offset, request_index, silence_s) + ) from exc errors = client.events.errors() if len(errors) > errors_before: raise RuntimeError(f"Seed-TTS Realtime TTS server error: {errors[-1]}") - new_audio_response_ids = [ - response_id - for response_id in client.events.response_ids[response_offset:] - if client.events.audio_bytes(response_id) - ] - if len(new_audio_response_ids) != 1: - raise RuntimeError( - f"Seed-TTS Realtime TTS turn {request_index} expected one audio response, " - f"got {len(new_audio_response_ids)}" - ) - response_id = new_audio_response_ids[0] + response_id = _seed_tts_turn_response_id(client.events, response_offset, request_index) timing = client.events.timing_summary( after_s=turn_started_at_s, input_committed_at_s=turn_started_at_s, @@ -2518,45 +2624,43 @@ async def async_request_openai_realtime_duplex( ) turn_transcripts.append(client.events.response_text(response_id)) await client.acknowledge_playback() - request_finished_at = time.perf_counter() - session_metrics = summarize_session_request_metrics( - turn_metrics, - session_id=session_id, - ) - await client.close_session(timeout_s=30.0) + await client.close_session(timeout_s=30.0) + request_finished_at = time.perf_counter() + session_metrics = summarize_session_request_metrics( + turn_metrics, + session_id=session_id, + ) - output.generated_text = " ".join(filter(None, turn_transcripts)) - from vllm_omni.clients.duplex import metric_mean + output.generated_text = " ".join(filter(None, turn_transcripts)) + from vllm_omni.clients.duplex import metric_mean - output.ttft = (metric_mean(session_metrics.get("ttft_ms")) or 0.0) / 1000.0 - output.audio_ttfp = (metric_mean(session_metrics.get("ttfp_ms")) or 0.0) / 1000.0 - output.audio_rtf = metric_mean(session_metrics.get("rtf")) or 0.0 - output.audio_duration = ( - sum(float(metric.get("audio_duration_ms") or 0.0) for metric in turn_metrics) / 1000.0 + output.ttft = (metric_mean(session_metrics.get("ttft_ms")) or 0.0) / 1000.0 + output.audio_ttfp = (metric_mean(session_metrics.get("ttfp_ms")) or 0.0) / 1000.0 + output.audio_rtf = metric_mean(session_metrics.get("rtf")) or 0.0 + output.audio_duration = sum(float(metric.get("audio_duration_ms") or 0.0) for metric in turn_metrics) / 1000.0 + output.audio_frames = int(output.audio_duration * _SEED_TTS_OUTPUT_SAMPLE_RATE_HZ) + output.latency = request_finished_at - output.start_time + output.tts_turn_pcm_bytes = turn_pcm_bytes + output.tts_output_pcm_bytes = b"".join(turn_pcm_bytes) + if bool((request_func_input.extra_body or {}).get("save_duplex_request_metrics")): + output.duplex_request_metrics = turn_metrics + output.duplex_session_metrics = session_metrics + output.output_tokens = sum( + int(stage0.get("output_token_count") or 0) + for timing in turn_timings + if isinstance((stage0 := timing.get("stage0_tokens")), dict) + ) + token_timing_measured = _apply_stage0_token_timings( + output, + [timing.get("stage0_tokens") for timing in turn_timings], + expected_output_tokens=output.output_tokens, + ) + if not token_timing_measured and output.output_tokens > 1: + logger.warning( + "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable", + session_id, ) - output.audio_frames = int(output.audio_duration * client.events.output_sample_rate_hz) - output.latency = request_finished_at - output.start_time - output.tts_turn_pcm_bytes = turn_pcm_bytes - output.tts_output_pcm_bytes = b"".join(turn_pcm_bytes) - if bool((request_func_input.extra_body or {}).get("save_duplex_request_metrics")): - output.duplex_request_metrics = turn_metrics - output.duplex_session_metrics = session_metrics - output.output_tokens = sum( - int(stage0.get("output_token_count") or 0) - for timing in turn_timings - if isinstance((stage0 := timing.get("stage0_tokens")), dict) - ) - token_timing_measured = _apply_stage0_token_timings( - output, - [timing.get("stage0_tokens") for timing in turn_timings], - expected_output_tokens=output.output_tokens, - ) - if not token_timing_measured and output.output_tokens > 1: - logger.warning( - "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable", - session_id, - ) - output.success = True + output.success = True except Exception: output.success = False output.error = traceback.format_exc() diff --git a/vllm_omni/clients/__init__.py b/vllm_omni/clients/__init__.py index 924582bf4..35f904778 100644 --- a/vllm_omni/clients/__init__.py +++ b/vllm_omni/clients/__init__.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Client-side libraries for vLLM-Omni serving APIs. -Modules under this package speak to a running vLLM-Omni server over the -network. They are for applications and tests; server runtime code +``vllm_omni.clients.duplex`` speaks to a running vLLM-Omni server over the +network; ``vllm_omni.clients.inline_duplex`` drives an in-process +``DuplexOmni`` handed in by the caller with the same client API. Modules here +import no server runtime code at module load, and server runtime code (``vllm_omni.engine``, ``vllm_omni.entrypoints``, ``vllm_omni.model_executor``) must never import from here. """ diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index a1d9fc1df..52266b64e 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -2,9 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Public asynchronous client for the vLLM-Omni full-duplex Realtime API. -Connects to ``/v1/realtime?duplex=1`` (the normative duplex contract) and -provides typed events, response demultiplexing, incremental playback -acking, and transparent session resume on transport drops. +:class:`DuplexClientBase` holds the transport-agnostic session client (typed +events, response demultiplexing, incremental playback acking, the session +handshake). :class:`DuplexClient` connects it to ``/v1/realtime?duplex=1`` +(the normative duplex contract) over WebSocket with transparent session +resume on transport drops; :class:`vllm_omni.clients.inline_duplex.InlineDuplexClient` +drives an in-process :class:`~vllm_omni.entrypoints.duplex_omni.DuplexOmni` +with exactly the same usage. Example:: @@ -35,10 +39,10 @@ import math import random import time import wave +from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import Protocol from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from uuid import uuid4 @@ -53,6 +57,7 @@ __all__ = [ "AudioDelta", "ConnectionResumed", "DuplexClient", + "DuplexClientBase", "DuplexClientError", "DuplexConnectionError", "DuplexEvent", @@ -229,7 +234,7 @@ class SessionConfig: idle_timeout_s: float | None = None extra_body: dict[str, object] = field(default_factory=dict) - def to_session_payload(self, *, model: str, session_id: str | None = None) -> dict[str, object]: + def to_session_payload(self, *, model: str) -> dict[str, object]: payload: dict[str, object] = { "model": model, "modalities": list(self.modalities), @@ -246,8 +251,6 @@ class SessionConfig: }, "turn_detection": self.turn_detection, } - if session_id: - payload["session_id"] = session_id if self.instructions is not None: payload["instructions"] = self.instructions if self.ref_audio is not None: @@ -368,11 +371,6 @@ class SessionCreated(DuplexEvent): value = self.raw.get("resume_token") return value if isinstance(value, str) else None - @property - def incarnation(self) -> int: - value = self.raw.get("incarnation") - return value if isinstance(value, int) else 0 - class SessionResumed(SessionCreated): pass @@ -598,79 +596,84 @@ class ResponseHandle: # Client -class WebSocketTransport(Protocol): +class WebSocketTransport(ABC): """Minimal WebSocket surface the client needs (satisfied by ``websockets``).""" + @abstractmethod async def send(self, data: str) -> None: ... + @abstractmethod async def recv(self) -> str | bytes: ... + @abstractmethod async def close(self) -> None: ... ConnectFn = Callable[[str], Awaitable["WebSocketTransport"]] -class DuplexClient: - """Async client for one duplex session over ``/v1/realtime?duplex=1``. +class DuplexClientBase(ABC): + """Transport-agnostic duplex session client. - ``session_id`` only names the session this client creates: entering the - client always performs the ``session.update`` handshake. Resuming an - existing session is supported only as automatic reconnect within the same - client instance (see ``ReconnectPolicy``); taking over a session from a - new client requires the wire-level ``session.resume`` handshake - (``resume_token``, ``incarnation``, ``last_received_server_event_seq``), - which this client does not expose. + Everything a caller touches (typed events, response demux, input helpers, + playback acks, the ``session.update`` handshake) lives here. A subclass + only supplies the transport: :meth:`_open` establishes it and must deliver + the server's ``session.created`` payload through :meth:`_dispatch`, + :meth:`_send_command` delivers one client event, :meth:`_teardown` + releases the transport. """ def __init__( self, - url: str, *, model: str, config: SessionConfig | None = None, - session_id: str | None = None, - reconnect: ReconnectPolicy | None = ReconnectPolicy(), - heartbeat_interval_s: float | None = 30.0, handshake_timeout_s: float = 30.0, - connect: ConnectFn | None = None, ) -> None: - self.url = url self.model = model self.config = config or SessionConfig() - self.session_id = session_id + #: Allocated by the server; ``None`` until ``session.created`` arrives. + self.session_id: str | None = None self.session_info: dict[str, object] = {} - self.incarnation = 0 - self.resume_token: str | None = None - self._reconnect = reconnect - self._heartbeat_interval_s = heartbeat_interval_s self._handshake_timeout_s = handshake_timeout_s - self._connect_fn: ConnectFn = connect or self._default_connect - self._ws: WebSocketTransport | None = None - self._reader_task: asyncio.Task[None] | None = None - self._heartbeat_task: asyncio.Task[None] | None = None self._subscribers: list[asyncio.Queue[DuplexEvent | _ClosedMarker]] = [] self._response_queue: asyncio.Queue[ResponseHandle | _ClosedMarker | _ErrorMarker] = asyncio.Queue( maxsize=_MAX_BUFFERED_EVENTS ) self._responses: dict[str, ResponseHandle] = {} - self._last_server_event_seq: int | None = None self._input_audio_end_ms = 0.0 self._closing = False self._closed = asyncio.Event() self._closed_marker: _ClosedMarker | None = None + # -- transport hooks ------------------------------------------------------ + + @abstractmethod + async def _open(self) -> None: + """Establish the transport and start delivering server events via ``_dispatch``. + + Called once from :meth:`__aenter__`; the base then sends the + ``session.update`` handshake and waits for ``session.created``. + """ + + @abstractmethod + async def _send_command(self, payload: dict[str, object]) -> None: + """Deliver one client event (already stamped with ``event_id``).""" + + @abstractmethod + async def _teardown(self) -> None: + """Release the transport; must be idempotent.""" + # -- lifecycle ---------------------------------------------------------- - async def __aenter__(self) -> DuplexClient: - self._ws = await self._connect_fn(self._target_url()) + async def __aenter__(self) -> DuplexClientBase: handshake_queue = self._add_subscriber() try: - self._reader_task = asyncio.create_task(self._read_loop(), name="duplex-client-reader") + await self._open() await self.send( { "type": "session.update", - "session": self.config.to_session_payload(model=self.model, session_id=self.session_id), + "session": self.config.to_session_payload(model=self.model), } ) created = await self._wait_on_queue( @@ -679,14 +682,12 @@ class DuplexClient: timeout_s=self._handshake_timeout_s, ) except BaseException: - await self._teardown() + await self._teardown_quietly() raise finally: self._remove_subscriber(handshake_queue) assert isinstance(created, SessionCreated) self._adopt_session(created) - if self._heartbeat_interval_s is not None: - self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="duplex-client-heartbeat") return self async def __aexit__(self, exc_type, exc, traceback) -> None: @@ -706,7 +707,7 @@ class DuplexClient: except DuplexClientError: pass finally: - await self._teardown() + await self._teardown_quietly() async def close(self, *, timeout_s: float = 20.0) -> None: """Send ``session.close`` and wait for the server to confirm.""" @@ -723,6 +724,12 @@ class DuplexClient: except asyncio.TimeoutError: self._finalize("close timed out", expected=True) + async def _teardown_quietly(self) -> None: + self._closing = True + if not self._closed.is_set(): + self._finalize("closed", expected=True) + await self._teardown() + # -- input ---------------------------------------------------------------- async def append_audio( @@ -925,53 +932,25 @@ class DuplexClient: payload = dict(event) payload.setdefault("event_id", f"evt_{uuid4().hex}") try: - await self._ws.send(json.dumps(payload)) + await self._send_command(payload) except asyncio.CancelledError: raise + except DuplexClientError: + raise except Exception as exc: raise DuplexConnectionError(f"send failed: {exc}") from exc return payload["event_id"] # -- internals --------------------------------------------------------------- - def _target_url(self) -> str: - parts = urlsplit(self.url) - path = parts.path if parts.path not in ("", "/") else "/v1/realtime" - query = dict(parse_qsl(parts.query, keep_blank_values=True)) - query.setdefault("duplex", "1") - query.setdefault("model", self.model) - # The client always opens with an explicit session.update (and resumes - # with an explicit session.resume). Autostart would race that handshake: - # with a model query param, the server creates a bare default session - # before reading the first client event, silently dropping ref_audio - # and the model-specific extra_body. Force it off even when the caller's - # URL carries autostart=1 — this client cannot operate over it. - query["autostart"] = "0" - return urlunsplit((parts.scheme or "ws", parts.netloc, path, urlencode(query), parts.fragment)) - - async def _default_connect(self, url: str) -> WebSocketTransport: - try: - import websockets - except ImportError as exc: # pragma: no cover - websockets is a pinned dep - raise DuplexConnectionError("The duplex client requires the 'websockets' package") from exc - try: - return await websockets.connect(url, max_size=_MAX_FRAME_BYTES) - except asyncio.CancelledError: - raise - except Exception as exc: - raise DuplexConnectionError(f"connect to {url} failed: {exc}") from exc - def _adopt_session(self, event: SessionCreated) -> None: - # A resume activation carries no nested session object; keep the - # session_info captured at the original handshake in that case. + # session.created and session.resumed both carry the engine's current + # session object; keep the last one seen. if event.session: self.session_info = event.session session_id = event.session.get("session_id") or event.session.get("id") or event.session_id if isinstance(session_id, str) and session_id: self.session_id = session_id - self.incarnation = event.incarnation - if event.resume_token: - self.resume_token = event.resume_token def _announce(self, handle: ResponseHandle) -> None: """Deliver a handle to responses() exactly once, decision known.""" @@ -1013,36 +992,18 @@ class DuplexClient: if item.type in types: return item - async def _read_loop(self) -> None: - while True: - try: - while True: - raw = await self._ws.recv() - if isinstance(raw, (bytes, bytearray)): - continue - try: - data = json.loads(raw) - except json.JSONDecodeError: - continue - if isinstance(data, dict): - await self._dispatch(data) - except asyncio.CancelledError: - raise - except Exception as exc: - if self._closing or self._closed.is_set(): - self._finalize("closed", expected=True) - return - if await self._attempt_resume(): - continue - self._finalize(f"connection lost: {exc}", expected=False) - return + def _accept_event(self, data: dict[str, object]) -> bool: + """Transport hook: return False to drop ``data`` before dispatch (duplicates).""" + del data + return True + + async def _after_dispatch(self, event: DuplexEvent, data: dict[str, object]) -> None: + """Transport hook run after an event was delivered to subscribers.""" + del event, data async def _dispatch(self, data: dict[str, object]) -> None: - seq = data.get("server_event_seq") - if isinstance(seq, int): - if self._last_server_event_seq is not None and seq <= self._last_server_event_seq: - return # duplicate delivered by resume replay - self._last_server_event_seq = seq + if not self._accept_event(data): + return event = wrap_event(data) if isinstance(event, SessionResumed): @@ -1109,14 +1070,130 @@ class DuplexClient: elif isinstance(event, SessionExpired): reason = data.get("reason") self._finalize(f"expired: {reason}" if reason else "expired", expected=False) - elif event.type == "session.resync_required": + + await self._after_dispatch(event, data) + + def _finalize(self, reason: str, *, expected: bool) -> None: + if self._closed.is_set(): + return + marker = _ClosedMarker(reason=reason, expected=expected) + self._closed_marker = marker + self._closed.set() + for handle in self._responses.values(): + handle._finish(None) + self._responses.clear() + for queue in (*self._subscribers, self._response_queue): + _put_drop_oldest(queue, marker) + + +class DuplexClient(DuplexClientBase): + """Async client for one duplex session over ``/v1/realtime?duplex=1``.""" + + def __init__( + self, + url: str, + *, + model: str, + config: SessionConfig | None = None, + reconnect: ReconnectPolicy | None = ReconnectPolicy(), + heartbeat_interval_s: float | None = 30.0, + handshake_timeout_s: float = 30.0, + connect: ConnectFn | None = None, + ) -> None: + super().__init__(model=model, config=config, handshake_timeout_s=handshake_timeout_s) + self.url = url + self.resume_token: str | None = None + self._reconnect = reconnect + self._heartbeat_interval_s = heartbeat_interval_s + self._connect_fn: ConnectFn = connect or self._default_connect + self._ws: WebSocketTransport | None = None + self._reader_task: asyncio.Task[None] | None = None + self._heartbeat_task: asyncio.Task[None] | None = None + self._last_server_event_seq: int | None = None + + # -- transport hooks ------------------------------------------------------ + + async def _open(self) -> None: + self._ws = await self._connect_fn(self._target_url()) + self._reader_task = asyncio.create_task(self._read_loop(), name="duplex-client-reader") + + async def __aenter__(self) -> DuplexClient: + await super().__aenter__() + if self._heartbeat_interval_s is not None: + self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="duplex-client-heartbeat") + return self + + async def _send_command(self, payload: dict[str, object]) -> None: + if self._ws is None: + raise DuplexConnectionError("send failed: transport is not connected") + await self._ws.send(json.dumps(payload)) + + async def _teardown(self) -> None: + for task in (self._heartbeat_task, self._reader_task): + if task is not None and not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + if self._ws is not None: + try: + await self._ws.close() + except Exception: # noqa: BLE001 + pass + + # -- internals --------------------------------------------------------------- + + def _target_url(self) -> str: + parts = urlsplit(self.url) + path = parts.path if parts.path not in ("", "/") else "/v1/realtime" + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query.setdefault("duplex", "1") + query.setdefault("model", self.model) + # The client always opens with an explicit session.update (and resumes + # with an explicit session.resume). Autostart would race that handshake: + # with a model query param, the server creates a bare default session + # before reading the first client event, silently dropping ref_audio + # and the model-specific extra_body. Force it off even when the caller's + # URL carries autostart=1 — this client cannot operate over it. + query["autostart"] = "0" + return urlunsplit((parts.scheme or "ws", parts.netloc, path, urlencode(query), parts.fragment)) + + async def _default_connect(self, url: str) -> WebSocketTransport: + try: + import websockets + except ImportError as exc: # pragma: no cover - websockets is a pinned dep + raise DuplexConnectionError("The duplex client requires the 'websockets' package") from exc + try: + return await websockets.connect(url, max_size=_MAX_FRAME_BYTES) + except asyncio.CancelledError: + raise + except Exception as exc: + raise DuplexConnectionError(f"connect to {url} failed: {exc}") from exc + + def _adopt_session(self, event: SessionCreated) -> None: + super()._adopt_session(event) + if event.resume_token: + self.resume_token = event.resume_token + + def _accept_event(self, data: dict[str, object]) -> bool: + seq = data.get("server_event_seq") + if isinstance(seq, int): + if self._last_server_event_seq is not None and seq <= self._last_server_event_seq: + return False # duplicate delivered by resume replay + self._last_server_event_seq = seq + return True + + async def _after_dispatch(self, event: DuplexEvent, data: dict[str, object]) -> None: + if event.type == "session.resync_required": # The server stopped journaling this session's events, so any # retained resume credential will be refused on the next # session.resume. Drop it now: a later transport failure then # finalizes immediately instead of burning reconnect attempts. self.resume_token = None - if isinstance(seq, int) and not self._closed.is_set(): + seq = data.get("server_event_seq") + if isinstance(seq, int) and not self._closed.is_set() and self._ws is not None: try: await self._ws.send(json.dumps({"type": "session.event_ack", "server_event_seq": seq})) except asyncio.CancelledError: @@ -1124,6 +1201,30 @@ class DuplexClient: except Exception: pass # the reader will observe the transport failure itself + async def _read_loop(self) -> None: + while True: + try: + while True: + raw = await self._ws.recv() + if isinstance(raw, (bytes, bytearray)): + continue + try: + data = json.loads(raw) + except json.JSONDecodeError: + continue + if isinstance(data, dict): + await self._dispatch(data) + except asyncio.CancelledError: + raise + except Exception as exc: + if self._closing or self._closed.is_set(): + self._finalize("closed", expected=True) + return + if await self._attempt_resume(): + continue + self._finalize(f"connection lost: {exc}", expected=False) + return + async def _attempt_resume(self) -> bool: policy = self._reconnect if policy is None or not self.resume_token or not self.session_id: @@ -1142,7 +1243,6 @@ class DuplexClient: { "type": "session.resume", "session_id": self.session_id, - "incarnation": self.incarnation, "resume_token": self.resume_token, "last_received_server_event_seq": self._last_server_event_seq or 0, } @@ -1162,7 +1262,10 @@ class DuplexClient: continue event_type = data.get("type") if event_type == "session.resumed": + previous = self._ws self._ws = ws + if previous is not None and previous is not ws: + await self._close_quietly(previous) await self._dispatch({"type": "connection.resumed", "attempt": attempt}) await self._dispatch(data) for entry in pending: @@ -1212,35 +1315,6 @@ class DuplexClient: except DuplexClientError: return - def _finalize(self, reason: str, *, expected: bool) -> None: - if self._closed.is_set(): - return - marker = _ClosedMarker(reason=reason, expected=expected) - self._closed_marker = marker - self._closed.set() - for handle in self._responses.values(): - handle._finish(None) - self._responses.clear() - for queue in (*self._subscribers, self._response_queue): - _put_drop_oldest(queue, marker) - - async def _teardown(self) -> None: - self._closing = True - if not self._closed.is_set(): - self._finalize("closed", expected=True) - for task in (self._heartbeat_task, self._reader_task): - if task is not None and not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): # noqa: BLE001 - pass - if self._ws is not None: - try: - await self._ws.close() - except Exception: # noqa: BLE001 - pass - # --------------------------------------------------------------------------- # Test/benchmark collector @@ -1368,7 +1442,7 @@ class EventCollector: self.response_ids: list[str] = [] self.output_sample_rate_hz = 24_000 - async def consume(self, client: DuplexClient) -> None: + async def consume(self, client: DuplexClientBase) -> None: """Subscribe to ``client`` and collect until the session ends.""" try: async for event in client.events(): @@ -1795,15 +1869,12 @@ def build_realtime_url( model: str | None, *, autostart: bool | None = None, - native_duplex: bool | None = None, - session_id: str | None = None, extra_query: dict[str, str] | None = None, ) -> str: """Add explicit duplex query parameters to a Realtime URL. :class:`DuplexClient` builds its own URL; this helper is for drivers that - speak the wire protocol directly. ``native_duplex`` sets the per-session - model-native opt-in query flag; other model-specific query flags ride in + speak the wire protocol directly. Model-specific query flags ride in ``extra_query``. ``http(s)`` URLs are rewritten to ``ws(s)``. """ parts = urlsplit(url) @@ -1815,14 +1886,10 @@ def build_realtime_url( query.setdefault("duplex", "1") if model: query.setdefault("model", model) - if native_duplex is not None: - query["native_duplex"] = "1" if native_duplex else "0" for key, value in (extra_query or {}).items(): query.setdefault(key, value) if autostart is not None: query.setdefault("autostart", "1" if autostart else "0") - if session_id: - query.setdefault("session_id", session_id) return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) @@ -1880,7 +1947,7 @@ def summarize_session_request_metrics( return summary -async def acknowledge_collected_playback(client: DuplexClient, collector: EventCollector) -> None: +async def acknowledge_collected_playback(client: DuplexClientBase, collector: EventCollector) -> None: """Ack playback of every collected response's audio (probe shorthand). A response that has produced no audio yet is still acked (0 ms) unless it diff --git a/vllm_omni/clients/inline_duplex.py b/vllm_omni/clients/inline_duplex.py new file mode 100644 index 000000000..79c09e014 --- /dev/null +++ b/vllm_omni/clients/inline_duplex.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""In-process duplex client: the :class:`DuplexClient` usage over a ``DuplexOmni``. + +Example:: + + from vllm_omni.clients.inline_duplex import InlineDuplexClient + from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config + from vllm_omni.entrypoints.duplex_omni import DuplexOmni + + omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True) + cfg = create_duplex_session_config(ref_audio=audio_data_url(wav_bytes)) + async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=cfg) as client: + await client.stream_pcm(pcm16) + await client.commit() + async for response in client.responses(): + async for chunk in response.audio(): + play(chunk) + break + +There is no transport to drop, so this client has no reconnect, resume, +heartbeat or event-acknowledgement machinery. The ``DuplexOmni`` instance is +passed in by the caller; this module never imports server runtime code at +import time. +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from vllm_omni.clients.duplex import DuplexClientBase, DuplexConnectionError, SessionConfig + +if TYPE_CHECKING: + from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle + +__all__ = ["InlineDuplexClient"] + +#: Acknowledgements are a websocket-journal concern; heartbeats are lease touches and go through. +_TRANSPORT_ONLY_EVENTS = frozenset({"session.event_ack"}) + + +class InlineDuplexClient(DuplexClientBase): + """Async client for one duplex session on an in-process :class:`DuplexOmni`.""" + + def __init__( + self, + omni: DuplexOmni, + *, + model: str, + config: SessionConfig | None = None, + handshake_timeout_s: float = 30.0, + ) -> None: + super().__init__(model=model, config=config, handshake_timeout_s=handshake_timeout_s) + self.omni = omni + self._handle: DuplexSessionHandle | None = None + self._pump_task: asyncio.Task[None] | None = None + self._handshake_pending = False + + # -- transport hooks ------------------------------------------------------ + + async def _open(self) -> None: + # The handle's first event is ``session.created``; the pump delivers + # it through ``_dispatch`` where the base handshake is waiting. + self._handle = await self.omni.open_session(self.config.to_session_payload(model=self.model)) + # The base class follows up with the ``session.update`` handshake; the + # open above already carried that session object, so it is not resent. + self._handshake_pending = True + self._pump_task = asyncio.create_task(self._pump_events(), name="inline-duplex-client-pump") + + def _input_defaults(self) -> RealtimeInputDefaults: + """Audio format / sample-rate / VAD defaults declared by this session's object.""" + from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + + session_payload = self.session_info or self.config.to_session_payload(model=self.model) + return RealtimeInputDefaults().with_session_payload(session_payload) + + async def _send_command(self, payload: dict[str, object]) -> None: + handle = self._handle + if handle is None or handle.closed: + raise DuplexConnectionError("send failed: session is not open") + event_type = payload.get("type") + if event_type == "session.update" and self._handshake_pending: + self._handshake_pending = False + return + if event_type == "session.close": + await handle.close() + return + if event_type in _TRANSPORT_ONLY_EVENTS: + return + # Lazy import keeps ``vllm_omni.clients`` free of runtime imports at + # module load; the caller already holds a live DuplexOmni. + from vllm_omni.engine.duplex.commands import DuplexCommandError, command_from_realtime + from vllm_omni.engine.duplex.events import error_event + + try: + command = command_from_realtime(payload, defaults=self._input_defaults()) + except DuplexCommandError as exc: + await self._dispatch( + error_event(exc.code, str(exc), event_id=exc.event_id or payload.get("event_id")).to_realtime() + ) + return + await handle.submit(command) + + async def _teardown(self) -> None: + pump = self._pump_task + if pump is not None and not pump.done(): + pump.cancel() + try: + await pump + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + handle = self._handle + if handle is not None and not handle.closed: + try: + await handle.close() + except Exception: # noqa: BLE001 + pass + + # -- internals --------------------------------------------------------------- + + async def _pump_events(self) -> None: + handle = self._handle + assert handle is not None + reason = "closed" + try: + async for event in handle.events(): + await self._dispatch(event.to_realtime()) + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 + reason = f"session failed: {exc}" + self._finalize(reason, expected=False) + return + self._finalize(reason, expected=True) diff --git a/vllm_omni/clients/minicpmo_4_5.py b/vllm_omni/clients/minicpmo_4_5.py index 4b25346d4..b43346324 100644 --- a/vllm_omni/clients/minicpmo_4_5.py +++ b/vllm_omni/clients/minicpmo_4_5.py @@ -13,7 +13,7 @@ __all__ = ["create_duplex_session_config"] def create_duplex_session_config(*, ref_audio: str | None = None, **overrides: object) -> SessionConfig: """Session preset matching the MiniCPM-o 4.5 native duplex deployment.""" - extra_body: dict[str, object] = {"native_duplex": True, "force_listen_count": 0} + extra_body: dict[str, object] = {"force_listen_count": 0} extra_body.update(overrides.pop("extra_body", {})) config = SessionConfig( ref_audio=ref_audio, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 9c9186c16..c96f514e8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -309,14 +309,17 @@ class PipelineConfig: diffusers_class_name: str | None = None diffusers_class_aliases: tuple[str, ...] = () endpoint_restrictions: tuple[EndpointRestriction, ...] = () - # Optional model-owned duplex planner loaded by the stable engine runtime. + # Dotted path of the model's ``DuplexModelPlugin``. A pipeline is a duplex + # model iff this is set: ``vllm-omni serve`` then always runs it through + # ``DuplexOmni`` (every served surface runs on a duplex session) and the engine hosts one + # ``DuplexOrchestrator`` with the plugin loaded. + duplex_plugin: str | None = None + # Legacy duplex wiring of the models that are not ported to the plugin + # framework yet (PersonaPlex, Nemotron VoiceChat). Nothing reads them: a + # pipeline that only declares these is served turn-based. Each field goes + # away with the follow-up PR that ports its model to ``duplex_plugin``. duplex_runtime_extension: str | None = None - # Optional model-owned Serving adapter loaded only when the Realtime duplex - # endpoint is enabled. Generic OpenAI modules must not select a model. duplex_serving_adapter: str | None = None - # Explicitly enable the stable duplex control mechanism. This is separate - # from the optional model extension because turn-commit-only deployments - # do not require a model planner. duplex_control_enabled: bool = False # Bundled deploy defaults for this concrete pipeline topology. The file is # loaded from vllm_omni/deploy; None uses DeployConfig defaults. @@ -510,6 +513,12 @@ class DuplexSessionRuntimeConfig: max_pending_input_bytes_per_session: int = 16 * 1024 * 1024 max_pending_turns_per_session: int = 4 max_sessions: int = 1 + # Unread by the plugin framework. It used to bound the per-session + # completed-append table that made a retried append RPC submit once; the + # framework carries appends as one-way commands on the runner's ordered + # mailbox, so there is no client-visible retry to deduplicate. The field + # stays so the deploy configs of the models that are not ported yet still + # parse, and goes away with the PR that ports the last of them. completed_append_cache_size: int = 256 server_vad_model_path: str | None = None # Startup warmup: run this many silent 80 ms-style frames through a diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 4e9fd5a77..ca772afbd 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -1,8 +1,8 @@ # MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. Every other -# duplex_session knob keeps its runtime default. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 9c3818813..9600e5959 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -4,8 +4,8 @@ # isolated from codec/vocoder contention while preserving stage-2 batching. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. Every other -# duplex_session knob keeps its runtime default. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 5eddbb6f4..aae531218 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -4,8 +4,8 @@ # Code2Wav run independently on GPUs 0, 1, and 2. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. Every other -# duplex_session knob keeps its runtime default. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index eacc069d1..00b1090f8 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -15,8 +15,9 @@ # so 4096 is the empirically-verified safe ceiling. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. The default -# duplex_session caps sessions at 1, matching Stage 0/1 max_num_seqs below. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1). The default duplex_session caps sessions at 1, +# matching Stage 0/1 max_num_seqs below. session_mode: duplex active_stream_window: 1 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 6850e6b29..3469d8fee 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1,498 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -""" -Async Omni Engine for vLLM-Omni multi-stage runtime. - -AsyncOmniEngine in the caller's thread is a thin proxy that communicates -with the Orchestrator (running in a background thread) via janus queues. -""" +"""AsyncOmniEngine: the turn-based engine (add_request / streaming updates / CFG companions / interaction).""" from __future__ import annotations -import asyncio -import concurrent.futures import copy -import queue import shutil -import threading import time -import uuid -import weakref from collections.abc import Mapping, Sequence from typing import Any, Literal, cast -import janus -from vllm import envs as vllm_envs from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest -from vllm.v1.engine.input_processor import InputProcessor -from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config -from vllm_omni.config.stage_config import ( - DuplexSessionRuntimeConfig, - PipelineConfig, - load_deploy_config, -) from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_engine_utils import ( - SHUTDOWN_ENQUEUE_TIMEOUT_S, - SHUTDOWN_JOIN_TIMEOUT_S, apply_omni_final_stage_metadata, - enqueue_orchestrator_shutdown, inject_global_id, - is_abort_transport_shutdown, - is_janus_sync_queue_shutdown, - shutdown_runtime_after_orchestrator, upgrade_to_omni_request, - weak_shutdown_async_omni_engine, -) -from vllm_omni.engine.duplex.control_client import DuplexControlClient -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - load_duplex_runtime_extension, - validate_duplex_runtime_extension, ) from vllm_omni.engine.messages import ( - AbortRequestMessage, - AbortResultMessage, AddCompanionRequestMessage, - CollectiveRPCRequestMessage, - CollectiveRPCResultMessage, - EngineQueueMessage, - ErrorMessage, InteractionMessage, - OutputMessage, StageSubmissionMessage, ) -from vllm_omni.engine.orchestrator import Orchestrator -from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.engine.omni_engine_base import OmniEngineBase, StageRuntimeInfo +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorBase from vllm_omni.engine.serialization import deserialize_additional_information -from vllm_omni.engine.stage_client import StageClient -from vllm_omni.engine.stage_init_utils import build_stage0_input_processor -from vllm_omni.engine.stage_pool import StagePool -from vllm_omni.engine.stage_runtime import ( - StageRuntimeInfo, - create_stage_runtime, -) -from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin -from vllm_omni.entrypoints.utils import parse_stage_overrides from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams -from vllm_omni.metrics.prometheus import OmniRequestCounter logger = init_logger(__name__) -_STARTUP_POLL_INTERVAL_S = 1.0 -_REQUEST_QUEUE_MAXSIZE = 256 -_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None] +class AsyncOmniEngine(OmniEngineBase): + """Turn-based engine used by ``Omni`` / ``AsyncOmni``.""" -def load_and_resolve_stage_configs( - model: str, - kwargs: dict[str, Any], - *, - trust_remote_code: bool | None, - deploy_config_path: str | None, - stage_overrides: Mapping[str, Mapping[str, Any]] | None, - strategy_config_path: str | None, -) -> OmniConfigResolution: - """Compatibility seam delegating to the single config resolver.""" - return resolve_omni_config( - model, - trust_remote_code=trust_remote_code, - cli_overrides=kwargs, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ) - - -class AsyncOmniEngine: - """Thin proxy that launches an Orchestrator in a background thread. - - All stage clients, input/output processors, and stage-to-stage transfer - logic live inside the Orchestrator coroutine (running in its own thread - with a dedicated asyncio event loop). This class communicates with it - via janus queues (sync side for callers, async side for orchestrator). - - Args: - model: Model name or path - init_timeout: Total timeout waiting for orchestrator startup (seconds). - stage_init_timeout: Timeout for stage initialization (seconds) - **kwargs: Additional arguments - """ - - # Class-level defaults so tests that bypass __init__ via object.__new__ - # don't AttributeError when stage-init / forward paths touch these attrs. - _log_stats: bool = False - _coordinator_runtime: Any = None - _transfer_emitter: Any = None - _prom_metrics: Any = None - _enable_orch_monitor: bool = False - # Lazily created by get_output_blocking_async(). - _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None - - def __init__( - self, - model: str, - stage_init_timeout: int = 300, - init_timeout: int = 600, - single_stage_mode: bool = False, - transfer_emitter: Any = None, - prom_metrics: Any = None, - log_stats: bool = False, - tokenizer: str | None = None, - trust_remote_code: bool | None = None, - **kwargs: Any, - ) -> None: - self.model = model - self.tokenizer = tokenizer - # Cached by get_diffusion_od_config(). - self._diffusion_od_config_view: Any = None - startup_timeout = int(init_timeout) - # Forwarded into Orchestrator so its _forward_to_next_stage path can - # emit per-edge transfer_tx_s / transfer_size_bytes histograms. - # Optional: when None, Orchestrator silently skips TX emit (existing - # RX path still works via OrchestratorAggregator). - self._transfer_emitter = transfer_emitter - self._prom_metrics = prom_metrics - # Drives upstream EngineCore + scheduler stats production. When False - # the engine skips SchedulerStats / IterationStats; the per-(stage, - # replica) vllm:* wrap stays registered but reads zero. Respects the - # --log-stats CLI flag set by the user via OmniBase. - self._log_stats = log_stats - self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False)) - - logger.info(f"[AsyncOmniEngine] Initializing with model {model}") - - # ------------------------------------------------------------------ # - # Single-stage mode detection # - # ------------------------------------------------------------------ # - # Single-stage mode is enabled when the caller explicitly passes # - # single_stage_mode=True, or when a stage_id is provided in the args. # - _stage_id_kwarg = kwargs.get("stage_id") - if isinstance(_stage_id_kwarg, int) and not single_stage_mode: - single_stage_mode = True - - self.single_stage_mode: bool = single_stage_mode - self._single_stage_id_filter: int | None = ( - int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None - ) - self._omni_master_address: str | None = kwargs.get("omni_master_address") - self._omni_master_port: int | None = kwargs.get("omni_master_port") - - # New omni-coordinator flags. Consumed only in single_stage_mode. - # ``omni_dp_size_local`` is process-local: each invocation (head and - # every headless) launches that many replicas for its own stage. - self._omni_dp_size_local: int = int(kwargs.get("omni_dp_size_local") or 1) - if self._omni_dp_size_local < 1: - raise ValueError(f"--omni-dp-size-local must be >= 1, got {self._omni_dp_size_local}") - self._omni_lb_policy: str = str(kwargs.get("omni_lb_policy") or "random") - self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) - if self._omni_heartbeat_timeout <= 0: - raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") - # Concurrent same-device stage init (admission + SH/EX phase locks). - # Sourced from the parallel_stage_init orchestrator/CLI arg (config, - # not an env var); default False preserves serial init. - self._parallel_stage_init: bool = bool(kwargs.get("parallel_stage_init") or False) - - if single_stage_mode: - logger.info( - "[AsyncOmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", - self._single_stage_id_filter, - self._omni_master_address, - self._omni_master_port, - ) - - # Keep the historical tuple return from _resolve_stage_configs while - # retaining the richer resolver result for pipeline-wide settings. - # Overrides of that private seam fall back to the factory below. - deploy_config_path = kwargs.get("deploy_config") - self._config_resolution: OmniConfigResolution | None = None - self.config_path, self.stage_configs = self._resolve_stage_configs( - model, - kwargs, - trust_remote_code=trust_remote_code, - ) - if self._config_resolution is None: - pipeline_config = StageConfigFactory.get_pipeline_config( - model=model, - trust_remote_code=bool(trust_remote_code), - deploy_config_path=deploy_config_path, - ) - self._set_pipeline_runtime_config(pipeline_config, self.config_path) - else: - self._set_pipeline_runtime_config( - self._config_resolution.pipeline_config, - self._config_resolution.config_path, - ) - - self.num_stages = len(self.stage_configs) - stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None - self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) - self.stage_pools: list[StagePool] = [] - self.stage_clients: list[StageClient] = [] # logical-stage view for external readers - self.input_processor: InputProcessor | None = None - self.prompt_transform_func: Any | None = None - self.prompt_expand_func: Any | None = None - self.supported_tasks: tuple[str, ...] = ("generate",) - self.default_sampling_params_list: list[OmniSamplingParams] = [] - self.stage_metadata: list[StageRuntimeInfo] = [] - # Janus queues are constructed eagerly here (not deferred to the - # orchestrator thread) so the master server's ROUTER thread always - # sees a non-None ``self.request_queue`` when on_register fires. - # ``async_q`` lazily binds to whatever event loop first awaits on - # it (the orchestrator loop), so cross-thread use stays correct. - self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) - self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() - self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() - self._shutdown_called = False - self._weak_finalizer: weakref.finalize | None = None - self._correlated_rpc_client: CorrelatedRpcClient | None = None - self._duplex_control_client: DuplexControlClient | None = None - self._running_counter = OmniRequestCounter() - self._engines_waiting_counter = OmniRequestCounter() - - logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") - - # Launch orchestrator background thread - startup_future: concurrent.futures.Future = concurrent.futures.Future() - - self.orchestrator_thread = threading.Thread( - target=self._bootstrap_orchestrator, - args=( - stage_init_timeout, - startup_future, - ), - daemon=True, - name="orchestrator", - ) - self.orchestrator_thread.start() - self._wait_for_orchestrator_init(startup_future, startup_timeout) - self._correlated_rpc_client = CorrelatedRpcClient( - self.request_queue.sync_q, - self.rpc_output_queue.sync_q, - ) - - # Stage runtime fields are assigned directly on self by the bootstrap thread. - self._weak_finalizer = weakref.finalize( - self, - weak_shutdown_async_omni_engine, - self.orchestrator_thread, - self.request_queue, - self.output_queue, - self.rpc_output_queue, - self._correlated_rpc_client, - ) - - logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") - - def get_diffusion_od_config(self) -> Any: - """Expose the diffusion ``model_class_name`` to client-side model-extras. - - The worker holds the full config; here we just resolve the pipeline class - name from the model config (cached). ``model_class_name`` may be ``None``. - """ - if self._diffusion_od_config_view is None: - from types import SimpleNamespace - - from vllm_omni.diffusion.data import resolve_model_class_name - from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata - - model_class_name = resolve_model_class_name(self.model) - metadata = get_diffusion_model_metadata(model_class_name) - self._diffusion_od_config_view = SimpleNamespace( - model_class_name=model_class_name, - supports_multimodal_inputs=metadata.supports_multimodal_inputs, - max_multimodal_image_inputs=metadata.max_multimodal_image_inputs, - supports_mixed_reference_inputs=metadata.supports_mixed_reference_inputs, - ) - return self._diffusion_od_config_view - - def _initialize_stages(self, stage_init_timeout: int) -> None: - """Initialize stage clients/processors via StageRuntime and assign to self.""" - self._runtime = create_stage_runtime( - stage_configs=self.stage_configs, - model=self.model, - config_path=self.config_path, - single_stage_mode=self.single_stage_mode, - stage_init_timeout=stage_init_timeout, - async_chunk=self.async_chunk, - tokenizer=self.tokenizer, - parallel_stage_init=self._parallel_stage_init, - single_stage_id_filter=self._single_stage_id_filter, - omni_master_address=self._omni_master_address, - omni_master_port=self._omni_master_port, - omni_dp_size_local=self._omni_dp_size_local, - omni_heartbeat_timeout=self._omni_heartbeat_timeout, - omni_lb_policy=self._omni_lb_policy, - request_queue=self.request_queue, - log_stats=self._log_stats, - ) - self._runtime.initialize() - - self.num_stages = len(self.stage_configs) - self.stage_pools = self._runtime.stage_pools - self.stage_clients = [ - cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None - ] - self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] - self.output_processors = [pool.output_processor for pool in self.stage_pools] - self.input_processor = ( - build_stage0_input_processor(self.stage_vllm_configs[0]) - if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None - else None - ) - self.prompt_transform_func = ( - getattr(self.stage_clients[0], "prompt_transform_func", None) if self.stage_clients else None - ) - self.prompt_expand_func = next( - ( - getattr(client, "prompt_expand_func", None) - for client in self.stage_clients - if getattr(client, "prompt_expand_func", None) is not None - ), - None, - ) - self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] - self.stage_metadata = [ - StageRuntimeInfo( - final_output=client.final_output, - final_output_type=client.final_output_type, - stage_type=client.stage_type, - model_stage=getattr(client, "model_stage", None), - ) - for client in self.stage_clients - ] - supported_tasks: set[str] = set() - if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): - supported_tasks.add("generate") - if any(meta.final_output_type == "audio" for meta in self.stage_metadata): - supported_tasks.add("speech") - self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) - - def _bootstrap_orchestrator( - self, - stage_init_timeout: int, - startup_future: concurrent.futures.Future, - ) -> None: - """Create loop, initialize stages, then run Orchestrator.""" - - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - async def _run_orchestrator() -> None: - self._initialize_stages(stage_init_timeout) - - pd_config = self._detect_pd_config() - - membership_controller = self._runtime.create_membership_controller() - duplex_runtime_extension = None - if self._duplex_control_enabled: - duplex_runtime_extension = load_duplex_runtime_extension( - getattr(self, "_duplex_runtime_extension_path", None) - ) - if duplex_runtime_extension is not None: - stage_clients = [] - for pool in self.stage_pools: - assert pool.stage_client is not None - stage_clients.append(pool.stage_client) - validate_duplex_runtime_extension( - duplex_runtime_extension, - sampling_defaults=tuple(client.default_sampling_params for client in stage_clients), - ) - - orchestrator = Orchestrator( - request_async_queue=self.request_queue.async_q, - output_async_queue=self.output_queue.async_q, - rpc_async_queue=self.rpc_output_queue.async_q, - stage_pools=self.stage_pools, - async_chunk=self.async_chunk, - pd_config=pd_config, - membership_controller=membership_controller, - running_counter=self._running_counter, - engines_waiting_counter=self._engines_waiting_counter, - transfer_emitter=self._transfer_emitter, - prom_metrics=self._prom_metrics, - log_stats=self._log_stats, - enable_orch_monitor=self._enable_orch_monitor, - duplex_runtime_extension=duplex_runtime_extension, - enable_duplex_control=self._duplex_control_enabled, - duplex_session_config=self.duplex_session_config, - ) - if not startup_future.done(): - startup_future.set_result(asyncio.get_running_loop()) - await orchestrator.run() - - try: - loop.run_until_complete(_run_orchestrator()) - except Exception as e: - if not startup_future.done(): - wrapped = RuntimeError(f"Orchestrator initialization failed: {e}") - wrapped.__cause__ = e - startup_future.set_exception(wrapped) - logger.exception("[AsyncOmniEngine] Orchestrator thread crashed") - error_text = str(e) or "Orchestrator thread crashed" - try: - error_msg = ErrorMessage(error=error_text, fatal=True) - if self.output_queue is not None: - self.output_queue.sync_q.put_nowait(error_msg) - if self.rpc_output_queue is not None: - self.rpc_output_queue.sync_q.put_nowait(error_msg) - except Exception: - pass - raise - finally: - try: - pending = [task for task in asyncio.all_tasks(loop) if not task.done()] - for task in pending: - task.cancel() - if pending: - loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) - loop.run_until_complete(loop.shutdown_asyncgens()) - if hasattr(loop, "shutdown_default_executor"): - loop.run_until_complete(loop.shutdown_default_executor()) - except Exception: - logger.exception("[AsyncOmniEngine] Failed during orchestrator loop cleanup") - finally: - asyncio.set_event_loop(None) - loop.close() - - def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None: - """ - Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process. - """ - deadline = time.monotonic() + startup_timeout - while True: - remaining = deadline - time.monotonic() - if remaining <= 0: - logger.warning( - "[AsyncOmniEngine] Orchestrator startup timed out after %ss. " - "Multi-stage deployments that initialize stages sequentially on one device " - "or load checkpoints from slow storage may need larger --init-timeout and " - "--stage-init-timeout values.", - startup_timeout, - ) - self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup timeout") - raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") - try: - startup_future.result( - timeout=min(remaining, _STARTUP_POLL_INTERVAL_S), - ) - break - except concurrent.futures.TimeoutError: - if not self.orchestrator_thread.is_alive(): - self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure") - if startup_future.done(): - startup_future.result() # re-raises the real exception - raise RuntimeError("Orchestrator thread died during startup") - except Exception: - self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure") - raise + def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: + return Orchestrator(**orchestrator_kwargs) # ---- request helpers ---- @@ -918,158 +465,6 @@ class AsyncOmniEngine: len(companions), ) - def _detect_pd_config(self) -> dict[str, Any] | None: - """Detect PD (Prefill-Decode) disaggregation config from stage_configs. - Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. - """ - pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs) - if pd_pair is None: - return None - prefill_idx, decode_idx = pd_pair - - # Extract bootstrap address from prefill stage engine_args - bootstrap_addr: str | None = None - try: - prefill_cfg = self.stage_configs[prefill_idx] - ea = getattr(prefill_cfg, "engine_args", None) - kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None - if kv_cfg is not None: - port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT - kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1" - bootstrap_addr = f"http://{kv_ip}:{port}" - except Exception as exc: - logger.warning("[AsyncOmniEngine] Could not extract PD bootstrap address: %s", exc) - - logger.info( - "[AsyncOmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s", - prefill_idx, - decode_idx, - bootstrap_addr, - ) - prefill_engine_id: str | None = None - try: - prefill_client = self.stage_clients[prefill_idx] - kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None) - prefill_engine_id = getattr(kv_cfg, "engine_id", None) - except Exception as exc: - logger.warning("[AsyncOmniEngine] Could not extract prefill engine_id: %s", exc) - - return { - "pd_pair": (prefill_idx, decode_idx), - "bootstrap_addr": bootstrap_addr, - "prefill_engine_id": prefill_engine_id, - } - - def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: - """Apply a strategy-derived ``omni_lb_policy`` to the engine. - - Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is - the engine default and is treated as "unset" (indistinguishable from no - flag), so a strategy value overrides it. If the user explicitly passed a - non-default policy that conflicts with the strategy-derived one, raise so - the mismatch is not silently ignored. - """ - if not derived: - return - explicit = kwargs.get("omni_lb_policy") - user_set = explicit is not None and str(explicit) != "random" - if user_set: - if str(explicit) != str(derived): - raise ValueError( - f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given " - f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. " - "Drop --omni-lb-policy to use the strategy value, or make them match." - ) - return - if self._omni_lb_policy != str(derived): - logger.info( - "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).", - derived, - self._omni_lb_policy, - ) - self._omni_lb_policy = str(derived) - - @staticmethod - def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]: - """Compatibility seam for the factory-owned diffusion fallback.""" - return StageConfigFactory.create_default_diffusion(kwargs) - - def _set_pipeline_runtime_config( - self, - pipeline_config: PipelineConfig | None, - config_path: str | None, - ) -> None: - """Initialize engine-wide settings resolved from pipeline metadata.""" - self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () - self._duplex_runtime_extension_path = ( - pipeline_config.duplex_runtime_extension if pipeline_config is not None else None - ) - self.duplex_serving_adapter_path = ( - pipeline_config.duplex_serving_adapter if pipeline_config is not None else None - ) - self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled) - self.duplex_session_config = DuplexSessionRuntimeConfig() - if config_path is not None: - self.duplex_session_config = load_deploy_config(config_path).duplex_session - - def _resolve_stage_configs( - self, - model: str, - kwargs: dict[str, Any], - *, - trust_remote_code: bool | None, - ) -> tuple[str, list[Any]]: - """Resolve stage configs and inject defaults shared by orchestrator/headless.""" - - for legacy_arg in ("stage_configs_path", "stage_configs"): - if legacy_arg in kwargs: - raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") - - # log_stats is captured by __init__; its CLI-only negative alias must - # not cross into per-stage structured config ownership validation. - kwargs.pop("disable_log_stats", None) - deploy_config_path = kwargs.pop("deploy_config", None) - strategy_config_path = kwargs.pop("strategy_config", None) - # CLI callers arrive pre-parsed; offline Python callers may use the - # JSON-string form documented in recipes. - stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None)) - - # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; - # stage configs know the field as ``streaming_output``. Mirror only a - # truthy value so the CLI's False default does not override deploy YAML. - if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: - kwargs["streaming_output"] = True - - resolution = cast( - _ConfigResolutionResult, - load_and_resolve_stage_configs( - model, - kwargs, - trust_remote_code=trust_remote_code, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ), - ) - if isinstance(resolution, OmniConfigResolution): - self._config_resolution = resolution - config_path = resolution.config_path - stage_configs = list(resolution.stage_configs) - strategy_lb_policy = resolution.omni_lb_policy - else: - # Compatibility for overrides of the historical tuple-returning - # seam. Production always receives OmniConfigResolution above. - config_path, stage_configs, strategy_lb_policy = resolution - - # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is - # an orchestrator-level knob (read once at construction), so apply it here - # rather than as a per-stage config field. - self._apply_strategy_lb_policy(strategy_lb_policy, kwargs) - - return cast(str, config_path), stage_configs - - # ==================== Public API ==================== - def add_request( self, request_id: str, @@ -1241,390 +636,6 @@ class AsyncOmniEngine: resumable=resumable, ) - def open_duplex_session( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Open an engine-level duplex session.""" - return self._get_duplex_control_client().open( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - runtime_config=runtime_config, - fence=fence, - timeout=timeout, - ) - - async def open_duplex_session_async( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for opening an engine-level duplex session.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.open_duplex_session( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - runtime_config=runtime_config, - fence=fence, - timeout=timeout, - ), - ) - - def append_duplex_input( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - expected_epoch: int | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Append input to an engine-level duplex session.""" - return self._get_duplex_control_client().append( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - expected_epoch=expected_epoch, - fence=fence, - timeout=timeout, - ) - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - expected_epoch: int | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for appending duplex input.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.append_duplex_input( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - expected_epoch=expected_epoch, - fence=fence, - timeout=timeout, - ), - ) - - def signal_duplex_turn( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Signal an engine-level duplex turn.""" - return self._get_duplex_control_client().signal( - session_id, - event=event, - fence=fence, - next_fence=next_fence, - session_config=session_config, - runtime_config=runtime_config, - timeout=timeout, - ) - - async def signal_duplex_turn_async( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for signaling a duplex turn.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.signal_duplex_turn( - session_id, - event=event, - fence=fence, - next_fence=next_fence, - session_config=session_config, - runtime_config=runtime_config, - timeout=timeout, - ), - ) - - def close_duplex_session( - self, - session_id: str, - *, - reason: str = "client_close", - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Close an engine-level duplex session.""" - return self._get_duplex_control_client().close( - session_id, - reason=reason, - fence=fence, - timeout=timeout, - ) - - def touch_duplex_session( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return self._get_duplex_control_client().touch( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ) - - async def touch_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = 10.0, - ) -> dict[str, object]: - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.touch_duplex_session( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ), - ) - - def resume_duplex_session( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return self._get_duplex_control_client().resume( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ) - - async def resume_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = 10.0, - ) -> dict[str, object]: - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.resume_duplex_session( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ), - ) - - def _get_duplex_control_client(self) -> DuplexControlClient: - client = getattr(self, "_duplex_control_client", None) - if client is None: - transport = getattr(self, "_correlated_rpc_client", None) - if transport is None: - raise RuntimeError("correlated RPC client is not initialized") - client = DuplexControlClient( - transport, - control_id_factory=lambda: uuid.uuid4().hex, - ) - self._duplex_control_client = client - return client - - async def close_duplex_session_async( - self, - session_id: str, - *, - reason: str = "client_close", - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for closing an engine-level duplex session.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.close_duplex_session(session_id, reason=reason, fence=fence, timeout=timeout), - ) - - def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: - """Read one output message from the Orchestrator output queue.""" - try: - return self.output_queue.sync_q.get(timeout=timeout) - except queue.Empty: - if not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return None - - async def try_get_output_async(self) -> EngineQueueMessage | None: - """Async read from the Orchestrator output queue.""" - try: - return self.output_queue.sync_q.get_nowait() - except queue.Empty: - if not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return None - - async def get_output_blocking_async(self, timeout: float = 1.0) -> EngineQueueMessage | None: - """Blocking-wait read from the Orchestrator output queue. - - Waits up to ``timeout`` seconds in a dedicated drain thread for the - next message (condition-variable wakeup instead of a poll cadence); - returns ``None`` on timeout so the caller keeps its liveness check, - mirroring ``try_get_output_async``'s contract. Used by the serving - final-output drain when ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` is on. - """ - executor = self._output_drain_executor - if executor is None: - executor = concurrent.futures.ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="omni-output-drain", - ) - self._output_drain_executor = executor - - sync_q = self.output_queue.sync_q - - def _drain_get() -> EngineQueueMessage | None: - # Exceptions are swallowed to a None sentinel: the queue may be - # closed mid-shutdown, and an exception left on an executor future - # after task cancellation would warn as never-retrieved. - try: - return sync_q.get(timeout=timeout) - except queue.Empty: - return None - except Exception: - return None - - loop = asyncio.get_running_loop() - msg = await loop.run_in_executor(executor, _drain_get) - if msg is None and not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return msg - - def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: - """Get cached metadata for a stage.""" - return self.stage_metadata[stage_id] - - def abort(self, request_ids: list[str]) -> None: - """Fire-and-forget abort: enqueue and return without waiting. - - Prefer :meth:`abort_async` when the caller needs acknowledgment that - stage aborts, binding release, and orchestrator request cleanup finished. - """ - if not request_ids or getattr(self, "_shutdown_called", False): - return - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") - try: - self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) - except Exception as exc: - if getattr(self, "_shutdown_called", False) and is_janus_sync_queue_shutdown(exc): - return - raise - - async def abort_async( - self, - request_ids: list[str], - timeout: float | None = None, - ) -> list[OutputMessage]: - """Abort requests and wait for orchestrator acknowledgment. - - Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the - :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and - raises if the orchestrator reports failure or times out. - - Returns: - Final-stage AR abort ``OutputMessage`` list carrying partial - tokens generated before abort (empty for diffusion / no OP state). - """ - if not request_ids or getattr(self, "_shutdown_called", False): - return [] - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") - transport = self._correlated_rpc_client - if transport is None: - raise RuntimeError("correlated RPC client is not initialized") - - rpc_id = uuid.uuid4().hex - msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id) - - def _wait() -> AbortResultMessage: - result_msg = transport.execute( - ("abort", rpc_id), - msg, - timeout=timeout, - timeout_message=f"abort timed out after {timeout} seconds", - block_on_submit=True, - ) - if not isinstance(result_msg, AbortResultMessage): - raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}") - return result_msg - - loop = asyncio.get_running_loop() - try: - result_msg = await loop.run_in_executor(None, _wait) - except Exception as exc: - if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): - return [] - raise - if not result_msg.success: - raise RuntimeError(result_msg.error or "abort failed") - return list(result_msg.abort_outputs or []) - def submit_interaction( self, request_id: str, @@ -1649,166 +660,6 @@ class AsyncOmniEngine: """Async interaction API.""" self.submit_interaction(request_id, interaction) - def collective_rpc( - self, - method: str, - timeout: float | None = None, - args: tuple[Any, ...] = (), - kwargs: dict[str, Any] | None = None, - stage_ids: list[int] | None = None, - ) -> list[Any]: - """Send a control RPC to the Orchestrator and wait for aggregated results. - This uses a dedicated RPC output queue so control-plane messages do not - race with the normal request output polling loop. - """ - rpc_id = uuid.uuid4().hex - msg = CollectiveRPCRequestMessage( - rpc_id=rpc_id, - method=method, - timeout=timeout, - args=tuple(args), - kwargs=kwargs or {}, - stage_ids=stage_ids, - ) - - transport = self._correlated_rpc_client - if transport is None: - raise RuntimeError("correlated RPC client is not initialized") - result_msg = transport.execute( - ("collective", rpc_id), - msg, - timeout=timeout, - timeout_message=f"collective_rpc timed out after {timeout} seconds", - block_on_submit=True, - ) - if not isinstance(result_msg, CollectiveRPCResultMessage): - raise RuntimeError(f"unexpected collective RPC result type: {type(result_msg).__name__}") - return list(result_msg.results) - - async def collective_rpc_async( - self, - method: str, - timeout: float | None = None, - args: tuple[Any, ...] = (), - kwargs: dict[str, Any] | None = None, - stage_ids: list[int] | None = None, - ) -> list[Any]: - """Async wrapper around collective_rpc().""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.collective_rpc( - method=method, - timeout=timeout, - args=args, - kwargs=kwargs, - stage_ids=stage_ids, - ), - ) - - def is_alive(self) -> bool: - """Whether the orchestrator thread is alive.""" - return bool(self.orchestrator_thread.is_alive()) - - def shutdown(self) -> None: - """Send shutdown message and wait for the Orchestrator thread to exit.""" - if getattr(self, "_shutdown_called", False): - return - self._shutdown_called = True - finalizer = getattr(self, "_weak_finalizer", None) - if finalizer is not None and finalizer.alive: - finalizer.detach() - - logger.info("[AsyncOmniEngine] Shutting down Orchestrator") - request_queue_closed = False - shutdown_enqueued = enqueue_orchestrator_shutdown( - self.request_queue, - timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S, - ) - if self.request_queue is not None and not shutdown_enqueued: - logger.error( - "[AsyncOmniEngine] Failed to enqueue orchestrator shutdown; " - "closing the request queue to wake the request handler" - ) - try: - self.request_queue.close() - request_queue_closed = True - except Exception: - logger.exception("[AsyncOmniEngine] Failed to close the request queue") - - if self._correlated_rpc_client is not None: - try: - self._correlated_rpc_client.close() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to close correlated RPC client") - - orchestrator_stopped = False - try: - if self.is_alive(): - self.orchestrator_thread.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S) - orchestrator_stopped = not self.is_alive() - if not orchestrator_stopped: - logger.error( - "[AsyncOmniEngine] Orchestrator did not stop within %.1f seconds; continuing cleanup", - SHUTDOWN_JOIN_TIMEOUT_S, - ) - except Exception: - logger.exception("[AsyncOmniEngine] Failed to join Orchestrator thread") - - for q in (self.request_queue, self.output_queue, self.rpc_output_queue): - try: - if not (q is self.request_queue and request_queue_closed): - q.close() - except Exception: - pass - - if self._output_drain_executor is not None: - # Any in-flight blocking get bails out within its ≤1 s timeout - # (or immediately via the queue close above), so don't wait. - self._output_drain_executor.shutdown(wait=False) - self._output_drain_executor = None - - if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped: - try: - self._runtime.shutdown() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to shutdown StageRuntime") - elif hasattr(self, "_runtime") and self._runtime is not None: - logger.warning("[AsyncOmniEngine] Deferring StageRuntime shutdown until the Orchestrator exits") - threading.Thread( - target=shutdown_runtime_after_orchestrator, - args=(self.orchestrator_thread, self._runtime), - daemon=True, - name="omni-stage-runtime-shutdown", - ).start() - - # ── Release CuMem allocator memory pool ────────────────────────────── - # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory - # Management) allocator holds model weights in a singleton memory pool - # that lives in the parent process. Killing the engine-core subprocess - # does NOT release this pool — the weights stay resident on the GPU - # and can cause CUDA OOM for subsequent engine instances (especially - # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). - # - # CuMemAllocator.sleep() is NOT idempotent — calling it on already- - # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator - # cuMemRelease (double-free of the memory handle). Use release_pools() - # instead, which is the designed cleanup path: it drops MemPool refs - # and lets the destructor/free path handle asleep entries correctly - # (returns a null handle so the C extension skips unmap/release). - try: - from vllm.device_allocator.cumem import CuMemAllocator, cumem_available - - if cumem_available: - allocator = CuMemAllocator.get_instance() - allocator.release_pools() - logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") - except Exception: - pass - - def _try_shutdown(self, *args, **kwargs) -> None: - try: - self.shutdown() - except Exception: - logger.exception(*args, **kwargs) +# ``StageRuntimeInfo`` moved to ``omni_engine_base``; keep the historical import path. +__all__ = ["AsyncOmniEngine", "StageRuntimeInfo"] diff --git a/vllm_omni/engine/duplex/__init__.py b/vllm_omni/engine/duplex/__init__.py index bd5ae0b3f..533e31e23 100644 --- a/vllm_omni/engine/duplex/__init__.py +++ b/vllm_omni/engine/duplex/__init__.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Engine-side helpers for the full-duplex runtime.""" +"""Engine-resident full-duplex sessions: typed commands/events, the merged session +state, the per-session runner, the session manager and the model plugin contract.""" diff --git a/vllm_omni/entrypoints/duplex/audio.py b/vllm_omni/engine/duplex/audio.py similarity index 96% rename from vllm_omni/entrypoints/duplex/audio.py rename to vllm_omni/engine/duplex/audio.py index 747293a1a..5f163826a 100644 --- a/vllm_omni/entrypoints/duplex/audio.py +++ b/vllm_omni/engine/duplex/audio.py @@ -11,6 +11,8 @@ import wave import numpy as np +from vllm_omni.utils.audio_resample import StreamingAudioResampler + try: from audioop import alaw2lin, lin2alaw, lin2ulaw, ulaw2lin except ImportError: # pragma: no cover - audioop is removed in newer Python. @@ -39,14 +41,10 @@ def validate_input_sample_rate_hz(sample_rate_hz: object) -> int: def resample_pcm16_mono(raw: bytes, *, source_rate_hz: int, target_rate_hz: int) -> bytes: if source_rate_hz <= 0 or target_rate_hz <= 0 or source_rate_hz == target_rate_hz: return raw - # Import lazily so the generic Duplex audio module does not pull in the - # OpenAI API server while the Duplex stack itself is being imported. - from vllm_omni.entrypoints.openai.audio_utils_mixin import StreamingAudioResampler - samples = np.frombuffer(raw, dtype=" StreamingAudioResampler._max_polyphase_factor: + if reduced_factor > StreamingAudioResampler.max_polyphase_factor: # The FIR size scales with the reduced ratio. Keep unusual client # rates bounded while common audio rates use the high-quality path. if samples.size <= 1: diff --git a/vllm_omni/engine/duplex/commands.py b/vllm_omni/engine/duplex/commands.py new file mode 100644 index 000000000..90d48d050 --- /dev/null +++ b/vllm_omni/engine/duplex/commands.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Typed commands accepted by a duplex session. + +``DuplexSessionHandle.submit()`` takes one of these; the websocket handler and +``InlineDuplexClient`` build them with :func:`command_from_realtime` from the +OpenAI Realtime client event vocabulary. Every command can also render the +session-internal mailbox payload (``payload()``), which is the dictionary +vocabulary the session runner bodies were written against. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from typing import TYPE_CHECKING, ClassVar + +if TYPE_CHECKING: + from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + + +class DuplexCommandError(ValueError): + """A client payload could not be turned into a command.""" + + def __init__(self, message: str, *, code: str = "bad_event", event_id: str | None = None) -> None: + super().__init__(message) + self.code = code + self.event_id = event_id + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexCommand: + #: Mailbox event type this command renders to (see ``payload()``). + type: ClassVar[str] = "" + #: Client correlation id (OpenAI ``event_id``), echoed on error events. + event_id: str | None = None + + def payload(self) -> dict[str, object]: + """Render the session-internal mailbox dictionary.""" + data: dict[str, object] = {"type": self.type} + for f in fields(self): + if f.name == "event_id": + if self.event_id is not None: + data["realtime_event_id"] = self.event_id + continue + value = getattr(self, f.name) + if value is None: + continue + if isinstance(value, tuple): + value = list(value) + elif isinstance(value, Mapping): + value = dict(value) + data[f.name] = value + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendAudio(DuplexCommand): + type: ClassVar[str] = "input_audio_buffer.append" + #: Raw audio bytes in ``format`` at ``sample_rate_hz`` (base64 only on the wire). + audio: bytes + format: str = "pcm16" + sample_rate_hz: int | None = None + is_speech: bool | None = None + video_frames: tuple[str, ...] = () + duration_ms: int | None = None + audio_end_ms: int | None = None + #: Model-neutral hints carried through from the wire (rms, vad, transcript hints ...). + hints: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + hints = data.pop("hints", None) + if isinstance(hints, Mapping): + # Hints are raw wire values; the typed fields went through + # ``build_append_audio``'s normalization, so they have to win. An + # unset typed field is simply absent here (``payload`` skips None), + # which is what lets a hint still carry it. Merging the other way + # round let a client's ``"is_speech": 0`` override the computed + # ``bool | None`` and silently miss the runner's silent-commit path. + merged: dict[str, object] = dict(hints) + merged.update(data) + data = merged + data["audio"] = base64.b64encode(self.audio).decode("ascii") + if not data.get("video_frames"): + data.pop("video_frames", None) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendText(DuplexCommand): + type: ClassVar[str] = "input.text.append" + text: str + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Commit(DuplexCommand): + type: ClassVar[str] = "input_audio_buffer.commit" + final: bool = True + #: ``None`` means "no explicit request": the runner decides on commit (auto-response + #: sessions answer on their own); ``True`` / ``False`` force it. + create_response: bool | None = None + is_speech: bool | None = None + #: Realtime conversation item created for this commit (wire correlation only). + realtime_item_id: str | None = None + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + if self.create_response is not None: + data["response_create"] = data.pop("create_response") + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateResponse(DuplexCommand): + type: ClassVar[str] = "response.create" + #: Raw Realtime ``response`` object (instructions, voice, modalities, ...). + options: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + options = data.pop("options", None) + if isinstance(options, Mapping): + data["response"] = dict(options) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearInput(DuplexCommand): + type: ClassVar[str] = "input_audio_buffer.clear" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelInput(DuplexCommand): + type: ClassVar[str] = "input.cancel" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelResponse(DuplexCommand): + type: ClassVar[str] = "response.cancel" + response_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BargeIn(DuplexCommand): + type: ClassVar[str] = "barge_in" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearOutputAudio(DuplexCommand): + type: ClassVar[str] = "output_audio_buffer.clear" + #: Explicit response to clear; ``None`` targets the active/last response. + response_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SignalTurn(DuplexCommand): + """Generic ``turn.signal`` (local turn transitions such as ``user_started``).""" + + type: ClassVar[str] = "turn.signal" + event: str + signal_payload: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + signal_payload = data.pop("signal_payload", None) + if isinstance(signal_payload, Mapping) and signal_payload: + data["payload"] = dict(signal_payload) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class UpdateSession(DuplexCommand): + """``session.update``: a Realtime ``session`` object patch.""" + + type: ClassVar[str] = "turn.signal" + patch: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "session.update" + data["payload"] = dict(data.pop("patch", {}) or {}) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AckPlayback(DuplexCommand): + type: ClassVar[str] = "playback.ack" + played_ms: int + committed_ms: int | None = None + response_id: str | None = None + item_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Heartbeat(DuplexCommand): + type: ClassVar[str] = "session.heartbeat" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateItem(DuplexCommand): + """``conversation.item.create`` (history injection or function-call output).""" + + type: ClassVar[str] = "turn.signal" + item: Mapping[str, object] + previous_item_id: str | None = None + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "conversation.item.create" + payload: dict[str, object] = {"item": dict(data.pop("item"))} + previous = data.pop("previous_item_id", None) + if previous is not None: + payload["previous_item_id"] = previous + data["payload"] = payload + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DeleteItem(DuplexCommand): + type: ClassVar[str] = "turn.signal" + item_id: str + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "conversation.item.delete" + data["payload"] = {"item_id": data.pop("item_id")} + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TruncateItem(DuplexCommand): + type: ClassVar[str] = "turn.signal" + item_id: str + audio_end_ms: int + content_index: int = 0 + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "conversation.item.truncate" + data["payload"] = { + "item_id": data.pop("item_id"), + "audio_end_ms": data.pop("audio_end_ms"), + "content_index": data.pop("content_index", 0), + } + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CloseSession(DuplexCommand): + """Graceful close requested through the command stream (``session.close``).""" + + type: ClassVar[str] = "session.close" + reason: str = "client_close" + + +#: Realtime client event types that map onto commands (wire vocabulary). +REALTIME_COMMAND_TYPES: frozenset[str] = frozenset( + { + "input_audio_buffer.append", + "input_audio_buffer.commit", + "input_audio_buffer.clear", + "output_audio_buffer.clear", + "response.create", + "response.cancel", + "conversation.item.create", + "conversation.item.delete", + "conversation.item.truncate", + "session.update", + "playback.ack", + "session.heartbeat", + "session.close", + "turn.signal", + "input.text.append", + "input.cancel", + "barge_in", + } +) + + +def command_from_realtime( + payload: Mapping[str, object], + *, + defaults: RealtimeInputDefaults | None = None, +) -> DuplexCommand: + """Translate one OpenAI Realtime client event into a command. + + Envelope concerns (``event_id`` acknowledgements, resume bookkeeping) are + handled by the transport; this function only validates and maps the + payload shape. ``defaults`` carries the audio format / sample rate / VAD + defaults the session declared, so every transport decodes an append the + same way. Raises :class:`DuplexCommandError` for malformed or unsupported + payloads. + """ + from vllm_omni.engine.duplex.realtime_commands import translate_realtime_command + + return translate_realtime_command(payload, defaults=defaults) + + +__all__ = [ + "REALTIME_COMMAND_TYPES", + "AckPlayback", + "AppendAudio", + "AppendText", + "BargeIn", + "CancelInput", + "CancelResponse", + "ClearInput", + "ClearOutputAudio", + "CloseSession", + "Commit", + "CreateItem", + "CreateResponse", + "DeleteItem", + "DuplexCommand", + "DuplexCommandError", + "Heartbeat", + "SignalTurn", + "TruncateItem", + "UpdateSession", + "command_from_realtime", +] diff --git a/vllm_omni/engine/duplex/config.py b/vllm_omni/engine/duplex/config.py new file mode 100644 index 000000000..aab0fb1c5 --- /dev/null +++ b/vllm_omni/engine/duplex/config.py @@ -0,0 +1,815 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session configuration, capabilities and value types of the duplex session model. + +These classes are the typed half of the public contract: ``DuplexOmni`` +accepts a ``DuplexSessionConfig`` (or the equivalent OpenAI Realtime +``session`` object via :meth:`DuplexSessionConfig.from_realtime`) and reports +``DuplexCapabilities``. They used to live in ``entrypoints/duplex/protocol.py``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, cast + + +class DuplexConfigError(ValueError): + """A Realtime session/response object was rejected while mapping it onto the duplex config.""" + + def __init__(self, message: str, *, code: str = "bad_event", param: str | None = None) -> None: + super().__init__(message) + self.code = code + self.param = param + + +class DuplexOverlapPolicy(str, Enum): + LISTEN_ONLY = "listen_only" + BARGE_IN_ON_SPEECH = "barge_in_on_speech" + + +class DuplexPlaybackCommitPolicy(str, Enum): + COMMIT_ALL_ON_DONE = "commit_all_on_done" + ACK_ONLY = "ack_only" + + +class DuplexSessionState(str, Enum): + OPEN = "open" + CLOSING = "closing" + CLOSED = "closed" + + +class DuplexTurnState(str, Enum): + IDLE = "idle" + USER_SPEAKING = "user_speaking" + USER_COMMITTED = "user_committed" + ASSISTANT_GENERATING = "assistant_generating" + ASSISTANT_PLAYING = "assistant_playing" + BARGE_IN = "barge_in" + + +class DuplexTurnEventType(str, Enum): + USER_STARTED = "user_started" + USER_COMMITTED = "user_committed" + ASSISTANT_STARTED = "assistant_started" + ASSISTANT_DONE = "assistant_done" + BARGE_IN = "barge_in" + PLAYBACK_ACK = "playback_ack" + TIMEOUT = "timeout" + CLOSE = "close" + + +@dataclass +class DuplexCapabilities: + """Runtime/model capabilities exposed by the duplex serving protocol. + + These are intentionally explicit so the serving layer does not assume all + duplex models support the same input append, rollback, or turn policy. + ``supports_core_kv_lease`` is reserved for scheduler-owned KV lifecycle; + model-owned decoder/TTS state must use ``supports_model_internal_state``. + ``supports_core_resumable_request`` means the scheduler can resume the + same request id across streaming updates, but it is not a KV lease by + itself. Realtime support means this endpoint can speak the native Realtime + event schema for the supported audio duplex paths while keeping model- or + scheduler-specific limits explicit in the capability payload. + """ + + supports_session_adapter: bool = True + supports_model_native_turn_policy: bool = False + supports_external_turn_signal: bool = True + supports_client_commit: bool = True + supports_barge_in: bool = True + supports_playback_ack: bool = True + supports_input_append: bool = False + supports_replace_latest_chunk: bool = True + supports_reencode_context: bool = True + supports_rollback_to_checkpoint: bool = False + supports_turn_commit_only: bool = True + supports_kv_lease: bool = False + supports_core_kv_lease: bool = False + supports_model_internal_state: bool = False + supports_stage_resumption: bool = False + supports_scheduler_native_append: bool = False + supports_core_resumable_request: bool = False + supports_stage_connector_handoff: bool = False + supports_independent_io_streams: bool = False + supports_realtime_endpoint: bool = False + supports_multi_session: bool = False + supports_multi_session_same_replica: bool = False + supports_session_lease: bool = False + supports_session_resume: bool = False + session_admission_mode: str = "serving_managed" + supports_audio_truncate: bool = False + #: Whether this model can serve ``/v1/chat/completions`` on a duplex + #: session. That needs one thing the protocol cannot assume: the model + #: honours ``DuplexSessionConfig.initial_user_text``, so a turn can be + #: seeded from text. A model-native model decides to speak from the audio + #: it hears, and silence is its signal *not* to take a turn, so without + #: seeding a text prompt has no representation and the turn never + #: completes. ``text_turn_priming_units`` is how many silence units such a + #: seeded turn must be given to generate on -- it still speaks per unit. + supports_chat_completions: bool = False + text_turn_priming_units: int = 0 + requires_model_runner_kv: bool = False + requires_native_stage_role: bool = False + adapter_patterns: list[str] = field(default_factory=lambda: ["chunk_group_append"]) + signal_sources: list[str] = field(default_factory=lambda: ["client_event", "server_policy", "model_native"]) + stage_handoff_transport: str | None = None + chunk_period_ms: int | None = 1000 + target_barge_in_latency_ms: int | None = 1000 + + def as_dict(self) -> dict[str, object]: + return { + "supports_session_adapter": self.supports_session_adapter, + "supports_model_native_turn_policy": self.supports_model_native_turn_policy, + "supports_external_turn_signal": self.supports_external_turn_signal, + "supports_client_commit": self.supports_client_commit, + "supports_barge_in": self.supports_barge_in, + "supports_playback_ack": self.supports_playback_ack, + "supports_input_append": self.supports_input_append, + "supports_replace_latest_chunk": self.supports_replace_latest_chunk, + "supports_reencode_context": self.supports_reencode_context, + "supports_rollback_to_checkpoint": self.supports_rollback_to_checkpoint, + "supports_turn_commit_only": self.supports_turn_commit_only, + "supports_kv_lease": self.supports_kv_lease, + "supports_core_kv_lease": self.supports_core_kv_lease, + "supports_model_internal_state": self.supports_model_internal_state, + "supports_stage_resumption": self.supports_stage_resumption, + "supports_scheduler_native_append": self.supports_scheduler_native_append, + "supports_core_resumable_request": self.supports_core_resumable_request, + "supports_stage_connector_handoff": self.supports_stage_connector_handoff, + "supports_independent_io_streams": self.supports_independent_io_streams, + "supports_realtime_endpoint": self.supports_realtime_endpoint, + "supports_multi_session": self.supports_multi_session, + "supports_multi_session_same_replica": self.supports_multi_session_same_replica, + "supports_session_lease": self.supports_session_lease, + "supports_session_resume": self.supports_session_resume, + "session_admission_mode": self.session_admission_mode, + "supports_audio_truncate": self.supports_audio_truncate, + "requires_model_runner_kv": self.requires_model_runner_kv, + "requires_native_stage_role": self.requires_native_stage_role, + # Every duplex model is model-native and appends audio chunks (§7 D4); + # kept on the wire as constants for clients that still read them. + "implementation_level": "model_native_duplex", + "adapter_patterns": self.adapter_patterns, + "input_modes": ["append_audio_chunk"], + "signal_sources": self.signal_sources, + "stage_handoff_transport": self.stage_handoff_transport, + "chunk_period_ms": self.chunk_period_ms, + "target_barge_in_latency_ms": self.target_barge_in_latency_ms, + } + + +@dataclass +class DuplexPlaybackCursor: + generated_ms: int = 0 + sent_ms: int = 0 + played_ms: int = 0 + committed_ms: int = 0 + + def acknowledge(self, played_ms: int, committed_ms: int | None = None) -> None: + self.played_ms = max(self.played_ms, max(0, int(played_ms))) + if committed_ms is None: + committed_ms = self.played_ms + self.committed_ms = max(self.committed_ms, max(0, int(committed_ms))) + + def truncate_committed(self, committed_ms: int) -> None: + self.committed_ms = max(0, min(max(self.sent_ms, self.generated_ms), int(committed_ms))) + + def as_dict(self) -> dict[str, int]: + return { + "generated_ms": self.generated_ms, + "sent_ms": self.sent_ms, + "played_ms": self.played_ms, + "committed_ms": self.committed_ms, + } + + def snapshot(self) -> DuplexPlaybackView: + return DuplexPlaybackView(**self.as_dict()) + + +@dataclass(frozen=True, slots=True) +class DuplexPlaybackView: + generated_ms: int = 0 + sent_ms: int = 0 + played_ms: int = 0 + committed_ms: int = 0 + + def as_dict(self) -> dict[str, int]: + return { + "generated_ms": self.generated_ms, + "sent_ms": self.sent_ms, + "played_ms": self.played_ms, + "committed_ms": self.committed_ms, + } + + +@dataclass +class DuplexAudioChunk: + data: str + format: str = "wav" + sample_rate_hz: int | None = None + + +@dataclass +class DuplexSessionConfig: + model: str | None = None + modalities: list[str] = field(default_factory=lambda: ["text", "audio"]) + instructions: str | None = None + #: Text the session treats as the user's opening turn. A caller that cannot + #: reach the model by speech (an HTTP request, a batch job) puts its prompt + #: here; ``supports_text_only_turn`` says whether the model can answer it. + initial_user_text: str | None = None + voice: str | None = None + ref_audio: str | None = None + response_format: str = "wav" + temperature: float | None = None + max_tokens: int | None = None + speed: float | None = None + use_tts_template: bool = True + idle_timeout_s: float = 300.0 + overlap_policy: str = DuplexOverlapPolicy.LISTEN_ONLY.value + overlap_short_ack_ms: int = 700 + overlap_barge_in_ms: int = 1200 + overlap_silence_rms: float = 0.003 + playback_commit_policy: str = DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value + extra_body: dict[str, object] = field(default_factory=dict) + + def as_dict(self) -> dict[str, object]: + return { + "model": self.model, + "modalities": list(self.modalities), + "instructions": self.instructions, + "initial_user_text": self.initial_user_text, + "voice": self.voice, + "ref_audio": self.ref_audio, + "response_format": self.response_format, + "temperature": self.temperature, + "max_tokens": self.max_tokens, + "speed": self.speed, + "use_tts_template": self.use_tts_template, + "idle_timeout_s": self.idle_timeout_s, + "overlap_policy": self.overlap_policy, + "overlap_short_ack_ms": self.overlap_short_ack_ms, + "overlap_barge_in_ms": self.overlap_barge_in_ms, + "overlap_silence_rms": self.overlap_silence_rms, + "playback_commit_policy": self.playback_commit_policy, + "extra_body": dict(self.extra_body), + } + + def normalized(self) -> DuplexSessionConfig: + """Apply the clamps / defaults ``from_event`` applies to a wire object (typed configs skip them).""" + return DuplexSessionConfig.from_event({"session": self.as_dict()}) + + @classmethod + def from_event(cls, event: dict[str, object]) -> DuplexSessionConfig: + payload = event.get("session") + if isinstance(payload, dict): + source = payload + else: + source = event + + config = cls() + if isinstance(source.get("model"), str): + config.model = source["model"] + if isinstance(source.get("instructions"), str): + config.instructions = source["instructions"] + if isinstance(source.get("initial_user_text"), str): + config.initial_user_text = source["initial_user_text"] + if isinstance(source.get("voice"), str): + config.voice = source["voice"] + if isinstance(source.get("ref_audio"), str): + config.ref_audio = source["ref_audio"] + if isinstance(source.get("response_format"), str): + config.response_format = source["response_format"] + if isinstance(source.get("use_tts_template"), bool): + config.use_tts_template = bool(source["use_tts_template"]) + if isinstance(source.get("temperature"), int | float): + config.temperature = float(source["temperature"]) + if isinstance(source.get("max_tokens"), int): + config.max_tokens = int(source["max_tokens"]) + if isinstance(source.get("speed"), int | float): + config.speed = float(source["speed"]) + if isinstance(source.get("idle_timeout_s"), int | float): + config.idle_timeout_s = float(source["idle_timeout_s"]) + if isinstance(source.get("overlap_policy"), str): + config.overlap_policy = cls._normalize_overlap_policy(source["overlap_policy"]) + if isinstance(source.get("overlap_short_ack_ms"), int | float): + config.overlap_short_ack_ms = max(0, int(source["overlap_short_ack_ms"])) + if isinstance(source.get("overlap_barge_in_ms"), int | float): + config.overlap_barge_in_ms = max(0, int(source["overlap_barge_in_ms"])) + if isinstance(source.get("overlap_silence_rms"), int | float): + config.overlap_silence_rms = max(0.0, float(source["overlap_silence_rms"])) + if isinstance(source.get("playback_commit_policy"), str): + config.playback_commit_policy = cls._normalize_playback_commit_policy(source["playback_commit_policy"]) + if isinstance(source.get("modalities"), list) and all(isinstance(x, str) for x in source["modalities"]): + config.modalities = list(source["modalities"]) + if isinstance(source.get("extra_body"), dict): + config.extra_body = dict(source["extra_body"]) + extra = config.extra_body + if isinstance(extra.get("overlap_policy"), str): + config.overlap_policy = cls._normalize_overlap_policy(extra["overlap_policy"]) + if isinstance(extra.get("overlap_short_ack_ms"), int | float): + config.overlap_short_ack_ms = max(0, int(extra["overlap_short_ack_ms"])) + if isinstance(extra.get("overlap_barge_in_ms"), int | float): + config.overlap_barge_in_ms = max(0, int(extra["overlap_barge_in_ms"])) + if isinstance(extra.get("overlap_silence_rms"), int | float): + config.overlap_silence_rms = max(0.0, float(extra["overlap_silence_rms"])) + if isinstance(extra.get("playback_commit_policy"), str): + config.playback_commit_policy = cls._normalize_playback_commit_policy(extra["playback_commit_policy"]) + return config + + @classmethod + def from_realtime( + cls, + session_payload: Mapping[str, object], + *, + model: str | None = None, + ) -> DuplexSessionConfig: + """Build a config from an OpenAI Realtime ``session`` object (session.update / open_session). + + Applies the wire defaults, validates audio formats and turn detection + (raising :class:`DuplexConfigError` with the same codes the Realtime + translator used), fills ``turn_detection`` defaults / derives + ``overlap_policy``, and stores the Realtime-only fields under the + ``realtime_*`` keys of ``extra_body`` exactly like the old + ``_session_create_from_realtime``. + """ + from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + duplex_response_format, + input_audio_transcription_config, + json_safe_realtime_payload, + realtime_max_output_tokens, + realtime_overlap_fields, + validate_realtime_session_audio_formats, + ) + from vllm_omni.engine.duplex.turn_detection import ( + normalize_turn_detection_session_payload, + validate_realtime_turn_detection, + ) + + payload: dict[str, object] = dict(session_payload) + format_error = validate_realtime_session_audio_formats(payload) + if format_error is not None: + raise DuplexConfigError(format_error, code="unsupported_audio_format") + turn_detection_error = validate_realtime_turn_detection(payload) + if turn_detection_error is not None: + raise DuplexConfigError(turn_detection_error, code="unsupported_turn_detection", param="turn_detection") + normalize_turn_detection_session_payload(payload) + defaults = RealtimeInputDefaults().with_session_payload(payload) + payload.update(realtime_overlap_fields(payload)) + + audio_config = payload.get("audio") + audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None + audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None + extra_body_payload = payload.get("extra_body") + extra_body: dict[str, object] = dict(extra_body_payload) if isinstance(extra_body_payload, dict) else {} + extra_body["realtime_session_payload"] = json_safe_realtime_payload(payload) + if isinstance(payload.get("tools"), list): + extra_body["realtime_tools"] = payload["tools"] + if isinstance(payload.get("tool_choice"), str | dict): + extra_body["realtime_tool_choice"] = payload["tool_choice"] + metadata = payload.get("metadata") + if isinstance(metadata, dict): + extra_body["realtime_metadata"] = dict(metadata) + include = payload.get("include") + if isinstance(include, list): + extra_body["realtime_include"] = list(include) + prompt = payload.get("prompt") + if isinstance(prompt, dict): + extra_body["realtime_prompt"] = dict(prompt) + transcription = input_audio_transcription_config(payload) + if isinstance(transcription, dict): + extra_body["realtime_input_audio_transcription"] = dict(transcription) + noise_reduction = payload.get("input_audio_noise_reduction") + if isinstance(noise_reduction, dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(noise_reduction) + audio_noise_reduction = audio_input.get("noise_reduction") if isinstance(audio_input, dict) else None + if isinstance(audio_noise_reduction, dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(audio_noise_reduction) + if isinstance(audio_config, dict): + extra_body["realtime_audio"] = dict(audio_config) + if isinstance(payload.get("tracing"), str | dict): + extra_body["realtime_tracing"] = payload["tracing"] + if isinstance(payload.get("turn_detection"), dict): + extra_body["realtime_turn_detection"] = dict(cast("Mapping[str, object]", payload["turn_detection"])) + elif "turn_detection" in payload and payload.get("turn_detection") is None: + extra_body["realtime_turn_detection"] = None + extra_body.setdefault("realtime_output_audio_format", defaults.output_audio_format) + voice = payload.get("voice") + if not isinstance(voice, str) and isinstance(audio_output, dict): + voice = audio_output.get("voice") + speed = payload.get("speed") + if not isinstance(speed, int | float) and isinstance(audio_output, dict): + speed = audio_output.get("speed") + resolved_model = payload.get("model") if isinstance(payload.get("model"), str) else model + session: dict[str, object] = { + "model": resolved_model, + "modalities": payload.get("modalities") or payload.get("output_modalities") or ["text", "audio"], + "instructions": payload.get("instructions"), + # Not a Realtime field, but a session one: a client that cannot + # reach the model by speech seeds its opening turn here. Without + # this a websocket client could only use the older + # ``extra_body.duplex_initial_user_text`` spelling. + "initial_user_text": payload.get("initial_user_text"), + "voice": voice, + "ref_audio": payload.get("ref_audio"), + "response_format": duplex_response_format(defaults.output_audio_format), + "temperature": payload.get("temperature"), + "max_tokens": realtime_max_output_tokens( + payload.get("max_response_output_tokens") + or payload.get("max_output_tokens") + or payload.get("max_tokens") + ), + "speed": speed, + "idle_timeout_s": payload.get("idle_timeout_s") or 300.0, + **realtime_overlap_fields(payload), + "extra_body": extra_body, + } + return cls.from_event({"session": session}) + + def apply_realtime_update( + self, + patch: Mapping[str, object], + *, + session_id: str | None = None, + audio_started: bool = False, + ) -> None: + """Apply a Realtime ``session.update`` patch in place. + + Raises :class:`DuplexConfigError` (``code`` in ``model_update_unsupported``, + ``voice_update_after_audio_unsupported``, ``ref_audio_update_unsupported``) + when the patch changes something a live session cannot change. + ``audio_started`` is ``playback.generated_ms > 0 or playback.sent_ms > 0``. + """ + from vllm_omni.engine.duplex.realtime_commands import ( + REALTIME_OUTPUT_AUDIO_FORMATS, + duplex_response_format, + input_audio_transcription_config, + json_safe_realtime_payload, + parse_realtime_audio_format, + realtime_max_output_tokens, + ) + + # Any: the patch is read with ``isinstance(payload.get(...))`` guards and consumed + # through ``payload[...]``; typing the values ``object`` would need a cast per read. + payload: dict[str, Any] = dict(patch) + model = payload.get("model") + audio_config = payload.get("audio") + audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None + audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None + voice = payload.get("voice") + if not isinstance(voice, str) and isinstance(audio_output, dict): + voice = audio_output.get("voice") + if isinstance(model, str) and self.model is not None and model != self.model: + raise DuplexConfigError( + "session.update cannot change model for an open realtime duplex session", + code="model_update_unsupported", + ) + if isinstance(model, str) and self.model is None: + self.model = model + if isinstance(voice, str) and audio_started: + raise DuplexConfigError( + "session.update cannot change voice after audio output has started", + code="voice_update_after_audio_unsupported", + ) + if isinstance(payload.get("ref_audio"), str): + raise DuplexConfigError( + "session.update cannot change ref_audio after the session is open", code="ref_audio_update_unsupported" + ) + if isinstance(payload.get("instructions"), str): + self.instructions = str(payload["instructions"]) + elif "instructions" in payload and payload.get("instructions") is None: + self.instructions = None + if isinstance(voice, str): + self.voice = str(voice) + elif "voice" in payload and payload.get("voice") is None: + self.voice = None + response_format = payload.get("output_audio_format") or payload.get("response_format") + if response_format is None and isinstance(audio_output, dict): + response_format = audio_output.get("format") + response_format, _ = parse_realtime_audio_format(response_format) + if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + self.response_format = duplex_response_format(response_format) + if isinstance(payload.get("temperature"), int | float): + self.temperature = float(payload["temperature"]) + speed = payload.get("speed") + if not isinstance(speed, int | float) and isinstance(audio_output, dict): + speed = audio_output.get("speed") + if isinstance(speed, int | float): + self.speed = float(speed) + max_tokens = ( + payload.get("max_response_output_tokens") + if "max_response_output_tokens" in payload + else payload.get("max_output_tokens") + if "max_output_tokens" in payload + else payload.get("max_tokens") + ) + if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: + self.max_tokens = realtime_max_output_tokens(max_tokens) + if isinstance(payload.get("overlap_policy"), str): + self.overlap_policy = self._normalize_overlap_policy(str(payload["overlap_policy"])) + if isinstance(payload.get("overlap_short_ack_ms"), int | float): + self.overlap_short_ack_ms = max(0, int(payload["overlap_short_ack_ms"])) + if isinstance(payload.get("overlap_barge_in_ms"), int | float): + self.overlap_barge_in_ms = max(0, int(payload["overlap_barge_in_ms"])) + if isinstance(payload.get("overlap_silence_rms"), int | float): + self.overlap_silence_rms = max(0.0, float(payload["overlap_silence_rms"])) + if isinstance(payload.get("playback_commit_policy"), str): + self.playback_commit_policy = self._normalize_playback_commit_policy(str(payload["playback_commit_policy"])) + modalities = payload.get("modalities") or payload.get("output_modalities") + if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): + self.modalities = list(modalities) + if isinstance(payload.get("extra_body"), dict): + self.extra_body.update(payload["extra_body"]) + extra = payload["extra_body"] + if isinstance(extra.get("overlap_policy"), str): + self.overlap_policy = self._normalize_overlap_policy(str(extra["overlap_policy"])) + if isinstance(extra.get("playback_commit_policy"), str): + self.playback_commit_policy = self._normalize_playback_commit_policy( + str(extra["playback_commit_policy"]) + ) + if isinstance(payload.get("tools"), list): + self.extra_body["realtime_tools"] = payload["tools"] + elif "tools" in payload and payload.get("tools") is None: + self.extra_body.pop("realtime_tools", None) + if isinstance(payload.get("tool_choice"), str | dict): + self.extra_body["realtime_tool_choice"] = payload["tool_choice"] + elif "tool_choice" in payload and payload.get("tool_choice") is None: + self.extra_body.pop("realtime_tool_choice", None) + if isinstance(payload.get("metadata"), dict): + self.extra_body["realtime_metadata"] = dict(payload["metadata"]) + elif "metadata" in payload and payload.get("metadata") is None: + self.extra_body.pop("realtime_metadata", None) + if isinstance(payload.get("include"), list): + self.extra_body["realtime_include"] = list(payload["include"]) + elif "include" in payload and payload.get("include") is None: + self.extra_body.pop("realtime_include", None) + if isinstance(payload.get("prompt"), dict): + self.extra_body["realtime_prompt"] = dict(payload["prompt"]) + elif "prompt" in payload and payload.get("prompt") is None: + self.extra_body.pop("realtime_prompt", None) + transcription = input_audio_transcription_config(payload) + if isinstance(transcription, dict): + self.extra_body["realtime_input_audio_transcription"] = dict(transcription) + elif "input_audio_transcription" in payload and payload.get("input_audio_transcription") is None: + self.extra_body.pop("realtime_input_audio_transcription", None) + if isinstance(payload.get("input_audio_noise_reduction"), dict): + self.extra_body["realtime_input_audio_noise_reduction"] = dict(payload["input_audio_noise_reduction"]) + elif "input_audio_noise_reduction" in payload and payload.get("input_audio_noise_reduction") is None: + self.extra_body.pop("realtime_input_audio_noise_reduction", None) + if isinstance(audio_input, dict) and isinstance(audio_input.get("noise_reduction"), dict): + self.extra_body["realtime_input_audio_noise_reduction"] = dict(audio_input["noise_reduction"]) + elif isinstance(audio_input, dict) and audio_input.get("noise_reduction") is None: + self.extra_body.pop("realtime_input_audio_noise_reduction", None) + if isinstance(payload.get("audio"), dict): + self.extra_body["realtime_audio"] = dict(payload["audio"]) + elif "audio" in payload and payload.get("audio") is None: + self.extra_body.pop("realtime_audio", None) + if isinstance(payload.get("tracing"), str | dict): + self.extra_body["realtime_tracing"] = payload["tracing"] + elif "tracing" in payload and payload.get("tracing") is None: + self.extra_body.pop("realtime_tracing", None) + if isinstance(payload.get("turn_detection"), dict): + self.extra_body["realtime_turn_detection"] = dict(payload["turn_detection"]) + elif "turn_detection" in payload and payload.get("turn_detection") is None: + self.extra_body["realtime_turn_detection"] = None + self.extra_body["realtime_session_payload"] = json_safe_realtime_payload(payload) + return None + + @staticmethod + def _normalize_overlap_policy(value: str) -> str: + normalized = value.strip().lower() + if normalized in {policy.value for policy in DuplexOverlapPolicy}: + return normalized + return DuplexOverlapPolicy.LISTEN_ONLY.value + + @staticmethod + def _normalize_playback_commit_policy(value: str) -> str: + normalized = value.strip().lower() + if normalized in {policy.value for policy in DuplexPlaybackCommitPolicy}: + return normalized + return DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value + + +@dataclass(frozen=True) +class ResponseCreateOptions: + instructions: str | None = None + voice: str | None = None + response_format: str | None = None + temperature: float | None = None + max_tokens: int | None = None + speed: float | None = None + modalities: tuple[str, ...] | None = None + extra_body: Mapping[str, object] = field(default_factory=dict) + + def apply_to(self, config: DuplexSessionConfig) -> None: + for field_name in ("instructions", "voice", "response_format", "temperature", "max_tokens", "speed"): + value = getattr(self, field_name) + if value is not None: + setattr(config, field_name, value) + if self.modalities is not None: + config.modalities = list(self.modalities) + config.extra_body.update(self.extra_body) + + @classmethod + def from_realtime( + cls, + response_payload: Mapping[str, object], + *, + private_runtime_config_keys: frozenset[str] = frozenset(), + ) -> ResponseCreateOptions: + """Parse an OpenAI Realtime ``response`` object into response-scoped options. + + Raises :class:`DuplexConfigError` (``code="unsupported_native_response_options"``) + for options a model-native duplex session cannot apply per response. + Private runtime keys in ``extra_body`` are dropped. + """ + from vllm_omni.engine.duplex.realtime_commands import ( + REALTIME_OUTPUT_AUDIO_FORMATS, + duplex_response_format, + parse_realtime_audio_format, + realtime_max_output_tokens, + ) + + payload: dict[str, object] = dict(response_payload) + audio_config = payload.get("audio") + audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None + nested_voice = audio_output.get("voice") if isinstance(audio_output, dict) else None + unsupported = ( + payload.get("instructions") is not None + or payload.get("voice") is not None + or nested_voice is not None + or payload.get("temperature") is not None + or any( + payload.get(field_name) is not None + for field_name in ("max_response_output_tokens", "max_output_tokens", "max_tokens") + ) + or payload.get("tools") is not None + or payload.get("tool_choice") is not None + ) + if unsupported: + raise DuplexConfigError( + "response.create options are not supported by a model-native duplex session", + code="unsupported_native_response_options", + ) + instructions = str(payload["instructions"]) if isinstance(payload.get("instructions"), str) else None + voice = payload.get("voice") + if not isinstance(voice, str) and isinstance(audio_output, dict): + voice = audio_output.get("voice") + voice = str(voice) if isinstance(voice, str) else None + response_format = payload.get("output_audio_format") or payload.get("response_format") + if response_format is None and isinstance(audio_output, dict): + response_format = audio_output.get("format") + response_format, _ = parse_realtime_audio_format(response_format) + if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + response_format = duplex_response_format(response_format) + else: + response_format = None + temperature = ( + float(cast("int | float", payload["temperature"])) + if isinstance(payload.get("temperature"), int | float) + else None + ) + speed = payload.get("speed") + if not isinstance(speed, int | float) and isinstance(audio_output, dict): + speed = audio_output.get("speed") + speed = float(speed) if isinstance(speed, int | float) else None + max_tokens = ( + payload.get("max_response_output_tokens") + if "max_response_output_tokens" in payload + else payload.get("max_output_tokens") + if "max_output_tokens" in payload + else payload.get("max_tokens") + ) + if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: + max_tokens = realtime_max_output_tokens(max_tokens) + else: + max_tokens = None + modalities = payload.get("modalities") or payload.get("output_modalities") + if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): + modalities = tuple(modalities) + else: + modalities = None + response_extra: dict[str, object] = {} + conversation = payload.get("conversation") + if isinstance(conversation, str): + response_extra["realtime_response_conversation"] = conversation + metadata = payload.get("metadata") + if isinstance(metadata, dict): + response_extra["realtime_response_metadata"] = dict(metadata) + prompt = payload.get("prompt") + if isinstance(prompt, dict): + response_extra["realtime_response_prompt"] = dict(prompt) + if isinstance(payload.get("tools"), list): + response_extra["realtime_response_tools"] = payload["tools"] + if isinstance(payload.get("tool_choice"), str | dict): + response_extra["realtime_response_tool_choice"] = payload["tool_choice"] + extra_body = payload.get("extra_body") + if isinstance(extra_body, dict): + response_extra.update( + (key, value) for key, value in extra_body.items() if key not in private_runtime_config_keys + ) + return cls( + instructions=instructions, + voice=voice, + response_format=response_format, + temperature=temperature, + max_tokens=max_tokens, + speed=speed, + modalities=modalities, + extra_body=response_extra, + ) + + +@dataclass +class DuplexCommittedInput: + message: dict[str, object] + turn_id: int + epoch: int + input_commit_seq: int + + +@dataclass +class DuplexAssistantAudioTextMark: + text_chars: int + audio_end_ms: int + + +def realtime_item_to_history_message(item: object) -> dict[str, object] | None: + """Convert a Realtime conversation item into a chat history message (None when empty).""" + if not isinstance(item, dict): + return None + role = item.get("role") + if role not in {"system", "user", "assistant"}: + return None + content = item.get("content") + if isinstance(content, str): + text = content.strip() + return {"role": role, "content": text} if text else None + if not isinstance(content, list): + return None + text_chunks: list[str] = [] + audio_chunks: list[dict[str, object]] = [] + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in {"input_text", "text", "output_text"} and isinstance(part.get("text"), str): + text_chunks.append(str(part["text"])) + elif part_type in {"input_audio", "audio"}: + audio = part.get("audio") or part.get("data") + fmt = part.get("format") if isinstance(part.get("format"), str) else "wav" + if isinstance(audio, str) and audio: + audio_chunks.append({"type": "audio_url", "audio_url": {"url": f"data:audio/{fmt};base64,{audio}"}}) + elif part_type in {"audio_transcript", "transcript"} and isinstance(part.get("text"), str): + text_chunks.append(str(part["text"])) + text = "".join(text_chunks).strip() + if audio_chunks: + content_items: list[dict[str, object]] = [] + if text: + content_items.append({"type": "text", "text": text}) + content_items.extend(audio_chunks) + return {"role": role, "content": content_items} + if text: + return {"role": role, "content": text} + return None + + +def realtime_max_output_tokens(value: object) -> int | None: + """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" + from vllm_omni.engine.duplex.realtime_commands import realtime_max_output_tokens as _impl + + return _impl(value) + + +def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: + """Return the ``input_audio_transcription`` object of a Realtime session payload.""" + from vllm_omni.engine.duplex.realtime_commands import input_audio_transcription_config as _impl + + return _impl(session_payload) + + +__all__ = [ + "DuplexAssistantAudioTextMark", + "DuplexAudioChunk", + "DuplexCapabilities", + "DuplexCommittedInput", + "DuplexConfigError", + "DuplexOverlapPolicy", + "DuplexPlaybackCommitPolicy", + "DuplexPlaybackCursor", + "DuplexPlaybackView", + "DuplexSessionConfig", + "DuplexSessionState", + "DuplexTurnEventType", + "DuplexTurnState", + "ResponseCreateOptions", + "input_audio_transcription_config", + "realtime_item_to_history_message", + "realtime_max_output_tokens", +] diff --git a/vllm_omni/engine/duplex/contracts.py b/vllm_omni/engine/duplex/contracts.py index 6af4fdea5..ed14f5fb8 100644 --- a/vllm_omni/engine/duplex/contracts.py +++ b/vllm_omni/engine/duplex/contracts.py @@ -1,102 +1,50 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Model-neutral contracts for the experimental duplex engine plugin. +"""Model-neutral value types shared by the duplex engine components. -This module contains only immutable data transfer objects and narrow protocols. -Duplex control algorithms, session implementations, model policy, and Realtime -serving remain in sibling experimental modules. +Immutable DTOs plus the ``DuplexStagePort`` base class that ``DuplexOrchestrator`` +implements for the session manager/runner. """ from __future__ import annotations import base64 -from collections.abc import Iterable, Mapping +from abc import ABC, abstractmethod +from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import Any, Protocol - -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.messages import EngineQueueMessage -class SessionMode(str, Enum): - TURN = "turn" - DUPLEX = "duplex" +@dataclass(frozen=True, slots=True) +class DuplexFence: + """Engine-internal session identity used for stage request ids and stale filtering.""" - -class DuplexInputMode(str, Enum): - APPEND_TOKENS = "append_tokens" - APPEND_AUDIO_CHUNK = "append_audio_chunk" - REPLACE_LATEST_CHUNK = "replace_latest_chunk" - REENCODE_CONTEXT = "reencode_context" - ROLLBACK_TO_CHECKPOINT = "rollback_to_checkpoint" - TURN_COMMIT_ONLY = "turn_commit_only" + session_id: str + epoch: int = 0 + turn_id: int = 0 class DuplexOutputAction(str, Enum): DIRECT_RESPONSE = "direct_response" -@dataclass -class DuplexRuntimeCapabilities: - input_modes: set[DuplexInputMode] = field(default_factory=lambda: {DuplexInputMode.TURN_COMMIT_ONLY}) - implementation_level: str = "serving_session_adapter" - - @dataclass(frozen=True) class DuplexAppendPlan: - prompt: dict[str, Any] + prompt: dict[str, object] @dataclass(frozen=True) class DuplexOutputDecision: action: DuplexOutputAction - metadata: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, object] = field(default_factory=dict) final_output_type: str = "text" def __post_init__(self) -> None: object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) -class DuplexRuntimeExtension(Protocol): - """Pure model policy invoked by the experimental duplex control plane.""" - - def configure_sampling_params( - self, - *, - runtime_config: dict[str, Any], - defaults: tuple[object, ...], - ) -> tuple[object, ...]: ... - - def plan_append( - self, - *, - request_id: str, - fence: DuplexFence, - session_config: dict[str, Any], - runtime_config: dict[str, Any], - seq: int, - turn_seq: int, - mode: DuplexInputMode, - payload: object, - final: bool, - sampling_params: object, - ) -> DuplexAppendPlan: ... - - def decide_output( - self, - *, - stage_id: int, - final_stage_id: int, - segment_finished: bool, - segment_token_ids: tuple[int, ...], - segment_output_metadata: dict[str, Any], - output: object, - ) -> DuplexOutputDecision | None: ... - - @dataclass(frozen=True) class DuplexRequestIdentity: session_id: str @@ -112,8 +60,8 @@ class DuplexStageRequestContext: final_stage_id: int config_generation: int sampling_params: tuple[object, ...] - session_config: Mapping[str, Any] = field(default_factory=dict) - runtime_config: Mapping[str, Any] = field(default_factory=dict) + session_config: Mapping[str, object] = field(default_factory=dict) + runtime_config: Mapping[str, object] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__(self, "sampling_params", tuple(self.sampling_params)) @@ -128,7 +76,7 @@ class DuplexStageRequestContext: @dataclass(frozen=True) class DuplexStageSubmission: context: DuplexStageRequestContext - prompt: Mapping[str, Any] + prompt: Mapping[str, object] already_submitted: bool def __post_init__(self) -> None: @@ -148,7 +96,7 @@ class DuplexOutputContext: final_stage_id: int segment_finished: bool segment_token_ids: tuple[int, ...] = () - segment_output_metadata: Mapping[str, Any] = field(default_factory=dict) + segment_output_metadata: Mapping[str, object] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__(self, "segment_token_ids", tuple(self.segment_token_ids)) @@ -159,53 +107,27 @@ class DuplexOutputContext: ) -class DuplexStagePort(Protocol): +class DuplexStagePort(ABC): + """Narrow stage-management surface the session runner/manager use (implemented by DuplexOrchestrator).""" + @property + @abstractmethod def stage_count(self) -> int: ... + @abstractmethod def sampling_defaults(self) -> tuple[object, ...]: ... + @abstractmethod def ensure_request(self, context: DuplexStageRequestContext) -> None: ... + @abstractmethod async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: ... + @abstractmethod async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: ... - -class DuplexControlPlanePort(Protocol): - @property - def sessions(self) -> object: ... - - def accepts(self, message: object) -> bool: ... - - def dispatch(self, message: object) -> None: ... - - async def shutdown(self) -> None: ... - - def close_sessions_for_request_ids(self, request_ids: list[str]) -> dict[str, list[str]]: ... - - def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: ... - - def session_for_identity(self, identity: DuplexRequestIdentity | None) -> object | None: ... - - def decide_output( - self, - stage_id: int, - output: object, - context: DuplexOutputContext | None, - ) -> DuplexOutputDecision | None: ... - - -class CorrelatedRpcTransport(Protocol): - def execute( - self, - key: tuple[str, str], - message: EngineQueueMessage, - *, - timeout: float | None, - timeout_message: str, - block_on_submit: bool = False, - ) -> EngineQueueMessage: ... + @abstractmethod + async def abort_requests(self, request_ids: list[str]) -> None: ... def duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]: @@ -232,20 +154,19 @@ def duplex_resource_request_id(fence: DuplexFence, role: str) -> str: ): raise ValueError(f"invalid duplex resource role: {role!r}") encoded_session_id = base64.urlsafe_b64encode(fence.session_id.encode("utf-8")).decode("ascii").rstrip("=") - return f"duplex-s.{encoded_session_id}.i.{fence.incarnation}.e.{fence.epoch}.r.{role}" + return f"duplex-s.{encoded_session_id}.e.{fence.epoch}.r.{role}" def duplex_resource_request_belongs_to_session(request_id: str, session_id: str) -> bool: """Return whether a current-format resource request belongs to a session.""" parts = request_id.split(".") - if len(parts) != 8 or parts[0] != "duplex-s" or parts[2] != "i" or parts[4] != "e" or parts[6] != "r": + if len(parts) != 6 or parts[0] != "duplex-s" or parts[2] != "e" or parts[4] != "r": return False try: int(parts[3]) - int(parts[5]) except ValueError: return False - role = parts[7] + role = parts[5] if not role or any( character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role ): @@ -255,21 +176,16 @@ def duplex_resource_request_belongs_to_session(request_id: str, session_id: str) __all__ = [ - "CorrelatedRpcTransport", + "DuplexFence", "DuplexAppendPlan", - "DuplexControlPlanePort", - "DuplexInputMode", "DuplexOutputAction", "DuplexOutputContext", "DuplexOutputDecision", "DuplexRequestIdentity", - "DuplexRuntimeCapabilities", - "DuplexRuntimeExtension", "DuplexStagePort", "DuplexStageRequestContext", "DuplexStageSubmission", "DuplexStageSubmissionResult", - "SessionMode", "duplex_data_plane_request_info", "duplex_resource_request_belongs_to_session", "duplex_resource_request_id", diff --git a/vllm_omni/engine/duplex/control_client.py b/vllm_omni/engine/duplex/control_client.py deleted file mode 100644 index 05ec46e49..000000000 --- a/vllm_omni/engine/duplex/control_client.py +++ /dev/null @@ -1,228 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import uuid -from collections.abc import Callable - -from vllm_omni.engine.duplex.contracts import CorrelatedRpcTransport -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexControlResultMessage, - DuplexFence, - OpenDuplexSessionMessage, - ResumeDuplexSessionMessage, - SignalDuplexTurnMessage, - TouchDuplexSessionMessage, -) - -DuplexControlMessage = ( - OpenDuplexSessionMessage - | AppendDuplexInputMessage - | SignalDuplexTurnMessage - | CloseDuplexSessionMessage - | TouchDuplexSessionMessage - | ResumeDuplexSessionMessage -) - - -class DuplexControlRequestError(RuntimeError): - """Typed client-side failure returned by the duplex control plane.""" - - def __init__(self, result: dict[str, object]) -> None: - error = result.get("error") - error_data = error if isinstance(error, dict) else {} - self.result = result - self.code = str(error_data.get("code") or "internal_error") - self.retryable = bool(error_data.get("retryable", False)) - accepted_fence = result.get("accepted_fence") - self.accepted_fence = accepted_fence if isinstance(accepted_fence, DuplexFence) else None - lease_generation = result.get("lease_generation") - self.lease_generation = lease_generation if isinstance(lease_generation, int) else None - super().__init__(f"duplex {result.get('operation')} failed: {result}") - - -class DuplexControlClient: - """Builds duplex commands on top of an engine-owned RPC transport.""" - - def __init__( - self, - transport: CorrelatedRpcTransport, - *, - control_id_factory: Callable[[], str] | None = None, - ) -> None: - self._transport = transport - self._control_id_factory = control_id_factory or (lambda: uuid.uuid4().hex) - - def execute(self, message: DuplexControlMessage, *, timeout: float | None) -> dict[str, object]: - result_message = self._transport.execute( - ("duplex", message.control_id), - message, - timeout=timeout, - timeout_message=f"duplex control timed out after {timeout} seconds", - ) - if not isinstance(result_message, DuplexControlResultMessage): - raise RuntimeError(f"unexpected duplex control result type: {type(result_message).__name__}") - if result_message.fence != message.fence: - raise RuntimeError( - f"duplex control fence mismatch: expected {message.fence!r}, got {result_message.fence!r}" - ) - result = { - "fence": result_message.fence, - "operation": result_message.operation, - "session_id": result_message.session_id, - "ok": result_message.ok, - "stage_results": list(result_message.stage_results), - "unsupported_count": result_message.unsupported_count, - "error_count": result_message.error_count, - "accepted_fence": result_message.accepted_fence, - "lease_generation": result_message.lease_generation, - "error": ( - { - "code": result_message.error.code, - "message": result_message.error.message, - "retryable": result_message.error.retryable, - } - if result_message.error is not None - else None - ), - } - if not result_message.ok: - raise DuplexControlRequestError(result) - return result - - def open( - self, - session_id: str, - *, - session_mode: str, - capabilities: dict[str, object] | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - OpenDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - session_mode=session_mode, - capabilities=dict(capabilities or {}), - session_config=dict(session_config or {}), - runtime_config=dict(runtime_config or {}), - ), - timeout=timeout, - ) - - def append( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None, - final: bool, - expected_epoch: int | None, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - if expected_epoch is not None and expected_epoch != fence.epoch: - raise ValueError("expected_epoch must match fence.epoch") - return self.execute( - AppendDuplexInputMessage( - control_id=self._control_id_factory(), - operation_id=operation_id, - fence=fence, - session_id=session_id, - expected_epoch=expected_epoch, - mode=mode, - payload=payload, - final=final, - ), - timeout=timeout, - ) - - def signal( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - SignalDuplexTurnMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - event=event, - next_fence=next_fence, - session_config=dict(session_config) if session_config is not None else None, - runtime_config=dict(runtime_config) if runtime_config is not None else None, - ), - timeout=timeout, - ) - - def close( - self, - session_id: str, - *, - reason: str, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - CloseDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - reason=reason, - ), - timeout=timeout, - ) - - def touch( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - TouchDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - activity=activity.value, - ), - timeout=timeout, - ) - - def resume( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - ResumeDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - expected_lease_generation=expected_lease_generation, - ), - timeout=timeout, - ) - - -__all__ = ["CorrelatedRpcTransport", "DuplexControlClient", "DuplexControlMessage"] diff --git a/vllm_omni/engine/duplex/control_plane.py b/vllm_omni/engine/duplex/control_plane.py deleted file mode 100644 index b425542b8..000000000 --- a/vllm_omni/engine/duplex/control_plane.py +++ /dev/null @@ -1,1063 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Orchestrator-side control plane for experimental duplex sessions. - -``Orchestrator`` creates this component only when duplex control is enabled -and supplies a narrow :class:`DuplexStagePort` implementation. Queue messages -are routed through :meth:`DuplexControlPlane.handle`; stage outputs are passed -to :meth:`DuplexControlPlane.decide_output`; request cleanup is reported back -through the control plane so session leases and stage bindings stay coherent. - -The module owns duplex session/control algorithms. It deliberately does not -own stage pools, request queues, or OpenAI Realtime protocol state. -""" - -from __future__ import annotations - -import asyncio -from collections.abc import Callable, Iterable -from dataclasses import dataclass -from typing import Any, Protocol - -from vllm.logger import init_logger - -from vllm_omni.engine.duplex.contracts import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexOutputContext, - DuplexOutputDecision, - DuplexRequestIdentity, - DuplexRuntimeCapabilities, - DuplexRuntimeExtension, - DuplexStagePort, - DuplexStageRequestContext, - DuplexStageSubmission, - DuplexStageSubmissionResult, - SessionMode, - duplex_resource_request_id, -) -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity, DuplexLeaseConfig -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexControlError, - DuplexControlResultMessage, - DuplexFence, - DuplexSessionLifecycleMessage, - OpenDuplexSessionMessage, - ResumeDuplexSessionMessage, - SignalDuplexTurnMessage, - TouchDuplexSessionMessage, -) -from vllm_omni.engine.duplex.session import ( - DuplexAppendReservation, - DuplexFenceMismatchError, - DuplexSessionExpiry, - DuplexSessionRuntimeManager, - DuplexSessionRuntimeState, -) - -logger = init_logger(__name__) - - -@dataclass(frozen=True) -class _PendingControlCleanup: - kind: str - session_id: str - fence: DuplexFence - submitted_request_ids: tuple[str, ...] - reserved_request_ids: tuple[str, ...] = () - - -@dataclass(frozen=True) -class _PendingSubmissionCleanup: - session_id: str - request_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class _PendingRequestCleanup: - session_id: str - fence: DuplexFence - lease_generation: int - request_ids: tuple[str, ...] - abort: bool - - -class DuplexResultSink(Protocol): - async def put(self, message: DuplexControlResultMessage) -> None: ... - - -class DuplexLifecycleSink(Protocol): - async def put(self, message: DuplexSessionLifecycleMessage) -> None: ... - - -class DuplexControlPlane: - _MESSAGE_TYPES = ( - OpenDuplexSessionMessage, - AppendDuplexInputMessage, - SignalDuplexTurnMessage, - CloseDuplexSessionMessage, - TouchDuplexSessionMessage, - ResumeDuplexSessionMessage, - ) - - def __init__( - self, - *, - extension: DuplexRuntimeExtension | None, - stage_port: DuplexStagePort, - result_sink: DuplexResultSink, - lifecycle_sink: DuplexLifecycleSink | None = None, - lease_config: DuplexLeaseConfig | None = None, - clock: Callable[[], float] | None = None, - max_sessions: int | None = None, - completed_append_limit: int = 256, - ) -> None: - self._extension = extension - self._stage_port = stage_port - self._result_sink = result_sink - self._lifecycle_sink = lifecycle_sink - self._lease_config = lease_config or DuplexLeaseConfig() - self._sessions = DuplexSessionRuntimeManager( - clock=clock, - max_sessions=max_sessions, - completed_append_limit=completed_append_limit, - ) - self._pending_expirations: dict[tuple[str, int], DuplexSessionExpiry] = {} - self._pending_control_cleanups: dict[tuple[str, str, int, int, int, int], _PendingControlCleanup] = {} - self._control_cleanup_tasks: dict[ - tuple[str, str, int, int, int, int], - asyncio.Task[None], - ] = {} - self._pending_submission_cleanups: dict[str, _PendingSubmissionCleanup] = {} - self._submission_cleanup_tasks: dict[str, asyncio.Task[None]] = {} - self._pending_request_cleanups: dict[tuple[str, int, int], _PendingRequestCleanup] = {} - self._request_cleanup_tasks: dict[tuple[str, int, int], asyncio.Task[None]] = {} - self._request_cleanups_in_progress: set[tuple[str, int, int]] = set() - self._session_control_tails: dict[str, asyncio.Task[None]] = {} - self._dispatched_control_tasks: set[asyncio.Task[None]] = set() - - @property - def sessions(self) -> DuplexSessionRuntimeManager: - return self._sessions - - @property - def pending_submission_cleanup_count(self) -> int: - return len(self._pending_submission_cleanups) - - def accepts(self, message: object) -> bool: - return isinstance(message, self._MESSAGE_TYPES) - - async def handle(self, message: object) -> None: - if isinstance(message, OpenDuplexSessionMessage): - await self.handle_open(message) - elif isinstance(message, AppendDuplexInputMessage): - await self.handle_append(message) - elif isinstance(message, SignalDuplexTurnMessage): - await self.handle_signal(message) - elif isinstance(message, CloseDuplexSessionMessage): - await self.handle_close(message) - elif isinstance(message, TouchDuplexSessionMessage): - await self.handle_touch(message) - elif isinstance(message, ResumeDuplexSessionMessage): - await self.handle_resume(message) - else: - raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") - - def dispatch(self, message: object) -> None: - """Schedule one control command without blocking unrelated sessions.""" - if not self.accepts(message): - raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") - session_id = message.session_id - predecessor = self._session_control_tails.get(session_id) - - async def run_ordered() -> None: - if predecessor is not None: - await predecessor - await self.handle(message) - - task = asyncio.create_task( - run_ordered(), - name=f"duplex-control-{session_id}-{message.control_id}", - ) - self._session_control_tails[session_id] = task - self._dispatched_control_tasks.add(task) - - def discard(completed: asyncio.Task[None]) -> None: - self._dispatched_control_tasks.discard(completed) - if self._session_control_tails.get(session_id) is completed: - self._session_control_tails.pop(session_id, None) - if not completed.cancelled(): - completed.exception() - - task.add_done_callback(discard) - - async def drain(self) -> None: - while self._dispatched_control_tasks: - await asyncio.gather(*tuple(self._dispatched_control_tasks)) - - async def shutdown(self) -> None: - tasks = tuple( - self._dispatched_control_tasks - | set(self._control_cleanup_tasks.values()) - | set(self._submission_cleanup_tasks.values()) - | set(self._request_cleanup_tasks.values()) - ) - for task in tasks: - task.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - - @staticmethod - def coerce_capabilities(raw: dict[str, object]) -> DuplexRuntimeCapabilities: - input_modes: set[DuplexInputMode] = set() - values = raw.get("input_modes") - if values is not None and not isinstance(values, list): - raise TypeError("duplex input_modes capability must be a list") - if isinstance(values, list): - for value in values: - try: - input_modes.add(DuplexInputMode(value)) - except (TypeError, ValueError) as exc: - raise ValueError(f"unknown duplex input mode: {value!r}") from exc - return DuplexRuntimeCapabilities( - input_modes=input_modes or {DuplexInputMode.TURN_COMMIT_ONLY}, - implementation_level=str(raw.get("implementation_level") or "serving_session_adapter"), - ) - - def sampling_params_for_config(self, runtime_config: dict[str, Any]) -> list[object]: - defaults = self._stage_port.sampling_defaults() - if self._extension is None: - return list(defaults) - configured = self._extension.configure_sampling_params( - runtime_config=runtime_config, - defaults=defaults, - ) - if not isinstance(configured, tuple): - raise TypeError("duplex runtime extension must return sampling parameters as a tuple") - if len(configured) != len(defaults): - raise ValueError("duplex runtime extension must return one sampling parameter per stage") - return list(configured) - - @staticmethod - def stage_request_id(fence: DuplexFence, *, stage_id: int) -> str: - return duplex_resource_request_id(fence, f"stage{stage_id}") - - def ensure_stage_request( - self, - session: DuplexSessionRuntimeState, - *, - stage_id: int, - fence: DuplexFence | None = None, - ) -> DuplexStageRequestContext | None: - if stage_id >= self._stage_port.stage_count: - return None - effective_fence = fence or session.fence - request_id = self.stage_request_id(effective_fence, stage_id=stage_id) - session.reserve_stage_request(stage_id, request_id, fence=effective_fence) - context = DuplexStageRequestContext( - request_id=request_id, - session_id=session.session_id, - fence=effective_fence, - stage_id=stage_id, - final_stage_id=self._stage_port.stage_count - 1, - config_generation=session.config_generation, - sampling_params=tuple(self.sampling_params_for_config(session.runtime_config)), - session_config=session.session_config, - runtime_config=session.runtime_config, - ) - self._stage_port.ensure_request(context) - return context - - async def handle_open(self, message: OpenDuplexSessionMessage) -> None: - session: DuplexSessionRuntimeState | None = None - try: - session_mode = SessionMode(message.session_mode) - capabilities = self.coerce_capabilities(message.capabilities) - native_append_modes = capabilities.input_modes - {DuplexInputMode.TURN_COMMIT_ONLY} - if native_append_modes and self._extension is None: - raise RuntimeError("duplex_runtime_extension_not_configured") - session = self.sessions.open_session( - message.fence, - capabilities=capabilities, - session_config=message.session_config, - runtime_config=message.runtime_config, - lease_config=self._lease_config, - ) - request_context = self.ensure_stage_request(session, stage_id=0) if self._extension is not None else None - await self.put_result( - message.control_id, - fence=message.fence, - operation="open", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "implementation_level": capabilities.implementation_level, - "data_plane_session": True, - "session_mode": session_mode.value, - "scheduler_request_context": request_context is not None, - "request_id": request_context.request_id if request_context is not None else None, - }, - } - ], - ) - except Exception as exc: - if self._control_error(exc).code == "resource_exhausted": - logger.info("open_duplex_session rejected: %s", exc) - else: - logger.exception("open_duplex_session failed: %s", exc) - if session is not None and self.sessions.get(message.session_id) is session: - reserved_request_ids = tuple(session.resource_request_ids()) - self.sessions.begin_close_session(session.fence, reason="open_rollback") - cleanup_key = self._cleanup_key("open_rollback", session.fence) - pending = _PendingControlCleanup( - kind="open_rollback", - session_id=session.session_id, - fence=session.fence, - reserved_request_ids=reserved_request_ids, - submitted_request_ids=(), - ) - self._pending_control_cleanups[cleanup_key] = pending - try: - await self._complete_control_cleanup(cleanup_key, pending) - except Exception as cleanup_exc: - logger.warning( - "duplex open rollback remains pending for session %s: %s", - session.session_id, - cleanup_exc, - ) - await self.put_result( - message.control_id, - fence=message.fence, - operation="open", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_append(self, message: AppendDuplexInputMessage) -> None: - session: DuplexSessionRuntimeState | None = None - lease_operation_id = f"append:{message.control_id}" - operation_started = False - try: - session = self.sessions.require(message.session_id) - await self._complete_pending_submission_cleanup(message.session_id) - if message.expected_epoch is not None and message.expected_epoch != message.fence.epoch: - raise ValueError("expected_epoch must match fence.epoch") - mode = DuplexInputMode(message.mode) - if message.operation_id is not None: - completed = session.completed_append( - message.operation_id, - fence=message.fence, - mode=mode, - final=message.final, - ) - if completed is not None: - session.touch(message.fence, DuplexLeaseActivity.APPEND) - await self.put_result( - message.control_id, - fence=message.fence, - operation="append", - session_id=message.session_id, - stage_results=completed, - ) - return - session.begin_operation(message.fence, lease_operation_id) - operation_started = True - reservation = session.prepare_append(mode=mode, fence=message.fence) - stage_results = await self.append_via_data_plane( - message, - session=session, - reservation=reservation, - mode=mode, - ) - if message.operation_id is not None: - session.record_completed_append( - message.operation_id, - fence=message.fence, - mode=mode, - final=message.final, - stage_results=stage_results, - ) - await self.put_result( - message.control_id, - fence=message.fence, - operation="append", - session_id=message.session_id, - stage_results=stage_results, - ) - except Exception as exc: - logger.exception("append_duplex_input failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="append", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - finally: - if session is not None and operation_started and lease_operation_id in session.lease.active_operations: - session.end_operation(session.fence, lease_operation_id) - - async def append_via_data_plane( - self, - message: AppendDuplexInputMessage, - *, - session: DuplexSessionRuntimeState, - reservation: DuplexAppendReservation, - mode: DuplexInputMode, - ) -> list[dict[str, object]]: - if self._extension is None and mode is DuplexInputMode.TURN_COMMIT_ONLY: - update = session.commit_append(reservation) - return [ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_append": False, - "seq": update.seq, - "turn_id": update.turn_id, - "turn_seq": update.turn_seq, - "mode": mode.value, - }, - } - ] - if self._stage_port.stage_count == 0: - return [ - { - "stage_id": -1, - "replica_id": -1, - "result": {"supported": False, "error": "duplex_data_plane_has_no_stage"}, - } - ] - if self._extension is None: - raise RuntimeError("duplex_runtime_extension_not_configured") - - stage_id = 0 - request_id = self.stage_request_id(message.fence, stage_id=stage_id) - existing_binding = session.stage_bindings.get(stage_id) - already_submitted = existing_binding is not None and existing_binding.request_id == request_id - request_context = self.ensure_stage_request( - session, - stage_id=stage_id, - fence=message.fence, - ) - if request_context is None: - raise RuntimeError("duplex_data_plane_has_no_stage") - append_plan = self._extension.plan_append( - request_id=request_id, - fence=message.fence, - session_config=dict(request_context.session_config), - runtime_config=dict(request_context.runtime_config), - seq=reservation.update.seq, - turn_seq=reservation.update.turn_seq, - mode=mode, - payload=message.payload, - final=message.final, - sampling_params=request_context.stage_sampling_params, - ) - if not isinstance(append_plan, DuplexAppendPlan): - raise TypeError("duplex runtime extension plan_append() must return DuplexAppendPlan") - submission = DuplexStageSubmission( - context=request_context, - prompt=append_plan.prompt, - already_submitted=already_submitted, - ) - submission_result = await self._stage_port.submit(submission) - try: - if submission_result.request_id != request_id or submission_result.stage_id != stage_id: - raise RuntimeError("duplex stage adapter returned a mismatched submission result") - update = session.commit_append(reservation) - session.bind_stage_request(stage_id, request_id, fence=message.fence) - except BaseException: - self._pending_submission_cleanups[session.session_id] = _PendingSubmissionCleanup( - session_id=session.session_id, - request_ids=(request_id,), - ) - try: - await self._complete_pending_submission_cleanup(session.session_id) - except Exception as cleanup_exc: - logger.warning( - "duplex append compensation remains pending for session %s: %s", - session.session_id, - cleanup_exc, - ) - raise - return [ - { - "stage_id": stage_id, - "replica_id": submission_result.replica_id, - "result": { - "supported": True, - "implementation_level": session.capabilities.implementation_level, - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": request_context.final_stage_id, - "seq": update.seq, - "turn_id": update.turn_id, - "response_seq": message.fence.response_seq, - "turn_seq": update.turn_seq, - "mode": mode.value, - "resumable": True, - }, - } - ] - - async def handle_signal(self, message: SignalDuplexTurnMessage) -> None: - try: - cancel_events = {"barge_in", "input.cancel", "response.cancel"} - if message.event not in {*cancel_events, "session.update"}: - raise ValueError(f"unsupported duplex runtime signal: {message.event}") - session = self.sessions.require(message.session_id) - await self._complete_pending_submission_cleanup(message.session_id) - effective_next_fence = message.next_fence - if message.event in cancel_events: - if effective_next_fence is None: - raise ValueError(f"{message.event} requires next_fence") - cleanup_key = self._cleanup_key("cancel", message.fence) - pending = self._pending_control_cleanups.get(cleanup_key) - if pending is None: - stale_request_ids = session.prepare_cancel_fence(message.fence, effective_next_fence) - pending = _PendingControlCleanup( - kind="cancel", - session_id=message.session_id, - fence=message.fence, - submitted_request_ids=tuple(stale_request_ids), - ) - self._pending_control_cleanups[cleanup_key] = pending - await self._complete_control_cleanup(cleanup_key, pending) - else: - session.accept_fence(message.fence) - if message.runtime_config is not None: - self.sampling_params_for_config(message.runtime_config) - session.replace_configs( - session_config=message.session_config, - runtime_config=message.runtime_config, - ) - session.touch(session.fence, DuplexLeaseActivity.SIGNAL) - await self.put_result( - message.control_id, - fence=message.fence, - operation="signal", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_signal": True, - "event": message.event, - "fence": message.fence, - "next_fence": effective_next_fence, - }, - } - ], - ) - except Exception as exc: - logger.exception("signal_duplex_turn failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="signal", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_close(self, message: CloseDuplexSessionMessage) -> None: - try: - await self._complete_pending_submission_cleanup(message.session_id) - session = self.sessions.get(message.session_id) - if session is None: - await self.put_result( - message.control_id, - fence=message.fence, - operation="close", - session_id=message.session_id, - stage_results=[], - ) - return - cleanup_key = self._cleanup_key("close", message.fence) - pending = self._pending_control_cleanups.get(cleanup_key) - if pending is None: - submitted = tuple(session.resource_request_ids(submitted=True)) - reserved = tuple(session.resource_request_ids(submitted=False)) - self.sessions.begin_close_session(message.fence, reason=message.reason) - pending = _PendingControlCleanup( - kind="close", - session_id=message.session_id, - fence=message.fence, - submitted_request_ids=submitted, - reserved_request_ids=reserved, - ) - self._pending_control_cleanups[cleanup_key] = pending - await self._complete_control_cleanup(cleanup_key, pending) - await self.put_result( - message.control_id, - fence=message.fence, - operation="close", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_close": True, - "reason": message.reason, - }, - } - ], - ) - except Exception as exc: - logger.exception("close_duplex_session failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="close", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_touch(self, message: TouchDuplexSessionMessage) -> None: - try: - session = self.sessions.require(message.session_id) - activity = DuplexLeaseActivity(message.activity) - if activity is DuplexLeaseActivity.DETACH: - session.detach(message.fence) - else: - session.touch(message.fence, activity) - await self.put_result( - message.control_id, - fence=message.fence, - operation="touch", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "activity": activity.value, - "lease_generation": session.lease.generation, - }, - } - ], - ) - except Exception as exc: - logger.exception("touch_duplex_session failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="touch", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_resume(self, message: ResumeDuplexSessionMessage) -> None: - try: - session = self.sessions.require(message.session_id) - generation = session.resume( - message.fence, - expected_lease_generation=message.expected_lease_generation, - ) - await self.put_result( - message.control_id, - fence=message.fence, - operation="resume", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "lease_generation": generation, - }, - } - ], - ) - except Exception as exc: - logger.exception("resume_duplex_session failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="resume", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def reap_expired(self, now: float | None = None) -> int: - completed = 0 - for session_id in list(self._pending_submission_cleanups): - try: - await self._complete_pending_submission_cleanup(session_id) - except Exception as exc: - logger.warning( - "duplex append compensation remains pending for session %s: %s", - session_id, - exc, - ) - for key, pending in list(self._pending_control_cleanups.items()): - try: - await self._complete_control_cleanup(key, pending) - except Exception as exc: - logger.warning( - "duplex %s cleanup remains pending for session %s: %s", - pending.kind, - pending.session_id, - exc, - ) - for key, pending in list(self._pending_request_cleanups.items()): - if key in self._request_cleanups_in_progress: - continue - try: - await self._complete_request_cleanup(key, pending) - except Exception as exc: - logger.warning( - "duplex request cleanup remains pending for session %s: %s", - pending.session_id, - exc, - ) - continue - completed += 1 - for item in self.sessions.collect_expired( - now, - excluded_session_ids=set(self._pending_submission_cleanups), - ): - self._pending_expirations.setdefault( - (item.session_id, item.lease_generation), - item, - ) - for key, item in list(self._pending_expirations.items()): - try: - if item.submitted_request_ids: - await self._stage_port.cleanup(list(item.submitted_request_ids), abort=True) - if item.reserved_request_ids: - await self._stage_port.cleanup(list(item.reserved_request_ids)) - if self._lifecycle_sink is not None: - await self._lifecycle_sink.put( - DuplexSessionLifecycleMessage( - fence=item.fence, - session_id=item.session_id, - event="expired", - reason=item.reason, - lease_generation=item.lease_generation, - submitted_request_ids=list(item.submitted_request_ids), - reserved_request_ids=list(item.reserved_request_ids), - ) - ) - except Exception as exc: - logger.warning( - "duplex expiry cleanup remains pending for session %s: %s", - item.session_id, - exc, - ) - continue - session = self.sessions.get(item.session_id) - if session is not None and session.lease.generation == item.lease_generation: - self.sessions.finalize_close_session(session) - self._pending_expirations.pop(key, None) - completed += 1 - return completed - - async def _complete_pending_submission_cleanup(self, session_id: str) -> None: - pending = self._pending_submission_cleanups.get(session_id) - if pending is None: - return - task = self._submission_cleanup_tasks.get(session_id) - if task is not None and task.done(): - self._submission_cleanup_tasks.pop(session_id, None) - task = None - if task is None: - task = asyncio.create_task( - self._run_submission_cleanup(pending), - name=f"duplex-append-compensation-{session_id}", - ) - self._submission_cleanup_tasks[session_id] = task - - def discard(completed: asyncio.Task[None]) -> None: - if self._submission_cleanup_tasks.get(session_id) is completed: - self._submission_cleanup_tasks.pop(session_id, None) - - task.add_done_callback(discard) - await asyncio.shield(task) - - async def _run_submission_cleanup(self, pending: _PendingSubmissionCleanup) -> None: - await self._stage_port.cleanup(list(pending.request_ids), abort=True) - session = self.sessions.get(pending.session_id) - if session is not None: - session.release_request_ids(list(pending.request_ids)) - if self._pending_submission_cleanups.get(pending.session_id) is pending: - self._pending_submission_cleanups.pop(pending.session_id, None) - - async def _complete_request_cleanup( - self, - key: tuple[str, int, int], - pending: _PendingRequestCleanup, - ) -> None: - task = self._request_cleanup_tasks.get(key) - if task is not None and task.done(): - self._request_cleanup_tasks.pop(key, None) - task = None - if task is None: - task = asyncio.create_task( - self._run_request_cleanup(key, pending), - name=f"duplex-request-cleanup-{pending.session_id}", - ) - self._request_cleanup_tasks[key] = task - - def discard(completed: asyncio.Task[None]) -> None: - if self._request_cleanup_tasks.get(key) is completed: - self._request_cleanup_tasks.pop(key, None) - - task.add_done_callback(discard) - await asyncio.shield(task) - - async def _run_request_cleanup( - self, - key: tuple[str, int, int], - pending: _PendingRequestCleanup, - ) -> None: - await self._stage_port.cleanup(list(pending.request_ids), abort=pending.abort) - session = self.sessions.get(pending.session_id) - if ( - session is not None - and session.fence.incarnation == pending.fence.incarnation - and session.lease.generation == pending.lease_generation - ): - self.sessions.finalize_close_session(session) - if self._pending_request_cleanups.get(key) is pending: - self._pending_request_cleanups.pop(key, None) - self._request_cleanups_in_progress.discard(key) - - @staticmethod - def _cleanup_key(kind: str, fence: DuplexFence) -> tuple[str, str, int, int, int, int]: - return (kind, fence.session_id, fence.incarnation, fence.epoch, fence.turn_id, fence.response_seq) - - async def _complete_control_cleanup( - self, - key: tuple[str, str, int, int, int, int], - pending: _PendingControlCleanup, - ) -> None: - task = self._control_cleanup_tasks.get(key) - if task is not None and task.done(): - self._control_cleanup_tasks.pop(key, None) - task = None - if task is None: - task = asyncio.create_task( - self._run_control_cleanup(key, pending), - name=f"duplex-{pending.kind}-cleanup-{pending.session_id}", - ) - self._control_cleanup_tasks[key] = task - - def discard(completed: asyncio.Task[None]) -> None: - if self._control_cleanup_tasks.get(key) is completed: - self._control_cleanup_tasks.pop(key, None) - - task.add_done_callback(discard) - await asyncio.shield(task) - - async def _run_control_cleanup( - self, - key: tuple[str, str, int, int, int, int], - pending: _PendingControlCleanup, - ) -> None: - if pending.submitted_request_ids: - await self._stage_port.cleanup(list(pending.submitted_request_ids), abort=True) - if pending.reserved_request_ids: - await self._stage_port.cleanup(list(pending.reserved_request_ids)) - session = self.sessions.get(pending.session_id) - if session is not None: - if pending.kind == "cancel": - session.release_fence(pending.fence) - elif pending.kind in {"close", "open_rollback"}: - self.sessions.finalize_close_session(session) - self._pending_control_cleanups.pop(key, None) - - @classmethod - def _iter_result_dicts(cls, result: object): - if isinstance(result, dict): - yield result - elif isinstance(result, list | tuple): - for item in result: - yield from cls._iter_result_dicts(item) - - @classmethod - def _result_counts(cls, stage_results: list[dict[str, object]]) -> tuple[int, int]: - unsupported_count = 0 - error_count = 0 - for item in stage_results: - for result in cls._iter_result_dicts(item.get("result")): - if result.get("supported") is False: - unsupported_count += 1 - if result.get("error"): - error_count += 1 - return unsupported_count, error_count - - async def put_result( - self, - control_id: str, - *, - fence: DuplexFence, - operation: str, - session_id: str, - stage_results: list[dict[str, object]], - error: BaseException | str | None = None, - ) -> None: - control_error = self._control_error(error) if error is not None else None - if error is not None: - stage_results = [ - { - "stage_id": -1, - "replica_id": -1, - "result": {"supported": False, "error": control_error.message}, - } - ] - unsupported_count, error_count = self._result_counts(stage_results) - session = self.sessions.get(session_id) - await self._result_sink.put( - DuplexControlResultMessage( - control_id=control_id, - fence=fence, - operation=operation, - session_id=session_id, - ok=error_count == 0 and unsupported_count == 0, - stage_results=stage_results, - unsupported_count=unsupported_count, - error_count=error_count, - error=control_error, - accepted_fence=session.fence if session is not None else None, - lease_generation=session.lease.generation if session is not None else None, - ) - ) - - @staticmethod - def _control_error(error: BaseException | str) -> DuplexControlError: - message = str(error) - if isinstance(error, DuplexFenceMismatchError): - code = "stale_fence" - retryable = False - elif "unknown duplex input mode" in message: - code = "invalid_capability" - retryable = False - elif "duplex_session_capacity_exhausted" in message: - code = "resource_exhausted" - retryable = True - elif isinstance(error, KeyError): - code = "not_found" - retryable = False - elif isinstance(error, (TypeError, ValueError)): - code = "invalid_argument" - retryable = False - elif isinstance(error, TimeoutError): - code = "timeout" - retryable = True - else: - code = "failed_precondition" - retryable = False - return DuplexControlError(code=code, message=message, retryable=retryable) - - def decide_output( - self, - stage_id: int, - output: object, - context: DuplexOutputContext | None, - ) -> DuplexOutputDecision | None: - if context is None or self._extension is None: - return None - session = self._sessions.get(context.identity.session_id) - if session is None: - return None - try: - session.touch(context.identity.fence, DuplexLeaseActivity.MODEL_OUTPUT) - except RuntimeError: - return None - decision = self._extension.decide_output( - stage_id=stage_id, - final_stage_id=context.final_stage_id, - segment_finished=context.segment_finished, - segment_token_ids=context.segment_token_ids, - segment_output_metadata=dict(context.segment_output_metadata), - output=output, - ) - if decision is not None and not isinstance(decision, DuplexOutputDecision): - raise TypeError("duplex runtime extension decide_output() must return DuplexOutputDecision or None") - return decision - - def session_for_identity(self, identity: DuplexRequestIdentity | None) -> DuplexSessionRuntimeState | None: - if identity is None: - return None - return self._sessions.get(identity.session_id) - - def close_sessions_for_request_ids( - self, - request_ids: list[str], - *, - abort: bool = False, - cleanup_in_progress: bool = False, - ) -> dict[str, list[str]]: - closed = self._sessions.close_sessions_for_request_ids(request_ids) - for session_id, stale_request_ids in closed.items(): - session = self._sessions.get(session_id) - if session is None: - continue - key = (session_id, session.fence.incarnation, session.lease.generation) - existing = self._pending_request_cleanups.get(key) - merged_request_ids = tuple( - dict.fromkeys( - [ - *(existing.request_ids if existing is not None else ()), - *stale_request_ids, - ] - ) - ) - self._pending_request_cleanups[key] = _PendingRequestCleanup( - session_id=session_id, - fence=session.fence, - lease_generation=session.lease.generation, - request_ids=merged_request_ids, - abort=abort or (existing.abort if existing is not None else False), - ) - if cleanup_in_progress: - self._request_cleanups_in_progress.add(key) - return closed - - def defer_request_cleanups(self, session_ids: Iterable[str]) -> None: - session_id_set = set(session_ids) - active_keys = {key for key in self._request_cleanups_in_progress if key[0] in session_id_set} - self._request_cleanups_in_progress.difference_update(active_keys) - - def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: - session_id_set = set(session_ids) - for key in list(self._pending_request_cleanups): - if key[0] not in session_id_set: - continue - self._pending_request_cleanups.pop(key, None) - self._request_cleanups_in_progress.discard(key) - self._sessions.finalize_closed_sessions(session_ids) - - -__all__ = [ - "DuplexControlPlane", - "DuplexOutputContext", - "DuplexRequestIdentity", - "DuplexResultSink", - "DuplexLifecycleSink", - "DuplexStagePort", - "DuplexStageRequestContext", - "DuplexStageSubmission", - "DuplexStageSubmissionResult", -] diff --git a/vllm_omni/engine/duplex/events.py b/vllm_omni/engine/duplex/events.py new file mode 100644 index 000000000..9c131a893 --- /dev/null +++ b/vllm_omni/engine/duplex/events.py @@ -0,0 +1,756 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Typed events emitted by a duplex session. + +Every public event is a frozen dataclass with explicit fields. The Realtime +wire JSON is *derived* from those fields by :meth:`DuplexEvent.to_realtime`, +which is pure: it never consults session state. The stateful part (response / +item ids, content-part bookkeeping) lives in +``vllm_omni.engine.duplex.realtime_events`` on the session runner and is +consumed when the events are *constructed*. + +``session_id`` / ``epoch`` are bound by +``DuplexSessionManager.emit`` when an event leaves the runner; producers +(projector, serving handshake helpers) build events with the defaults. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from typing import ClassVar, cast +from uuid import uuid4 + + +def new_event_id() -> str: + return f"event_{uuid4().hex}" + + +#: OpenAI Realtime ``error.type`` for each internal error code. +REALTIME_ERROR_TYPES_BY_CODE: dict[str, str] = { + "bad_event": "invalid_request_error", + "bad_audio": "invalid_request_error", + "config_timeout": "invalid_request_error", + "invalid_json": "invalid_request_error", + "event_too_large": "invalid_request_error", + "unknown_event": "invalid_request_error", + "internal_error": "server_error", + "runtime_append_failed": "server_error", + "runtime_append_task_failed": "server_error", + "runtime_signal_failed": "server_error", + "runtime_abort_failed": "server_error", + "runtime_data_plane_stream_failed": "server_error", + "runtime_data_plane_text_without_audio": "server_error", + "resource_exhausted": "rate_limit_error", + "session_exists": "invalid_request_error", + "session_closed": "invalid_request_error", + "unknown_session": "invalid_request_error", + "invalid_duplex_runtime_config": "invalid_request_error", + "instructions_update_unsupported": "invalid_request_error", + "persona_update_unsupported": "invalid_request_error", + "voice_update_unsupported": "invalid_request_error", + "unsupported_nemotron_duplex_mode": "invalid_request_error", + "unsupported_native_response_options": "invalid_request_error", + "runtime_touch_failed": "server_error", + "engine_error": "server_error", + "input_backpressure": "rate_limit_error", + "response_already_active": "invalid_request_error", + "response_not_active": "invalid_request_error", + "response_create_without_input": "invalid_request_error", + "text_only_turn_unsupported": "invalid_request_error", + "commit_aborted": "server_error", + "input_audio_buffer_empty": "invalid_request_error", + "missing_item_id": "invalid_request_error", + "item_not_found": "invalid_request_error", + "playback_item_mismatch": "invalid_request_error", + "playback_item_not_found": "invalid_request_error", + "playback_ack_too_late": "invalid_request_error", + "unsupported_audio_format": "invalid_request_error", + "unsupported_turn_detection": "invalid_request_error", + "unsupported_ref_audio_path": "invalid_request_error", + "ref_audio_required": "invalid_request_error", + "model_update_unsupported": "invalid_request_error", + "voice_update_after_audio_unsupported": "invalid_request_error", + "ref_audio_update_unsupported": "invalid_request_error", + "native_text_append_unsupported": "invalid_request_error", + "invalid_video_frames": "invalid_request_error", + "invalid_function_call_output": "invalid_request_error", + "server_vad_unavailable": "server_error", +} + +_IDENTITY_FIELDS = frozenset({"session_id", "epoch", "event_id"}) + + +def _wire_value(value: object) -> object: + if isinstance(value, Mapping): + return {str(k): _wire_value(v) for k, v in value.items()} + if isinstance(value, tuple | list): + return [_wire_value(v) for v in value] + return value + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexEvent: + """Base of every public session event.""" + + #: Realtime ``type`` this class renders to. + wire_type: ClassVar[str] = "" + #: Fields that are omitted from the wire object when ``None``. + optional_wire_fields: ClassVar[frozenset[str]] = frozenset() + + session_id: str = "" + epoch: int | None = None + #: Server event id (OpenAI ``event_id``); generated at construction. + event_id: str = field(default_factory=new_event_id) + + # ---- identity ---- + + @property + def type(self) -> str: + return self.wire_type + + @property + def is_terminal(self) -> bool: + return False + + # ---- generic accessors (None when the class has no such field) ---- + + @property + def response_id(self) -> str | None: + return None + + @property + def item_id(self) -> str | None: + return None + + @property + def text(self) -> str | None: + return None + + @property + def audio(self) -> bytes | None: + return None + + # ---- wire rendering (pure) ---- + + def _wire_fields(self) -> dict[str, object]: + data: dict[str, object] = {} + for f in fields(self): + if f.name in _IDENTITY_FIELDS: + continue + value = getattr(self, f.name) + if value is None and f.name in self.optional_wire_fields: + continue + data[f.name] = _wire_value(value) + return data + + def to_realtime(self) -> dict[str, object]: + """The Realtime wire JSON object for this event (derived from the fields).""" + return {"type": self.wire_type, "event_id": self.event_id, **self._wire_fields()} + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"{type(self).__name__}(type={self.wire_type!r}, session_id={self.session_id!r}, epoch={self.epoch!r})" + + +# ---- session lifecycle ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionCreated(DuplexEvent): + wire_type = "session.created" + optional_wire_fields = frozenset({"attachment_generation", "resume_token"}) + + session: Mapping[str, object] = field(default_factory=dict) + #: Transport resume credentials (set by the websocket handler when resume is supported). + attachment_generation: int | None = None + resume_token: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionUpdated(DuplexEvent): + wire_type = "session.updated" + + session: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionResumed(SessionCreated): + wire_type = "session.resumed" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionReplaced(DuplexEvent): + """A newer connection took the session over (sent to the replaced socket).""" + + wire_type = "session.replaced" + + attachment_generation: int = 0 + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionResyncRequired(DuplexEvent): + """The replay journal cannot bridge the gap; the client must start over.""" + + wire_type = "session.resync_required" + + reason: str = "journal_gap" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionClosed(DuplexEvent): + wire_type = "session.closed" + + reason: str = "closed" + #: The session-internal close event (kept for clients that inspect it). + details: Mapping[str, object] = field(default_factory=dict) + + @property + def is_terminal(self) -> bool: + return True + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "reason": self.reason, "event": _wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionExpired(SessionClosed): + wire_type = "session.expired" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "reason": self.reason} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionHeartbeatAck(DuplexEvent): + wire_type = "session.heartbeat_ack" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TurnEvent(DuplexEvent): + """Turn-state transition (``turn.event``).""" + + wire_type = "turn.event" + + event: str = "" + turn_state: str = "" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "event": self.event, "turn_state": self.turn_state, "epoch": self.epoch} + + +# ---- response lifecycle ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ResponseEvent(DuplexEvent): + """Events addressed to one response (and usually one output item / content part).""" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + content_index: int = 0 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseCreated(DuplexEvent): + wire_type = "response.created" + + response_id: str | None = None # type: ignore[assignment] + response: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Listen(DuplexEvent): + """The model decided to keep listening (no spoken response for this turn).""" + + wire_type = "response.listen" + optional_wire_fields = frozenset({"response_id"}) + + response_id: str | None = None # type: ignore[assignment] + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + response: dict[str, object] = { + "object": "realtime.response", + "status": "listening", + "metadata": _wire_value(self.details), + } + data: dict[str, object] = {"session_id": self.session_id, "epoch": self.epoch, "response": response} + if self.response_id: + response["id"] = self.response_id + data["response_id"] = self.response_id + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Speak(_ResponseEvent): + wire_type = "response.speak" + + metadata: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OverlapDecision(DuplexEvent): + wire_type = "overlap.decision" + + policy: str | None = None + action: str | None = None + reason: str | None = None + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "session_id": self.session_id, + "epoch": self.epoch, + "policy": self.policy, + "action": self.action, + "reason": self.reason, + "metadata": _wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputItemAdded(DuplexEvent): + wire_type = "response.output_item.added" + + response_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputItemDone(OutputItemAdded): + wire_type = "response.output_item.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ContentPartAdded(_ResponseEvent): + wire_type = "response.content_part.added" + + part: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ContentPartDone(ContentPartAdded): + wire_type = "response.content_part.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDelta(_ResponseEvent): + wire_type = "response.output_audio.delta" + optional_wire_fields = frozenset({"sample_rate_hz", "metadata"}) + + #: Base64 audio in ``format``. + delta: str = "" + format: str = "pcm16" + sample_rate_hz: int | None = None + metadata: Mapping[str, object] | None = None + + @property + def audio(self) -> bytes | None: + try: + return base64.b64decode(self.delta) + except (ValueError, TypeError): + return None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDone(_ResponseEvent): + wire_type = "response.output_audio.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TranscriptDelta(_ResponseEvent): + wire_type = "response.output_audio_transcript.delta" + + delta: str = "" + + @property + def text(self) -> str | None: + return self.delta + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TranscriptDone(_ResponseEvent): + wire_type = "response.output_audio_transcript.done" + + transcript: str = "" + + @property + def text(self) -> str | None: + return self.transcript + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TextDelta(_ResponseEvent): + wire_type = "response.output_text.delta" + + delta: str = "" + + @property + def text(self) -> str | None: + return self.delta + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TextDone(_ResponseEvent): + wire_type = "response.output_text.done" + + text: str = "" # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FunctionCallArgumentsDelta(DuplexEvent): + wire_type = "response.function_call_arguments.delta" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + call_id: str = "" + delta: str = "" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FunctionCallArgumentsDone(DuplexEvent): + wire_type = "response.function_call_arguments.done" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + call_id: str = "" + arguments: str = "" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseDone(DuplexEvent): + wire_type = "response.done" + + response_id: str | None = None # type: ignore[assignment] + response: Mapping[str, object] = field(default_factory=dict) + + @property + def status(self) -> str | None: + value = self.response.get("status") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RateLimitsUpdated(DuplexEvent): + wire_type = "rate_limits.updated" + + rate_limits: tuple[Mapping[str, object], ...] = () + + +# ---- input buffer / conversation ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCommitted(DuplexEvent): + wire_type = "input_audio_buffer.committed" + + previous_item_id: str | None = None + item_id: str | None = None # type: ignore[assignment] + #: The session-internal ``input.committed`` event. + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "previous_item_id": self.previous_item_id, + "item_id": self.item_id, + "event": _wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCleared(DuplexEvent): + wire_type = "input_audio_buffer.cleared" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpeechStarted(DuplexEvent): + wire_type = "input_audio_buffer.speech_started" + + audio_start_ms: int = 0 + item_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpeechStopped(DuplexEvent): + wire_type = "input_audio_buffer.speech_stopped" + + audio_end_ms: int = 0 + item_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ItemEvent(DuplexEvent): + previous_item_id: str | None = None + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemAdded(_ItemEvent): + wire_type = "conversation.item.added" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemCreated(_ItemEvent): + wire_type = "conversation.item.created" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDone(_ItemEvent): + wire_type = "conversation.item.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemRetrieved(DuplexEvent): + wire_type = "conversation.item.retrieved" + + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDeleted(DuplexEvent): + wire_type = "conversation.item.deleted" + + item_id: str | None = None # type: ignore[assignment] + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return {"item_id": self.item_id, "event": _wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemTruncated(DuplexEvent): + wire_type = "conversation.item.truncated" + + item_id: str | None = None # type: ignore[assignment] + content_index: int = 0 + audio_end_ms: int = 0 + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "item_id": self.item_id, + "content_index": self.content_index, + "audio_end_ms": self.audio_end_ms, + "event": _wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputTranscriptionCompleted(DuplexEvent): + wire_type = "conversation.item.input_audio_transcription.completed" + + item_id: str | None = None # type: ignore[assignment] + content_index: int = 0 + transcript: str = "" + + @property + def text(self) -> str | None: + return self.transcript + + +# ---- playback / control ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputAudioCleared(DuplexEvent): + wire_type = "output_audio_buffer.cleared" + + response_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PlaybackAcknowledged(DuplexEvent): + wire_type = "playback.acknowledged" + + #: The session-internal ``playback.acknowledged`` event (cursor, committed_ms ...). + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return {"event": _wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ErrorEvent(DuplexEvent): + wire_type = "error" + + code: str = "internal_error" + message: str = "" + #: Client ``event_id`` this error answers (OpenAI ``error.event_id``). + related_event_id: str | None = None + param: str | None = None + #: Extra keys merged into the wire ``error`` object (``retryable`` ...). + extra: Mapping[str, object] = field(default_factory=dict) + + @property + def error_type(self) -> str: + return REALTIME_ERROR_TYPES_BY_CODE.get(self.code, "invalid_request_error") + + @property + def error(self) -> dict[str, object]: + error: dict[str, object] = {"type": self.error_type, "code": self.code, "message": self.message} + if self.related_event_id: + error["event_id"] = self.related_event_id + if self.param: + error["param"] = self.param + error.update(cast("Mapping[str, object]", _wire_value(self.extra))) + return error + + @property + def text(self) -> str | None: + return self.message + + def _wire_fields(self) -> dict[str, object]: + return {"error": self.error} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexRawEvent(DuplexEvent): + """A session-internal event with no dedicated Realtime type (``duplex.``).""" + + wire_type = "duplex.raw" + + internal_type: str = "" + details: Mapping[str, object] = field(default_factory=dict) + + @property + def type(self) -> str: + return f"duplex.{self.internal_type}" + + def to_realtime(self) -> dict[str, object]: + return {"type": self.type, "event_id": self.event_id, "event": _wire_value(self.details)} + + +def error_event( + code: str, + message: str, + *, + event_id: object | None = None, + param: object | None = None, + extra: Mapping[str, object] | None = None, +) -> ErrorEvent: + """Build the error event for an internal code; ``event_id`` is the *client* event id.""" + return ErrorEvent( + code=code, + message=message, + related_event_id=event_id if isinstance(event_id, str) and event_id else None, + param=param if isinstance(param, str) and param else None, + extra=dict(extra or {}), + ) + + +#: Internal events that terminate a response/session and must never be dropped as stale. +DOMAIN_TERMINAL_EVENTS = frozenset( + { + "response.done", + "response.listen", + "audio.cancelled", + "input.cancelled", + "session.closed", + } +) + +#: Internal streaming model-output events subject to epoch/close stale filtering. +MODEL_OUTPUT_EVENTS = frozenset( + { + "response.created", + "response.listen", + "response.speak", + "response.output_item.added", + "response.content_part.added", + "response.output_audio.delta", + "response.output_audio.done", + "response.output_text.delta", + "response.output_text.done", + "response.text.delta", + "response.text.done", + "response.message", + "response.output_item.done", + "response.content_part.done", + "response.done", + } +) + + +__all__ = [ + "DOMAIN_TERMINAL_EVENTS", + "MODEL_OUTPUT_EVENTS", + "REALTIME_ERROR_TYPES_BY_CODE", + "AudioDelta", + "AudioDone", + "ContentPartAdded", + "ContentPartDone", + "DuplexEvent", + "DuplexRawEvent", + "ErrorEvent", + "FunctionCallArgumentsDelta", + "FunctionCallArgumentsDone", + "InputCleared", + "InputCommitted", + "InputTranscriptionCompleted", + "ItemAdded", + "ItemCreated", + "ItemDeleted", + "ItemDone", + "ItemRetrieved", + "ItemTruncated", + "Listen", + "OutputAudioCleared", + "OutputItemAdded", + "OutputItemDone", + "OverlapDecision", + "PlaybackAcknowledged", + "RateLimitsUpdated", + "ResponseCreated", + "ResponseDone", + "SessionClosed", + "SessionCreated", + "SessionExpired", + "SessionHeartbeatAck", + "SessionReplaced", + "SessionResumed", + "SessionResyncRequired", + "SessionUpdated", + "Speak", + "SpeechStarted", + "SpeechStopped", + "TextDelta", + "TextDone", + "TranscriptDelta", + "TranscriptDone", + "TurnEvent", + "error_event", + "new_event_id", +] diff --git a/vllm_omni/engine/duplex/messages.py b/vllm_omni/engine/duplex/messages.py index 2c49d573f..ad264e0c5 100644 --- a/vllm_omni/engine/duplex/messages.py +++ b/vllm_omni/engine/duplex/messages.py @@ -1,130 +1,114 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Queue contracts owned by the experimental duplex control plane.""" +"""Queue envelopes exchanged between ``DuplexOmniEngine`` and ``DuplexOrchestrator``. + +Sessions live inside the engine (``DuplexSessionRunner`` on the orchestrator +loop). The API layer only opens/closes/resumes/touches sessions through +correlated RPC and pushes ``DuplexCommand`` objects one-way; every session +output travels back as a ``DuplexSessionEventMessage``. +""" from __future__ import annotations -from dataclasses import dataclass -from typing import Literal +from typing import TYPE_CHECKING, Literal from vllm_omni.engine.messages import EngineQueueMessage +if TYPE_CHECKING: + from vllm_omni.engine.duplex.commands import DuplexCommand + from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig + from vllm_omni.engine.duplex.events import DuplexEvent -@dataclass(frozen=True, slots=True) -class DuplexFence: - """Identity fence carried by experimental duplex control messages.""" - session_id: str - epoch: int = 0 - turn_id: int = 0 - response_seq: int = 0 - incarnation: int = 0 +class DuplexSessionError(RuntimeError): + """A session control operation was rejected by the engine.""" + + def __init__( + self, + message: str, + *, + code: str = "internal_error", + retryable: bool = False, + session_id: str | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.retryable = retryable + self.session_id = session_id class OpenDuplexSessionMessage(EngineQueueMessage, kw_only=True): type: Literal["open_duplex_session"] = "open_duplex_session" control_id: str - fence: DuplexFence session_id: str - session_mode: str = "duplex" - capabilities: dict[str, object] - session_config: dict[str, object] | None = None - runtime_config: dict[str, object] | None = None - - -class AppendDuplexInputMessage(EngineQueueMessage, kw_only=True): - type: Literal["append_duplex_input"] = "append_duplex_input" - control_id: str - operation_id: str | None = None - fence: DuplexFence - session_id: str - expected_epoch: int | None = None - mode: str - payload: object - final: bool = False - - -class SignalDuplexTurnMessage(EngineQueueMessage, kw_only=True): - type: Literal["signal_duplex_turn"] = "signal_duplex_turn" - control_id: str - fence: DuplexFence - session_id: str - event: str - next_fence: DuplexFence | None = None - session_config: dict[str, object] | None = None - runtime_config: dict[str, object] | None = None + #: Already normalized by ``DuplexOmni``; the queue is in-process, so the object crosses as is. + session_config: DuplexSessionConfig class CloseDuplexSessionMessage(EngineQueueMessage, kw_only=True): type: Literal["close_duplex_session"] = "close_duplex_session" control_id: str - fence: DuplexFence session_id: str reason: str = "client_close" -class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True): - type: Literal["touch_duplex_session"] = "touch_duplex_session" - control_id: str - fence: DuplexFence - session_id: str - activity: str - - class ResumeDuplexSessionMessage(EngineQueueMessage, kw_only=True): type: Literal["resume_duplex_session"] = "resume_duplex_session" control_id: str - fence: DuplexFence session_id: str expected_lease_generation: int -class DuplexSessionLifecycleMessage(EngineQueueMessage, kw_only=True): - type: Literal["duplex_session_lifecycle"] = "duplex_session_lifecycle" - fence: DuplexFence +class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True): + type: Literal["touch_duplex_session"] = "touch_duplex_session" + control_id: str session_id: str - event: str - reason: str - lease_generation: int - submitted_request_ids: list[str] - reserved_request_ids: list[str] + activity: str -class DuplexControlError(EngineQueueMessage, kw_only=True): - code: str - message: str - retryable: bool = False +class DuplexSessionCommandMessage(EngineQueueMessage, kw_only=True): + """One-way session command; rejections come back as ``ErrorEvent``.""" + + type: Literal["duplex_session_command"] = "duplex_session_command" + session_id: str + command: DuplexCommand class DuplexControlResultMessage(EngineQueueMessage, kw_only=True): type: Literal["duplex_control_result"] = "duplex_control_result" control_id: str - fence: DuplexFence operation: str session_id: str ok: bool - stage_results: list[dict[str, object]] - unsupported_count: int = 0 - error_count: int = 0 - error: DuplexControlError | None = None - accepted_fence: DuplexFence | None = None lease_generation: int | None = None + capabilities: DuplexCapabilities | None = None + #: The Realtime ``session`` object of the session (wire payload, already a dict). + public_session: dict[str, object] | None = None + error_code: str | None = None + error_message: str | None = None + error_retryable: bool = False @property def rpc_correlation_key(self) -> tuple[str, str]: return ("duplex", self.control_id) +class DuplexSessionEventMessage(EngineQueueMessage, kw_only=True): + """A typed session event on the engine output queue.""" + + type: Literal["duplex_session_event"] = "duplex_session_event" + session_id: str + event: DuplexEvent + + __all__ = [ - "AppendDuplexInputMessage", "CloseDuplexSessionMessage", - "DuplexControlError", "DuplexControlResultMessage", - "DuplexFence", - "DuplexSessionLifecycleMessage", + "DuplexSessionCommandMessage", + "DuplexSessionError", + "DuplexSessionEventMessage", "OpenDuplexSessionMessage", "ResumeDuplexSessionMessage", - "SignalDuplexTurnMessage", "TouchDuplexSessionMessage", ] diff --git a/vllm_omni/engine/duplex/plugin.py b/vllm_omni/engine/duplex/plugin.py new file mode 100644 index 000000000..130d5204d --- /dev/null +++ b/vllm_omni/engine/duplex/plugin.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Model plugin contract for full-duplex models. + +One ``DuplexModelPlugin`` subclass per model binds what used to be two dotted +paths (the engine ``DuplexRuntimeExtension`` and the serving +``ServingRuntimeAdapter``). Everything runs engine-side now, so the plugin is +loaded once by ``DuplexOmniEngine`` and handed to ``DuplexOrchestrator`` / +``DuplexSessionManager``. +""" + +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Mapping +from importlib import import_module +from typing import TYPE_CHECKING + +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexOutputDecision, +) + +if TYPE_CHECKING: + from vllm.config import ModelConfig + + +class DuplexRuntimeConfigError(ValueError): + """A model plugin rejected client-visible runtime configuration.""" + + def __init__(self, message: str, *, code: str = "invalid_duplex_runtime_config") -> None: + super().__init__(message) + self.code = code + + +def reject_changed_runtime_value( + new_value: object, + current_value: object, + *, + message: str, + code: str, + error_cls: type[DuplexRuntimeConfigError] = DuplexRuntimeConfigError, +) -> None: + if new_value != current_value: + raise error_cls(message, code=code) + + +class PcmAppendReservation(ABC): + operation_id: str + payload: dict[str, object] | None + + @property + @abstractmethod + def active(self) -> bool: ... + + @property + @abstractmethod + def byte_count(self) -> int: ... + + @abstractmethod + def commit(self) -> None: ... + + @abstractmethod + def rollback(self) -> None: ... + + +class PcmAppendBuffer(ABC): + @property + @abstractmethod + def pending_byte_count(self) -> int: ... + + @abstractmethod + def clear(self) -> None: ... + + @abstractmethod + def clear_force_listen(self) -> None: ... + + @abstractmethod + def has_pending(self) -> bool: ... + + @abstractmethod + def has_reserved(self) -> bool: ... + + @abstractmethod + def prepare_append( + self, + payload: dict[str, object], + *, + operation_id: str, + chunk_period_ms: int, + allow_emit: bool, + ) -> PcmAppendReservation | None: ... + + @abstractmethod + def prepare_commit( + self, + *, + operation_id: str, + chunk_period_ms: int, + ) -> PcmAppendReservation: ... + + @abstractmethod + def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: ... + + +class DuplexModelSessionState(ABC): + """Model-owned per-session state; owned by the session runner (one per session).""" + + audio_buffer: PcmAppendBuffer + input_since_commit: bool + speech_since_commit: bool + context_locked: bool + committed_audio_payload: dict[str, object] | None + committed_audio_operation_id: str | None + committed_audio_reserved_bytes: int + deferred_response_create: bool + deferred_precreate_response: bool + continuation_owner_id: str | None + continuation_units: int + pending_silence_task: asyncio.Task[bool] | None + pending_silence_owner_id: str | None + + @abstractmethod + def retain_committed_audio( + self, + payload: dict[str, object], + *, + operation_id: str | None, + reserved_bytes: int = 0, + ) -> None: ... + + @abstractmethod + def clear_committed_audio(self) -> int: ... + + @abstractmethod + def clear_continuation(self) -> None: ... + + +class DuplexDataPlane(ABC): + """Projects raw stage outputs of one model into internal duplex events.""" + + @abstractmethod + def begin_request(self, request_id: str) -> None: ... + + @abstractmethod + def is_terminal(self, request_id: str | None) -> bool: ... + + @abstractmethod + def mark_terminal(self, request_id: str) -> None: ... + + @abstractmethod + def close_stream(self, request_id: str) -> None: ... + + @abstractmethod + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: ... + + @abstractmethod + def project(self, result: object, *, context: object | None = None) -> Iterable[dict[str, object]]: ... + + +EncodeAudio = Callable[[object, int, str, float | None], str | None] + + +class DuplexModelPlugin(ABC): + """Everything vLLM-Omni needs to know about one full-duplex model. + + Engine policy (sampling params, append planning, output decisions) and + session policy (capabilities, runtime configuration, per-session state, + data-plane projection) live on the same object so a mismatch between the + two halves is impossible by construction. + """ + + plugin_id: str = "" + private_runtime_config_keys: frozenset[str] = frozenset() + #: Samples per silence unit the runner appends to keep a model turn going. + silence_continuation_samples: int = 16000 + data_plane: DuplexDataPlane + + def __init__(self, encode_audio: EncodeAudio) -> None: + # Constructor-only: concrete plugins hand the encoder to their data plane. + del encode_audio + + # ---- engine policy (was DuplexRuntimeExtension) ---- + + @abstractmethod + def configure_sampling_params( + self, + *, + runtime_config: dict[str, object], + defaults: tuple[object, ...], + ) -> tuple[object, ...]: ... + + @abstractmethod + def plan_append( + self, + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, object], + runtime_config: dict[str, object], + seq: int, + turn_seq: int, + payload: object, + final: bool, + sampling_params: object, + ) -> DuplexAppendPlan: ... + + @abstractmethod + def decide_output( + self, + *, + stage_id: int, + final_stage_id: int, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: dict[str, object], + output: object, + ) -> DuplexOutputDecision | None: ... + + # ---- session policy (was ServingRuntimeAdapter) ---- + + @abstractmethod + def create_session_state(self) -> DuplexModelSessionState: ... + + @abstractmethod + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: ... + + @abstractmethod + def validate_client_extra_body(self, extra_body: object) -> None: ... + + @abstractmethod + async def prepare_runtime_config( + self, config: DuplexSessionConfig, *, model_config: ModelConfig | None + ) -> dict[str, object]: ... + + @abstractmethod + def runtime_config_for_update( + self, + config: DuplexSessionConfig, + current: Mapping[str, object], + ) -> dict[str, object]: ... + + @abstractmethod + def data_plane_context( + self, + *, + epoch: int, + turn_id: int, + active_response_turn_id: int | None, + active_response_id: str | None, + auto_responds: bool, + response_format: str, + speed: float | None, + modalities: tuple[str, ...], + ) -> object: ... + + # Optional hook: build the runtime config patch for a function-call output + # item. Plugins without tools keep the default (no change). + def runtime_config_for_function_output( + self, + config: DuplexSessionConfig, + current: Mapping[str, object], + item: Mapping[str, object], + ) -> dict[str, object] | None: + del config, current, item + return None + + +def load_duplex_plugin(path: str, encode_audio: EncodeAudio) -> DuplexModelPlugin: + module_name, separator, attribute_name = path.rpartition(".") + if not separator: + raise ValueError(f"Invalid duplex plugin path: {path!r}") + plugin_type = getattr(import_module(module_name), attribute_name) + plugin = plugin_type(encode_audio) + if not isinstance(plugin, DuplexModelPlugin): + raise TypeError(f"{path!r} is not a DuplexModelPlugin") + if not plugin.plugin_id: + raise TypeError("Duplex plugin must declare plugin_id") + if not isinstance(getattr(plugin, "data_plane", None), DuplexDataPlane): + raise TypeError("Duplex plugin must declare a DuplexDataPlane as data_plane") + return plugin + + +def validate_duplex_plugin_sampling(plugin: DuplexModelPlugin, *, sampling_defaults: tuple[object, ...]) -> None: + """Fail fast when the plugin cannot produce one sampling parameter per stage.""" + configured = plugin.configure_sampling_params(runtime_config={}, defaults=sampling_defaults) + if not isinstance(configured, tuple): + raise TypeError("Duplex plugin must return sampling parameters as a tuple") + if len(configured) != len(sampling_defaults): + raise ValueError("Duplex plugin must return one sampling parameter per stage") + for stage_id, (value, default) in enumerate(zip(configured, sampling_defaults, strict=True)): + if default is not None and not isinstance(value, type(default)): + raise TypeError( + "Duplex plugin sampling parameter type mismatch " + f"for stage {stage_id}: expected {type(default).__name__}, got {type(value).__name__}" + ) + + +def payload_turn_id(payload: object) -> int | None: + if not isinstance(payload, Mapping): + return None + return coerce_int(payload.get("duplex_turn_id", payload.get("model_turn_id"))) + + +def coerce_int(value: object) -> int | None: + try: + return int(value) if value is not None else None + except (TypeError, ValueError): + return None + + +__all__ = [ + "DuplexDataPlane", + "DuplexModelPlugin", + "DuplexModelSessionState", + "DuplexRuntimeConfigError", + "EncodeAudio", + "PcmAppendBuffer", + "PcmAppendReservation", + "coerce_int", + "load_duplex_plugin", + "payload_turn_id", + "reject_changed_runtime_value", + "validate_duplex_plugin_sampling", +] diff --git a/vllm_omni/engine/duplex/realtime_commands.py b/vllm_omni/engine/duplex/realtime_commands.py new file mode 100644 index 000000000..370364017 --- /dev/null +++ b/vllm_omni/engine/duplex/realtime_commands.py @@ -0,0 +1,887 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI Realtime client events -> :class:`DuplexCommand` (stateless). + +This is the mapping half of the former ``entrypoints/duplex/realtime_input.py`` +translator. Everything that needed per-session state in the old translator +(input-buffer emptiness for commits, response-id fallbacks for cancels, +conversation-item lookups, VAD) is resolved by the session runner through the +helpers on :class:`~vllm_omni.engine.duplex.realtime_events.RealtimeProjectionState`; +the commands produced here carry the raw client intent only. +""" + +from __future__ import annotations + +import base64 +import binascii +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Any, cast +from uuid import uuid4 + +import numpy as np + +from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate +from vllm_omni.engine.duplex.commands import ( + AckPlayback, + AppendAudio, + AppendText, + BargeIn, + CancelInput, + CancelResponse, + ClearInput, + ClearOutputAudio, + CloseSession, + Commit, + CreateItem, + CreateResponse, + DeleteItem, + DuplexCommand, + DuplexCommandError, + Heartbeat, + SignalTurn, + TruncateItem, + UpdateSession, +) + +REALTIME_INPUT_AUDIO_FORMATS = { + "pcm16", + "pcm_s16le", + "s16le", + "pcm_f32le", + "g711_ulaw", + "g711_alaw", +} +REALTIME_OUTPUT_AUDIO_FORMATS = { + "pcm16", + "pcm_s16le", + "s16le", + "wav", + "pcm", + "g711_ulaw", + "g711_alaw", +} + +#: Wire hint keys copied verbatim from a client append onto the internal payload. +REALTIME_INPUT_HINT_KEYS = ( + "duration_ms", + "audio_duration_ms", + "audio_start_ms", + "audio_end_ms", + "is_speech", + "speech", + "speech_probability", + "vad", + "overlap_action", + "overlap", + "force_barge_in", + "force_listen", + "text", + "transcript", +) + + +@dataclass(frozen=True, slots=True) +class RealtimeInputDefaults: + """Session-level wire defaults an append may omit (derived from the session object).""" + + input_audio_format: str = "pcm16" + input_sample_rate_hz: int = 16000 + output_audio_format: str = "pcm16" + output_sample_rate_hz: int | None = None + overlap_silence_rms: float = 0.003 + + def with_session_payload(self, session_payload: Mapping[str, object]) -> RealtimeInputDefaults: + """Return defaults updated from a Realtime ``session`` object (session.update).""" + values = apply_realtime_session_defaults(self, session_payload) + return values + + +# ---- audio format helpers ---- + + +def parse_realtime_audio_format(raw_format: object) -> tuple[object, int | None]: + def normalize_format(fmt: str) -> str: + normalized = fmt.lower() + if normalized in {"audio/pcm", "pcm"}: + return "pcm16" + if normalized in {"audio/wav", "wav"}: + return "wav" + if normalized in {"audio/pcm16", "pcm16", "pcm_s16le", "s16le"}: + return "pcm16" + if normalized in {"audio/pcm_f32le", "pcm_f32le", "f32le"}: + return "pcm_f32le" + if normalized in {"audio/g711_ulaw", "g711_ulaw", "g711-ulaw", "ulaw", "mulaw"}: + return "g711_ulaw" + if normalized in {"audio/g711_alaw", "g711_alaw", "g711-alaw", "alaw"}: + return "g711_alaw" + return fmt + + if isinstance(raw_format, str): + return normalize_format(raw_format), None + if not isinstance(raw_format, dict): + return raw_format, None + rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") + sample_rate_hz = int(rate) if isinstance(rate, int | float) and rate > 0 else None + fmt = raw_format.get("type") or raw_format.get("format") + if not isinstance(fmt, str): + return raw_format, sample_rate_hz + return normalize_format(fmt), sample_rate_hz + + +def duplex_response_format(realtime_format: str) -> str: + normalized = realtime_format.lower() + if normalized in {"pcm16", "pcm_s16le", "s16le"}: + return "pcm" + if normalized in {"g711_ulaw", "g711_alaw"}: + return "pcm" + if normalized in {"wav", "pcm"}: + return normalized + return "wav" + + +def realtime_output_format(duplex_format: object) -> str: + if isinstance(duplex_format, str) and duplex_format.lower() in {"g711_ulaw", "g711_alaw"}: + return duplex_format.lower() + if isinstance(duplex_format, str) and duplex_format.lower() == "pcm": + return "pcm16" + return str(duplex_format or "wav") + + +def is_supported_realtime_input_format(fmt: object) -> bool: + return isinstance(fmt, str) and fmt.lower() in REALTIME_INPUT_AUDIO_FORMATS + + +def realtime_audio_format_object(fmt: object, *, sample_rate_hz: int | None = None) -> dict[str, object]: + if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le", "pcm"}: + payload: dict[str, object] = {"type": "audio/pcm"} + elif isinstance(fmt, str) and fmt.lower() == "pcm_f32le": + payload = {"type": "audio/pcm_f32le"} + elif isinstance(fmt, str) and fmt.lower() == "g711_ulaw": + payload = {"type": "audio/g711_ulaw"} + elif isinstance(fmt, str) and fmt.lower() == "g711_alaw": + payload = {"type": "audio/g711_alaw"} + else: + payload = {"type": "audio/wav"} + if sample_rate_hz is not None: + payload["rate"] = int(sample_rate_hz) + return payload + + +def validate_realtime_session_audio_formats(session_payload: Mapping[str, object]) -> str | None: + audio_config = session_payload.get("audio") + input_format: object = session_payload.get("input_audio_format") + if input_format is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + input_format = audio_input.get("format") + parsed_input, _ = parse_realtime_audio_format(input_format) + if input_format is not None and not ( + isinstance(parsed_input, str) and parsed_input.lower() in REALTIME_INPUT_AUDIO_FORMATS + ): + return f"Unsupported input_audio_format: {input_format}" + + output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + parsed_output, _ = parse_realtime_audio_format(output_format) + if output_format is not None and not ( + isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS + ): + return f"Unsupported output_audio_format: {output_format}" + return None + + +def validate_realtime_response_audio_formats(response_payload: Mapping[str, object]) -> str | None: + output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") + audio_config = response_payload.get("audio") + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + parsed_output, _ = parse_realtime_audio_format(output_format) + if output_format is not None and not ( + isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS + ): + return f"Unsupported output_audio_format: {output_format}" + return None + + +def validate_conversation_item_audio_formats(item: object) -> str | None: + if not isinstance(item, dict): + return None + content = item.get("content") + if not isinstance(content, list): + return None + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in {"input_audio", "audio"}: + continue + raw_format = part.get("format") + parsed_format, _ = parse_realtime_audio_format(raw_format) + if raw_format is not None and not ( + isinstance(parsed_format, str) and parsed_format.lower() in REALTIME_INPUT_AUDIO_FORMATS + ): + return f"Unsupported input_audio format in conversation.item.create: {raw_format}" + return None + + +def apply_realtime_session_defaults( + defaults: RealtimeInputDefaults, + session_payload: Mapping[str, object], +) -> RealtimeInputDefaults: + """Derive the wire defaults a Realtime ``session`` object declares.""" + input_format: object = session_payload.get("input_audio_format") + audio_config = session_payload.get("audio") + if input_format is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + input_format = audio_input.get("format") + input_format, input_rate = parse_realtime_audio_format(input_format) + # Any: heterogeneous field values splatted into ``dataclasses.replace`` below. + updates: dict[str, Any] = {} + if isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS: + updates["input_audio_format"] = input_format + output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") + output_rate_raw: object | None = None + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") + output_format, output_rate = parse_realtime_audio_format(output_format) + if output_rate is None and isinstance(output_rate_raw, int | float) and output_rate_raw > 0: + output_rate = int(output_rate_raw) + if isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + updates["output_audio_format"] = realtime_output_format(output_format) + sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") + if sample_rate is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") + if sample_rate is None: + sample_rate = input_rate + if isinstance(sample_rate, int | float) and sample_rate > 0: + updates["input_sample_rate_hz"] = int(sample_rate) + if isinstance(output_rate, int | float) and output_rate > 0: + updates["output_sample_rate_hz"] = int(output_rate) + overlap_fields = realtime_overlap_fields(session_payload) + overlap_silence_rms = overlap_fields.get("overlap_silence_rms") + if isinstance(overlap_silence_rms, int | float): + updates["overlap_silence_rms"] = max(0.0, float(overlap_silence_rms)) + return replace(defaults, **updates) if updates else defaults + + +def realtime_overlap_fields(session_payload: Mapping[str, object]) -> dict[str, object]: + fields: dict[str, object] = {} + if isinstance(session_payload.get("overlap_policy"), str): + fields["overlap_policy"] = session_payload["overlap_policy"] + for key in ("overlap_short_ack_ms", "overlap_barge_in_ms", "overlap_silence_rms"): + value = session_payload.get(key) + if isinstance(value, int | float): + fields[key] = value + if isinstance(session_payload.get("playback_commit_policy"), str): + fields["playback_commit_policy"] = session_payload["playback_commit_policy"] + return fields + + +def json_safe_realtime_payload(payload: Mapping[str, object]) -> dict[str, object]: + clean: dict[str, object] = {} + for key, value in payload.items(): + if key == "extra_body": + continue + if isinstance(value, str | int | float | bool) or value is None: + clean[key] = value + elif isinstance(value, dict): + clean[key] = json_safe_realtime_payload(value) + elif isinstance(value, list): + clean[key] = [ + (json_safe_realtime_payload(item) if isinstance(item, dict) else item) + for item in value + if isinstance(item, str | int | float | bool | dict) or item is None + ] + return clean + + +def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: + transcription = session_payload.get("input_audio_transcription") + if isinstance(transcription, dict): + return transcription + audio_config = session_payload.get("audio") + if not isinstance(audio_config, dict): + return None + audio_input = audio_config.get("input") + if not isinstance(audio_input, dict): + return None + transcription = audio_input.get("transcription") + return transcription if isinstance(transcription, dict) else None + + +def realtime_max_output_tokens(value: object) -> int | None: + """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" + if value is None: + return None + if isinstance(value, str) and value.strip().lower() in {"inf", "infinity", "unlimited"}: + return None + if isinstance(value, int) and value > 0: + return int(value) + return None + + +# ---- speech hints ---- + + +def input_explicitly_non_speech(event: Mapping[str, object]) -> bool: + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return not value + vad = event.get("vad") + if isinstance(vad, dict): + value = vad.get("is_speech") + if isinstance(value, bool): + return not value + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) < 0.5 + probability = event.get("speech_probability") + return isinstance(probability, int | float) and float(probability) < 0.5 + + +def input_looks_like_speech( + event: Mapping[str, object], + *, + audio: object, + fmt: object, + overlap_silence_rms: float, +) -> bool: + if input_explicitly_non_speech(event): + return False + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return value + vad = event.get("vad") + if isinstance(vad, dict): + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) >= 0.5 + probability = event.get("speech_probability") + if isinstance(probability, int | float): + return float(probability) >= 0.5 + if fmt != "pcm_f32le": + return True + if isinstance(audio, bytes | bytearray): + raw = bytes(audio) + elif isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return True + else: + return True + if len(raw) < 4 or len(raw) % 4 != 0: + return True + samples = np.frombuffer(raw, dtype=np.float32) + if samples.size == 0: + return False + rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) + threshold = event.get("overlap_silence_rms") + if not isinstance(threshold, int | float): + vad = event.get("vad") + if isinstance(vad, dict): + threshold = vad.get("silence_rms") + silence_rms = float(threshold) if isinstance(threshold, int | float) else overlap_silence_rms + return rms >= max(0.0, silence_rms) + + +def copy_realtime_input_hints(source: Mapping[str, object], target: dict[str, object]) -> None: + for key in REALTIME_INPUT_HINT_KEYS: + if key in source: + target[key] = source[key] + + +def validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: + """Validate omni-duplex camera frames on input_audio_buffer.append. + + Wire contract matches the official MiniCPM-o duplex loop: one base + base64 JPEG per ~1 s audio chunk, optionally followed by that unit's + stacked composite tiling the sub-frames captured inside it (at most 2 + images either way). A caller-supplied ``max_slice_nums`` is rejected rather + than silently ignored: slicing is Stage 0's decision here, and Stage 0 + already applies the official HD suggestion for a stacked unit + (``max_slice_nums=[2, 1]``). The wire simply does not let the client choose + it. + """ + if max_slice_nums not in (None, 1): + return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" + if not isinstance(video_frames, list): + return "video_frames must be a list of base64-encoded images" + frames = [frame for frame in video_frames if frame is not None] + if len(frames) > 2: + return "video_frames carries more than 2 frames for one append; send ~1 frame per 1 s chunk" + for frame in frames: + if not isinstance(frame, str) or not frame: + return "video_frames entries must be non-empty base64 strings" + if len(frame) > 4_000_000: + return "video_frames entry exceeds 4MB base64; reduce capture resolution or JPEG quality" + try: + header = base64.b64decode(frame[:64] + "=" * (-len(frame[:64]) % 4)) + except (binascii.Error, ValueError): + return "video_frames entries must be valid base64" + if not (header.startswith(b"\xff\xd8") or header.startswith(b"\x89PNG")): + return "video_frames entries must be JPEG or PNG images" + return None + + +# ---- conversation items ---- + + +def normalize_conversation_item(item: Mapping[str, object]) -> dict[str, object]: + normalized = dict(item) + normalized.setdefault("id", f"item_{uuid4().hex}") + normalized.setdefault("object", "realtime.item") + normalized.setdefault("type", "message") + normalized.setdefault("status", "completed") + if "role" not in normalized and normalized.get("type") == "message": + normalized["role"] = "user" + if not isinstance(normalized.get("content"), list): + normalized["content"] = [] + return normalized + + +def text_chars_for_audio_ms_from_marks( + audio_end_ms: int, + text_len: int, + marks: list[object], + *, + final_ms: object | None = None, +) -> int: + if text_len <= 0: + return 0 + clean_marks: list[tuple[int, int]] = [] + for mark in marks: + if not isinstance(mark, dict): + continue + raw_text_chars = mark.get("text_chars") + raw_audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) + if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): + continue + clean_marks.append((max(0, int(raw_audio_end_ms)), min(text_len, max(0, int(raw_text_chars))))) + if not clean_marks: + return 0 if audio_end_ms <= 0 else text_len + clean_marks.sort(key=lambda item: item[0]) + audio_end_ms = max(0, int(audio_end_ms)) + if audio_end_ms <= 0: + return 0 + previous_ms = 0 + previous_chars = 0 + for mark_ms, mark_chars in clean_marks: + mark_ms = max(previous_ms, mark_ms) + mark_chars = max(previous_chars, min(text_len, mark_chars)) + if audio_end_ms <= mark_ms: + if mark_ms <= previous_ms: + return mark_chars + ratio = (audio_end_ms - previous_ms) / max(1, mark_ms - previous_ms) + return int(previous_chars + (mark_chars - previous_chars) * max(0.0, min(1.0, ratio))) + previous_ms = mark_ms + previous_chars = mark_chars + if isinstance(final_ms, int | float) and int(final_ms) > previous_ms: + if audio_end_ms >= int(final_ms): + return text_len + ratio = (audio_end_ms - previous_ms) / max(1, int(final_ms) - previous_ms) + return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) + return text_len if audio_end_ms >= previous_ms else previous_chars + + +def truncate_realtime_item_content(item: dict[str, object], *, content_index: int, audio_end_ms: int) -> None: + content = item.get("content") + if not isinstance(content, list) or not content: + return + index = max(0, int(content_index)) + if index >= len(content): + return + part = content[index] + if not isinstance(part, dict): + return + transcript = part.get("transcript") + if not isinstance(transcript, str) or not transcript: + return + marks = part.get("audio_text_marks") + if isinstance(marks, list): + keep_chars = text_chars_for_audio_ms_from_marks( + audio_end_ms, + len(transcript), + marks, + final_ms=part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms"), + ) + part["transcript"] = transcript[:keep_chars].rstrip() + return + duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") + if isinstance(duration_ms, int | float) and duration_ms > 0: + keep_chars = int(len(transcript) * max(0.0, min(1.0, int(audio_end_ms) / float(duration_ms)))) + part["transcript"] = transcript[:keep_chars].rstrip() + elif audio_end_ms <= 0: + part["transcript"] = "" + + +def validate_realtime_item_truncate(item: Mapping[str, object], *, content_index: int, audio_end_ms: int) -> str | None: + if item.get("type") != "message" or item.get("role") != "assistant": + return "conversation.item.truncate only supports assistant message items" + if audio_end_ms < 0: + return "conversation.item.truncate requires non-negative audio_end_ms" + content = item.get("content") + if not isinstance(content, list) or not content: + return None + index = max(0, int(content_index)) + if index >= len(content): + return f"conversation.item.truncate content_index out of range: {content_index}" + part = content[index] + if not isinstance(part, dict): + return "conversation.item.truncate target content part is invalid" + if part.get("type") not in {"audio", "output_audio"}: + return "conversation.item.truncate target content part must be audio" + duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") + if isinstance(duration_ms, int | float) and int(duration_ms) >= 0 and audio_end_ms > int(duration_ms): + return "conversation.item.truncate audio_end_ms exceeds item audio duration" + return None + + +def input_transcript_from_item(item: Mapping[str, object]) -> str: + content = item.get("content") + if not isinstance(content, list): + return "" + parts: list[str] = [] + for part in content: + if not isinstance(part, dict): + continue + for key in ("transcript", "text"): + value = part.get(key) + if isinstance(value, str) and value: + if part.get("type") in {"input_audio", "audio", "audio_transcript", "transcript"}: + parts.append(value) + break + return "".join(parts).strip() + + +# ---- append / item audio conversion ---- + + +def build_append_audio( + event: Mapping[str, object], + *, + defaults: RealtimeInputDefaults, + hints_source: Mapping[str, object] | None = None, +) -> AppendAudio: + """Validate, convert (to 16 kHz ``pcm_f32le`` base64) and pack one audio append. + + Speech classification without VAD (``is_speech`` on the command) uses the + explicit wire hints and the RMS fallback exactly like the old translator. + """ + event_id = cast("str", event.get("event_id")) if isinstance(event.get("event_id"), str) else None + audio = event.get("audio") or event.get("delta") + fmt, format_rate = parse_realtime_audio_format( + event.get("format") or event.get("input_audio_format") or defaults.input_audio_format + ) + sample_rate_hz = ( + event.get("sample_rate_hz") or event.get("sample_rate") or format_rate or defaults.input_sample_rate_hz + ) + if not is_supported_realtime_input_format(fmt): + raise DuplexCommandError( + f"Unsupported input_audio_format: {fmt}", + code="unsupported_audio_format", + event_id=event_id, + ) + try: + audio, fmt, sample_rate_hz = convert_input_audio_with_rate( + audio, + fmt, + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, + ) + except ValueError as exc: + raise DuplexCommandError(str(exc), code="bad_event", event_id=event_id) from exc + hints: dict[str, object] = {} + if hints_source is not None: + copy_realtime_input_hints(hints_source, hints) + copy_realtime_input_hints(event, hints) + looks_like_speech = input_looks_like_speech( + {**(hints_source or {}), **event}, + audio=audio, + fmt=fmt, + overlap_silence_rms=defaults.overlap_silence_rms, + ) + video_frames: tuple[str, ...] = () + raw_frames = event.get("video_frames") + if raw_frames is not None: + frames_error = validate_realtime_video_frames(raw_frames, event.get("max_slice_nums")) + if frames_error is not None: + raise DuplexCommandError(frames_error, code="invalid_video_frames", event_id=event_id) + if isinstance(raw_frames, list): + video_frames = tuple(frame for frame in raw_frames if isinstance(frame, str) and frame) + duration_ms = hints.get("duration_ms", hints.get("audio_duration_ms")) + audio_end_ms = hints.get("audio_end_ms") + try: + audio_bytes = base64.b64decode(audio, validate=True) if isinstance(audio, str) and audio else b"" + except (binascii.Error, ValueError) as exc: + raise DuplexCommandError("input audio is not valid base64", code="bad_audio", event_id=event_id) from exc + return AppendAudio( + event_id=event_id, + audio=audio_bytes, + format=str(fmt), + sample_rate_hz=int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else None, + is_speech=looks_like_speech, + video_frames=video_frames, + duration_ms=int(duration_ms) if isinstance(duration_ms, int | float) else None, + audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else None, + hints=hints, + ) + + +# ---- translation ---- + + +def translate_realtime_command( + payload: Mapping[str, object], + *, + defaults: RealtimeInputDefaults | None = None, +) -> DuplexCommand: + """Map one OpenAI Realtime client event onto a :class:`DuplexCommand`. + + Raises :class:`DuplexCommandError` for malformed or unsupported payloads. + ``session.resume`` and ``session.event_ack`` are transport concerns and are + rejected with ``code="unknown_event"``. + """ + defaults = defaults or RealtimeInputDefaults() + event_type = payload.get("type") + event_id = cast("str", payload.get("event_id")) if isinstance(payload.get("event_id"), str) else None + if not isinstance(event_type, str): + raise DuplexCommandError("Duplex event missing string type", code="bad_event", event_id=event_id) + + if event_type == "session.update": + session = payload.get("session") + session_payload: Mapping[str, object] = session if isinstance(session, dict) else payload + format_error = validate_realtime_session_audio_formats(session_payload) + if format_error is not None: + raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) + from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection + + turn_detection_error = validate_realtime_turn_detection(session_payload) + if turn_detection_error is not None: + raise DuplexCommandError(turn_detection_error, code="unsupported_turn_detection", event_id=event_id) + return UpdateSession(event_id=event_id, patch=dict(session_payload)) + + if event_type == "conversation.item.create": + item = payload.get("item") + format_error = validate_conversation_item_audio_formats(item) + if format_error is not None: + raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) + if not isinstance(item, dict): + raise DuplexCommandError("conversation.item.create requires item", code="bad_event", event_id=event_id) + previous_item_id = payload.get("previous_item_id") + return CreateItem( + event_id=event_id, + item=normalize_conversation_item(item), + previous_item_id=previous_item_id if isinstance(previous_item_id, str) else None, + ) + + if event_type == "conversation.item.delete": + item_id = payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + raise DuplexCommandError( + "conversation.item.delete requires item_id", code="missing_item_id", event_id=event_id + ) + return DeleteItem(event_id=event_id, item_id=item_id) + + if event_type == "conversation.item.truncate": + item_id = payload.get("item_id") + audio_end_ms = payload.get("audio_end_ms") + content_index = payload.get("content_index", 0) + if not isinstance(item_id, str) or not item_id: + raise DuplexCommandError( + "conversation.item.truncate requires item_id", code="missing_item_id", event_id=event_id + ) + if not isinstance(audio_end_ms, int | float): + raise DuplexCommandError( + "conversation.item.truncate requires numeric audio_end_ms", code="bad_event", event_id=event_id + ) + return TruncateItem( + event_id=event_id, + item_id=item_id, + audio_end_ms=int(audio_end_ms), + content_index=int(content_index) if isinstance(content_index, int | float) else 0, + ) + + if event_type == "input_audio_buffer.append": + return build_append_audio(dict(payload), defaults=defaults) + + if event_type in {"input_audio_buffer.commit", "input.commit"}: + final = payload.get("final", True) + create_response = payload.get("response_create", payload.get("create_response")) + is_speech = payload.get("is_speech") + return Commit( + event_id=event_id, + final=bool(final) if isinstance(final, bool) else True, + create_response=( + bool(create_response) + if isinstance(create_response, bool) + else (True if event_type == "input.commit" and create_response is None else None) + ), + is_speech=is_speech if isinstance(is_speech, bool) else None, + ) + + if event_type == "input_audio_buffer.clear": + return ClearInput(event_id=event_id) + + if event_type == "output_audio_buffer.clear": + response_id = payload.get("response_id") + return ClearOutputAudio( + event_id=event_id, + response_id=response_id if isinstance(response_id, str) and response_id else None, + ) + + if event_type == "response.cancel": + response_id = payload.get("response_id") + return CancelResponse( + event_id=event_id, + response_id=response_id if isinstance(response_id, str) and response_id else None, + ) + + if event_type == "response.create": + response_payload = payload.get("response") + if isinstance(response_payload, dict): + format_error = validate_realtime_response_audio_formats(response_payload) + if format_error is not None: + raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) + return CreateResponse( + event_id=event_id, + options=dict(response_payload) if isinstance(response_payload, dict) else {}, + ) + + if event_type in {"playback.ack", "audio.playback_ack"}: + played_ms = payload.get("played_ms") + if not isinstance(played_ms, int | float): + raise DuplexCommandError("playback.ack requires numeric played_ms", code="bad_event", event_id=event_id) + committed_ms = payload.get("committed_ms") + response_id = payload.get("response_id") + item_id = payload.get("item_id") + return AckPlayback( + event_id=event_id, + played_ms=int(played_ms), + committed_ms=int(committed_ms) if isinstance(committed_ms, int | float) else None, + response_id=response_id if isinstance(response_id, str) and response_id else None, + item_id=item_id if isinstance(item_id, str) and item_id else None, + ) + + if event_type == "session.heartbeat": + return Heartbeat(event_id=event_id) + + if event_type == "conversation.item.retrieve": + # The projected conversation items live engine-side; the runner answers. + return SignalTurn(event_id=event_id, event="conversation.item.retrieve", signal_payload=dict(payload)) + + if event_type in {"session.close", "close", "close_session"}: + reason = payload.get("reason") + return CloseSession(event_id=event_id, reason=reason if isinstance(reason, str) and reason else "client_close") + + if event_type in {"input.text.append", "input_text.append", "push_text"}: + text = payload.get("text") + if not isinstance(text, str): + raise DuplexCommandError("input.text.append requires text", code="bad_event", event_id=event_id) + return AppendText(event_id=event_id, text=text) + + if event_type == "input.cancel": + return CancelInput(event_id=event_id) + + if event_type == "barge_in": + return BargeIn(event_id=event_id) + + if event_type in {"turn.signal", "signal_turn"}: + signal_event = payload.get("event") + if not isinstance(signal_event, str) or not signal_event: + raise DuplexCommandError("turn.signal requires event", code="bad_event", event_id=event_id) + signal_payload = payload.get("payload") + signal_payload = dict(signal_payload) if isinstance(signal_payload, dict) else {} + if signal_event == "input.cancel": + return CancelInput(event_id=event_id) + if signal_event == "barge_in": + return BargeIn(event_id=event_id) + if signal_event == "response.cancel": + response_id = signal_payload.get("response_id", payload.get("response_id")) + return CancelResponse( + event_id=event_id, + response_id=response_id if isinstance(response_id, str) and response_id else None, + ) + if signal_event == "session.update": + return UpdateSession(event_id=event_id, patch=signal_payload) + if signal_event == "conversation.item.create": + item = signal_payload.get("item") + if not isinstance(item, dict): + raise DuplexCommandError("conversation.item.create requires item", code="bad_event", event_id=event_id) + previous_item_id = signal_payload.get("previous_item_id") + return CreateItem( + event_id=event_id, + item=normalize_conversation_item(item), + previous_item_id=previous_item_id if isinstance(previous_item_id, str) else None, + ) + if signal_event == "conversation.item.delete": + item_id = signal_payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + raise DuplexCommandError( + "conversation.item.delete requires item_id", code="missing_item_id", event_id=event_id + ) + return DeleteItem(event_id=event_id, item_id=item_id) + if signal_event == "conversation.item.truncate": + item_id = signal_payload.get("item_id") + audio_end_ms = signal_payload.get("audio_end_ms") + content_index = signal_payload.get("content_index", 0) + if not isinstance(item_id, str) or not item_id or not isinstance(audio_end_ms, int | float): + raise DuplexCommandError( + "conversation.item.truncate requires item_id and numeric audio_end_ms", + code="bad_event", + event_id=event_id, + ) + return TruncateItem( + event_id=event_id, + item_id=item_id, + audio_end_ms=int(audio_end_ms), + content_index=int(content_index) if isinstance(content_index, int | float) else 0, + ) + return SignalTurn(event_id=event_id, event=signal_event, signal_payload=signal_payload) + + raise DuplexCommandError(f"Unknown duplex event type: {event_type}", code="unknown_event", event_id=event_id) + + +__all__ = [ + "REALTIME_INPUT_AUDIO_FORMATS", + "REALTIME_INPUT_HINT_KEYS", + "REALTIME_OUTPUT_AUDIO_FORMATS", + "RealtimeInputDefaults", + "apply_realtime_session_defaults", + "build_append_audio", + "copy_realtime_input_hints", + "duplex_response_format", + "input_audio_transcription_config", + "input_explicitly_non_speech", + "input_looks_like_speech", + "input_transcript_from_item", + "is_supported_realtime_input_format", + "json_safe_realtime_payload", + "normalize_conversation_item", + "parse_realtime_audio_format", + "realtime_audio_format_object", + "realtime_max_output_tokens", + "realtime_output_format", + "realtime_overlap_fields", + "text_chars_for_audio_ms_from_marks", + "translate_realtime_command", + "truncate_realtime_item_content", + "validate_conversation_item_audio_formats", + "validate_realtime_item_truncate", + "validate_realtime_response_audio_formats", + "validate_realtime_session_audio_formats", + "validate_realtime_video_frames", +] diff --git a/vllm_omni/engine/duplex/realtime_events.py b/vllm_omni/engine/duplex/realtime_events.py new file mode 100644 index 000000000..91359f297 --- /dev/null +++ b/vllm_omni/engine/duplex/realtime_events.py @@ -0,0 +1,1526 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session-internal duplex events -> typed :class:`DuplexEvent` objects (stateful projection). + +This is the former ``entrypoints/duplex/realtime_output.py`` projector plus the +projection-relevant fields of ``realtime_state.py``, now owned by the session +runner through :class:`RealtimeProjectionState`. The projection consumes the +state (response / item ids, content-part bookkeeping) when it *constructs* +events; rendering an event to wire JSON (``event.to_realtime()``) is pure and +lives on the event classes in ``vllm_omni.engine.duplex.events``. + +Besides the output projection (:func:`project_internal_event`) the state also +carries the input-side bookkeeping the old input translator kept (input buffer +flags, conversation items, response-id fallbacks); the ``resolve_*`` / +``note_*`` helpers give the runner the same behaviour for the corresponding +commands. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, cast +from uuid import uuid4 + +from vllm_omni.engine.duplex.audio import convert_output_audio +from vllm_omni.engine.duplex.commands import ( + AppendAudio, + CancelResponse, + ClearOutputAudio, + Commit, + CreateItem, + DeleteItem, + TruncateItem, +) +from vllm_omni.engine.duplex.events import ( + AudioDelta, + AudioDone, + ContentPartAdded, + ContentPartDone, + DuplexEvent, + DuplexRawEvent, + FunctionCallArgumentsDelta, + FunctionCallArgumentsDone, + InputCleared, + InputCommitted, + InputTranscriptionCompleted, + ItemAdded, + ItemCreated, + ItemDeleted, + ItemDone, + ItemRetrieved, + ItemTruncated, + Listen, + OutputAudioCleared, + OutputItemAdded, + OutputItemDone, + RateLimitsUpdated, + ResponseCreated, + ResponseDone, + SessionClosed, + SessionCreated, + SessionReplaced, + SessionResumed, + SessionResyncRequired, + SessionUpdated, + Speak, + SpeechStarted, + SpeechStopped, + TextDelta, + TextDone, + TranscriptDelta, + TranscriptDone, + error_event, +) +from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + apply_realtime_session_defaults, + build_append_audio, + copy_realtime_input_hints, + input_looks_like_speech, + input_transcript_from_item, + parse_realtime_audio_format, + realtime_audio_format_object, + realtime_output_format, + truncate_realtime_item_content, + validate_realtime_item_truncate, +) + +if TYPE_CHECKING: + # Runtime import pulls in the VAD backend; the projection only needs the type. + from vllm_omni.engine.duplex.turn_detection import TurnDetectionResult + + +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _int_or(value: object, default: int = 0) -> int: + return int(value) if isinstance(value, int | float) else default + + +@dataclass(slots=True) +class _ResponseProjection: + item_id: str + transcript_parts: list[str] = field(default_factory=list) + text_parts: list[str] = field(default_factory=list) + audio_duration_ms: int | None = None + audio_text_marks: list[dict[str, int]] = field(default_factory=list) + audio_delta_emitted: bool = False + audio_done_emitted: bool = False + audio_part_added: bool = False + audio_part_done: bool = False + text_part_added: bool = False + text_part_done: bool = False + output_text_done: bool = False + output_item_done: bool = False + conversation_item_done: bool = False + speak_emitted: bool = False + done_emitted: bool = False + + @property + def transcript(self) -> str: + return "".join(self.transcript_parts) + + @property + def text(self) -> str: + return "".join(self.text_parts) + + +@dataclass(slots=True) +class RealtimeProjectionState: + """Per-session Realtime projection state (lives on the engine-side runner).""" + + session_id: str + model: str | None = None + default_payload: Mapping[str, object] | None = None + #: Wire defaults (input/output formats, rates, silence rms) declared by the session object. + defaults: RealtimeInputDefaults = field(default_factory=RealtimeInputDefaults) + #: The first ``session.created`` also emits ``session.updated`` (client sent session.update to open). + initial_session_update: bool = False + # ---- response / item projection ---- + response_states: dict[str | int, _ResponseProjection] = field(default_factory=dict) + item_truncation_cursors: dict[str, tuple[int, int]] = field(default_factory=dict) + active_response_id: str | None = None + last_response_id: str | None = None + conversation_items: dict[str, dict[str, object]] = field(default_factory=dict) + last_conversation_item_id: str | None = None + pending_commit_item_ids: list[str] = field(default_factory=list) + # ---- input buffer projection (was on the input translator) ---- + input_speech_started: bool = False + active_input_item_id: str | None = None + input_audio_buffer_has_audio: bool = False + input_audio_buffer_had_non_speech: bool = False + input_audio_buffer_transcript_parts: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.default_payload is not None: + self.apply_session_defaults(self.default_payload) + + # ---- session defaults ---- + + def apply_session_defaults(self, session_payload: Mapping[str, object]) -> None: + self.defaults = apply_realtime_session_defaults(self.defaults, session_payload) + + @property + def input_audio_format(self) -> str: + return self.defaults.input_audio_format + + @property + def input_sample_rate_hz(self) -> int: + return self.defaults.input_sample_rate_hz + + @property + def output_audio_format(self) -> str: + return self.defaults.output_audio_format + + @property + def output_sample_rate_hz(self) -> int | None: + return self.defaults.output_sample_rate_hz + + @property + def overlap_silence_rms(self) -> float: + return self.defaults.overlap_silence_rms + + +# ---- response projection helpers ---- + + +def _response_state( + state: RealtimeProjectionState, + response_id: object, + *, + event: Mapping[str, object] | None = None, + create: bool = True, +) -> _ResponseProjection | None: + if not (isinstance(response_id, str) and response_id): + # An event without a response id has no projection to attach to; a + # key derived from the transient event object would alias later + # anonymous events (and leak an entry per terminal). + return None + key = response_id + item_id = f"item_{response_id}" + projection = state.response_states.get(key) + if projection is None and create: + projection = _ResponseProjection(item_id=item_id) + state.response_states[key] = projection + return projection + + +def response_is_done(state: RealtimeProjectionState, response_id: object) -> bool: + projection = _response_state(state, response_id, create=False) + return projection is not None and projection.done_emitted + + +def _response_item_id(state: RealtimeProjectionState, response_id: object) -> str: + projection = _response_state(state, response_id) + if projection is not None: + return projection.item_id + return f"item_{uuid4().hex}" + + +def _response_content_part(*, transcript: str = "") -> dict[str, object]: + return {"type": "audio", "transcript": transcript} + + +def _response_text_content_part(*, text: str = "") -> dict[str, object]: + return {"type": "text", "text": text} + + +def _response_item_content_part( + *, + transcript: str = "", + audio_duration_ms: int | None = None, + audio_text_marks: list[dict[str, int]] | None = None, +) -> dict[str, object]: + part: dict[str, object] = {"type": "output_audio", "transcript": transcript} + if audio_duration_ms is not None: + part["audio_duration_ms"] = int(audio_duration_ms) + if audio_text_marks: + part["audio_text_marks"] = [dict(mark) for mark in audio_text_marks] + return part + + +def _response_item_text_content_part(*, text: str = "") -> dict[str, object]: + return {"type": "output_text", "text": text} + + +def _previous_item_id(state: RealtimeProjectionState, item_id: str) -> str | None: + previous: str | None = None + for known_id in state.conversation_items: + if known_id == item_id: + return previous + previous = known_id + if state.last_conversation_item_id == item_id: + return previous + return state.last_conversation_item_id + + +def _conversation_item_added_events(state: RealtimeProjectionState, item: dict[str, object]) -> list[DuplexEvent]: + item_id = item.get("id") + explicit_previous_item_id = item.pop("_previous_item_id", None) + previous_item_id = ( + explicit_previous_item_id if isinstance(explicit_previous_item_id, str) else state.last_conversation_item_id + ) + if isinstance(item_id, str) and item_id: + state.last_conversation_item_id = item_id + return [ + ItemAdded(previous_item_id=previous_item_id, item=item), + ItemCreated(previous_item_id=previous_item_id, item=item), + ] + + +def _conversation_item_done_event(state: RealtimeProjectionState, item: dict[str, object]) -> ItemDone: + item_id = item.get("id") + previous_item_id = _previous_item_id(state, item_id) if isinstance(item_id, str) else None + if isinstance(item_id, str) and item_id: + state.conversation_items[item_id] = item + state.last_conversation_item_id = item_id + return ItemDone(previous_item_id=previous_item_id, item=item) + + +def _remove_conversation_item(state: RealtimeProjectionState, item_id: str) -> bool: + removed = state.conversation_items.pop(item_id, None) is not None + if state.last_conversation_item_id == item_id: + remaining_ids = list(state.conversation_items) + state.last_conversation_item_id = remaining_ids[-1] if remaining_ids else None + state.item_truncation_cursors.pop(item_id, None) + return removed + + +def _apply_pending_item_truncation(state: RealtimeProjectionState, item: dict[str, object]) -> None: + item_id = item.get("id") + if not isinstance(item_id, str) or not item_id: + return + cursor = state.item_truncation_cursors.get(item_id) + if cursor is None: + return + content_index, audio_end_ms = cursor + truncate_realtime_item_content(item, content_index=content_index, audio_end_ms=audio_end_ms) + + +def _response_done_output_item( + state: RealtimeProjectionState, response_id: object, *, status: str +) -> dict[str, object]: + item_id = _response_item_id(state, response_id) + projection = _response_state(state, response_id) + transcript = projection.transcript if projection is not None else "" + text = projection.text if projection is not None else "" + audio_duration_ms = projection.audio_duration_ms if projection is not None else None + audio_text_marks = projection.audio_text_marks if projection is not None else None + content: list[dict[str, object]] = [] + if (projection is not None and projection.audio_part_added) or transcript or audio_duration_ms is not None: + content.append( + _response_item_content_part( + transcript=transcript, + audio_duration_ms=audio_duration_ms, + audio_text_marks=audio_text_marks, + ) + ) + if text: + content.append(_response_item_text_content_part(text=text)) + item = { + "id": item_id, + "object": "realtime.item", + "type": "message", + "role": "assistant", + "status": status, + "content": content, + } + _apply_pending_item_truncation(state, item) + return item + + +def _refresh_in_progress_response_item(state: RealtimeProjectionState, response_id: object) -> None: + if not isinstance(response_id, str) or not response_id: + return + item_id = _response_item_id(state, response_id) + item = state.conversation_items.get(item_id) + if not isinstance(item, dict): + return + content = item.get("content") + if not isinstance(content, list): + content = [] + item["content"] = content + projection = _response_state(state, response_id) + if projection is None: + return + transcript = projection.transcript + audio_duration_ms = projection.audio_duration_ms + audio_text_marks = projection.audio_text_marks + has_audio = projection.audio_part_added or bool(transcript) or audio_duration_ms is not None + if has_audio: + audio_part = _response_item_content_part( + transcript=transcript, + audio_duration_ms=audio_duration_ms, + audio_text_marks=audio_text_marks, + ) + if content and isinstance(content[0], dict) and content[0].get("type") in {"audio", "output_audio"}: + content[0] = audio_part + else: + content.insert(0, audio_part) + text = projection.text + if text: + text_index = ( + 1 if content and isinstance(content[0], dict) and content[0].get("type") in {"audio", "output_audio"} else 0 + ) + text_part = _response_item_text_content_part(text=text) + if ( + len(content) > text_index + and isinstance(content[text_index], dict) + and content[text_index].get("type") in {"text", "output_text"} + ): + content[text_index] = text_part + else: + content.insert(text_index, text_part) + _apply_pending_item_truncation(state, item) + + +def _append_response_transcript(state: RealtimeProjectionState, response_id: object, text: str) -> None: + projection = _response_state(state, response_id) + if projection is None or not text: + return + projection.transcript_parts.append(text) + + +def _ensure_response_text_part_added(state: RealtimeProjectionState, response_id: object) -> list[DuplexEvent]: + projection = _response_state(state, response_id) + if projection is None or projection.text_part_added: + return [] + projection.text_part_added = True + return [ + ContentPartAdded( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + content_index=1 if projection.audio_part_added else 0, + part=_response_text_content_part(), + ) + ] + + +def _ensure_response_audio_part_added(state: RealtimeProjectionState, response_id: object) -> list[DuplexEvent]: + projection = _response_state(state, response_id) + if projection is None or projection.audio_part_added: + return [] + projection.audio_part_added = True + return [ + ContentPartAdded( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + part=_response_content_part(), + ) + ] + + +def _remember_response_audio_metadata( + state: RealtimeProjectionState, response_id: object, event: Mapping[str, object] +) -> None: + projection = _response_state(state, response_id) + if projection is None: + return + duration = event.get("audio_duration_ms") + playback = event.get("playback") + if not isinstance(duration, int | float) and isinstance(playback, dict): + duration = playback.get("sent_ms") or playback.get("generated_ms") + if isinstance(duration, int | float): + projection.audio_duration_ms = max(projection.audio_duration_ms or 0, int(duration)) + marks = event.get("audio_text_marks") + if not isinstance(marks, list): + return + clean_marks: list[dict[str, int]] = [] + for mark in marks: + if not isinstance(mark, dict): + continue + text_chars = mark.get("text_chars") + audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) + if not isinstance(text_chars, int | float) or not isinstance(audio_end_ms, int | float): + continue + clean_marks.append({"text_chars": max(0, int(text_chars)), "audio_end_ms": max(0, int(audio_end_ms))}) + if clean_marks: + merged = list(projection.audio_text_marks) + merged.extend(clean_marks) + deduped: dict[tuple[int, int], dict[str, int]] = {} + for mark in merged: + deduped[(int(mark["audio_end_ms"]), int(mark["text_chars"]))] = mark + projection.audio_text_marks = sorted( + deduped.values(), key=lambda mark: (mark["audio_end_ms"], mark["text_chars"]) + ) + + +def _response_created_event(event: Mapping[str, object]) -> ResponseCreated: + response_id = event.get("response_id") + metadata = event.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + metadata = {**metadata, "duplex_event": dict(event)} + return ResponseCreated( + response_id=_str_or_none(response_id), + response={ + "id": response_id, + "object": "realtime.response", + "status": "in_progress", + "status_details": None, + "output": [], + "modalities": event.get("modalities") or ["audio", "text"], + "metadata": metadata, + }, + ) + + +def _response_speak_metadata(event: Mapping[str, object]) -> dict[str, object]: + return {key: event[key] for key in ("session_id", "epoch", "model_speak") if key in event} + + +def _realtime_audio_delta_events( + state: RealtimeProjectionState, + event: Mapping[str, object], + response_id: object, + audio: str, +) -> list[DuplexEvent]: + item_id = _response_item_id(state, response_id) + fmt, format_rate = parse_realtime_audio_format(event.get("format", "wav")) + source_fmt = realtime_output_format(fmt) + source_sample_rate_hz = event.get("sample_rate_hz") or format_rate + target_sample_rate_hz = ( + state.output_sample_rate_hz if state.output_audio_format in {"g711_ulaw", "g711_alaw"} else None + ) + audio, fmt, converted_sample_rate_hz = convert_output_audio( + audio, + source_fmt=source_fmt, + target_fmt=state.output_audio_format, + source_sample_rate_hz=(int(source_sample_rate_hz) if isinstance(source_sample_rate_hz, int | float) else None), + target_sample_rate_hz=target_sample_rate_hz, + ) + sample_rate_hz = ( + converted_sample_rate_hz + if fmt in {"g711_ulaw", "g711_alaw"} + else event.get("sample_rate_hz") or format_rate or state.output_sample_rate_hz + ) + metadata: dict[str, object] = {} + for key in ("session_id", "epoch", "model_speak", "end_of_turn", "playback", "vllm_omni"): + if key in event: + metadata[key] = event[key] + duration_ms = event.get("audio_duration_ms") + if isinstance(duration_ms, int | float): + metadata["audio_duration_ms"] = int(duration_ms) + marks = event.get("audio_text_marks") + if isinstance(marks, list): + metadata["audio_text_marks"] = marks + projection = _response_state(state, response_id) + if projection is not None: + projection.audio_delta_emitted = True + _remember_response_audio_metadata(state, response_id, event) + events: list[DuplexEvent] = [] + if projection is not None and not projection.speak_emitted and metadata.get("model_speak") is True: + projection.speak_emitted = True + events.append( + Speak(response_id=_str_or_none(response_id), item_id=item_id, metadata=_response_speak_metadata(event)) + ) + events.append( + AudioDelta( + response_id=_str_or_none(response_id), + item_id=item_id, + delta=str(audio), + format=str(fmt), + sample_rate_hz=int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else None, + metadata=metadata or None, + ) + ) + return events + + +def _realtime_audio_done_events( + state: RealtimeProjectionState, event: Mapping[str, object], response_id: object +) -> list[DuplexEvent]: + item_id = _response_item_id(state, response_id) + projection = _response_state(state, response_id, event=event) + transcript = projection.transcript if projection is not None else "" + events: list[DuplexEvent] = [] + if ( + isinstance(response_id, str) + and projection is not None + and not projection.audio_delta_emitted + and not transcript + ): + return events + if projection is None or not projection.audio_done_emitted: + if projection is not None: + projection.audio_done_emitted = True + events.append(AudioDone(response_id=_str_or_none(response_id), item_id=item_id)) + if transcript: + events.append(TranscriptDone(response_id=_str_or_none(response_id), item_id=item_id, transcript=transcript)) + return events + + +def _realtime_response_done_event( + state: RealtimeProjectionState, + event: Mapping[str, object], + *, + projection: _ResponseProjection | None = None, + status: str = "completed", + status_details: dict[str, object] | None = None, +) -> ResponseDone | None: + response_id = event.get("response_id") + if projection is None: + projection = _response_state(state, response_id, event=event) + if projection is not None: + if projection.done_emitted: + return None + projection.done_emitted = True + return ResponseDone( + response_id=_str_or_none(response_id), + response={ + "id": response_id, + "object": "realtime.response", + "status": status, + "status_details": status_details, + "output": [_response_done_output_item(state, response_id, status=status)], + "metadata": dict(event), + }, + ) + + +def _realtime_response_terminal_events( + state: RealtimeProjectionState, + event: Mapping[str, object], + response_id: object, + *, + status: str = "completed", + status_details: dict[str, object] | None = None, +) -> list[DuplexEvent]: + item_id = _response_item_id(state, response_id) + projection = _response_state(state, response_id, event=event) + transcript = projection.transcript if projection is not None else "" + rid = _str_or_none(response_id) + events: list[DuplexEvent] = [] + if projection is not None and projection.audio_part_added and not projection.audio_part_done: + projection.audio_part_done = True + events.append( + ContentPartDone(response_id=rid, item_id=item_id, part=_response_content_part(transcript=transcript)) + ) + if projection is not None and projection.text_parts and not projection.output_text_done: + projection.output_text_done = True + events.append( + TextDone( + response_id=rid, + item_id=item_id, + content_index=1 if projection.audio_part_added else 0, + text=projection.text, + ) + ) + if projection is not None and projection.text_parts and not projection.text_part_done: + projection.text_part_done = True + events.append( + ContentPartDone( + response_id=rid, + item_id=item_id, + content_index=1 if projection.audio_part_added else 0, + part=_response_text_content_part(text=projection.text), + ) + ) + if projection is None or not projection.output_item_done: + if projection is not None: + projection.output_item_done = True + item = _response_done_output_item(state, response_id, status=status) + events.append(OutputItemDone(response_id=rid, item=item)) + if projection is None or not projection.conversation_item_done: + if projection is not None: + projection.conversation_item_done = True + events.append(_conversation_item_done_event(state, item)) + done_event = _realtime_response_done_event( + state, + {**event, "response_id": response_id}, + projection=projection, + status=status, + status_details=status_details, + ) + if done_event is not None: + events.append(done_event) + # vLLM-Omni does not expose a Realtime-specific quota budget; emit the + # terminal event with an empty list so clients that sequence on it work. + events.append(RateLimitsUpdated()) + if isinstance(response_id, str) and response_id == state.active_response_id: + state.active_response_id = None + return events + + +def _user_item_content_from_duplex_message(message: object) -> list[dict[str, object]]: + if not isinstance(message, dict): + return [{"type": "input_audio", "transcript": ""}] + message_transcript = message.get("transcript") + content = message.get("content") + if isinstance(content, str): + return [{"type": "input_text", "text": content}] + if not isinstance(content, list): + return [ + {"type": "input_audio", "transcript": message_transcript if isinstance(message_transcript, str) else ""} + ] + parts: list[dict[str, object]] = [] + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in {"text", "input_text"} and isinstance(part.get("text"), str): + parts.append({"type": "input_text", "text": part["text"]}) + continue + if part_type == "audio_url": + transcript = ( + part.get("transcript") + if isinstance(part.get("transcript"), str) + else message_transcript + if isinstance(message_transcript, str) + else "" + ) + parts.append({"type": "input_audio", "transcript": transcript}) + continue + if part_type in {"audio", "input_audio"}: + transcript = part.get("transcript") if isinstance(part.get("transcript"), str) else "" + parts.append({"type": "input_audio", "transcript": transcript}) + return parts or [{"type": "input_audio", "transcript": ""}] + + +def _input_audio_transcription_completed_event( + item_id: str, item: Mapping[str, object] +) -> InputTranscriptionCompleted | None: + content = item.get("content") + if not isinstance(content, list): + return None + for index, part in enumerate(content): + if not isinstance(part, dict) or part.get("type") != "input_audio": + continue + transcript = part.get("transcript") + if isinstance(transcript, str) and transcript: + return InputTranscriptionCompleted(item_id=item_id, content_index=index, transcript=transcript) + return None + + +def _pop_pending_commit_item_id(state: RealtimeProjectionState) -> str: + if state.pending_commit_item_ids: + return state.pending_commit_item_ids.pop(0) + return f"item_{uuid4().hex}" + + +def realtime_session_payload(state: RealtimeProjectionState, session: object) -> dict[str, object]: + """Fill the Realtime ``session`` object defaults for session.created/updated.""" + payload = dict(session) if isinstance(session, Mapping) else {} + payload.setdefault("object", "realtime.session") + payload.setdefault("type", "realtime") + payload.setdefault("id", payload.get("id") or state.session_id) + payload.setdefault("model", payload.get("model") or state.model) + payload.setdefault("input_audio_format", state.input_audio_format) + payload.setdefault("output_audio_format", state.output_audio_format) + payload.setdefault("modalities", payload.get("modalities") or ["text", "audio"]) + payload.setdefault("output_modalities", payload.get("output_modalities") or payload.get("modalities")) + payload.setdefault( + "audio", + { + "input": { + "format": realtime_audio_format_object( + state.input_audio_format, sample_rate_hz=state.input_sample_rate_hz + ), + "sample_rate_hz": state.input_sample_rate_hz, + }, + "output": { + "format": realtime_audio_format_object( + state.output_audio_format, sample_rate_hz=state.output_sample_rate_hz + ), + }, + }, + ) + payload.setdefault("turn_detection", payload.get("turn_detection")) + payload.setdefault("input_audio_transcription", payload.get("input_audio_transcription")) + payload.setdefault("tracing", payload.get("tracing")) + return payload + + +# ---- main projection ---- + + +def project_internal_event(state: RealtimeProjectionState, event: Mapping[str, object]) -> list[DuplexEvent]: + """Project one session-internal event onto 0..n typed public events.""" + return _project(state, dict(event)) + + +def _project(state: RealtimeProjectionState, event: dict[str, object]) -> list[DuplexEvent]: + event_type = event.get("type") + if event_type == "session.created": + session = realtime_session_payload(state, event.get("session")) + events: list[DuplexEvent] = [ + SessionCreated( + session=session, + attachment_generation=( + cast("int", event["attachment_generation"]) + if isinstance(event.get("attachment_generation"), int) + else None + ), + resume_token=_str_or_none(event.get("resume_token")), + ) + ] + if state.initial_session_update: + events.append(SessionUpdated(session=session)) + state.initial_session_update = False + return events + if event_type == "session.updated": + return [SessionUpdated(session=realtime_session_payload(state, event.get("session")))] + if event_type == "session.resumed": + return [ + SessionResumed( + session=realtime_session_payload(state, event.get("session")), + attachment_generation=( + cast("int", event["attachment_generation"]) + if isinstance(event.get("attachment_generation"), int) + else None + ), + resume_token=_str_or_none(event.get("resume_token")), + ) + ] + if event_type == "session.replaced": + return [SessionReplaced(attachment_generation=_int_or(event.get("attachment_generation")))] + if event_type == "session.resync_required": + return [SessionResyncRequired(reason=str(event.get("reason") or "journal_gap"))] + if event_type == "session.closed": + return [SessionClosed(reason=str(event.get("reason") or "closed"), details=event)] + if event_type == "response.created": + response_id = event.get("response_id") + if isinstance(response_id, str) and response_id: + state.active_response_id = response_id + state.last_response_id = response_id + item_id = _response_item_id(state, response_id) + modalities = event.get("modalities") + has_audio_modality = not isinstance(modalities, list) or "audio" in modalities + item = { + "id": item_id, + "object": "realtime.item", + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + } + state.conversation_items[item_id] = item + events = [ + _response_created_event(event), + *_conversation_item_added_events(state, item), + OutputItemAdded(response_id=_str_or_none(response_id), item=item), + ] + if has_audio_modality: + events.extend(_ensure_response_audio_part_added(state, response_id)) + return events + if event_type == "response.listen": + return [Listen(response_id=_str_or_none(event.get("response_id")), details=event)] + if event_type == "response.speak": + response_id = event.get("response_id") + projection = _response_state(state, response_id) + if projection is not None: + if projection.speak_emitted: + return [] + projection.speak_emitted = True + return [ + Speak( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + metadata=_response_speak_metadata(event), + ) + ] + if event_type == "response.output_audio.delta": + response_id = event.get("response_id") + audio = event.get("audio", "") + events = [] + if isinstance(audio, str) and audio: + events.extend(_ensure_response_audio_part_added(state, response_id)) + events.extend(_realtime_audio_delta_events(state, event, response_id, audio)) + _refresh_in_progress_response_item(state, response_id) + text = event.get("text") + has_text = isinstance(text, str) and bool(text) + if has_text: + _append_response_transcript(state, response_id, cast("str", text)) + _refresh_in_progress_response_item(state, response_id) + # Keep the audio.delta + transcript.delta pair invariant even for + # text-less units so clients that treat the pair as unit-complete work. + if has_text or (isinstance(audio, str) and bool(audio)): + events.append( + TranscriptDelta( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + delta=cast("str", text) if has_text else "", + ) + ) + if event.get("end_of_turn") is True: + events.extend(_realtime_audio_done_events(state, event, response_id)) + events.extend( + _realtime_response_terminal_events( + state, + event, + response_id, + status="completed", + status_details={"type": "completed", "reason": event.get("finish_reason") or "stop"}, + ) + ) + return events + if event_type == "response.text.delta": + response_id = event.get("response_id") + text = event.get("delta", "") + projection = _response_state(state, response_id) + if projection is not None and isinstance(text, str) and text: + projection.text_parts.append(text) + _refresh_in_progress_response_item(state, response_id) + events = _ensure_response_text_part_added(state, response_id) + events.append( + TextDelta( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + content_index=1 if projection is not None and projection.audio_part_added else 0, + delta=str(text) if isinstance(text, str) else "", + ) + ) + return events + if event_type == "response.done": + response_id = event.get("response_id") + status = cast("str", event.get("status")) if isinstance(event.get("status"), str) else "completed" + status_details = ( + cast("dict[str, object]", event.get("status_details")) + if isinstance(event.get("status_details"), dict) + else None + ) + return [ + *_realtime_audio_done_events(state, event, response_id), + *_realtime_response_terminal_events( + state, event, response_id, status=status, status_details=status_details + ), + ] + if event_type == "input.committed": + event_item_id = event.get("realtime_item_id") + item_id = ( + event_item_id if isinstance(event_item_id, str) and event_item_id else _pop_pending_commit_item_id(state) + ) + item = state.conversation_items.get(item_id) + events = [] + if item is None: + message = event.get("message") + no_response = event.get("no_response") is True + is_speech = event.get("is_speech") + item = { + "id": item_id, + "object": "realtime.item", + "type": "message", + "role": "user", + "status": "completed", + "content": ( + [{"type": "input_audio", "transcript": "", "is_speech": False}] + if no_response and is_speech is False + else _user_item_content_from_duplex_message(message) + ), + } + state.conversation_items[item_id] = item + events.extend(_conversation_item_added_events(state, item)) + item["status"] = "completed" + events.append( + InputCommitted(previous_item_id=_previous_item_id(state, item_id), item_id=item_id, details=event) + ) + transcription_event = _input_audio_transcription_completed_event(item_id, item) + if transcription_event is not None: + events.append(transcription_event) + events.append(_conversation_item_done_event(state, item)) + return events + if event_type == "input.cancelled": + return [InputCleared()] + if event_type == "audio.cancelled": + response_id = event.get("response_id") + events = [] + if event.get("reason") == "output_audio_buffer_clear": + if not isinstance(response_id, str) or not response_id: + response_id = state.active_response_id or state.last_response_id + events.append(OutputAudioCleared(response_id=_str_or_none(response_id))) + if not isinstance(response_id, str) or not response_id: + return events + elif not isinstance(response_id, str) or not response_id: + response_id = state.active_response_id + if not isinstance(response_id, str) or not response_id: + return events + if response_is_done(state, response_id): + return events + committed_ms = event.get("committed_ms") + if isinstance(committed_ms, int | float): + item_id = _response_item_id(state, response_id) + committed_audio_ms = max(0, int(committed_ms)) + state.item_truncation_cursors[item_id] = (0, committed_audio_ms) + item = state.conversation_items.get(item_id) + if item is not None: + truncate_realtime_item_content(item, content_index=0, audio_end_ms=committed_audio_ms) + events.extend(_realtime_audio_done_events(state, event, response_id)) + events.extend( + _realtime_response_terminal_events( + state, + event, + response_id, + status="cancelled", + status_details={"type": "cancelled", "reason": event.get("reason") or "client_cancelled"}, + ) + ) + if isinstance(response_id, str) and response_id == state.active_response_id: + state.active_response_id = None + return events + if event_type == "conversation.item.created": + item = event.get("item") + if isinstance(item, dict) and isinstance(item.get("id"), str): + item_id = str(item["id"]) + already_known = item_id in state.conversation_items + state.conversation_items[item_id] = item + if already_known: + if item.get("status") == "completed": + return [_conversation_item_done_event(state, item)] + return [] + events = _conversation_item_added_events(state, item) + if item.get("status") == "completed": + events.append(_conversation_item_done_event(state, item)) + return events + return [ItemCreated(item=item if isinstance(item, Mapping) else {})] + if event_type == "conversation.item.deleted": + item_id = event.get("item_id") + if isinstance(item_id, str): + _remove_conversation_item(state, item_id) + return [ItemDeleted(item_id=_str_or_none(item_id), details=event)] + if event_type == "conversation.item.truncated": + item_id = event.get("item_id") + audio_end_ms = event.get("audio_end_ms") + content_index = event.get("content_index", 0) + if isinstance(item_id, str): + item = state.conversation_items.get(item_id) + if item is not None: + truncate_realtime_item_content( + item, + content_index=_int_or(content_index), + audio_end_ms=_int_or(audio_end_ms), + ) + return [ + ItemTruncated( + item_id=_str_or_none(item_id), + content_index=_int_or(content_index), + audio_end_ms=_int_or(audio_end_ms), + details=event, + ) + ] + if event_type == "conversation.item.retrieved": + item = event.get("item") + return [ItemRetrieved(item=item if isinstance(item, Mapping) else {})] + if event_type == "response.output_item.done": + response_id = event.get("response_id") + return _realtime_response_terminal_events( + state, + event, + response_id, + status="completed", + status_details={"type": "completed", "reason": event.get("finish_reason") or "stop"}, + ) + if event_type == "function_call.done": + return _function_call_done_events(state, event) + if event_type == "input_audio_buffer.speech_started": + return [ + SpeechStarted( + audio_start_ms=_int_or(event.get("audio_start_ms")), item_id=_str_or_none(event.get("item_id")) + ) + ] + if event_type == "input_audio_buffer.speech_stopped": + return [ + SpeechStopped(audio_end_ms=_int_or(event.get("audio_end_ms")), item_id=_str_or_none(event.get("item_id"))) + ] + return [DuplexRawEvent(internal_type=str(event_type), details=event)] + + +def _function_call_done_events(state: RealtimeProjectionState, event: dict[str, object]) -> list[DuplexEvent]: + call_id = event.get("call_id") + name = event.get("name") + arguments = event.get("arguments", "") + if not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name: + return [DuplexRawEvent(internal_type="function_call.done", details=event)] + if not isinstance(arguments, str): + arguments = str(arguments) + response_id = f"resp_{uuid4().hex}" + item_id = f"item_{uuid4().hex}" + item = { + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "name": name, + "call_id": call_id, + "arguments": arguments, + } + state.conversation_items[item_id] = item + response = { + "id": response_id, + "object": "realtime.response", + "status": "completed", + "status_details": {"type": "completed", "reason": "completed"}, + "output": [dict(item)], + "metadata": {}, + } + return [ + ResponseCreated(response_id=response_id, response={**response, "status": "in_progress", "output": []}), + *_conversation_item_added_events(state, item), + OutputItemAdded(response_id=response_id, item=dict(item)), + FunctionCallArgumentsDelta(response_id=response_id, item_id=item_id, call_id=call_id, delta=arguments), + FunctionCallArgumentsDone(response_id=response_id, item_id=item_id, call_id=call_id, arguments=arguments), + OutputItemDone(response_id=response_id, item=dict(item)), + _conversation_item_done_event(state, item), + ResponseDone(response_id=response_id, response=response), + ] + + +# ---- input-side helpers for the runner (were on the input translator) ---- + + +def register_user_item( + state: RealtimeProjectionState, + item: Mapping[str, object], + *, + previous_item_id: str | None = None, +) -> list[DuplexEvent]: + """Record a client-created user message item and return its ack events. + + Replaces the old translator's ``_send_realtime_input_ack``: call it when + handling a ``CreateItem`` whose item role is ``user`` (before the internal + ``conversation.item.created`` is emitted, which then projects to ``done``). + """ + normalized = dict(item) + if isinstance(previous_item_id, str): + normalized["_previous_item_id"] = previous_item_id + if normalized.get("role") != "user": + return [] + state.conversation_items[str(normalized["id"])] = normalized + return _conversation_item_added_events(state, normalized) + + +def retrieve_item_events(state: RealtimeProjectionState, payload: Mapping[str, object]) -> list[DuplexEvent]: + """Answer ``conversation.item.retrieve`` from the projected conversation items.""" + event_id = payload.get("event_id") + item_id = payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + return [error_event("missing_item_id", "conversation.item.retrieve requires item_id", event_id=event_id)] + item = state.conversation_items.get(item_id) + if item is None: + return [error_event("item_not_found", f"Conversation item not found: {item_id}", event_id=event_id)] + return [ItemRetrieved(item=item)] + + +def emit_input_speech_started(state: RealtimeProjectionState, audio_start_ms: object = 0) -> list[DuplexEvent]: + if state.input_speech_started: + return [] + state.input_speech_started = True + if state.active_input_item_id is None: + state.active_input_item_id = f"item_{uuid4().hex}" + return [SpeechStarted(audio_start_ms=_int_or(audio_start_ms), item_id=state.active_input_item_id)] + + +def emit_input_speech_stopped( + state: RealtimeProjectionState, *, item_id: str, audio_end_ms: object = 0 +) -> list[DuplexEvent]: + if not state.input_speech_started: + return [] + state.input_speech_started = False + return [SpeechStopped(audio_end_ms=_int_or(audio_end_ms), item_id=item_id)] + + +def _remember_input_transcript_hint(state: RealtimeProjectionState, event: Mapping[str, object]) -> None: + transcript = event.get("transcript") + if not isinstance(transcript, str): + transcript = event.get("text") if isinstance(event.get("text"), str) else None + if not isinstance(transcript, str): + hints = event.get("hints") + if isinstance(hints, dict): + transcript = hints.get("transcript") + if not isinstance(transcript, str): + transcript = hints.get("text") if isinstance(hints.get("text"), str) else None + if isinstance(transcript, str) and transcript: + parts = state.input_audio_buffer_transcript_parts + if parts and parts[-1] == transcript: + return + parts.append(transcript) + + +def _consume_input_transcript_hint(state: RealtimeProjectionState) -> str: + transcript = "".join(state.input_audio_buffer_transcript_parts).strip() + state.input_audio_buffer_transcript_parts.clear() + return transcript + + +def note_input_append( + state: RealtimeProjectionState, + payload: dict[str, object], + *, + vad_result: TurnDetectionResult | None = None, +) -> list[DuplexEvent]: + """Update the input-buffer projection for one appended chunk; returns typed events. + + ``payload`` is the internal ``input_audio_buffer.append`` dictionary + (``AppendAudio.payload()``), optionally after ``apply_turn_detection_result``; + ``vad_result`` is the ``TurnDetectionResult`` when server VAD is active. The + returned events are ``input_audio_buffer.speech_started`` / + ``speech_stopped`` exactly as the old translator produced them. + """ + audio = payload.get("audio") + looks_like_speech = bool(vad_result.is_speech) if vad_result is not None else payload.get("is_speech") is not False + has_audio = isinstance(audio, str) and bool(audio) + state.input_audio_buffer_has_audio = state.input_audio_buffer_has_audio or (looks_like_speech and has_audio) + state.input_audio_buffer_had_non_speech = state.input_audio_buffer_had_non_speech or ( + not looks_like_speech and has_audio + ) + events: list[DuplexEvent] = [] + stop_ms: object = payload.get("audio_end_ms", payload.get("audio_ms", 0)) + if vad_result is not None: + if vad_result.speech_stopped and vad_result.audio_end_ms is not None: + stop_ms = vad_result.audio_end_ms + if vad_result.speech_stopped and vad_result.speech_active: + events.extend( + emit_input_speech_stopped( + state, + item_id=state.active_input_item_id or f"item_{uuid4().hex}", + audio_end_ms=stop_ms, + ) + ) + if vad_result.speech_started: + start_ms = ( + vad_result.audio_start_ms if vad_result.audio_start_ms is not None else payload.get("audio_start_ms", 0) + ) + events.extend(emit_input_speech_started(state, start_ms)) + if vad_result.speech_stopped and not vad_result.speech_active: + events.extend( + emit_input_speech_stopped( + state, + item_id=state.active_input_item_id or f"item_{uuid4().hex}", + audio_end_ms=stop_ms, + ) + ) + elif looks_like_speech: + events.extend(emit_input_speech_started(state, payload.get("audio_start_ms", 0))) + if looks_like_speech: + _remember_input_transcript_hint(state, payload) + return events + + +def discard_pending_input_audio(state: RealtimeProjectionState, audio_end_ms: int | None = None) -> list[DuplexEvent]: + """Drop Realtime input-buffer state that was consumed as overlap (returns typed events).""" + events: list[DuplexEvent] = [] + if state.input_speech_started and state.active_input_item_id is not None: + events.append(SpeechStopped(audio_end_ms=max(0, int(audio_end_ms or 0)), item_id=state.active_input_item_id)) + state.input_speech_started = False + state.active_input_item_id = None + state.input_audio_buffer_has_audio = False + state.input_audio_buffer_had_non_speech = False + state.input_audio_buffer_transcript_parts.clear() + return events + + +def clear_input_buffer(state: RealtimeProjectionState) -> None: + """Input-buffer projection reset for ``input_audio_buffer.clear``.""" + state.input_speech_started = False + state.active_input_item_id = None + state.input_audio_buffer_has_audio = False + state.input_audio_buffer_had_non_speech = False + state.input_audio_buffer_transcript_parts.clear() + + +@dataclass(frozen=True, slots=True) +class ResolvedCommit: + """What a ``Commit`` means given the projected input buffer.""" + + #: Internal ``input_audio_buffer.commit`` payload to run, or None when rejected. + payload: dict[str, object] | None + #: Events to emit before running (speech_stopped) or instead (error). + events: list[DuplexEvent] + #: True when the VAD detector should be reset. + reset_vad: bool = False + + +def resolve_commit(state: RealtimeProjectionState, command: Commit) -> ResolvedCommit: + """Apply the old translator's commit rules (empty buffer, non-speech commit, item id).""" + if command.realtime_item_id is None and not state.input_audio_buffer_has_audio: + if state.input_audio_buffer_had_non_speech: + state.input_audio_buffer_had_non_speech = False + state.active_input_item_id = None + payload: dict[str, object] = { + "type": "input_audio_buffer.commit", + "final": command.final, + "response_create": False, + "is_speech": False, + } + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedCommit(payload=payload, events=[], reset_vad=True) + return ResolvedCommit( + payload=None, + events=[ + error_event( + "input_audio_buffer_empty", + "input_audio_buffer.commit requires a non-empty input audio buffer", + event_id=command.event_id, + ) + ], + ) + item_id = command.realtime_item_id or state.active_input_item_id or f"item_{uuid4().hex}" + state.pending_commit_item_ids.append(item_id) + events = emit_input_speech_stopped(state, item_id=item_id, audio_end_ms=0) + transcript = _consume_input_transcript_hint(state) + state.active_input_item_id = None + state.input_audio_buffer_has_audio = False + state.input_audio_buffer_had_non_speech = False + payload = { + "type": "input_audio_buffer.commit", + "final": command.final, + "realtime_item_id": item_id, + "response_create": bool(command.create_response) if command.create_response is not None else False, + } + if command.is_speech is not None: + payload["is_speech"] = command.is_speech + if transcript: + payload["transcript"] = transcript + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedCommit(payload=payload, events=events, reset_vad=True) + + +@dataclass(frozen=True, slots=True) +class ResolvedControl: + """Result of resolving a cancel/clear/delete/truncate against the projection.""" + + #: Internal payload(s) to run in order (empty when nothing to do). + payloads: list[dict[str, object]] + #: Events to emit instead (already-done response ack, errors). + events: list[DuplexEvent] + + +def resolve_cancel_response(state: RealtimeProjectionState, command: CancelResponse) -> ResolvedControl: + payload: dict[str, object] = {"type": "response.cancel", "reason": "response.cancel"} + response_id = command.response_id or state.active_response_id or state.last_response_id + if response_is_done(state, response_id): + return ResolvedControl( + payloads=[], + events=[ + error_event( + "response_not_active", f"Response is already complete: {response_id}", event_id=command.event_id + ) + ], + ) + if isinstance(response_id, str) and response_id: + payload["response_id"] = response_id + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedControl(payloads=[payload], events=[]) + + +def resolve_clear_output_audio(state: RealtimeProjectionState, command: ClearOutputAudio) -> ResolvedControl: + payload: dict[str, object] = {"type": "output_audio_buffer.clear", "reason": "output_audio_buffer.clear"} + response_id = command.response_id or state.active_response_id or state.last_response_id + if response_is_done(state, response_id): + return ResolvedControl(payloads=[], events=[OutputAudioCleared(response_id=_str_or_none(response_id))]) + if isinstance(response_id, str) and response_id: + payload["response_id"] = response_id + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedControl(payloads=[payload], events=[]) + + +def resolve_delete_item(state: RealtimeProjectionState, command: DeleteItem) -> ResolvedControl: + if command.item_id not in state.conversation_items: + return ResolvedControl( + payloads=[], + events=[ + error_event( + "item_not_found", f"Conversation item not found: {command.item_id}", event_id=command.event_id + ) + ], + ) + _remove_conversation_item(state, command.item_id) + return ResolvedControl(payloads=[command.payload()], events=[]) + + +def resolve_truncate_item(state: RealtimeProjectionState, command: TruncateItem) -> ResolvedControl: + item = state.conversation_items.get(command.item_id) + if not isinstance(item, dict): + return ResolvedControl( + payloads=[], + events=[ + error_event( + "item_not_found", f"Conversation item not found: {command.item_id}", event_id=command.event_id + ) + ], + ) + truncate_error = validate_realtime_item_truncate( + item, content_index=command.content_index, audio_end_ms=command.audio_end_ms + ) + if truncate_error is not None: + return ResolvedControl( + payloads=[], events=[error_event("bad_event", truncate_error, event_id=command.event_id)] + ) + state.item_truncation_cursors[command.item_id] = (command.content_index, command.audio_end_ms) + truncate_realtime_item_content(item, content_index=command.content_index, audio_end_ms=command.audio_end_ms) + ack_payload: dict[str, object] = { + "type": "playback.ack", + "item_id": command.item_id, + "committed_ms": int(command.audio_end_ms), + "played_ms": int(command.audio_end_ms), + "truncate": True, + } + return ResolvedControl(payloads=[command.payload(), ack_payload], events=[]) + + +def _duplicate_function_call_output(state: RealtimeProjectionState, call_id: object) -> tuple[bool, bool]: + matching_call = any( + known.get("type") == "function_call" and known.get("call_id") == call_id + for known in state.conversation_items.values() + ) + duplicate_output = any( + known.get("type") == "function_call_output" and known.get("call_id") == call_id + for known in state.conversation_items.values() + ) + return matching_call, duplicate_output + + +def resolve_create_item(state: RealtimeProjectionState, command: CreateItem) -> ResolvedControl: + """Expand ``conversation.item.create`` the way the old translator did. + + Assistant/system/function items and text-only user items become one + ``turn.signal conversation.item.create`` payload. User items carrying audio + become one ``input_audio_buffer.append`` per audio part followed by an + ``input_audio_buffer.commit`` (no response). The user-item ack events are + returned in ``events`` and must be emitted *before* running the payloads. + """ + item = dict(command.item) + if isinstance(command.previous_item_id, str): + item["_previous_item_id"] = command.previous_item_id + item_id = str(item["id"]) + item_type = item.get("type") + role = item.get("role") + if item_type == "function_call_output": + call_id = item.get("call_id") + output = item.get("output") + matching_call, duplicate_output = _duplicate_function_call_output(state, call_id) + if not isinstance(call_id, str) or not call_id or not matching_call: + return ResolvedControl( + payloads=[], + events=[ + error_event( + "invalid_function_call_output", + "function_call_output requires the call_id of a completed function call", + event_id=command.event_id, + param="item.call_id", + ) + ], + ) + if not isinstance(output, str) or duplicate_output: + return ResolvedControl( + payloads=[], + events=[ + error_event( + "invalid_function_call_output", + ( + "function_call_output requires a string output" + if not isinstance(output, str) + else f"function_call_output already exists for call_id {call_id}" + ), + event_id=command.event_id, + param="item.output", + ) + ], + ) + ack_events = register_user_item(state, item) if item_type == "message" and role == "user" else [] + signal_payload = { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": {"item": item}, + } + if command.event_id is not None: + signal_payload["realtime_event_id"] = command.event_id + if item_type != "message" or role in {"assistant", "system"}: + return ResolvedControl(payloads=[signal_payload], events=ack_events) + content = item.get("content") + if not isinstance(content, list): + return ResolvedControl(payloads=[], events=ack_events) + text_chunks: list[str] = [] + audio_payloads: list[dict[str, object]] = [] + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") in {"input_text", "text"} and isinstance(part.get("text"), str): + text_chunks.append(str(part["text"])) + if part.get("type") in {"input_audio", "audio"}: + audio = part.get("audio") or part.get("data") + if not isinstance(audio, str) or not audio: + continue + try: + append = build_append_audio( + { + "audio": audio, + "format": part.get("format"), + "sample_rate_hz": part.get("sample_rate_hz") or part.get("sample_rate"), + "event_id": command.event_id, + }, + defaults=state.defaults, + hints_source={**item, **part}, + ) + except Exception as exc: # DuplexCommandError from validation/conversion + code = getattr(exc, "code", "bad_event") + if code == "unsupported_audio_format": + continue + return ResolvedControl( + payloads=[], + events=[error_event(code, str(exc), event_id=command.event_id, param="sample_rate_hz")], + ) + speech_hints = {**item, **part} + if not input_looks_like_speech( + speech_hints, audio=append.audio, fmt=append.format, overlap_silence_rms=state.overlap_silence_rms + ): + continue + state.input_speech_started = True + payload = append.payload() + copy_realtime_input_hints(part, payload) + copy_realtime_input_hints(item, payload) + audio_payloads.append(payload) + if audio_payloads: + state.input_audio_buffer_has_audio = True + transcript = input_transcript_from_item(item) + commit_payload: dict[str, object] = { + "type": "input_audio_buffer.commit", + "final": True, + "realtime_item_id": item_id, + "response_create": False, + } + if transcript: + commit_payload["transcript"] = transcript + return ResolvedControl(payloads=[*audio_payloads, commit_payload], events=ack_events) + if not text_chunks: + return ResolvedControl(payloads=[], events=ack_events) + text_item = dict(item) + text_item["status"] = "completed" + signal_payload["payload"] = {"item": text_item} + return ResolvedControl(payloads=[signal_payload], events=ack_events) + + +def append_audio_payload(command: AppendAudio) -> dict[str, object]: + """Internal ``input_audio_buffer.append`` payload for a command (convenience for the runner).""" + return command.payload() + + +__all__ = [ + "RealtimeProjectionState", + "ResolvedCommit", + "ResolvedControl", + "append_audio_payload", + "clear_input_buffer", + "discard_pending_input_audio", + "emit_input_speech_started", + "emit_input_speech_stopped", + "note_input_append", + "project_internal_event", + "realtime_session_payload", + "register_user_item", + "resolve_cancel_response", + "resolve_clear_output_audio", + "resolve_commit", + "resolve_create_item", + "resolve_delete_item", + "resolve_truncate_item", + "response_is_done", + "retrieve_item_events", +] diff --git a/vllm_omni/engine/duplex/runtime.py b/vllm_omni/engine/duplex/runtime.py deleted file mode 100644 index dede0e0a6..000000000 --- a/vllm_omni/engine/duplex/runtime.py +++ /dev/null @@ -1,108 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Experimental duplex runtime loading and compatibility exports.""" - -from __future__ import annotations - -from importlib import import_module - -from vllm_omni.engine.duplex.contracts import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexOutputAction, - DuplexOutputDecision, - DuplexRuntimeCapabilities, - DuplexRuntimeExtension, - SessionMode, - duplex_data_plane_request_info, - duplex_resource_request_belongs_to_session, - duplex_resource_request_id, -) - - -def load_duplex_runtime_extension(path: str | None) -> DuplexRuntimeExtension | None: - if not path: - return None - module_name, separator, attribute_name = path.rpartition(".") - if not separator: - raise ValueError(f"Invalid duplex runtime extension path: {path!r}") - extension_type = getattr(import_module(module_name), attribute_name) - return validate_duplex_runtime_extension(extension_type()) - - -def validate_duplex_runtime_extension( - extension: object, - *, - sampling_defaults: tuple[object, ...] | None = None, -) -> DuplexRuntimeExtension: - required_methods = ( - "configure_sampling_params", - "plan_append", - "decide_output", - ) - missing = [name for name in required_methods if not callable(getattr(extension, name, None))] - if missing: - raise TypeError(f"Duplex runtime extension is missing callable method(s): {', '.join(missing)}") - typed_extension = extension # type: ignore[assignment] - if sampling_defaults is not None: - configured = typed_extension.configure_sampling_params( - runtime_config={}, - defaults=sampling_defaults, - ) - if not isinstance(configured, tuple): - raise TypeError("Duplex runtime extension must return sampling parameters as a tuple") - if len(configured) != len(sampling_defaults): - raise ValueError("Duplex runtime extension must return one sampling parameter per stage") - for stage_id, (value, default) in enumerate(zip(configured, sampling_defaults, strict=True)): - if default is not None and not isinstance(value, type(default)): - raise TypeError( - "Duplex runtime extension sampling parameter type mismatch " - f"for stage {stage_id}: expected {type(default).__name__}, got {type(value).__name__}" - ) - return typed_extension - - -from vllm_omni.engine.duplex.lease import ( # noqa: E402, F401 - DuplexLeaseActivity, - DuplexLeaseConfig, - DuplexLeaseState, - DuplexSessionExpiry, -) -from vllm_omni.engine.duplex.session import ( # noqa: E402, F401 - DuplexAppendReservation, - DuplexCompletedAppend, - DuplexFenceMismatchError, - DuplexInputAppend, - DuplexRequestResource, - DuplexSessionRuntimeManager, - DuplexSessionRuntimeState, - DuplexStageBinding, -) - -__all__ = [ - "DuplexAppendPlan", - "DuplexAppendReservation", - "DuplexCompletedAppend", - "DuplexFenceMismatchError", - "DuplexInputAppend", - "DuplexInputMode", - "DuplexLeaseActivity", - "DuplexLeaseConfig", - "DuplexLeaseState", - "DuplexOutputAction", - "DuplexOutputDecision", - "DuplexRequestResource", - "DuplexRuntimeCapabilities", - "DuplexRuntimeExtension", - "DuplexSessionExpiry", - "DuplexSessionRuntimeManager", - "DuplexSessionRuntimeState", - "DuplexStageBinding", - "SessionMode", - "duplex_data_plane_request_info", - "duplex_resource_request_belongs_to_session", - "duplex_resource_request_id", - "load_duplex_runtime_extension", - "validate_duplex_runtime_extension", -] diff --git a/vllm_omni/engine/duplex/session.py b/vllm_omni/engine/duplex/session.py deleted file mode 100644 index b7ad51504..000000000 --- a/vllm_omni/engine/duplex/session.py +++ /dev/null @@ -1,534 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import time -from collections import OrderedDict -from collections.abc import Callable, Iterable -from dataclasses import dataclass, field -from typing import Any - -from vllm_omni.engine.duplex.contracts import ( - DuplexInputMode, - DuplexRuntimeCapabilities, -) -from vllm_omni.engine.duplex.lease import ( - DuplexLeaseActivity, - DuplexLeaseConfig, - DuplexLeaseState, - DuplexSessionExpiry, -) -from vllm_omni.engine.duplex.messages import DuplexFence - - -def _default_capabilities() -> DuplexRuntimeCapabilities: - return DuplexRuntimeCapabilities() - - -class DuplexFenceMismatchError(RuntimeError): - def __init__(self, expected: DuplexFence, actual: DuplexFence) -> None: - super().__init__(f"duplex fence mismatch: expected {expected!r}, got {actual!r}") - self.expected = expected - self.actual = actual - - -@dataclass -class DuplexStageBinding: - request_id: str - fence: DuplexFence - - -@dataclass -class DuplexRequestResource: - stage_id: int - request_id: str - fence: DuplexFence - submitted: bool = False - - -@dataclass -class DuplexInputAppend: - seq: int - turn_seq: int - turn_id: int - - -@dataclass(frozen=True) -class DuplexAppendReservation: - fence: DuplexFence - mode: DuplexInputMode - base_fence: DuplexFence - base_input_seq: int - base_input_turn_seq: int - base_append_turn_key: tuple[int, int, int] | None - update: DuplexInputAppend - - -@dataclass(frozen=True) -class DuplexCompletedAppend: - fence: DuplexFence - mode: DuplexInputMode - final: bool - stage_results: tuple[dict[str, object], ...] - - -@dataclass -class DuplexSessionRuntimeState: - """Engine resource handles associated with a core-owned identity fence.""" - - fence: DuplexFence - lease: DuplexLeaseState - _clock: Callable[[], float] = field(repr=False) - capabilities: DuplexRuntimeCapabilities = field(default_factory=_default_capabilities) - session_config: dict[str, Any] = field(default_factory=dict) - runtime_config: dict[str, Any] = field(default_factory=dict) - config_generation: int = 0 - stage_bindings: dict[int, DuplexStageBinding] = field(default_factory=dict) - request_resources: dict[tuple[int, str], DuplexRequestResource] = field(default_factory=dict) - input_seq: int = 0 - input_turn_seq: int = 0 - _append_turn_key: tuple[int, int, int] | None = None - completed_append_limit: int = 256 - completed_appends: OrderedDict[str, DuplexCompletedAppend] = field(default_factory=OrderedDict) - - def __post_init__(self) -> None: - if self.completed_append_limit <= 0: - raise ValueError("completed_append_limit must be positive") - - @property - def session_id(self) -> str: - return self.fence.session_id - - @property - def epoch(self) -> int: - return self.fence.epoch - - @property - def turn_id(self) -> int: - return self.fence.turn_id - - def _validate_fence(self, fence: DuplexFence) -> None: - if fence.session_id != self.session_id or fence.incarnation != self.fence.incarnation: - raise DuplexFenceMismatchError(self.fence, fence) - current = self.fence - if fence.epoch < current.epoch or ( - fence.epoch == current.epoch - and (fence.turn_id < current.turn_id or fence.response_seq < current.response_seq) - ): - raise DuplexFenceMismatchError(current, fence) - - def accept_fence(self, fence: DuplexFence) -> None: - self._validate_fence(fence) - if fence.epoch != self.fence.epoch: - self.input_seq = 0 - self.input_turn_seq = 0 - self._append_turn_key = None - self.completed_appends.clear() - self.fence = fence - - def touch(self, fence: DuplexFence, activity: DuplexLeaseActivity) -> None: - self._validate_fence(fence) - self.lease.touch(self._clock(), activity) - - def detach(self, fence: DuplexFence) -> None: - self._validate_fence(fence) - self.lease.detach(self._clock()) - - def resume(self, fence: DuplexFence, *, expected_lease_generation: int) -> int: - self._validate_fence(fence) - return self.lease.resume( - self._clock(), - expected_generation=expected_lease_generation, - ) - - def begin_operation(self, fence: DuplexFence, operation_id: str) -> None: - self._validate_fence(fence) - self.lease.begin_operation(self._clock(), operation_id) - - def end_operation(self, fence: DuplexFence, operation_id: str) -> None: - self._validate_fence(fence) - self.lease.end_operation(self._clock(), operation_id) - - def replace_session_config(self, session_config: dict[str, Any]) -> None: - self.session_config = dict(session_config) - self.config_generation += 1 - - def replace_runtime_config(self, runtime_config: dict[str, Any]) -> None: - self.runtime_config = dict(runtime_config) - self.config_generation += 1 - - def replace_configs( - self, - *, - session_config: dict[str, Any] | None = None, - runtime_config: dict[str, Any] | None = None, - ) -> None: - """Atomically publish one validated configuration generation.""" - if session_config is None and runtime_config is None: - return - next_session_config = self.session_config if session_config is None else dict(session_config) - next_runtime_config = self.runtime_config if runtime_config is None else dict(runtime_config) - self.session_config = next_session_config - self.runtime_config = next_runtime_config - self.config_generation += 1 - - def reserve_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: - self._validate_fence(fence) - resource_key = (stage_id, request_id) - existing = self.request_resources.get(resource_key) - if existing is not None: - if ( - existing.fence.session_id != fence.session_id - or existing.fence.incarnation != fence.incarnation - or existing.fence.epoch != fence.epoch - ): - raise ValueError(f"Duplex request resource already reserved with different identity: {request_id}") - return - self.request_resources[resource_key] = DuplexRequestResource( - stage_id=stage_id, - request_id=request_id, - fence=fence, - ) - - def bind_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: - self.reserve_stage_request(stage_id, request_id, fence=fence) - self.accept_fence(fence) - resource = self.request_resources[(stage_id, request_id)] - resource.fence = fence - resource.submitted = True - self.stage_bindings[stage_id] = DuplexStageBinding(request_id=request_id, fence=fence) - - def stage_request_ids(self, fence: DuplexFence | None = None) -> list[str]: - return [ - binding.request_id for binding in self.stage_bindings.values() if fence is None or binding.fence == fence - ] - - def resource_request_ids( - self, - fence: DuplexFence | None = None, - *, - submitted: bool | None = None, - ) -> list[str]: - return list( - dict.fromkeys( - resource.request_id - for resource in self.request_resources.values() - if (fence is None or resource.fence == fence) and (submitted is None or resource.submitted is submitted) - ) - ) - - def release_request_ids(self, request_ids: list[str]) -> None: - released = set(request_ids) - if not released: - return - self.stage_bindings = { - stage_id: binding for stage_id, binding in self.stage_bindings.items() if binding.request_id not in released - } - self.request_resources = { - resource_key: resource - for resource_key, resource in self.request_resources.items() - if resource.request_id not in released - } - - def prepare_append(self, *, mode: DuplexInputMode, fence: DuplexFence) -> DuplexAppendReservation: - if mode not in self.capabilities.input_modes: - raise ValueError(f"Duplex input mode {mode.value!r} is not supported by session {self.session_id}") - self._validate_fence(fence) - input_seq = 0 if fence.epoch != self.fence.epoch else self.input_seq - input_turn_seq = 0 if fence.epoch != self.fence.epoch else self.input_turn_seq - append_turn_key = None if fence.epoch != self.fence.epoch else self._append_turn_key - turn_key = (fence.epoch, fence.turn_id, fence.response_seq) - turn_seq = input_turn_seq + 1 if turn_key == append_turn_key else 1 - return DuplexAppendReservation( - fence=fence, - mode=mode, - base_fence=self.fence, - base_input_seq=self.input_seq, - base_input_turn_seq=self.input_turn_seq, - base_append_turn_key=self._append_turn_key, - update=DuplexInputAppend( - seq=input_seq + 1, - turn_seq=turn_seq, - turn_id=fence.turn_id, - ), - ) - - def commit_append(self, reservation: DuplexAppendReservation) -> DuplexInputAppend: - if ( - self.fence != reservation.base_fence - or self.input_seq != reservation.base_input_seq - or self.input_turn_seq != reservation.base_input_turn_seq - or self._append_turn_key != reservation.base_append_turn_key - ): - raise RuntimeError("duplex append reservation is stale") - self.accept_fence(reservation.fence) - self.input_seq = reservation.update.seq - self.input_turn_seq = reservation.update.turn_seq - self._append_turn_key = ( - reservation.fence.epoch, - reservation.fence.turn_id, - reservation.fence.response_seq, - ) - return reservation.update - - def completed_append( - self, - operation_id: str, - *, - fence: DuplexFence, - mode: DuplexInputMode, - final: bool, - ) -> list[dict[str, object]] | None: - self._validate_fence(fence) - completed = self.completed_appends.get(operation_id) - if completed is None: - return None - if completed.fence != fence or completed.mode != mode or completed.final != final: - raise ValueError(f"duplex append operation {operation_id!r} was reused with different metadata") - return [dict(result) for result in completed.stage_results] - - def record_completed_append( - self, - operation_id: str, - *, - fence: DuplexFence, - mode: DuplexInputMode, - final: bool, - stage_results: list[dict[str, object]], - ) -> None: - self.completed_appends[operation_id] = DuplexCompletedAppend( - fence=fence, - mode=mode, - final=final, - stage_results=tuple(dict(result) for result in stage_results), - ) - self.completed_appends.move_to_end(operation_id) - while len(self.completed_appends) > self.completed_append_limit: - self.completed_appends.popitem(last=False) - - def append_input(self, *, mode: DuplexInputMode, fence: DuplexFence) -> DuplexInputAppend: - return self.commit_append(self.prepare_append(mode=mode, fence=fence)) - - def release_fence(self, fence: DuplexFence) -> list[str]: - stale = self.resource_request_ids(fence) - self.stage_bindings = { - stage_id: binding for stage_id, binding in self.stage_bindings.items() if binding.fence != fence - } - self.request_resources = { - resource_key: resource - for resource_key, resource in self.request_resources.items() - if resource.fence != fence - } - return stale - - def cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: - stale = self.prepare_cancel_fence(cancelled_fence, next_fence) - self.release_fence(cancelled_fence) - return stale - - def prepare_cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: - """Advance the cancellation fence without dropping cleanup records.""" - if cancelled_fence.session_id != self.session_id or cancelled_fence.incarnation != self.fence.incarnation: - raise DuplexFenceMismatchError(self.fence, cancelled_fence) - if ( - next_fence.session_id != self.session_id - or next_fence.incarnation != self.fence.incarnation - or next_fence.epoch <= cancelled_fence.epoch - ): - raise DuplexFenceMismatchError(cancelled_fence, next_fence) - current_key = (self.fence.epoch, self.fence.turn_id, self.fence.response_seq) - cancelled_key = (cancelled_fence.epoch, cancelled_fence.turn_id, cancelled_fence.response_seq) - next_key = (next_fence.epoch, next_fence.turn_id, next_fence.response_seq) - if cancelled_key > current_key: - raise DuplexFenceMismatchError(self.fence, cancelled_fence) - if next_key > current_key: - self.accept_fence(next_fence) - return self.resource_request_ids(cancelled_fence) - - def begin_close(self, fence: DuplexFence, *, reason: str) -> bool: - """Make close irreversible while retaining resources for cleanup retry.""" - self.accept_fence(fence) - if self.lease.terminal_reason is not None: - return True - return self.lease.mark_terminal(reason) - - def finalize_close(self) -> list[str]: - return self.close() - - def close(self, fence: DuplexFence | None = None) -> list[str]: - if fence is not None: - self.accept_fence(fence) - stale = self.resource_request_ids() - self.stage_bindings.clear() - self.request_resources.clear() - return stale - - def terminate(self, fence: DuplexFence, *, reason: str) -> bool: - if not self.begin_close(fence, reason=reason): - return False - self.finalize_close() - return True - - -class DuplexSessionRuntimeManager: - def __init__( - self, - *, - clock: Callable[[], float] | None = None, - max_sessions: int | None = None, - completed_append_limit: int = 256, - ) -> None: - if max_sessions is not None and max_sessions <= 0: - raise ValueError("max_sessions must be positive or null") - if completed_append_limit <= 0: - raise ValueError("completed_append_limit must be positive") - self._sessions: dict[str, DuplexSessionRuntimeState] = {} - self._clock = clock or time.monotonic - self._max_sessions = max_sessions - self._completed_append_limit = completed_append_limit - - def open_session( - self, - fence: DuplexFence, - *, - capabilities: DuplexRuntimeCapabilities | None = None, - session_config: dict[str, Any] | None = None, - runtime_config: dict[str, Any] | None = None, - lease_config: DuplexLeaseConfig | None = None, - ) -> DuplexSessionRuntimeState: - if not isinstance(fence, DuplexFence): - raise TypeError("open_session requires DuplexFence") - if fence.session_id in self._sessions: - raise ValueError(f"Duplex session already exists: {fence.session_id}") - if self._max_sessions is not None and len(self._sessions) >= self._max_sessions: - raise RuntimeError(f"duplex_session_capacity_exhausted: limit={self._max_sessions}") - session = DuplexSessionRuntimeState( - fence=fence, - lease=DuplexLeaseState( - config=lease_config or DuplexLeaseConfig(), - generation=0, - last_activity=self._clock(), - ), - _clock=self._clock, - capabilities=capabilities or _default_capabilities(), - session_config=dict(session_config or {}), - runtime_config=dict(runtime_config or {}), - completed_append_limit=self._completed_append_limit, - ) - self._sessions[fence.session_id] = session - return session - - def get(self, session_id: str) -> DuplexSessionRuntimeState | None: - return self._sessions.get(session_id) - - def require(self, session_id: str) -> DuplexSessionRuntimeState: - session = self.get(session_id) - if session is None: - raise KeyError(f"Unknown duplex session: {session_id}") - return session - - def close_session( - self, - fence: DuplexFence, - *, - reason: str = "explicit_close", - ) -> DuplexSessionRuntimeState | None: - if not isinstance(fence, DuplexFence): - raise TypeError("close_session requires DuplexFence") - session = self._sessions.get(fence.session_id) - if session is not None: - if session.lease.terminal_reason is not None: - return None - if not session.terminate(fence, reason=reason): - return None - if self._sessions.get(fence.session_id) is session: - self._sessions.pop(fence.session_id) - return session - - def begin_close_session( - self, - fence: DuplexFence, - *, - reason: str = "explicit_close", - ) -> DuplexSessionRuntimeState | None: - if not isinstance(fence, DuplexFence): - raise TypeError("begin_close_session requires DuplexFence") - session = self._sessions.get(fence.session_id) - if session is not None: - session.begin_close(fence, reason=reason) - return session - - def finalize_close_session(self, session: DuplexSessionRuntimeState) -> None: - session.finalize_close() - if self._sessions.get(session.session_id) is session: - self._sessions.pop(session.session_id, None) - - def collect_expired( - self, - now: float | None = None, - *, - excluded_session_ids: set[str] | None = None, - ) -> list[DuplexSessionExpiry]: - effective_now = self._clock() if now is None else now - excluded = excluded_session_ids or set() - expired: list[DuplexSessionExpiry] = [] - for session_id, session in list(self._sessions.items()): - if session_id in excluded: - continue - if session.lease.disconnect_grace_expired(effective_now): - reason = "disconnect_grace_expired" - elif session.lease.idle_expired(effective_now): - reason = "idle_ttl_expired" - else: - continue - submitted = tuple(session.resource_request_ids(submitted=True)) - reserved = tuple(session.resource_request_ids(submitted=False)) - if not session.begin_close(session.fence, reason=reason): - continue - expired.append( - DuplexSessionExpiry( - session_id=session_id, - fence=session.fence, - lease_generation=session.lease.generation, - reason=reason, - submitted_request_ids=submitted, - reserved_request_ids=reserved, - ) - ) - return expired - - def close_sessions_for_request_ids(self, request_ids: list[str]) -> dict[str, list[str]]: - request_id_set = set(request_ids) - closed: dict[str, list[str]] = {} - for session_id, session in list(self._sessions.items()): - stale = session.resource_request_ids() - if request_id_set.isdisjoint(stale): - continue - if not session.begin_close(session.fence, reason="request_cleanup"): - continue - closed[session_id] = stale - return closed - - def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: - for session_id in session_ids: - session = self._sessions.get(session_id) - if session is None or session.lease.terminal_reason is None: - continue - self.finalize_close_session(session) - - -__all__ = [ - "DuplexAppendReservation", - "DuplexCompletedAppend", - "DuplexFenceMismatchError", - "DuplexInputAppend", - "DuplexLeaseActivity", - "DuplexLeaseConfig", - "DuplexLeaseState", - "DuplexRequestResource", - "DuplexSessionExpiry", - "DuplexSessionRuntimeManager", - "DuplexSessionRuntimeState", - "DuplexStageBinding", -] diff --git a/vllm_omni/engine/duplex/session/__init__.py b/vllm_omni/engine/duplex/session/__init__.py new file mode 100644 index 000000000..a062ebf0e --- /dev/null +++ b/vllm_omni/engine/duplex/session/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""One engine-resident duplex session, and everything that runs it. + +``DuplexEngineSession`` (``engine_session``) is the state -- the input, +response, playback and conversation ledgers, the lease and the identity fence. +``DuplexSessionRunner`` (``runner``) owns one of those on the orchestrator loop +and is the only writer of it; ``DuplexSessionManager`` (``manager``) owns the +runners and the admission slots. + +The rest of the package is what the runner was split into, and each module +names the thing it owns rather than a layer: + +* ``context`` -- the state a runner and its components share, and the only + services a component may ask the runner for. +* ``emitter`` -- everything the session sends: the Realtime projection, the + epoch filter, the domain effects a terminal event applies. +* ``model_channel`` -- the model side: submitting an append, projecting the + stage output that comes back, continuing a turn with silence. +* ``control`` -- server VAD and the events that reconfigure the session. +* ``append_task`` -- one append in flight, and the rollback its failure owes. +* ``helpers`` -- the small pure reads and payload builders over a session. +* ``overlap_policy`` / ``commit_policy`` / ``playback_ledger`` -- the decisions + a session makes about speech that arrives while the model is speaking, what a + commit should do, and what the client has actually played. +* ``lease`` -- the idle/activity lease that decides when a session expires. + +Imports here are the package's public surface; a module of this package is fair +game for anything inside it, but code outside ``engine.duplex`` should reach for +these three names. +""" + +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner + +__all__ = ["DuplexEngineSession", "DuplexSessionManager", "DuplexSessionRunner"] diff --git a/vllm_omni/engine/duplex/session/append_task.py b/vllm_omni/engine/duplex/session/append_task.py new file mode 100644 index 000000000..724dc7f82 --- /dev/null +++ b/vllm_omni/engine/duplex/session/append_task.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""One append in flight, and what its failure owes the session. + +An append is not a call: it is a task queued behind the previous append, which +may be cancelled, may find the session closed before it starts, and on any of +those paths has to give back what it reserved --- the PCM bytes it took from the +input budget, the committed audio it was going to consume, and the response it +precreated. + +That compensation is the reason this is an object. The five steps of an append +all need the same eleven values, so as closures over ``_start_append`` they read +as one function with five entry points and no way to test any of them. Named +fields make the captured state explicit and the compensation a method. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from dataclasses import dataclass + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DuplexSessionState, DuplexTurnEventType +from vllm_omni.engine.duplex.plugin import PcmAppendReservation +from vllm_omni.engine.duplex.session import helpers +from vllm_omni.engine.duplex.session.context import DuplexSessionContext +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.model_channel import ModelChannel + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class AppendAttempt: + """One queued append to the model, with the rollback each failure path owes.""" + + ctx: DuplexSessionContext + out: SessionEmitter + model: ModelChannel + #: Marks the session closing before the close coroutine gets to run. + fail_session: Callable[[str], None] + + payload: dict[str, object] + epoch: int + request_id: str + final: bool + pcm_reservation: PcmAppendReservation | None + operation_id: str | None + #: Committed audio this append consumes; released once it is safely submitted. + retained_committed_payload: dict[str, object] | None + #: Response reserved before submission, to be failed if submission does not happen. + precreated_response_id: str | None + #: Last chance to call the append off, checked once the predecessor is done. + before_append: Callable[[], bool] | None + + # ------------------------------------------------------------------ # + # Compensation # + # ------------------------------------------------------------------ # + + def discard_retained_audio(self) -> None: + """Release the committed audio this append was going to consume. + + Only if it is still the session's: a later commit may have replaced it, + and that one belongs to the append that will carry it. + """ + model_state = self.ctx.model_state + if self.retained_committed_payload is not None and ( + model_state.committed_audio_payload is self.retained_committed_payload + ): + self.ctx.session.release_input_bytes(model_state.clear_committed_audio()) + + def abandon(self) -> None: + """Give back everything this append reserved but never used.""" + if self.pcm_reservation is not None: + self.pcm_reservation.rollback() + self.discard_retained_audio() + + def release_on_failure(self, done: asyncio.Task[bool]) -> None: + """Done-callback: a cancelled or failed append must not hold the audio.""" + if done.cancelled(): + self.discard_retained_audio() + return + try: + append_ok = done.result() + except Exception: + append_ok = False + if not append_ok: + self.discard_retained_audio() + + def clear_pending_silence(self, done: asyncio.Task[bool]) -> None: + """Done-callback of a silence continuation: it is no longer pending.""" + model_state = self.ctx.model_state + if model_state.pending_silence_task is done: + model_state.pending_silence_task = None + model_state.pending_silence_owner_id = None + + # ------------------------------------------------------------------ # + # Running # + # ------------------------------------------------------------------ # + + async def run_in_wire_order(self, predecessor: asyncio.Task[bool] | None) -> bool: + """Wait for the previous append, then submit this one if it still applies.""" + if predecessor is not None: + try: + predecessor_ok = await predecessor + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + predecessor_ok = False + except Exception: + predecessor_ok = False + if not predecessor_ok: + self.abandon() + return False + run = self.ctx.run + if run.closing or run.runtime_closed or self.ctx.session.state != DuplexSessionState.OPEN: + self.abandon() + return False + if self.before_append is not None and not self.before_append(): + self.abandon() + # Called off, not failed: the chain behind it still runs. + return True + if self.pcm_reservation is not None and not self.pcm_reservation.active: + self.discard_retained_audio() + return False + return await self._submit() + + async def _submit(self) -> bool: + session = self.ctx.session + model_state = self.ctx.model_state + try: + append_ok, emitted_response = await self.model.append_runtime_input( + self.payload, + operation_id=( + self.pcm_reservation.operation_id if self.pcm_reservation is not None else self.operation_id + ), + final=self.final, + expected_epoch=self.epoch, + ) + if append_ok: + model_state.context_locked = True + if self.pcm_reservation is not None: + self.pcm_reservation.commit() + session.release_input_bytes(self.pcm_reservation.byte_count) + self.discard_retained_audio() + else: + self.abandon() + if not append_ok and self._precreated_response_still_active(): + self._fail_precreated_response() + if not append_ok and session.state == DuplexSessionState.CLOSED: + self.ctx.run.runtime_closed = True + return False + if not emitted_response and session.epoch == self.epoch: + if session.active_request_id == helpers.stage0_request_id(session, self.epoch): + session.clear_request(self.request_id) + if self.final: + self.out.emit_events([session.signal_turn(DuplexTurnEventType.USER_STARTED.value)]) + return append_ok + except asyncio.CancelledError: + if self.pcm_reservation is not None: + self.pcm_reservation.rollback() + raise + except Exception as exc: + self.abandon() + logger.exception("Native duplex append task failed: %s", exc) + self.model.send_runtime_error("runtime_append_task_failed", exc) + if session.state != DuplexSessionState.CLOSED: + self.fail_session("runtime_append_task_failed") + return False + + def _precreated_response_still_active(self) -> bool: + return ( + self.precreated_response_id is not None + and self.ctx.session.active_response_id == self.precreated_response_id + ) + + def _fail_precreated_response(self) -> None: + session = self.ctx.session + session.end_response(commit_text=False) + self.out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": self.precreated_response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": "runtime_append_failed"}, + "playback": session.playback.as_dict(), + } + ) diff --git a/vllm_omni/entrypoints/duplex/commit_policy.py b/vllm_omni/engine/duplex/session/commit_policy.py similarity index 91% rename from vllm_omni/entrypoints/duplex/commit_policy.py rename to vllm_omni/engine/duplex/session/commit_policy.py index 868b98c98..0c07a6fa3 100644 --- a/vllm_omni/entrypoints/duplex/commit_policy.py +++ b/vllm_omni/engine/duplex/session/commit_policy.py @@ -21,7 +21,7 @@ class CommitSnapshot: speech_since_commit: bool active_response_id: str | None overlap_speech_ms: int - native_response_in_progress: bool + response_in_progress: bool playback_active: bool @@ -31,7 +31,7 @@ def decide_commit_action(snapshot: CommitSnapshot) -> CommitAction: if snapshot.speech_since_commit: return CommitAction.START_AUTO_RESPONSE return CommitAction.COMMIT_ONLY - if snapshot.native_response_in_progress and snapshot.overlap_speech_ms > 0: + if snapshot.response_in_progress and snapshot.overlap_speech_ms > 0: return CommitAction.DEFER_ACTIVE_RESPONSE if snapshot.playback_active and snapshot.overlap_speech_ms > 0: return CommitAction.DEFER_ACTIVE_RESPONSE diff --git a/vllm_omni/engine/duplex/session/context.py b/vllm_omni/engine/duplex/session/context.py new file mode 100644 index 000000000..1e7b861bc --- /dev/null +++ b/vllm_omni/engine/duplex/session/context.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared state for :class:`DuplexSessionRunner` and its components. + +The runner is a state machine that several components read and write. Splitting +it into modules only helps if the shared part stops being implicit: a component +holding a back-reference to the runner is a mixin with extra steps, and one +holding its own copy of a flag diverges from the others. + +So the mutable part is named. ``DuplexRunState`` is the small set of flags that +more than one component touches -- ``grep "\\.run\\."`` finds every mutation -- +and ``DuplexSessionContext`` is the read-mostly collaborator bundle everything +is constructed with. Anything a component needs from the runner that is neither +of those is infrastructure, and goes through :class:`RunnerServices`. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol, TypeVar + +if TYPE_CHECKING: + from vllm.outputs import RequestOutput + + from vllm_omni.engine.duplex.contracts import ( + DuplexOutputContext, + DuplexOutputDecision, + DuplexStagePort, + ) + from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, DuplexModelSessionState + from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + from vllm_omni.engine.duplex.session.manager import DuplexSessionManager + from vllm_omni.metrics.stats import StageRequestStats + +_OffloadT = TypeVar("_OffloadT") + + +@dataclass(slots=True) +class DuplexRunState: + """Mutable per-run flags shared by the runner and its components. + + Every field here is written by one component and read by another; that is + the criterion for being in this object rather than private to a component. + """ + + #: An irreversible close has begun: commands and control ops are refused. + closing: bool = False + #: ``session.closed`` / ``session.expired`` already left the runner. + closed_emitted: bool = False + #: The terminal event is deferred to the manager, which emits it after the + #: stage cleanup so it also means "the admission slot is free again". + closed_deferred: bool = False + close_reason: str | None = None + #: The model runtime reported a close; further data-plane work is pointless. + runtime_closed: bool = False + #: Request id of the resumable data-plane stream currently bound to the session. + stream_request_id: str | None = None + + +class RunnerServices(Protocol): + """The only things a component may ask the runner for. + + Deliberately two methods. A component that needs more than task scheduling + from the runner is reaching for orchestration that belongs in the runner. + """ + + def spawn(self, coro: Awaitable[None], *, name: str) -> None: + """Run ``coro`` as a tracked background task on the session's loop.""" + ... + + async def offload(self, fn: Callable[..., _OffloadT], *args: object, **kwargs: object) -> _OffloadT: + """Run a blocking call off the orchestrator loop.""" + ... + + +@dataclass(slots=True) +class DuplexSessionContext: + """Collaborators one session's components are built with. + + Read-mostly: the session and the model state are mutated through their own + APIs, not by rebinding these fields. The mutable flags live in ``run``. + """ + + session: DuplexEngineSession + model_state: DuplexModelSessionState + plugin: DuplexModelPlugin + stage_port: DuplexStagePort + manager: DuplexSessionManager + tasks: DuplexSessionTasks + run: DuplexRunState + services: RunnerServices + + +# --------------------------------------------------------------------------- # +# Task bookkeeping # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True, slots=True) +class DuplexAppendTaskMeta: + epoch: int + final: bool + response_bound: bool + + +@dataclass +class DuplexSessionTasks: + """Tracked task handles of one session (append tail, active response, pending silence).""" + + append_tasks: dict[asyncio.Task[bool], DuplexAppendTaskMeta] = field(default_factory=dict) + append_tail: asyncio.Task[bool] | None = None + active_response_task: asyncio.Task[None] | None = None + + def track_append_task( + self, + task: asyncio.Task[bool], + *, + epoch: int, + final: bool, + response_bound: bool, + ) -> None: + self.append_tasks[task] = DuplexAppendTaskMeta(epoch, final, response_bound) + task.add_done_callback(self.append_tasks.pop) + + def has_response_bound_append_tasks(self) -> bool: + return any(meta.response_bound for meta in self.append_tasks.values()) + + async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: + tasks = [task for task, meta in self.append_tasks.items() if not response_bound_only or meta.response_bound] + if not tasks: + return False + cancelled_tail = self.append_tail if self.append_tail in tasks else None + for task in tasks: + task.cancel() + try: + await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) + except TimeoutError: + pass + if cancelled_tail is not None and self.append_tail is cancelled_tail: + self.append_tail = None + return True + + +# --------------------------------------------------------------------------- # +# Stage output (the mailbox item the orchestrator hands the runner) # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True, slots=True) +class StageOutput: + """One stage result, queued for the session that owns the request.""" + + stage_id: int + output: RequestOutput + metrics: StageRequestStats | None + request_id: str + context: DuplexOutputContext + decision: DuplexOutputDecision | None diff --git a/vllm_omni/engine/duplex/session/control.py b/vllm_omni/engine/duplex/session/control.py new file mode 100644 index 000000000..2d46154fe --- /dev/null +++ b/vllm_omni/engine/duplex/session/control.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Control of one duplex session: server VAD, and the events that reconfigure it. + +Server turn detection and ``session.update`` are the same concern seen from two +sides. The detector is per-session mutable state, and ``session.update`` is the +event that replaces it mid-flight -- along with the session config, the runtime +config and the sampling policy, each of which has to be validated against a +candidate before the live session adopts it. The rest of ``turn.signal`` +(conversation items, local turn markers) travels the same path and is small, so +it lives here rather than in a module of its own. + +The one thing this does not own is ordering against the data plane: a +``session.update`` must not overtake an append still in flight, so the runner +supplies ``wait_for_append_tail``. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from copy import deepcopy + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DuplexConfigError, realtime_item_to_history_message +from vllm_omni.engine.duplex.plugin import DuplexRuntimeConfigError +from vllm_omni.engine.duplex.session import helpers +from vllm_omni.engine.duplex.session.context import DuplexSessionContext +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.model_channel import ModelChannel +from vllm_omni.engine.duplex.turn_detection import ( + PendingTurnDetectionUpdate, + ServerTurnDetector, + ServerVADUnavailableError, + SileroVADBackendProvider, + TurnDetectionConfig, + TurnDetectionResult, + apply_turn_detection_result, +) + +logger = init_logger(__name__) + + +class SessionControl: + """Server VAD and the control events that reconfigure one session.""" + + def __init__( + self, + ctx: DuplexSessionContext, + out: SessionEmitter, + model: ModelChannel, + *, + wait_for_append_tail: Callable[[], Awaitable[bool]], + ) -> None: + self._ctx = ctx + self._out = out + self._model = model + self._wait_for_append_tail = wait_for_append_tail + self._config: TurnDetectionConfig | None = None + self._detector: ServerTurnDetector | None = None + + def init_turn_detection(self) -> None: + turn_detection = self._ctx.session.config.extra_body.get("realtime_turn_detection") + if not isinstance(turn_detection, dict) or turn_detection.get("type") != "server_vad": + self._config = None + self._detector = None + return + try: + config = TurnDetectionConfig.from_realtime(turn_detection) + self._config = config + self._detector = config.build_detector(self._vad_backend_provider()) + except Exception as exc: + logger.warning("Duplex session %s: turn detection disabled: %s", self._ctx.session.session_id, exc) + self._config = None + self._detector = None + + def _vad_backend_provider(self) -> SileroVADBackendProvider | None: + """The engine-wide Silero backend, so one model serves every session. + + Held by the manager rather than built here: the ONNX backend is shared, + and ``duplex_session.server_vad_model_path`` is a deploy-level setting. + """ + return getattr(self._ctx.manager, "vad_backend_provider", None) + + def reset_vad(self) -> None: + """Drop the detector's speech state (a barge-in or a clear starts a new turn).""" + if self._detector is not None: + self._detector.reset() + + async def run_turn_detection(self, event: dict[str, object]) -> TurnDetectionResult | None: + detector = self._detector + if detector is None: + return None + audio = event.get("audio") + if not isinstance(audio, str) or not audio: + return None + fmt = event.get("format") if isinstance(event.get("format"), str) else "pcm_f32le" + sample_rate_hz = event.get("sample_rate_hz") + try: + result = await self._ctx.services.offload( + detector.process, + audio, + fmt=fmt, + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int) else None, + audio_end_ms=event.get("audio_end_ms") if isinstance(event.get("audio_end_ms"), int) else None, + ) + except ServerVADUnavailableError as exc: + self._out.emit_error("server_vad_unavailable", str(exc)) + self._detector = None + return None + except ValueError as exc: + self._out.emit_error("bad_audio", str(exc)) + return None + apply_turn_detection_result(event, result) + return result + + # ------------------------------------------------------------------ # + # Control events # + # ------------------------------------------------------------------ # + + async def on_turn_signal(self, event: dict[str, object]) -> None: + session = self._ctx.session + turn_event = event.get("event") + realtime_event_id = event.get("realtime_event_id") + if not isinstance(turn_event, str): + self._out.emit_error("bad_event", "turn.signal requires event") + return + if turn_event == "barge_in" and not session.capabilities.supports_barge_in: + self._out.emit_events([helpers.barge_in_unsupported_error()]) + return + if turn_event == "session.update": + payload = event.get("payload") + if not isinstance(payload, dict): + self._out.emit_error( + "bad_event", "session.update requires a session payload", event_id=realtime_event_id + ) + return + await self.on_session_update( + payload, + realtime_event_id=realtime_event_id if isinstance(realtime_event_id, str) else None, + ) + return + if turn_event == "conversation.item.create": + await self._on_conversation_item_create(event) + return + if turn_event == "conversation.item.delete": + payload = event.get("payload") + item_id = payload.get("item_id") if isinstance(payload, dict) else None + deleted = session.delete_history_item(item_id) if isinstance(item_id, str) else False + self._out.emit( + { + "type": "conversation.item.deleted", + "session_id": session.session_id, + "item_id": item_id, + "deleted": deleted, + } + ) + return + if turn_event == "conversation.item.truncate": + payload = event.get("payload") + item_id = payload.get("item_id") if isinstance(payload, dict) else None + audio_end_ms = payload.get("audio_end_ms") if isinstance(payload, dict) else None + truncated = ( + session.truncate_history_item( + item_id, + audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, + hard=True, + ) + if isinstance(item_id, str) + else False + ) + self._out.emit( + { + "type": "conversation.item.truncated", + "session_id": session.session_id, + "item_id": item_id, + "content_index": (payload.get("content_index", 0) if isinstance(payload, dict) else 0), + "audio_end_ms": audio_end_ms, + "truncated": truncated, + } + ) + return + self._out.emit_events([session.signal_turn(turn_event, event)]) + + async def on_session_update(self, payload: dict[str, object], *, realtime_event_id: str | None) -> None: + session = self._ctx.session + model_state = self._ctx.model_state + pending_turn_detection: PendingTurnDetectionUpdate | None = None + + def reject_update() -> None: + if pending_turn_detection is not None: + pending_turn_detection.reject() + + try: + pending_turn_detection = PendingTurnDetectionUpdate.prepare( + payload, backend_provider=self._vad_backend_provider() + ) + except Exception as exc: + self._out.emit_error("unsupported_turn_detection", str(exc), event_id=realtime_event_id) + return + if not await self._wait_for_append_tail(): + self._out.emit_error( + "session_update_aborted", + "session.update was not applied because the preceding append failed", + event_id=realtime_event_id, + ) + reject_update() + return + try: + self._ctx.plugin.validate_client_extra_body(payload.get("extra_body")) + except DuplexRuntimeConfigError as exc: + self._out.emit_error(exc.code, str(exc), event_id=realtime_event_id) + reject_update() + return + candidate_config = deepcopy(session.config) + audio_started = session.playback.generated_ms > 0 or session.playback.sent_ms > 0 + try: + candidate_config.apply_realtime_update( + payload, + session_id=session.session_id, + audio_started=audio_started, + ) + except DuplexConfigError as exc: + self._out.emit_error(exc.code, str(exc) or "session.update was rejected", event_id=realtime_event_id) + reject_update() + return + if candidate_config.instructions != session.config.instructions and model_state.context_locked: + self._out.emit_error( + "instructions_update_unsupported", + "session.update cannot change instructions after the native duplex context is initialized", + event_id=realtime_event_id, + ) + reject_update() + return + requests_audio = any(str(modality).lower() == "audio" for modality in candidate_config.modalities) + if ( + requests_audio + and "ref_audio_data" not in session.runtime_config + and getattr(self._ctx.plugin, "requires_ref_audio", False) + ): + self._out.emit_error( + "ref_audio_required", "Native duplex audio output requires ref_audio", event_id=realtime_event_id + ) + reject_update() + return + try: + candidate_runtime_config = self._ctx.plugin.runtime_config_for_update( + candidate_config, + dict(session.runtime_config), + ) + # Validate the sampling policy for the candidate before adopting it. + self._ctx.plugin.configure_sampling_params( + runtime_config=dict(candidate_runtime_config), + defaults=tuple(self._ctx.stage_port.sampling_defaults()), + ) + except (DuplexRuntimeConfigError, DuplexConfigError) as exc: + self._out.emit_error(exc.code, str(exc), event_id=realtime_event_id) + reject_update() + return + except Exception as exc: + self._out.emit_error("runtime_signal_failed", str(exc), event_id=realtime_event_id) + reject_update() + return + session.replace_config(candidate_config) + session.replace_runtime_config(candidate_runtime_config) + try: + session.touch_lease(DuplexLeaseActivity.SIGNAL) + except Exception: + pass + if pending_turn_detection is not None: + self._config, self._detector = pending_turn_detection.commit(self._detector) + projector = self._out.require_projector() + projector.apply_session_defaults(payload) + self._out.emit({"type": "session.updated", "session": session.as_public_dict()}) + + async def _on_conversation_item_create(self, event: dict[str, object]) -> None: + session = self._ctx.session + payload = event.get("payload") + item = payload.get("item") if isinstance(payload, dict) else None + item_type = item.get("type") if isinstance(item, dict) else None + if item_type == "function_call_output" and isinstance(item, dict): + if not await self._wait_for_append_tail(): + return + try: + candidate_runtime_config = self._ctx.plugin.runtime_config_for_function_output( + session.config, + dict(session.runtime_config), + item, + ) + except DuplexRuntimeConfigError as exc: + self._out.emit_error(exc.code, str(exc)) + return + if candidate_runtime_config is not None: + session.replace_runtime_config(candidate_runtime_config) + self._out.emit( + { + "type": "conversation.item.created", + "session_id": session.session_id, + "item": item, + "created": True, + } + ) + self._ctx.services.spawn( + self._model.maybe_continue_response( + expected_epoch=session.epoch, + expected_model_turn_id=session.turn_id, + ), + name="duplex-continue", + ) + return + message = realtime_item_to_history_message(item) + item_id = item.get("id") if isinstance(item, dict) else None + if message is not None: + session.append_history_message(message) + session.register_history_item(item_id if isinstance(item_id, str) else None, message) + if message.get("role") == "user": + # A later response.create may answer this without any audio. + session.notify_new_user_item() + self._out.emit( + { + "type": "conversation.item.created", + "session_id": session.session_id, + "item": item, + "created": message is not None, + } + ) diff --git a/vllm_omni/engine/duplex/session/emitter.py b/vllm_omni/engine/duplex/session/emitter.py new file mode 100644 index 000000000..508bc7ef1 --- /dev/null +++ b/vllm_omni/engine/duplex/session/emitter.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Everything one duplex session sends to its client. + +Emission is the hub of the session runner: near enough every path through the +runner ends in an event, which is why this is a collaborator rather than a +module of functions. It owns three things the rest of the runner should not +have to think about --- the Realtime projection state, the epoch filter that +drops output from a superseded turn, and the domain transitions an outbound +terminal event applies to the session before it is projected. + +The one thing it cannot own is what happens *after* a terminal promotes +deferred overlap audio: that re-enters the mailbox of the runner, so the +runner injects it as ``promote_deferred_overlap``. +""" + +from __future__ import annotations + +from collections.abc import Callable + +from vllm_omni.engine.duplex.config import DuplexSessionState +from vllm_omni.engine.duplex.events import ( + DOMAIN_TERMINAL_EVENTS, + MODEL_OUTPUT_EVENTS, + DuplexEvent, + error_event, +) +from vllm_omni.engine.duplex.realtime_events import ( + RealtimeProjectionState, + discard_pending_input_audio, + project_internal_event, +) +from vllm_omni.engine.duplex.session import overlap_policy +from vllm_omni.engine.duplex.session.context import DuplexSessionContext + + +class SessionEmitter: + """Outbound events for one session.""" + + def __init__( + self, + ctx: DuplexSessionContext, + *, + promote_deferred_overlap: Callable[[dict[str, object], bool], None], + ) -> None: + self._ctx = ctx + self._promote_deferred_overlap = promote_deferred_overlap + self._projector: RealtimeProjectionState | None = ctx.session.projector + + # ------------------------------------------------------------------ # + # Projection state # + # ------------------------------------------------------------------ # + + @property + def projector(self) -> RealtimeProjectionState | None: + """The Realtime projection state, or ``None`` before anything needs it.""" + return self._projector + + @projector.setter + def projector(self, value: RealtimeProjectionState | None) -> None: + self._projector = value + self._ctx.session.projector = value + + def require_projector(self) -> RealtimeProjectionState: + """The projection state, created on first use. + + A session torn down before ``start()`` never builds one, so the lazy + path here is reachable rather than defensive. + """ + projector = self._projector + if projector is None: + session = self._ctx.session + projector = RealtimeProjectionState( + session_id=session.session_id, + model=session.config.model, + initial_session_update=True, + ) + self.projector = projector + return projector + + # ------------------------------------------------------------------ # + # Sending # + # ------------------------------------------------------------------ # + + def emit_events(self, events: list[DuplexEvent]) -> None: + for event in events: + self._ctx.manager.emit(self._ctx.session, event) + + def emit_error( + self, + code: str, + message: str, + *, + event_id: object | None = None, + retryable: bool | None = None, + ) -> None: + """Send one typed ``error`` event (``event_id`` is the client event it answers).""" + extra = {} if retryable is None else {"retryable": retryable} + self.emit_events([error_event(code, message, event_id=event_id, extra=extra)]) + + def emit(self, payload: dict[str, object]) -> None: + """Apply the domain effects of an internal event, then project it to typed events and send them. + + Only events with domain effects (response / cancel / close terminals) or + Realtime projection state (response items, content parts) still travel as + internal dictionaries; stateless events are constructed typed at the + emit site. + """ + accepted, deferred_overlap_payload = self._apply_outbound_session_event(payload) + if not accepted: + return + self.emit_events(project_internal_event(self.require_projector(), payload)) + if deferred_overlap_payload is not None and not self._ctx.run.closing: + precreate_response = self._ctx.model_state.deferred_precreate_response + self._ctx.model_state.deferred_precreate_response = False + self._promote_deferred_overlap(deferred_overlap_payload, precreate_response) + + # ------------------------------------------------------------------ # + # Filtering and domain effects # + # ------------------------------------------------------------------ # + + def auto_responds(self) -> bool: + """Whether the session answers committed input without a ``response.create``.""" + extra = getattr(self._ctx.session.config, "extra_body", None) + if not isinstance(extra, dict): + return False + return extra.get("auto_response") is True or extra.get("full_duplex") is True + + def is_stale_model_output(self, payload: dict[str, object]) -> bool: + """Whether ``payload`` belongs to a turn the session has already moved past.""" + session = self._ctx.session + event_type = payload.get("type") + if event_type in DOMAIN_TERMINAL_EVENTS: + return False + if event_type not in MODEL_OUTPUT_EVENTS: + return False + if self._ctx.run.closing and event_type != "response.listen": + return True + if session.state == DuplexSessionState.CLOSED and event_type != "response.listen": + return True + epoch = payload.get("epoch") + return isinstance(epoch, int) and epoch != session.epoch + + def _apply_outbound_session_event(self, payload: dict[str, object]) -> tuple[bool, dict[str, object] | None]: + """Apply domain transitions before an event is projected (moved from serving).""" + session = self._ctx.session + model_state = self._ctx.model_state + payload_type = payload.get("type") + is_terminal = payload_type in DOMAIN_TERMINAL_EVENTS + if is_terminal: + payload_epoch = payload.get("epoch") + if isinstance(payload_epoch, int) and payload_epoch != session.epoch: + return False, None + if payload_type in {"response.done", "response.listen"} and ( + self._ctx.run.closing or session.state == DuplexSessionState.CLOSED + ): + return False, None + elif self.is_stale_model_output(payload): + return False, None + + if payload_type == "session.closed": + self._ctx.run.close_reason = self._ctx.run.close_reason or str(payload.get("reason") or "closed") + session.mark_closing() + + if not is_terminal: + return True, None + + terminal_status = payload.get("status") + terminal_status_details = payload.get("status_details") + if terminal_status is None and isinstance(terminal_status_details, dict): + terminal_status = terminal_status_details.get("type") + response_terminal = payload_type == "response.done" or ( + payload_type == "response.listen" and session.active_response_id is not None + ) + can_promote_overlap = response_terminal and terminal_status not in {"cancelled", "failed"} + deferred_overlap_payload: dict[str, object] | None = None + continuous_input_crosses_terminal = ( + can_promote_overlap + and self.auto_responds() + and model_state.input_since_commit + and not model_state.deferred_response_create + ) + if continuous_input_crosses_terminal: + session.reset_overlap_speech() + return True, None + realtime_input_still_open = ( + can_promote_overlap and model_state.input_since_commit and not model_state.deferred_response_create + ) + if realtime_input_still_open: + session.reset_overlap_speech() + return True, None + if can_promote_overlap and session.overlap_speech_ms > 0: + has_deferred_overlap = ( + model_state.audio_buffer.has_pending() or model_state.committed_audio_payload is not None + ) + should_promote_overlap = ( + session.state == DuplexSessionState.OPEN + and has_deferred_overlap + and session.overlap_speech_ms > session.config.overlap_short_ack_ms + ) + if should_promote_overlap: + flushed_reserved_bytes = model_state.audio_buffer.pending_byte_count + deferred_overlap_payload = model_state.audio_buffer.flush( + chunk_period_ms=session.capabilities.chunk_period_ms or 1000 + ) + if model_state.committed_audio_payload is not None: + if deferred_overlap_payload is not None: + deferred_overlap_payload = overlap_policy.merge_audio_payloads( + model_state.committed_audio_payload, + deferred_overlap_payload, + ) + else: + deferred_overlap_payload = model_state.committed_audio_payload + if self.auto_responds() and deferred_overlap_payload is not None: + deferred_overlap_payload = dict(deferred_overlap_payload) + deferred_overlap_payload["force_listen"] = False + if deferred_overlap_payload is not None: + model_state.retain_committed_audio( + deferred_overlap_payload, + operation_id=model_state.committed_audio_operation_id, + reserved_bytes=flushed_reserved_bytes, + ) + model_state.input_since_commit = deferred_overlap_payload is not None + # Realtime path: defer the promoted response to the next terminal. + model_state.deferred_response_create = True + model_state.deferred_precreate_response = False + deferred_overlap_payload = None + else: + had_pending_overlap_audio = model_state.audio_buffer.has_pending() + model_state.audio_buffer.clear() + model_state.input_since_commit = False + model_state.speech_since_commit = False + if had_pending_overlap_audio and self._projector is not None: + self.emit_events(discard_pending_input_audio(self._projector, session.overlap_speech_ms)) + if payload_type in {"audio.cancelled", "input.cancelled", "session.closed"}: + session.release_input_bytes(model_state.clear_committed_audio()) + + session.reset_overlap_speech() + if ( + can_promote_overlap + and model_state.deferred_response_create + and model_state.committed_audio_payload is not None + ): + deferred_overlap_payload = model_state.committed_audio_payload + model_state.deferred_response_create = False + model_state.input_since_commit = False + model_state.speech_since_commit = False + return True, deferred_overlap_payload diff --git a/vllm_omni/entrypoints/duplex/protocol.py b/vllm_omni/engine/duplex/session/engine_session.py similarity index 56% rename from vllm_omni/entrypoints/duplex/protocol.py rename to vllm_omni/engine/duplex/session/engine_session.py index 042180838..bb5c38515 100644 --- a/vllm_omni/entrypoints/duplex/protocol.py +++ b/vllm_omni/engine/duplex/session/engine_session.py @@ -1,411 +1,64 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The one duplex session state, owned by the engine-side ``DuplexSessionRunner``. + +``DuplexEngineSession`` is the single session object: the input / response / +playback / conversation ledgers, the lease, the identity fence, the stage +request resources, the append sequencing, the model plugin's per-session state +(``model_state``) and the Realtime projection state (``projector``). There is +no cross-boundary fence protocol; ``session.fence`` is derived from the +session's own epoch/turn and ``accepted_fence`` is the monotonic high-water +mark of fences accepted for stage requests (``sync_fence()`` publishes the +current identity there). +""" + from __future__ import annotations import copy -from collections.abc import Mapping +import time +from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from enum import Enum from types import MappingProxyType +from typing import TYPE_CHECKING from uuid import uuid4 -import numpy as np - -from vllm_omni.entrypoints.duplex.audio import ( - encode_float32_mono_wav_base64, +from vllm_omni.engine.duplex.config import ( + DuplexAssistantAudioTextMark, + DuplexCapabilities, + DuplexCommittedInput, + DuplexPlaybackCommitPolicy, + DuplexPlaybackCursor, + DuplexPlaybackView, + DuplexSessionConfig, + DuplexSessionState, + DuplexTurnEventType, + DuplexTurnState, + ResponseCreateOptions, ) -from vllm_omni.entrypoints.duplex.server_vad import ( - ServerVADConfig, - parse_session_turn_detection, +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.engine.duplex.events import TurnEvent +from vllm_omni.engine.duplex.session.lease import ( + DuplexLeaseActivity, + DuplexLeaseConfig, + DuplexLeaseState, ) - -class DuplexOverlapPolicy(str, Enum): - LISTEN_ONLY = "listen_only" - BARGE_IN_ON_SPEECH = "barge_in_on_speech" +if TYPE_CHECKING: + from vllm_omni.engine.duplex.plugin import DuplexModelSessionState + from vllm_omni.engine.duplex.realtime_events import RealtimeProjectionState -class DuplexPlaybackCommitPolicy(str, Enum): - COMMIT_ALL_ON_DONE = "commit_all_on_done" - ACK_ONLY = "ack_only" - - -class DuplexSessionState(str, Enum): - OPEN = "open" - CLOSING = "closing" - CLOSED = "closed" - - -# Canonical extra_body key for the per-session opt-in to a model-native -# duplex runtime, plus the deprecated model-prefixed spelling it replaced. -# Client payloads may still use the legacy key; it is folded into the -# canonical one at ingestion so everything downstream (and every echo in -# session.created/session.updated) sees only NATIVE_DUPLEX_KEY. -NATIVE_DUPLEX_KEY = "native_duplex" -LEGACY_NATIVE_DUPLEX_KEY = "minicpmo45_native_duplex" - - -def normalize_native_duplex_key(extra_body: dict[str, object]) -> dict[str, object]: - """Fold the deprecated alias into the canonical key (canonical wins).""" - if LEGACY_NATIVE_DUPLEX_KEY in extra_body: - legacy = extra_body.pop(LEGACY_NATIVE_DUPLEX_KEY) - extra_body.setdefault(NATIVE_DUPLEX_KEY, legacy) - return extra_body - - -def native_duplex_opt_in(extra_body: Mapping[str, object]) -> object: - """Read the opt-in flag, accepting the deprecated alias for raw dicts.""" - if NATIVE_DUPLEX_KEY in extra_body: - return extra_body[NATIVE_DUPLEX_KEY] - return extra_body.get(LEGACY_NATIVE_DUPLEX_KEY) - - -class DuplexTurnState(str, Enum): - IDLE = "idle" - USER_SPEAKING = "user_speaking" - USER_COMMITTED = "user_committed" - ASSISTANT_GENERATING = "assistant_generating" - ASSISTANT_PLAYING = "assistant_playing" - BARGE_IN = "barge_in" - - -class DuplexTurnEventType(str, Enum): - USER_STARTED = "user_started" - USER_COMMITTED = "user_committed" - ASSISTANT_STARTED = "assistant_started" - ASSISTANT_DONE = "assistant_done" - BARGE_IN = "barge_in" - PLAYBACK_ACK = "playback_ack" - TIMEOUT = "timeout" - CLOSE = "close" - - -@dataclass -class DuplexCapabilities: - """Runtime/model capabilities exposed by the duplex serving protocol. - - These are intentionally explicit so the serving layer does not assume all - duplex models support the same input append, rollback, or turn policy. - ``supports_core_kv_lease`` is reserved for scheduler-owned KV lifecycle; - model-owned decoder/TTS state must use ``supports_model_internal_state``. - ``supports_core_resumable_request`` means the scheduler can resume the - same request id across streaming updates, but it is not a KV lease by - itself. Realtime support means this endpoint can speak the native Realtime - event schema for the supported audio duplex paths while keeping model- or - scheduler-specific limits explicit in the capability payload. - """ - - supports_session_adapter: bool = True - supports_model_native_turn_policy: bool = False - supports_external_turn_signal: bool = True - supports_client_commit: bool = True - supports_barge_in: bool = True - supports_playback_ack: bool = True - supports_input_append: bool = False - supports_replace_latest_chunk: bool = True - supports_reencode_context: bool = True - supports_rollback_to_checkpoint: bool = False - supports_turn_commit_only: bool = True - supports_kv_lease: bool = False - supports_core_kv_lease: bool = False - supports_model_internal_state: bool = False - supports_stage_resumption: bool = False - supports_scheduler_native_append: bool = False - supports_core_resumable_request: bool = False - supports_stage_connector_handoff: bool = False - supports_independent_io_streams: bool = False - supports_realtime_endpoint: bool = False - supports_multi_session: bool = False - supports_multi_session_same_replica: bool = False - supports_session_lease: bool = False - supports_session_resume: bool = False - session_admission_mode: str = "serving_managed" - supports_audio_truncate: bool = False - requires_model_runner_kv: bool = False - requires_native_stage_role: bool = False - implementation_level: str = "serving_session_adapter" - adapter_patterns: list[str] = field(default_factory=lambda: ["chunk_group_append"]) - input_modes: list[str] = field(default_factory=lambda: ["turn_commit_only", "reencode_context"]) - signal_sources: list[str] = field(default_factory=lambda: ["client_event", "server_policy", "model_native"]) - stage_handoff_transport: str | None = None - chunk_period_ms: int | None = 1000 - target_barge_in_latency_ms: int | None = 1000 - - def as_dict(self) -> dict[str, object]: - return { - "supports_session_adapter": self.supports_session_adapter, - "supports_model_native_turn_policy": self.supports_model_native_turn_policy, - "supports_external_turn_signal": self.supports_external_turn_signal, - "supports_client_commit": self.supports_client_commit, - "supports_barge_in": self.supports_barge_in, - "supports_playback_ack": self.supports_playback_ack, - "supports_input_append": self.supports_input_append, - "supports_replace_latest_chunk": self.supports_replace_latest_chunk, - "supports_reencode_context": self.supports_reencode_context, - "supports_rollback_to_checkpoint": self.supports_rollback_to_checkpoint, - "supports_turn_commit_only": self.supports_turn_commit_only, - "supports_kv_lease": self.supports_kv_lease, - "supports_core_kv_lease": self.supports_core_kv_lease, - "supports_model_internal_state": self.supports_model_internal_state, - "supports_stage_resumption": self.supports_stage_resumption, - "supports_scheduler_native_append": self.supports_scheduler_native_append, - "supports_core_resumable_request": self.supports_core_resumable_request, - "supports_stage_connector_handoff": self.supports_stage_connector_handoff, - "supports_independent_io_streams": self.supports_independent_io_streams, - "supports_realtime_endpoint": self.supports_realtime_endpoint, - "supports_multi_session": self.supports_multi_session, - "supports_multi_session_same_replica": self.supports_multi_session_same_replica, - "supports_session_lease": self.supports_session_lease, - "supports_session_resume": self.supports_session_resume, - "session_admission_mode": self.session_admission_mode, - "supports_audio_truncate": self.supports_audio_truncate, - "requires_model_runner_kv": self.requires_model_runner_kv, - "requires_native_stage_role": self.requires_native_stage_role, - "implementation_level": self.implementation_level, - "adapter_patterns": self.adapter_patterns, - "input_modes": self.input_modes, - "signal_sources": self.signal_sources, - "stage_handoff_transport": self.stage_handoff_transport, - "chunk_period_ms": self.chunk_period_ms, - "target_barge_in_latency_ms": self.target_barge_in_latency_ms, - } - - -@dataclass -class DuplexPlaybackCursor: - generated_ms: int = 0 - sent_ms: int = 0 - played_ms: int = 0 - committed_ms: int = 0 - generated_frames: int = 0 - sample_rate_hz: int | None = None - - def acknowledge(self, played_ms: int, committed_ms: int | None = None) -> None: - self.played_ms = max(self.played_ms, max(0, int(played_ms))) - if committed_ms is None: - committed_ms = self.played_ms - self.committed_ms = max(self.committed_ms, max(0, int(committed_ms))) - - def truncate_committed(self, committed_ms: int) -> None: - self.committed_ms = max(0, min(max(self.sent_ms, self.generated_ms), int(committed_ms))) - - def as_dict(self) -> dict[str, int]: - return { - "generated_ms": self.generated_ms, - "sent_ms": self.sent_ms, - "played_ms": self.played_ms, - "committed_ms": self.committed_ms, - } - - def snapshot(self) -> DuplexPlaybackView: - return DuplexPlaybackView(**self.as_dict()) - - -@dataclass(frozen=True, slots=True) -class DuplexPlaybackView: - generated_ms: int = 0 - sent_ms: int = 0 - played_ms: int = 0 - committed_ms: int = 0 - - def as_dict(self) -> dict[str, int]: - return { - "generated_ms": self.generated_ms, - "sent_ms": self.sent_ms, - "played_ms": self.played_ms, - "committed_ms": self.committed_ms, - } - - -@dataclass -class DuplexAudioChunk: - data: str - format: str = "wav" - sample_rate_hz: int | None = None - - -@dataclass -class DuplexSessionConfig: - model: str | None = None - modalities: list[str] = field(default_factory=lambda: ["text", "audio"]) - instructions: str | None = None - voice: str | None = None - ref_audio: str | None = None - response_format: str = "wav" - temperature: float | None = None - max_tokens: int | None = None - speed: float | None = None - use_tts_template: bool = True - idle_timeout_s: float = 300.0 - overlap_policy: str = DuplexOverlapPolicy.LISTEN_ONLY.value - overlap_short_ack_ms: int = 700 - overlap_barge_in_ms: int = 1200 - overlap_silence_rms: float = 0.003 - playback_commit_policy: str = DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value - turn_detection_configured: bool = False - server_vad: ServerVADConfig | None = None - extra_body: dict[str, object] = field(default_factory=dict) - - def as_dict(self) -> dict[str, object]: - payload: dict[str, object] = { - "model": self.model, - "modalities": list(self.modalities), - "instructions": self.instructions, - "voice": self.voice, - "ref_audio": self.ref_audio, - "response_format": self.response_format, - "temperature": self.temperature, - "max_tokens": self.max_tokens, - "speed": self.speed, - "use_tts_template": self.use_tts_template, - "idle_timeout_s": self.idle_timeout_s, - "overlap_policy": self.overlap_policy, - "overlap_short_ack_ms": self.overlap_short_ack_ms, - "overlap_barge_in_ms": self.overlap_barge_in_ms, - "overlap_silence_rms": self.overlap_silence_rms, - "playback_commit_policy": self.playback_commit_policy, - "extra_body": dict(self.extra_body), - } - if self.turn_detection_configured: - payload["turn_detection"] = self.server_vad.as_dict() if self.server_vad is not None else None - return payload - - @classmethod - def from_event(cls, event: dict[str, object]) -> DuplexSessionConfig: - payload = event.get("session") - if isinstance(payload, dict): - source = payload - else: - source = event - - config = cls() - if isinstance(source.get("model"), str): - config.model = source["model"] - if isinstance(source.get("instructions"), str): - config.instructions = source["instructions"] - if isinstance(source.get("voice"), str): - config.voice = source["voice"] - if isinstance(source.get("ref_audio"), str): - config.ref_audio = source["ref_audio"] - if isinstance(source.get("response_format"), str): - config.response_format = source["response_format"] - if isinstance(source.get("use_tts_template"), bool): - config.use_tts_template = bool(source["use_tts_template"]) - if isinstance(source.get("temperature"), int | float): - config.temperature = float(source["temperature"]) - if isinstance(source.get("max_tokens"), int): - config.max_tokens = int(source["max_tokens"]) - if isinstance(source.get("speed"), int | float): - config.speed = float(source["speed"]) - if isinstance(source.get("idle_timeout_s"), int | float): - config.idle_timeout_s = float(source["idle_timeout_s"]) - if isinstance(source.get("overlap_policy"), str): - config.overlap_policy = cls._normalize_overlap_policy(source["overlap_policy"]) - if isinstance(source.get("overlap_short_ack_ms"), int | float): - config.overlap_short_ack_ms = max(0, int(source["overlap_short_ack_ms"])) - if isinstance(source.get("overlap_barge_in_ms"), int | float): - config.overlap_barge_in_ms = max(0, int(source["overlap_barge_in_ms"])) - if isinstance(source.get("overlap_silence_rms"), int | float): - config.overlap_silence_rms = max(0.0, float(source["overlap_silence_rms"])) - if isinstance(source.get("playback_commit_policy"), str): - config.playback_commit_policy = cls._normalize_playback_commit_policy(source["playback_commit_policy"]) - if isinstance(source.get("modalities"), list) and all(isinstance(x, str) for x in source["modalities"]): - config.modalities = list(source["modalities"]) - turn_detection_configured, server_vad = parse_session_turn_detection(source) - if isinstance(source.get("extra_body"), dict): - config.extra_body = normalize_native_duplex_key(dict(source["extra_body"])) - extra = config.extra_body - raw_realtime_session = extra.get("realtime_session_payload") - if not turn_detection_configured and isinstance(raw_realtime_session, dict): - turn_detection_configured, server_vad = parse_session_turn_detection(raw_realtime_session) - if isinstance(extra.get("overlap_policy"), str): - config.overlap_policy = cls._normalize_overlap_policy(extra["overlap_policy"]) - if isinstance(extra.get("overlap_short_ack_ms"), int | float): - config.overlap_short_ack_ms = max(0, int(extra["overlap_short_ack_ms"])) - if isinstance(extra.get("overlap_barge_in_ms"), int | float): - config.overlap_barge_in_ms = max(0, int(extra["overlap_barge_in_ms"])) - if isinstance(extra.get("overlap_silence_rms"), int | float): - config.overlap_silence_rms = max(0.0, float(extra["overlap_silence_rms"])) - if isinstance(extra.get("playback_commit_policy"), str): - config.playback_commit_policy = cls._normalize_playback_commit_policy(extra["playback_commit_policy"]) - if turn_detection_configured: - config.turn_detection_configured = True - config.server_vad = server_vad - return config - - @staticmethod - def _normalize_overlap_policy(value: str) -> str: - normalized = value.strip().lower() - if normalized in {policy.value for policy in DuplexOverlapPolicy}: - return normalized - return DuplexOverlapPolicy.LISTEN_ONLY.value - - @staticmethod - def _normalize_playback_commit_policy(value: str) -> str: - normalized = value.strip().lower() - if normalized in {policy.value for policy in DuplexPlaybackCommitPolicy}: - return normalized - return DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value - - -@dataclass(frozen=True) -class ResponseCreateOptions: - instructions: str | None = None - voice: str | None = None - response_format: str | None = None - temperature: float | None = None - max_tokens: int | None = None - speed: float | None = None - modalities: tuple[str, ...] | None = None - extra_body: Mapping[str, object] = field(default_factory=dict) - - def apply_to(self, config: DuplexSessionConfig) -> None: - for field_name in ("instructions", "voice", "response_format", "temperature", "max_tokens", "speed"): - value = getattr(self, field_name) - if value is not None: - setattr(config, field_name, value) - if self.modalities is not None: - config.modalities = list(self.modalities) - config.extra_body.update(self.extra_body) - - -@dataclass -class DuplexCommittedInput: - message: dict[str, object] - turn_id: int - epoch: int - input_commit_seq: int - - -@dataclass -class PendingServerVADTurn: - committed: DuplexCommittedInput - create_response: bool - item_id: str - - -@dataclass -class DuplexAssistantAudioTextMark: - text_chars: int - audio_end_ms: int +def _default_lease() -> DuplexLeaseState: + return DuplexLeaseState(config=DuplexLeaseConfig(), generation=0, last_activity=time.monotonic()) @dataclass class InputBufferState: commit_seq: int = 0 - pending_text: list[str] = field(default_factory=list) - pending_audio: list[DuplexAudioChunk] = field(default_factory=list) overlap_speech_ms: int = 0 reserved_input_bytes: int = 0 pending_turns: int = 0 - turn_detection_config_locked: bool = False - vad_prefix_audio: list[np.ndarray] = field(default_factory=list) - vad_prefix_samples: int = 0 - vad_utterance_audio: list[np.ndarray] = field(default_factory=list) - vad_speech_active: bool = False - active_server_vad_item_id: str | None = None - pending_server_vad_turn: PendingServerVADTurn | None = None @dataclass @@ -413,10 +66,6 @@ class ResponseState: active_request_id: str | None = None active_response_id: str | None = None active_response_turn_id: int | None = None - request_started_at_s_by_turn: dict[int, float] = field(default_factory=dict) - active_response_request_started_at_s: float | None = None - active_response_ttft_ms: float | None = None - active_response_ttfp_ms: float | None = None active_response_input_commit_seq: int | None = None active_response_awaits_input_commit: bool = False last_response_id: str | None = None @@ -439,6 +88,10 @@ class PlaybackLedger: @dataclass class ConversationHistory: messages: list[dict[str, object]] = field(default_factory=list) + #: User items added by ``conversation.item.create`` that no response has + #: answered yet. A turn can start from these alone, the way the Realtime + #: protocol allows, without any audio having been committed. + unanswered_user_items: int = 0 item_ids: dict[str, dict[str, object]] = field(default_factory=dict) history_item_placeholders: dict[str, dict[str, object]] = field(default_factory=dict) item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) @@ -454,21 +107,264 @@ class ConversationHistory: ) +class DuplexFenceMismatchError(RuntimeError): + def __init__(self, expected: DuplexFence, actual: DuplexFence) -> None: + super().__init__(f"duplex fence mismatch: expected {expected!r}, got {actual!r}") + self.expected = expected + self.actual = actual + + @dataclass -class DuplexSession: +class DuplexRequestResource: + stage_id: int + request_id: str + fence: DuplexFence + submitted: bool = False + + +@dataclass +class DuplexInputAppend: + seq: int + turn_seq: int + turn_id: int + + +@dataclass(frozen=True) +class DuplexAppendReservation: + fence: DuplexFence + base_fence: DuplexFence + base_input_seq: int + base_input_turn_seq: int + base_append_turn_key: tuple[int, int, int] | None + update: DuplexInputAppend + + +@dataclass +class DuplexEngineSession: + """The one session state (see module docstring). + + Owned and mutated only by its ``DuplexSessionRunner`` on the orchestrator loop. + """ + session_id: str config: DuplexSessionConfig capabilities: DuplexCapabilities = field(default_factory=DuplexCapabilities) - incarnation: int = 0 state: DuplexSessionState = DuplexSessionState.OPEN turn_state: DuplexTurnState = DuplexTurnState.IDLE epoch: int = 0 turn_id: int = 0 + lease: DuplexLeaseState = field(default_factory=_default_lease, repr=False) + _clock: Callable[[], float] = field(default=time.monotonic, repr=False) + _runtime_config: dict[str, object] = field(default_factory=dict, repr=False) + #: Bumped on every published session / runtime config change. + config_generation: int = 0 + #: Stage request ids reserved or submitted for this session, keyed by ``(stage_id, request_id)``. + request_resources: dict[tuple[int, str], DuplexRequestResource] = field(default_factory=dict, repr=False) + #: Highest fence accepted for a stage request (monotonic; reset of the append + #: sequence happens when its epoch advances). + accepted_fence: DuplexFence = field(default=None, repr=False) # type: ignore[assignment] + input_seq: int = 0 + input_turn_seq: int = 0 + _append_turn_key: tuple[int, int, int] | None = field(default=None, repr=False) _input: InputBufferState = field(default_factory=InputBufferState, repr=False) _response: ResponseState = field(default_factory=ResponseState, repr=False) + #: Stage snapshots that arrived while no response was open. A stage whose + #: output feeds the next stage rather than the client reports its metrics + #: before the response those tokens end up in exists, so they are held here + #: and folded into the first response that opens after them. + _pending_stage_metrics: list[dict[str, dict[str, object]]] = field(default_factory=list, repr=False) _playback: PlaybackLedger = field(default_factory=PlaybackLedger, repr=False) _conversation: ConversationHistory = field(default_factory=ConversationHistory, repr=False) - _runtime_config: dict[str, object] = field(default_factory=dict, repr=False) + model_state: DuplexModelSessionState | None = field(default=None, repr=False) + projector: RealtimeProjectionState | None = field(default=None, repr=False) + created_monotonic: float = field(default_factory=time.monotonic) + + def __post_init__(self) -> None: + if self.accepted_fence is None: + self.accepted_fence = self.fence + else: + self.accept_fence(self.fence) + if self.config.initial_user_text: + # Seeded text is user input like any other: it waits for a response. + self.notify_new_user_item() + + # ---- identity / fence ---- + + @property + def fence(self) -> DuplexFence: + return DuplexFence(self.session_id, epoch=self.epoch, turn_id=self.turn_id) + + def sync_fence(self) -> DuplexFence: + """Publish the current (epoch, turn_id) identity as the accepted fence.""" + fence = self.fence + self.accept_fence(fence) + return fence + + def _validate_fence(self, fence: DuplexFence) -> None: + current = self.accepted_fence + if fence.session_id != self.session_id: + raise DuplexFenceMismatchError(current, fence) + if fence.epoch < current.epoch or (fence.epoch == current.epoch and fence.turn_id < current.turn_id): + raise DuplexFenceMismatchError(current, fence) + + def accept_fence(self, fence: DuplexFence) -> None: + self._validate_fence(fence) + if fence.epoch != self.accepted_fence.epoch: + self.input_seq = 0 + self.input_turn_seq = 0 + self._append_turn_key = None + self.accepted_fence = fence + + # ---- lease ---- + + @property + def lease_generation(self) -> int: + return self.lease.generation + + def touch_lease(self, activity: DuplexLeaseActivity) -> None: + self._validate_fence(self.fence) + self.lease.touch(self._clock(), activity) + + def detach_lease(self) -> None: + self._validate_fence(self.fence) + self.lease.detach(self._clock()) + + def resume_lease(self, *, expected_lease_generation: int) -> int: + self._validate_fence(self.fence) + return self.lease.resume(self._clock(), expected_generation=expected_lease_generation) + + def begin_lease_operation(self, fence: DuplexFence, operation_id: str) -> None: + self._validate_fence(fence) + self.lease.begin_operation(self._clock(), operation_id) + + def end_lease_operation(self, operation_id: str) -> None: + self._validate_fence(self.fence) + self.lease.end_operation(self._clock(), operation_id) + + # ---- stage request resources ---- + + def reserve_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: + self._validate_fence(fence) + resource_key = (stage_id, request_id) + existing = self.request_resources.get(resource_key) + if existing is not None: + if existing.fence.session_id != fence.session_id or existing.fence.epoch != fence.epoch: + raise ValueError(f"Duplex request resource already reserved with different identity: {request_id}") + return + self.request_resources[resource_key] = DuplexRequestResource( + stage_id=stage_id, + request_id=request_id, + fence=fence, + ) + + def bind_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: + self.reserve_stage_request(stage_id, request_id, fence=fence) + self.accept_fence(fence) + resource = self.request_resources[(stage_id, request_id)] + resource.fence = fence + resource.submitted = True + + def stage_request_submitted(self, stage_id: int, request_id: str) -> bool: + resource = self.request_resources.get((stage_id, request_id)) + return resource is not None and resource.submitted + + def resource_request_ids( + self, + fence: DuplexFence | None = None, + *, + submitted: bool | None = None, + ) -> list[str]: + return list( + dict.fromkeys( + resource.request_id + for resource in self.request_resources.values() + if (fence is None or resource.fence == fence) and (submitted is None or resource.submitted is submitted) + ) + ) + + def release_all_requests(self) -> list[str]: + """Drop every reserved/submitted stage request; returns their ids for stage cleanup.""" + stale = self.resource_request_ids() + self.request_resources.clear() + return stale + + # ---- append sequencing ---- + + def prepare_append(self, fence: DuplexFence) -> DuplexAppendReservation: + self._validate_fence(fence) + same_epoch = fence.epoch == self.accepted_fence.epoch + input_seq = self.input_seq if same_epoch else 0 + input_turn_seq = self.input_turn_seq if same_epoch else 0 + append_turn_key = self._append_turn_key if same_epoch else None + turn_key = (fence.epoch, fence.turn_id, 0) + turn_seq = input_turn_seq + 1 if turn_key == append_turn_key else 1 + return DuplexAppendReservation( + fence=fence, + base_fence=self.accepted_fence, + base_input_seq=self.input_seq, + base_input_turn_seq=self.input_turn_seq, + base_append_turn_key=self._append_turn_key, + update=DuplexInputAppend( + seq=input_seq + 1, + turn_seq=turn_seq, + turn_id=fence.turn_id, + ), + ) + + def commit_append(self, reservation: DuplexAppendReservation) -> DuplexInputAppend: + if ( + self.accepted_fence != reservation.base_fence + or self.input_seq != reservation.base_input_seq + or self.input_turn_seq != reservation.base_input_turn_seq + or self._append_turn_key != reservation.base_append_turn_key + ): + raise RuntimeError("duplex append reservation is stale") + self.accept_fence(reservation.fence) + self.input_seq = reservation.update.seq + self.input_turn_seq = reservation.update.turn_seq + self._append_turn_key = (reservation.fence.epoch, reservation.fence.turn_id, 0) + return reservation.update + + # ---- cancel / close of stage resources ---- + + def release_fence(self, fence: DuplexFence) -> list[str]: + stale = self.resource_request_ids(fence) + self.request_resources = { + resource_key: resource + for resource_key, resource in self.request_resources.items() + if resource.fence != fence + } + return stale + + def cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: + stale = self.prepare_cancel_fence(cancelled_fence, next_fence) + self.release_fence(cancelled_fence) + return stale + + def prepare_cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: + """Advance the cancellation fence without dropping cleanup records.""" + current = self.accepted_fence + if cancelled_fence.session_id != self.session_id: + raise DuplexFenceMismatchError(current, cancelled_fence) + if next_fence.session_id != self.session_id or next_fence.epoch <= cancelled_fence.epoch: + raise DuplexFenceMismatchError(cancelled_fence, next_fence) + current_key = (current.epoch, current.turn_id) + cancelled_key = (cancelled_fence.epoch, cancelled_fence.turn_id) + next_key = (next_fence.epoch, next_fence.turn_id) + if cancelled_key > current_key: + raise DuplexFenceMismatchError(current, cancelled_fence) + if next_key > current_key: + self.accept_fence(next_fence) + return self.resource_request_ids(cancelled_fence) + + def begin_close(self, *, reason: str) -> bool: + """Make close irreversible while retaining stage resources for cleanup retry.""" + self.accept_fence(self.fence) + if self.lease.terminal_reason is not None: + return True + return self.lease.mark_terminal(reason) + + # ---- configuration ---- @property def response_config(self) -> DuplexSessionConfig: @@ -486,19 +382,12 @@ class DuplexSession: def replace_runtime_config(self, runtime_config: Mapping[str, object]) -> None: self._runtime_config = dict(runtime_config) + self.config_generation += 1 @property def input_commit_seq(self) -> int: return self._input.commit_seq - @property - def pending_text(self) -> tuple[str, ...]: - return tuple(self._input.pending_text) - - @property - def pending_audio(self) -> tuple[DuplexAudioChunk, ...]: - return tuple(self._input.pending_audio) - @property def history(self) -> tuple[dict[str, object], ...]: placeholders = {id(message) for message in self._conversation.history_item_placeholders.values()} @@ -620,10 +509,6 @@ class DuplexSession: def playback(self) -> DuplexPlaybackView: return self._playback.current.snapshot() - @property - def response_playbacks(self) -> Mapping[str, DuplexPlaybackView]: - return MappingProxyType({key: value.snapshot() for key, value in self._playback.by_response.items()}) - @property def history_item_ids(self) -> Mapping[str, dict[str, object]]: return MappingProxyType({key: dict(value) for key, value in self._conversation.item_ids.items()}) @@ -637,10 +522,11 @@ class DuplexSession: return MappingProxyType(dict(self._conversation.pending_truncations_ms)) def replace_config(self, config: DuplexSessionConfig) -> None: + previous_seed = self.config.initial_user_text self.config = config - - def replace_capabilities(self, capabilities: DuplexCapabilities) -> None: - self.capabilities = capabilities + self.config_generation += 1 + if config.initial_user_text and config.initial_user_text != previous_seed: + self.notify_new_user_item() def transition_turn(self, state: DuplexTurnState) -> None: self.turn_state = state @@ -660,45 +546,6 @@ class DuplexSession: def bind_response_turn(self, turn_id: int | None) -> None: self._response.active_response_turn_id = turn_id - def mark_model_turn_request_started(self, turn_id: int, started_at_s: float) -> None: - """Record the latest native request start that can own one model turn.""" - turn_id = int(turn_id) - started_at_s = float(started_at_s) - if self.active_response_turn_id == turn_id: - if self._response.active_response_request_started_at_s is None: - self._response.active_response_request_started_at_s = started_at_s - return - self._response.request_started_at_s_by_turn[turn_id] = started_at_s - - def mark_response_first_outputs( - self, - *, - observed_at_s: float, - has_text: bool, - has_audio: bool, - ) -> dict[str, object]: - """Return server-monotonic TTF metrics observed for the active response.""" - started_at_s = self._response.active_response_request_started_at_s - if started_at_s is None: - return {} - elapsed_ms = max(0.0, (float(observed_at_s) - started_at_s) * 1000.0) - if has_text and self._response.active_response_ttft_ms is None: - self._response.active_response_ttft_ms = elapsed_ms - if has_audio and self._response.active_response_ttfp_ms is None: - self._response.active_response_ttfp_ms = elapsed_ms - metrics: dict[str, object] = { - "source": "server_monotonic_request_start", - "measurement_origin": { - "ttft": "native model-turn request execution start to first non-empty text output", - "ttfp": "native model-turn request execution start to first audio output", - }, - } - if self._response.active_response_ttft_ms is not None: - metrics["ttft_ms"] = self._response.active_response_ttft_ms - if self._response.active_response_ttfp_ms is not None: - metrics["ttfp_ms"] = self._response.active_response_ttfp_ms - return metrics - def active_response_accepts_model_turn(self, turn_id: int | None) -> bool: if self._response.active_response_id is None: return False @@ -708,19 +555,19 @@ class DuplexSession: return active_turn_id is None or int(turn_id) == active_turn_id def append_history_message(self, message: dict[str, object]) -> None: - response_id = self.active_response_id - if response_id is not None: - self.reserve_history_item(f"item_{response_id}") self._conversation.messages.append(message) - def stage_pending_history_item(self, item_id: str, message: dict[str, object]) -> None: - self._conversation.pending_item_ids[item_id] = dict(message) + def notify_new_user_item(self) -> None: + """Record a user item that a later ``response.create`` may answer.""" + self._conversation.unanswered_user_items += 1 - def append_text(self, text: str) -> None: - if not text: - return - self._input.pending_text.append(text) - self.turn_state = DuplexTurnState.USER_SPEAKING + def unanswered_user_items(self) -> int: + """How many user items are waiting for a response.""" + return self._conversation.unanswered_user_items + + def reset_unanswered_user_items(self) -> None: + """A turn has started: the pending items are now its input.""" + self._conversation.unanswered_user_items = 0 @property def pending_input_bytes(self) -> int: @@ -737,13 +584,6 @@ class DuplexSession: self._input.reserved_input_bytes += size return True - @property - def turn_detection_config_locked(self) -> bool: - return self._input.turn_detection_config_locked - - def lock_turn_detection_config(self) -> None: - self._input.turn_detection_config_locked = True - def release_input_bytes(self, size: int) -> None: self._input.reserved_input_bytes = max(0, self._input.reserved_input_bytes - max(0, int(size))) @@ -759,200 +599,18 @@ class DuplexSession: def release_pending_turn(self) -> None: self._input.pending_turns = max(0, self._input.pending_turns - 1) - def clear_pending_turn_reservations(self) -> None: - self._input.pending_turns = 0 - - def append_audio(self, data: str, *, fmt: str = "wav", sample_rate_hz: int | None = None) -> None: - if not data: - return - self._input.pending_audio.append(DuplexAudioChunk(data=data, format=fmt, sample_rate_hz=sample_rate_hz)) - self.turn_state = DuplexTurnState.USER_SPEAKING - - def append_server_vad_frame( - self, - samples: np.ndarray, - *, - speech_started: bool, - speech_stopped: bool, - prefix_samples: int, - ) -> int: - """Store one detector frame while keeping commit-eligible audio session-owned. - - Returns the number of raw float32 bytes released when the rolling - pre-speech prefix is trimmed. - """ - # Session-owned audio must not retain the batch backing a detector frame. - frame = np.array(samples, dtype=np.float32, copy=True).reshape(-1) - released_samples = 0 - if not self._input.vad_speech_active: - if speech_started: - self._input.vad_utterance_audio = [*self._input.vad_prefix_audio, frame] - self._input.vad_prefix_audio = [] - self._input.vad_prefix_samples = 0 - self._input.vad_speech_active = True - else: - self._input.vad_prefix_audio.append(frame) - self._input.vad_prefix_samples += frame.size - while self._input.vad_prefix_samples > prefix_samples and self._input.vad_prefix_audio: - overflow = self._input.vad_prefix_samples - prefix_samples - first = self._input.vad_prefix_audio[0] - if first.size <= overflow: - self._input.vad_prefix_audio.pop(0) - self._input.vad_prefix_samples -= first.size - released_samples += first.size - else: - self._input.vad_prefix_audio[0] = first[overflow:].copy() - self._input.vad_prefix_samples -= overflow - released_samples += overflow - else: - self._input.vad_utterance_audio.append(frame) - - if speech_stopped: - self._input.vad_speech_active = False - if self._input.vad_speech_active or speech_started: - self.turn_state = DuplexTurnState.USER_SPEAKING - return released_samples * np.dtype(np.float32).itemsize - - def stage_server_vad_audio_for_commit(self) -> bool: - if not self._input.vad_utterance_audio: - return False - samples = np.concatenate(self._input.vad_utterance_audio) - encoded = encode_float32_mono_wav_base64(samples, sample_rate_hz=16_000) - self._input.pending_audio.append(DuplexAudioChunk(data=encoded, format="wav", sample_rate_hz=16_000)) - self._input.vad_utterance_audio.clear() - self._input.vad_speech_active = False - return True - - def clear_server_vad_audio(self) -> None: - self._input.vad_prefix_audio.clear() - self._input.vad_prefix_samples = 0 - self._input.vad_utterance_audio.clear() - self._input.vad_speech_active = False - self._input.active_server_vad_item_id = None - - @property - def active_server_vad_item_id(self) -> str | None: - return self._input.active_server_vad_item_id - - @property - def server_vad_utterance_bytes(self) -> int: - return sum(frame.nbytes for frame in self._input.vad_utterance_audio) - - def discard_uncommitted_server_vad_utterance(self) -> int: - released = self.server_vad_utterance_bytes - self._input.vad_utterance_audio.clear() - self._input.vad_speech_active = False - self._input.active_server_vad_item_id = None - self._input.pending_text.clear() - return released - - def begin_server_vad_speech(self, item_id: str) -> None: - self._input.active_server_vad_item_id = item_id - - def finish_server_vad_speech(self) -> str | None: - item_id = self._input.active_server_vad_item_id - self._input.active_server_vad_item_id = None - return item_id - def mark_user_input_activity(self) -> None: self.turn_state = DuplexTurnState.USER_SPEAKING - def mark_assistant_generating(self) -> None: - self.turn_state = DuplexTurnState.ASSISTANT_GENERATING - def cancel_pending_input(self) -> dict[str, int]: - cancelled = { - "text_chunks": len(self._input.pending_text), - "audio_chunks": len(self._input.pending_audio), - } - self._input.pending_text.clear() - self._input.pending_audio.clear() + """Drop every pending reservation; the PCM buffer itself is model state cleared by the runner.""" + cancelled = {"text_chunks": 0, "audio_chunks": 0} self._input.reserved_input_bytes = 0 self._input.pending_turns = 0 - self.clear_server_vad_audio() self.turn_state = DuplexTurnState.IDLE return cancelled - def commit_user_input( - self, - *, - append_history: bool = True, - retain_input_bytes: int = 0, - ) -> DuplexCommittedInput | None: - text = "".join(self._input.pending_text).strip() - audio_chunks = list(self._input.pending_audio) - if not text and not audio_chunks: - return None - - content: str | list[dict[str, object]] - if audio_chunks: - content_items: list[dict[str, object]] = [] - if text: - content_items.append({"type": "text", "text": text}) - for chunk in audio_chunks: - content_items.append( - { - "type": "audio_url", - "audio_url": { - "url": f"data:audio/{chunk.format};base64,{chunk.data}", - }, - } - ) - content = content_items - else: - content = text - - self._bind_active_response_to_input_commit(self._input.commit_seq + 1) - self._input.commit_seq += 1 - message = {"role": "user", "content": content} - if append_history: - self._conversation.messages.append(message) - self._input.pending_text.clear() - self._input.pending_audio.clear() - self._input.reserved_input_bytes = max(0, int(retain_input_bytes)) - self.turn_state = DuplexTurnState.USER_COMMITTED - return DuplexCommittedInput( - message=message, - turn_id=self.turn_id, - epoch=self.epoch, - input_commit_seq=self._input.commit_seq, - ) - - @property - def pending_server_vad_turn(self) -> PendingServerVADTurn | None: - return self._input.pending_server_vad_turn - - def stage_server_vad_turn( - self, - committed: DuplexCommittedInput, - *, - create_response: bool, - item_id: str, - ) -> None: - if self._input.pending_server_vad_turn is not None: - raise RuntimeError("A Server VAD turn is already pending") - response_id = self.active_response_id - if response_id is not None: - # A pending user turn advances input_commit_seq before the active - # response completes. Reserve its current history position first. - self.reserve_history_item(f"item_{response_id}") - self._input.pending_server_vad_turn = PendingServerVADTurn( - committed=committed, - create_response=create_response, - item_id=item_id, - ) - - def activate_server_vad_turn(self) -> PendingServerVADTurn | None: - pending = self._input.pending_server_vad_turn - if pending is None: - return None - self._input.pending_server_vad_turn = None - self._conversation.messages.append(pending.committed.message) - self.register_history_item(pending.item_id, pending.committed.message) - self.turn_state = DuplexTurnState.USER_COMMITTED - return pending - - def commit_native_audio_input( + def commit_audio_input( self, *, transcript: str | None = None, @@ -970,8 +628,6 @@ class DuplexSession: if transcript: message["transcript"] = transcript self._conversation.messages.append(message) - self._input.pending_text.clear() - self._input.pending_audio.clear() self.turn_state = DuplexTurnState.USER_COMMITTED return DuplexCommittedInput( message=message, @@ -983,13 +639,9 @@ class DuplexSession: def complete_model_turn(self, turn_id: int) -> None: """Advance the model-owned output identity after its terminal signal.""" completed_turn_id = int(turn_id) - self._response.request_started_at_s_by_turn = { - pending_turn_id: started_at_s - for pending_turn_id, started_at_s in self._response.request_started_at_s_by_turn.items() - if pending_turn_id > completed_turn_id - } if completed_turn_id >= self.turn_id: self.turn_id = completed_turn_id + 1 + self.sync_fence() def _bind_active_response_to_input_commit(self, input_commit_seq: int) -> None: if self.active_response_id is not None and self._response.active_response_awaits_input_commit: @@ -1024,12 +676,6 @@ class DuplexSession: response_id = f"resp-{self.session_id}-{self.epoch}-{uuid4().hex[:8]}" self._response.active_response_id = response_id self._response.active_response_turn_id = self.turn_id if turn_id is None else int(turn_id) - self._response.active_response_request_started_at_s = self._response.request_started_at_s_by_turn.pop( - self._response.active_response_turn_id, - None, - ) - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = self.input_commit_seq self._response.active_response_awaits_input_commit = self.turn_state == DuplexTurnState.USER_SPEAKING self._response.last_response_id = response_id @@ -1048,11 +694,39 @@ class DuplexSession: self._response.stage_metric_tpot_weighted_ms.clear() self._response.stage_metric_tpot_weight.clear() + def stash_stage_metrics(self, stage_metrics: Mapping[object, object] | None) -> None: + """Hold a stage snapshot until a response exists to attribute it to. + + A stage that hands its output to the next stage instead of the client + (stage 0 feeding the TTS stage) reports its token metrics before the + response carrying those tokens is created. Dropping them would leave + every response without an engine-side token count; holding them keeps + the count whole, at the cost of attributing a turn the model never + spoke to the next response it does speak. + """ + if not isinstance(stage_metrics, Mapping): + return + snapshot = { + str(stage_id): dict(values) for stage_id, values in stage_metrics.items() if isinstance(values, Mapping) + } + if not snapshot: + return + if self.active_response_id is not None: + self.accumulate_response_stage_metrics(snapshot) + return + self._pending_stage_metrics.append(snapshot) + def accumulate_response_stage_metrics( self, stage_metrics: Mapping[object, object] | None, ) -> dict[str, dict[str, object]]: - if self.active_response_id is None or not isinstance(stage_metrics, Mapping): + if self.active_response_id is None: + return copy.deepcopy(self._response.stage_metrics) + if self._pending_stage_metrics: + pending, self._pending_stage_metrics = self._pending_stage_metrics, [] + for held in pending: + self.accumulate_response_stage_metrics(held) + if not isinstance(stage_metrics, Mapping): return copy.deepcopy(self._response.stage_metrics) additive_fields = ( @@ -1111,22 +785,17 @@ class DuplexSession: tpot_ms = raw_values.get("vllm_tpot_ms") token_count = raw_values.get("num_tokens_out") - if isinstance(tpot_ms, int | float) and not isinstance(tpot_ms, bool) and tpot_ms > 0: - weight = ( - max(int(token_count) - 1, 0) - if isinstance(token_count, int | float) and not isinstance(token_count, bool) - else 1 + if isinstance(tpot_ms, int | float) and tpot_ms > 0: + weight = max(int(token_count) - 1, 1) if isinstance(token_count, int | float) else 1 + self._response.stage_metric_tpot_weighted_ms[stage_id] = ( + self._response.stage_metric_tpot_weighted_ms.get(stage_id, 0.0) + float(tpot_ms) * weight + ) + self._response.stage_metric_tpot_weight[stage_id] = ( + self._response.stage_metric_tpot_weight.get(stage_id, 0) + weight + ) + current["vllm_tpot_ms"] = self._response.stage_metric_tpot_weighted_ms[stage_id] / float( + self._response.stage_metric_tpot_weight[stage_id] ) - if weight > 0: - self._response.stage_metric_tpot_weighted_ms[stage_id] = ( - self._response.stage_metric_tpot_weighted_ms.get(stage_id, 0.0) + float(tpot_ms) * weight - ) - self._response.stage_metric_tpot_weight[stage_id] = ( - self._response.stage_metric_tpot_weight.get(stage_id, 0) + weight - ) - current["vllm_tpot_ms"] = self._response.stage_metric_tpot_weighted_ms[stage_id] / float( - self._response.stage_metric_tpot_weight[stage_id] - ) for name, value in raw_values.items(): if name not in handled_fields: @@ -1164,54 +833,25 @@ class DuplexSession: if text: self._response.assistant_text_buffer.append(text) - def record_generated_audio(self, response_id: str, *, frame_count: int, sample_rate_hz: int) -> int | None: - """Accumulate source frames, rounding only the cumulative duration. - - A late chunk cannot change a different response's playback cursor. - This records encoded audio, not transport delivery or client playback. - """ - if response_id != self.active_response_id: - return None - if type(frame_count) is not int or frame_count < 0: - raise ValueError("Audio frame_count must be a non-negative integer") - if type(sample_rate_hz) is not int or sample_rate_hz <= 0: - raise ValueError("Audio sample_rate_hz must be a positive integer") - playback = self._playback.current - if playback.sample_rate_hz is not None and playback.sample_rate_hz != sample_rate_hz: - raise ValueError("Audio sample rate changed within a response") - playback.sample_rate_hz = sample_rate_hz - playback.generated_frames += frame_count - playback.generated_ms = playback.generated_frames * 1000 // sample_rate_hz - return playback.generated_ms - def mark_audio_sent( self, duration_ms: int | None = None, *, - response_id: str | None = None, text_chars: int | None = None, audio_text_marks: list[dict[str, object]] | None = None, ) -> None: - """Record cumulative audio accepted by the output path. - - ``sent_ms`` includes queued/journaled output; only client acknowledgements - establish ``played_ms``. It is not confirmation of socket delivery. - """ - playback = self._playback.current if response_id is None else self._playback.by_response.get(response_id) - if playback is None: - return - is_current = response_id is None or response_id == self.active_response_id + playback = self._playback.current if duration_ms is not None: playback.generated_ms = max(playback.generated_ms, duration_ms) playback.sent_ms = max(playback.sent_ms, duration_ms) - if is_current and text_chars is not None and text_chars >= 0: + if text_chars is not None and text_chars >= 0: self._response.assistant_audio_text_marks.append( DuplexAssistantAudioTextMark( text_chars=int(text_chars), audio_end_ms=max(0, int(duration_ms)), ) ) - if is_current and audio_text_marks: + if audio_text_marks: for raw_mark in audio_text_marks: if not isinstance(raw_mark, dict): continue @@ -1225,8 +865,7 @@ class DuplexSession: audio_end_ms=max(0, int(raw_audio_end_ms)), ) ) - if is_current: - self.turn_state = DuplexTurnState.ASSISTANT_PLAYING + self.turn_state = DuplexTurnState.ASSISTANT_PLAYING def _playback_cursor_for_response(self, response_id: str | None = None) -> DuplexPlaybackCursor: if response_id is None: @@ -1284,14 +923,7 @@ class DuplexSession: response_input_commit_seq = self._response.active_response_input_commit_seq if response_input_commit_seq is None: response_input_commit_seq = self.input_commit_seq - response_item_id = f"item_{response_id}" if response_id is not None else None - response_history_is_reserved = response_item_id is not None and ( - response_item_id in self._conversation.item_ids - or response_item_id in self._conversation.history_item_placeholders - ) - response_history_is_late = ( - self.input_commit_seq > response_input_commit_seq and not response_history_is_reserved - ) + response_history_is_late = self.input_commit_seq > response_input_commit_seq assistant_text = "".join(self._response.assistant_text_buffer).strip() message = None if assistant_text: @@ -1312,8 +944,9 @@ class DuplexSession: committed_text = "" if commit_text and committed_text: message = {"role": "assistant", "content": committed_text} - if response_item_id is not None and response_history_is_reserved: - self._store_history_item_message(response_item_id, message) + item_id = f"item_{response_id}" if response_id is not None else None + if item_id is not None and item_id in self._conversation.history_item_placeholders: + self._store_history_item_message(item_id, message) else: self._conversation.messages.append(message) effective_playback_policy = playback_commit_policy or self.config.playback_commit_policy @@ -1324,16 +957,13 @@ class DuplexSession: and effective_playback_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value ): self.register_history_item(f"item_{response_id}", None) - elif response_id is not None and (not assistant_text or not commit_text): + elif response_id is not None and not assistant_text: self._discard_history_item_placeholder(f"item_{response_id}") self._response.assistant_text_buffer.clear() if not preserve_request: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None - self._response.active_response_request_started_at_s = None - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self._clear_response_metrics() @@ -1386,11 +1016,6 @@ class DuplexSession: ) def delete_history_item(self, item_id: str) -> bool: - removed_pending_turn = False - pending_server_vad_turn = self._input.pending_server_vad_turn - if pending_server_vad_turn is not None and pending_server_vad_turn.item_id == item_id: - self._input.pending_server_vad_turn = None - removed_pending_turn = True response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None message = self._conversation.item_ids.pop(item_id, None) self._conversation.item_audio_text_marks.pop(item_id, None) @@ -1403,7 +1028,7 @@ class DuplexSession: if response_id is not None: self._conversation.assistant_response_snapshots.pop(response_id, None) if message is None: - return removed_pending_turn or pending is not None or removed_placeholder + return pending is not None or removed_placeholder try: self._conversation.messages.remove(message) except ValueError: @@ -1644,17 +1269,12 @@ class DuplexSession: def barge_in(self) -> int: self.epoch += 1 - self._input.pending_text.clear() - self._input.pending_audio.clear() + self.sync_fence() self._response.assistant_text_buffer.clear() self._response.assistant_audio_text_marks.clear() self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None - self._response.request_started_at_s_by_turn.clear() - self._response.active_response_request_started_at_s = None - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self._clear_response_metrics() @@ -1670,17 +1290,35 @@ class DuplexSession: self.state = DuplexSessionState.CLOSED self.turn_state = DuplexTurnState.IDLE self._response.active_response_turn_id = None - self._response.request_started_at_s_by_turn.clear() - self._response.active_response_request_started_at_s = None - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False - self.clear_server_vad_audio() - self._input.pending_server_vad_turn = None self._clear_response_metrics() self._restore_response_config() + def signal_turn(self, event_type: str, payload: Mapping[str, object] | None = None) -> TurnEvent: + """Apply one external turn signal and return the typed ``turn.event``.""" + payload = payload or {} + if event_type == DuplexTurnEventType.USER_STARTED.value: + self.transition_turn(DuplexTurnState.USER_SPEAKING) + elif event_type == DuplexTurnEventType.USER_COMMITTED.value: + self.transition_turn(DuplexTurnState.USER_COMMITTED) + elif event_type == DuplexTurnEventType.ASSISTANT_STARTED.value: + self.transition_turn(DuplexTurnState.ASSISTANT_GENERATING) + elif event_type == DuplexTurnEventType.ASSISTANT_DONE.value: + self.transition_turn(DuplexTurnState.IDLE) + elif event_type == DuplexTurnEventType.PLAYBACK_ACK.value: + played_ms = int(payload.get("played_ms", 0) or 0) + committed_ms = payload.get("committed_ms") + self.acknowledge_playback( + played_ms, + int(committed_ms) if isinstance(committed_ms, int | float) else None, + ) + elif event_type == DuplexTurnEventType.BARGE_IN.value: + self.transition_turn(DuplexTurnState.BARGE_IN) + elif event_type in {DuplexTurnEventType.CLOSE.value, DuplexTurnEventType.TIMEOUT.value}: + self.transition_session(DuplexSessionState.CLOSING) + return TurnEvent(event=event_type, turn_state=self.turn_state.value) + def as_public_dict(self) -> dict[str, object]: payload: dict[str, object] = { "id": self.session_id, @@ -1708,8 +1346,6 @@ class DuplexSession: "playback": self.playback.as_dict(), "capabilities": self.capabilities.as_dict(), } - if self.config.turn_detection_configured: - payload["turn_detection"] = self.config.server_vad.as_dict() if self.config.server_vad is not None else None if isinstance(self.config.extra_body.get("realtime_tools"), list): payload["tools"] = self.config.extra_body["realtime_tools"] if isinstance(self.config.extra_body.get("realtime_tool_choice"), str | dict): @@ -1738,73 +1374,14 @@ class DuplexSession: return payload -class DuplexTurnController: - """Interaction controller that accepts signals from multiple sources.""" - - def signal( - self, - session: DuplexSession, - event_type: str, - payload: dict[str, object] | None = None, - ) -> dict[str, object]: - payload = payload or {} - if event_type == DuplexTurnEventType.USER_STARTED.value: - session.transition_turn(DuplexTurnState.USER_SPEAKING) - elif event_type == DuplexTurnEventType.USER_COMMITTED.value: - session.transition_turn(DuplexTurnState.USER_COMMITTED) - elif event_type == DuplexTurnEventType.ASSISTANT_STARTED.value: - session.transition_turn(DuplexTurnState.ASSISTANT_GENERATING) - elif event_type == DuplexTurnEventType.ASSISTANT_DONE.value: - session.transition_turn(DuplexTurnState.IDLE) - elif event_type == DuplexTurnEventType.PLAYBACK_ACK.value: - played_ms = int(payload.get("played_ms", 0) or 0) - committed_ms = payload.get("committed_ms") - session.acknowledge_playback( - played_ms, - int(committed_ms) if isinstance(committed_ms, int | float) else None, - ) - elif event_type == DuplexTurnEventType.BARGE_IN.value: - session.transition_turn(DuplexTurnState.BARGE_IN) - elif event_type in {DuplexTurnEventType.CLOSE.value, DuplexTurnEventType.TIMEOUT.value}: - session.transition_session(DuplexSessionState.CLOSING) - return { - "type": "turn.event", - "session_id": session.session_id, - "event": event_type, - "turn_state": session.turn_state.value, - "epoch": session.epoch, - } - - -class DuplexSessionRegistry: - def __init__(self, capabilities: DuplexCapabilities | None = None) -> None: - self._capabilities = capabilities or DuplexCapabilities() - self._sessions: dict[str, DuplexSession] = {} - self._next_incarnation_by_session_id: dict[str, int] = {} - - def create(self, config: DuplexSessionConfig | None = None, session_id: str | None = None) -> DuplexSession: - sid = session_id or f"duplex-{uuid4().hex}" - if sid in self._sessions: - raise ValueError(f"Duplex session already exists: {sid}") - incarnation = self._next_incarnation_by_session_id.get(sid, 0) - self._next_incarnation_by_session_id[sid] = incarnation + 1 - session = DuplexSession( - session_id=sid, - config=config or DuplexSessionConfig(), - capabilities=self._capabilities, - incarnation=incarnation, - ) - self._sessions[sid] = session - return session - - def active_count(self) -> int: - return len(self._sessions) - - def get(self, session_id: str) -> DuplexSession | None: - return self._sessions.get(session_id) - - def close(self, session_id: str) -> DuplexSession | None: - session = self._sessions.pop(session_id, None) - if session is not None: - session.close() - return session +__all__ = [ + "ConversationHistory", + "DuplexAppendReservation", + "DuplexEngineSession", + "DuplexFenceMismatchError", + "DuplexInputAppend", + "DuplexRequestResource", + "InputBufferState", + "PlaybackLedger", + "ResponseState", +] diff --git a/vllm_omni/engine/duplex/session/helpers.py b/vllm_omni/engine/duplex/session/helpers.py new file mode 100644 index 000000000..f8b705e6d --- /dev/null +++ b/vllm_omni/engine/duplex/session/helpers.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Small reads and payload builders over one duplex session. + +None of these hold state or decide anything: they answer a question about the +session, or shape one of the wire payloads its events carry. They were methods +on the runner only because that is where the session reference lived, which +made every component that needed one a reason to reach back into the runner. + +Functions, not a component, because there is nothing here to own. +""" + +from __future__ import annotations + +import binascii +from typing import TYPE_CHECKING + +import pybase64 as base64 + +from vllm_omni.engine.duplex.config import DuplexPlaybackCommitPolicy +from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id +from vllm_omni.engine.duplex.events import ErrorEvent, OverlapDecision, error_event + +if TYPE_CHECKING: + from collections.abc import Mapping + + from vllm_omni.engine.duplex.session.context import DuplexSessionTasks + from vllm_omni.engine.duplex.session.engine_session import DuplexCommittedInput, DuplexEngineSession + + +# --------------------------------------------------------------------------- # +# Questions about the session # +# --------------------------------------------------------------------------- # + + +def stage0_request_id(session: DuplexEngineSession, epoch: int) -> str: + return duplex_resource_request_id(DuplexFence(session.session_id, epoch=epoch), "stage0") + + +def response_in_progress(session: DuplexEngineSession, tasks: DuplexSessionTasks) -> bool: + """Whether the model still owns the turn. + + Broader than ``active_response_id``: audio already sent but not yet played + back, and an append still in flight, both mean the turn is not free. + """ + if session.active_response_id is not None: + return True + if assistant_playback_active(session): + return True + active_task = tasks.active_response_task + if active_task is not None and not active_task.done(): + return True + return tasks.has_response_bound_append_tasks() + + +def assistant_playback_active(session: DuplexEngineSession) -> bool: + """Whether the client is still playing audio the session has sent.""" + return ( + session.config.playback_commit_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value + and session.playback.sent_ms > session.playback.committed_ms + ) + + +def audio_payload_size_bytes(payload: Mapping[str, object]) -> int: + audio = payload.get("audio") or payload.get("data") + if not isinstance(audio, str): + return 0 + try: + return len(base64.b64decode(audio, validate=True)) + except (ValueError, binascii.Error): + return 0 + + +# --------------------------------------------------------------------------- # +# Turn transitions # +# --------------------------------------------------------------------------- # + + +def advance_barge_in_epoch(session: DuplexEngineSession) -> tuple[int, dict[str, int]]: + """Start a new epoch for a barge-in, returning it with the playback it cut off.""" + old_playback = session.playback.as_dict() + new_epoch = session.barge_in() + session.clear_playback_cursor() + return new_epoch, old_playback + + +def commit_played_response_history( + session: DuplexEngineSession, + response_id: str | None, + committed_ms: int, +) -> None: + """Truncate an interrupted response in history to what the client actually heard.""" + if not response_id or committed_ms < 0: + return + session.truncate_history_item( + f"item_{response_id}", + audio_end_ms=committed_ms, + playback=session.playback_for_response(response_id), + ) + + +def commit_audio_input( + session: DuplexEngineSession, + *, + realtime_item_id: object | None = None, + transcript: object | None = None, + turn_id: int | None = None, +) -> DuplexCommittedInput: + clean_transcript = transcript.strip() if isinstance(transcript, str) else None + committed = session.commit_audio_input( + transcript=clean_transcript or None, + turn_id=turn_id, + ) + if isinstance(realtime_item_id, str) and realtime_item_id: + session.register_history_item(realtime_item_id, committed.message) + return committed + + +# --------------------------------------------------------------------------- # +# Payloads # +# --------------------------------------------------------------------------- # + + +def barge_in_unsupported_error() -> ErrorEvent: + return error_event("barge_in_unsupported", "Barge-in is not supported by this duplex model") + + +def overlap_decision_event(session: DuplexEngineSession, decision: dict[str, object]) -> OverlapDecision: + details: dict[str, object] = { + "type": "overlap.decision", + "session_id": session.session_id, + "epoch": session.epoch, + "policy": session.config.overlap_policy, + **decision, + } + action = decision.get("action") + reason = decision.get("reason") + return OverlapDecision( + policy=session.config.overlap_policy, + action=action if isinstance(action, str) else None, + reason=reason if isinstance(reason, str) else None, + details=details, + ) + + +def input_committed_payload( + session: DuplexEngineSession, + committed: DuplexCommittedInput, + *, + realtime_item_id: object | None = None, +) -> dict[str, object]: + payload: dict[str, object] = { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": committed.turn_id, + "epoch": committed.epoch, + "history_len": len(session.history), + "message": committed.message, + } + if isinstance(realtime_item_id, str) and realtime_item_id: + payload["realtime_item_id"] = realtime_item_id + return payload + + +def audio_committed_payload( + session: DuplexEngineSession, + *, + committed: DuplexCommittedInput | None = None, + realtime_item_id: object | None = None, + transcript: object | None = None, +) -> dict[str, object]: + message = committed.message if committed is not None else None + if not isinstance(message, dict): + input_audio_part: dict[str, object] = { + "type": "audio_url", + "audio_url": {"url": "native-duplex:input-audio"}, + } + if isinstance(transcript, str) and transcript: + input_audio_part["transcript"] = transcript + message = {"role": "user", "content": [input_audio_part]} + payload: dict[str, object] = { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": committed.turn_id if committed is not None else session.turn_id, + "epoch": committed.epoch if committed is not None else session.epoch, + "history_len": len(session.history), + "native_audio": True, + "message": message, + } + if isinstance(transcript, str) and transcript: + payload["transcript"] = transcript + if isinstance(realtime_item_id, str) and realtime_item_id: + payload["realtime_item_id"] = realtime_item_id + return payload diff --git a/vllm_omni/engine/duplex/lease.py b/vllm_omni/engine/duplex/session/lease.py similarity index 92% rename from vllm_omni/engine/duplex/lease.py rename to vllm_omni/engine/duplex/session/lease.py index c38faacd1..aea204039 100644 --- a/vllm_omni/engine/duplex/lease.py +++ b/vllm_omni/engine/duplex/session/lease.py @@ -8,8 +8,6 @@ from __future__ import annotations from dataclasses import dataclass, field from enum import Enum -from vllm_omni.engine.duplex.messages import DuplexFence - class DuplexLeaseActivity(str, Enum): APPEND = "append" @@ -111,19 +109,8 @@ class DuplexLeaseState: return True -@dataclass(frozen=True) -class DuplexSessionExpiry: - session_id: str - fence: DuplexFence - lease_generation: int - reason: str - submitted_request_ids: tuple[str, ...] - reserved_request_ids: tuple[str, ...] - - __all__ = [ "DuplexLeaseActivity", "DuplexLeaseConfig", "DuplexLeaseState", - "DuplexSessionExpiry", ] diff --git a/vllm_omni/engine/duplex/session/manager.py b/vllm_omni/engine/duplex/session/manager.py new file mode 100644 index 000000000..0b7a7a07e --- /dev/null +++ b/vllm_omni/engine/duplex/session/manager.py @@ -0,0 +1,872 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Engine-side owner of duplex sessions: admission, lease reaping, command dispatch. + +``DuplexOrchestrator`` hosts exactly one manager. The whole session +(``DuplexEngineSession``) lives here, owned by one ``DuplexSessionRunner`` per +session on the orchestrator loop; nothing outside the engine holds session +state. + +Control operations (open / close / resume / touch) arrive as correlated RPC +messages and answer through ``result_sink``; session commands arrive one-way +as ``DuplexSessionCommandMessage`` and are pushed onto the runner's ordered +mailbox after backpressure admission; everything a session emits leaves +through ``output_sink`` as ``DuplexSessionEventMessage``. +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import time +from collections.abc import Awaitable, Callable, Iterable +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.commands import AppendAudio, Commit, DuplexCommand +from vllm_omni.engine.duplex.contracts import ( + DuplexFence, + DuplexStagePort, + DuplexStageRequestContext, + duplex_resource_request_belongs_to_session, + duplex_resource_request_id, +) +from vllm_omni.engine.duplex.events import DuplexEvent, ErrorEvent, SessionClosed, SessionExpired, error_event +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionError, + DuplexSessionEventMessage, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, DuplexRuntimeConfigError, validate_duplex_plugin_sampling +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity, DuplexLeaseConfig, DuplexLeaseState +from vllm_omni.engine.duplex.turn_detection import SileroVADBackendProvider + +if TYPE_CHECKING: + import janus + from vllm.config import ModelConfig + + from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig + from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner + from vllm_omni.engine.messages import EngineQueueMessage + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _PendingRequestCleanup: + session_id: str + lease_generation: int + request_ids: tuple[str, ...] + abort: bool + + +@dataclass(frozen=True) +class _PendingSessionCleanup: + """A close/expiry whose stage cleanup has not succeeded yet (retried by the reaper).""" + + kind: str + session_id: str + lease_generation: int + reason: str + submitted_request_ids: tuple[str, ...] + reserved_request_ids: tuple[str, ...] + + +class DuplexSessionManager: + _CONTROL_TYPES = ( + OpenDuplexSessionMessage, + CloseDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, + ) + _MESSAGE_TYPES = (*_CONTROL_TYPES, DuplexSessionCommandMessage) + + def __init__( + self, + *, + plugin: DuplexModelPlugin, + stage_port: DuplexStagePort, + output_sink: janus.AsyncQueue[EngineQueueMessage], + result_sink: janus.AsyncQueue[EngineQueueMessage], + runtime_config: DuplexSessionRuntimeConfig, + model_config: ModelConfig | None, + clock: Callable[[], float] | None = None, + executor: concurrent.futures.ThreadPoolExecutor | None = None, + ) -> None: + validate_duplex_plugin_sampling(plugin, sampling_defaults=tuple(stage_port.sampling_defaults())) + self.plugin = plugin + self.stage_port = stage_port + self.model_config = model_config + self.runtime_config = runtime_config + self._output_sink = output_sink + self._result_sink = result_sink + self._clock = clock or time.monotonic + self.executor = executor or concurrent.futures.ThreadPoolExecutor( + max_workers=4, + thread_name_prefix="duplex-session", + ) + self._owns_executor = executor is None + self._lease_config = DuplexLeaseConfig( + idle_ttl_s=runtime_config.idle_ttl_s, + disconnect_grace_s=runtime_config.disconnect_grace_s, + ) + #: One Silero backend for the whole engine: the ONNX graph is stateless + #: per call, so every session's detector can share it. Resolved lazily, + #: on the first session that actually turns server VAD on. + self.vad_backend_provider = SileroVADBackendProvider( + model_path=getattr(runtime_config, "server_vad_model_path", None) + ) + self.runners: dict[str, DuplexSessionRunner] = {} + #: Sessions whose close/expiry began but whose stage cleanup has not finalized; + #: they keep their admission slot until finalized. + self._closing: dict[str, _PendingSessionCleanup] = {} + #: Session ids whose ``open`` passed the capacity check and is still awaiting the plugin. + self._admitting: set[str] = set() + #: Opens whose caller gave up (control RPC timeout / cancellation) before + #: admission finished. The open completes anyway -- it is already past + #: the capacity check -- so it must tear itself down on arrival instead + #: of holding a slot and a stage resource until idle expiry. + self._abandoned_opens: set[str] = set() + self._request_index: dict[str, str] = {} + self._session_control_tails: dict[str, asyncio.Task[None]] = {} + self._dispatched_control_tasks: set[asyncio.Task[None]] = set() + self._pending_request_cleanups: dict[tuple[str, int], _PendingRequestCleanup] = {} + self._request_cleanup_tasks: dict[tuple[str, int], asyncio.Task[None]] = {} + self._request_cleanups_in_progress: set[tuple[str, int]] = set() + #: Sessions retained for cleanup retry after their runner is gone. + self._session_snapshots: dict[str, DuplexEngineSession] = {} + + # ------------------------------------------------------------------ # + # Views # + # ------------------------------------------------------------------ # + + def active_count(self) -> int: + return len(self.runners) + + def get(self, session_id: str) -> DuplexEngineSession | None: + runner = self.runners.get(session_id) + return runner.session if runner is not None else None + + def sessions(self) -> dict[str, DuplexEngineSession]: + return {session_id: runner.session for session_id, runner in self.runners.items()} + + def runner_for_request_id(self, request_id: str) -> DuplexSessionRunner | None: + session_id = self._request_index.get(request_id) + if session_id is None: + for candidate_id in self.runners: + if duplex_resource_request_belongs_to_session(request_id, candidate_id): + session_id = candidate_id + break + if session_id is None: + return None + return self.runners.get(session_id) + + def register_request(self, request_id: str, session_id: str) -> None: + self._request_index[request_id] = session_id + + def unregister_request(self, request_id: str) -> None: + self._request_index.pop(request_id, None) + + def _unregister_session_requests(self, session_id: str) -> None: + for request_id in [rid for rid, sid in self._request_index.items() if sid == session_id]: + self._request_index.pop(request_id, None) + + # ------------------------------------------------------------------ # + # Message routing # + # ------------------------------------------------------------------ # + + def accepts(self, message: object) -> bool: + return isinstance(message, self._MESSAGE_TYPES) + + def dispatch(self, message: object) -> None: + """Route one engine request-queue message without blocking the request handler.""" + if isinstance(message, DuplexSessionCommandMessage): + self._dispatch_command(message) + return + if not isinstance(message, self._CONTROL_TYPES): + raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") + session_id = message.session_id + + async def handle_message() -> None: + await self.handle(message) + + self._run_control(session_id, f"duplex-control-{session_id}-{message.control_id}", handle_message) + + def _run_control( + self, + session_id: str, + name: str, + operation: Callable[[], Awaitable[None]], + ) -> asyncio.Task[None]: + """Run one control operation after the session's previous one (per-session order).""" + predecessor = self._session_control_tails.get(session_id) + + async def run_ordered() -> None: + if predecessor is not None: + try: + await predecessor + except BaseException: + pass + await operation() + + task = asyncio.create_task(run_ordered(), name=name) + self._session_control_tails[session_id] = task + self._track_task(task) + + def discard(completed: asyncio.Task[None]) -> None: + if self._session_control_tails.get(session_id) is completed: + self._session_control_tails.pop(session_id, None) + + task.add_done_callback(discard) + return task + + def _track_task(self, task: asyncio.Task[None]) -> None: + """Keep a background task alive until it finishes and log an escaped exception.""" + self._dispatched_control_tasks.add(task) + + def done(completed: asyncio.Task[None]) -> None: + self._dispatched_control_tasks.discard(completed) + if not completed.cancelled() and completed.exception() is not None: + logger.error("duplex control task %s failed: %r", completed.get_name(), completed.exception()) + + task.add_done_callback(done) + + async def handle(self, message: object) -> None: + if isinstance(message, OpenDuplexSessionMessage): + await self.open(message) + elif isinstance(message, CloseDuplexSessionMessage): + await self.close(message) + elif isinstance(message, ResumeDuplexSessionMessage): + await self.resume(message) + elif isinstance(message, TouchDuplexSessionMessage): + await self.touch(message) + elif isinstance(message, DuplexSessionCommandMessage): + self._dispatch_command(message) + else: + raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") + + def _dispatch_command(self, message: DuplexSessionCommandMessage) -> None: + """Admit one command: identity check, backpressure, then the runner's ordered mailbox. + + Backpressure contract shared with the runner: a ``Commit`` that passes + admission holds one pending-turn reservation and an ``AppendAudio`` + holds its wire-size byte reservation until the runner dequeues the + command (the runner then re-reserves the decoded size around its PCM + reservation), so the mailbox never holds more bytes than the session + limit even while the worker is busy. + """ + runner = self.runners.get(message.session_id) + command = message.command + if runner is None: + self._emit_raw( + message.session_id, + self._error_event( + "unknown_session", + f"Unknown or closed duplex session: {message.session_id}", + command=command, + ), + ) + return + session = runner.session + if runner.closing: + self.emit( + session, + self._error_event( + "session_closed", f"Duplex session is closing: {session.session_id}", command=command + ), + ) + return + if isinstance(command, AppendAudio): + limit = int(self.runtime_config.max_pending_input_bytes_per_session) + if not session.reserve_input_bytes(len(command.audio), limit=limit): + self.emit( + session, + self._error_event( + "input_backpressure", + "Duplex session has too many pending input bytes", + command=command, + ), + ) + return + elif isinstance(command, Commit): + if not session.reserve_pending_turn(limit=int(self.runtime_config.max_pending_turns_per_session)): + self.emit( + session, + self._error_event( + "input_backpressure", + "Duplex session has too many pending input turns", + command=command, + ), + ) + return + runner.submit(command) + + @staticmethod + def _error_event(code: str, message: str, *, command: DuplexCommand | None = None) -> ErrorEvent: + return error_event(code, message, event_id=command.event_id if command is not None else None) + + # ------------------------------------------------------------------ # + # Emission # + # ------------------------------------------------------------------ # + + def emit(self, session: DuplexEngineSession, event: DuplexEvent) -> None: + """Bind the session identity to one typed event and push it to the engine output queue.""" + self._emit_raw(session.session_id, event, epoch=session.epoch) + + def _emit_raw( + self, + session_id: str, + event: DuplexEvent, + *, + epoch: int | None = None, + ) -> None: + event = replace(event, session_id=session_id, epoch=epoch) + message = DuplexSessionEventMessage(session_id=session_id, event=event) + put_nowait = getattr(self._output_sink, "put_nowait", None) + if callable(put_nowait): + put_nowait(message) + else: # pragma: no cover - sink without put_nowait + asyncio.ensure_future(self._output_sink.put(message)) + + async def _put_result( + self, + message: ( + OpenDuplexSessionMessage + | CloseDuplexSessionMessage + | ResumeDuplexSessionMessage + | TouchDuplexSessionMessage + ), + *, + operation: str, + ok: bool, + session: DuplexEngineSession | None = None, + error: BaseException | None = None, + ) -> None: + error_code, error_message, error_retryable = ( + self._control_error(error) if error is not None else (None, None, False) + ) + result = DuplexControlResultMessage( + control_id=message.control_id, + operation=operation, + session_id=message.session_id, + ok=ok, + lease_generation=session.lease_generation if session is not None else None, + capabilities=session.capabilities if session is not None and ok else None, + public_session=session.as_public_dict() if session is not None and ok else None, + error_code=error_code, + error_message=error_message, + error_retryable=error_retryable, + ) + await self._result_sink.put(result) + + @staticmethod + def _control_error(error: BaseException) -> tuple[str, str, bool]: + """Map an exception to ``(code, message, retryable)`` for a control result.""" + message = str(error) + retryable = False + if isinstance(error, DuplexSessionError): + code = error.code + retryable = error.retryable + elif isinstance(error, DuplexRuntimeConfigError): + code = error.code + elif isinstance(error, DuplexFenceMismatchError): + code = "stale_fence" + elif isinstance(error, KeyError): + code = "not_found" + elif isinstance(error, TypeError | ValueError): + code = "invalid_argument" + elif isinstance(error, TimeoutError): + code = "timeout" + retryable = True + else: + code = "failed_precondition" + return code, message, retryable + + # ------------------------------------------------------------------ # + # Sampling / stage request helpers shared with the runner # + # ------------------------------------------------------------------ # + + def sampling_params_for(self, session: DuplexEngineSession) -> tuple[object, ...]: + defaults = tuple(self.stage_port.sampling_defaults()) + configured = self.plugin.configure_sampling_params( + runtime_config=dict(session.runtime_config), + defaults=defaults, + ) + if not isinstance(configured, tuple): + raise TypeError("duplex plugin must return sampling parameters as a tuple") + if len(configured) != len(defaults): + raise ValueError("duplex plugin must return one sampling parameter per stage") + return configured + + @staticmethod + def stage_request_id(fence: DuplexFence, *, stage_id: int) -> str: + return duplex_resource_request_id(fence, f"stage{stage_id}") + + def ensure_stage_request( + self, + session: DuplexEngineSession, + *, + stage_id: int, + fence: DuplexFence | None = None, + ) -> DuplexStageRequestContext | None: + """Reserve the session's stage request id and register it with the stage port.""" + if stage_id >= self.stage_port.stage_count: + return None + effective_fence = fence or session.fence + request_id = self.stage_request_id(effective_fence, stage_id=stage_id) + session.reserve_stage_request(stage_id, request_id, fence=effective_fence) + context = DuplexStageRequestContext( + request_id=request_id, + session_id=session.session_id, + fence=effective_fence, + stage_id=stage_id, + final_stage_id=self.stage_port.stage_count - 1, + config_generation=session.config_generation, + sampling_params=self.sampling_params_for(session), + session_config=session.config.as_dict(), + runtime_config=session.runtime_config, + ) + self.stage_port.ensure_request(context) + self.register_request(request_id, session.session_id) + return context + + # ------------------------------------------------------------------ # + # Control operations # + # ------------------------------------------------------------------ # + + def _admission_count(self) -> int: + return len(set(self.runners) | set(self._closing) | self._admitting) + + async def open(self, message: OpenDuplexSessionMessage) -> None: + from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner + + session_id = message.session_id + session: DuplexEngineSession | None = None + runner: DuplexSessionRunner | None = None + # Only the open that put the id in ``_admitting`` may take it out again: + # a duplicate open must not release the slot the first one is holding. + holds_admission_slot = False + try: + if session_id in self.runners or session_id in self._closing or session_id in self._admitting: + raise DuplexSessionError(f"Duplex session already exists: {session_id}", code="session_exists") + max_sessions = int(self.runtime_config.max_sessions) + if self._admission_count() >= max_sessions: + raise DuplexSessionError( + f"duplex_session_capacity_exhausted: limit={max_sessions}", + code="resource_exhausted", + retryable=True, + ) + # Hold the slot across the plugin await: opens of different + # sessions run concurrently and must not all pass the check above. + self._admitting.add(session_id) + holds_admission_slot = True + config = message.session_config + self.plugin.validate_client_extra_body(config.extra_body) + try: + runtime_config = await self.plugin.prepare_runtime_config(config, model_config=self.model_config) + except DuplexRuntimeConfigError: + raise + except ValueError as exc: + raise DuplexRuntimeConfigError(str(exc)) from exc + capabilities = self.plugin.capabilities(max_sessions=max_sessions) + session = DuplexEngineSession( + session_id=session_id, + config=config, + capabilities=capabilities, + lease=DuplexLeaseState(config=self._lease_config, generation=0, last_activity=self._clock()), + _clock=self._clock, + _runtime_config=dict(runtime_config), + ) + # Validates the plugin's sampling policy for this runtime config before admission. + self.sampling_params_for(session) + session.model_state = self.plugin.create_session_state() + runner = DuplexSessionRunner( + session=session, + plugin=self.plugin, + stage_port=self.stage_port, + manager=self, + model_config=self.model_config, + ) + self.runners[session_id] = runner + # Reserve the Stage0 request resource atomically with admission. + self.ensure_stage_request(session, stage_id=0) + runner.start() + await self._put_result(message, operation="open", ok=True, session=session) + if session_id in self._abandoned_opens: + # The caller timed out while we were admitting; hand the slot + # and the stage resource straight back. + self._abandoned_opens.discard(session_id) + await self._close_runner(runner, kind="close", reason="open_abandoned") + except Exception as exc: + error_code, _, _ = self._control_error(exc) + if error_code in {"resource_exhausted", "session_exists"}: + logger.info("open_duplex_session rejected: %s", exc) + else: + logger.exception("open_duplex_session failed: %s", exc) + if runner is not None and self.runners.get(session_id) is runner: + self.runners.pop(session_id, None) + try: + await runner.shutdown() + except Exception: + logger.exception("duplex open rollback: runner shutdown failed for %s", session_id) + if session is not None: + reserved = session.release_all_requests() + if reserved: + try: + await self.stage_port.cleanup(list(reserved)) + except Exception: + logger.warning("duplex open rollback: request cleanup pending for %s", session_id) + self._unregister_session_requests(session_id) + await self._put_result(message, operation="open", ok=False, error=exc) + finally: + if holds_admission_slot: + self._admitting.discard(session_id) + self._abandoned_opens.discard(session_id) + + def _require_runner(self, session_id: str) -> DuplexSessionRunner: + runner = self.runners.get(session_id) + if runner is None: + raise DuplexSessionError(f"Unknown duplex session: {session_id}", code="unknown_session") + if runner.closing: + raise DuplexSessionError(f"Duplex session is closing: {session_id}", code="session_closed") + return runner + + async def close(self, message: CloseDuplexSessionMessage) -> None: + session_id = message.session_id + runner = self.runners.get(session_id) + if runner is None: + if session_id in self._admitting: + # The open is still awaiting the plugin. It cannot be cancelled + # from here, so leave a note it checks before returning. + self._abandoned_opens.add(session_id) + # Already closing, or closed and forgotten: ids are never reused + # (D16), so a repeated close is idempotent rather than an error. + await self._put_result(message, operation="close", ok=True) + return + session = runner.session + try: + await self._close_runner(runner, kind="close", reason=message.reason) + except Exception as exc: + # The session is closed either way; only the stage cleanup is + # outstanding and the reaper retries it while the slot stays held. + logger.warning("close_duplex_session: stage cleanup for %s remains pending: %s", session_id, exc) + await self._put_result(message, operation="close", ok=True, session=session) + + def close_from_runner(self, runner: DuplexSessionRunner, reason: str) -> None: + """Finish a close the runner started from its command stream (``session.close``). + + Runs on the orchestrator loop as a tracked task (the runner's worker + cannot await its own teardown): stage requests are aborted, the + admission slot is released and ``session.closed`` is emitted last. + """ + session_id = runner.session.session_id + if self.runners.get(session_id) is not runner: + return + + async def close_runner() -> None: + if self.runners.get(session_id) is runner: + await self._close_runner(runner, kind="close", reason=reason) + + # Chained on the session's control tail so a control RPC issued after + # the wire close observes the closed session, never a half-torn one. + self._run_control(session_id, f"duplex-close-{session_id}", close_runner) + + async def _close_runner(self, runner: DuplexSessionRunner, *, kind: str, reason: str) -> None: + """Begin an irreversible close, tear down the runner, then release stage resources. + + The session keeps its admission slot (``_closing``) until stage cleanup + succeeded; a failed cleanup is retried by the reaper. For an explicit + close the ``session.closed`` event is emitted only after the cleanup + attempt, so a client that sees it can open a replacement session at + once (the runner used to emit it before the stage requests were + aborted, which let a prompt reopen hit ``resource_exhausted``). + """ + session = runner.session + submitted = tuple(session.resource_request_ids(submitted=True)) + reserved = tuple(session.resource_request_ids(submitted=False)) + session.begin_close(reason=reason) + pending = _PendingSessionCleanup( + kind=kind, + session_id=session.session_id, + lease_generation=session.lease_generation, + reason=reason, + submitted_request_ids=submitted, + reserved_request_ids=reserved, + ) + self._closing[session.session_id] = pending + if self.runners.get(session.session_id) is runner: + self.runners.pop(session.session_id, None) + emit_after_cleanup = kind in {"close", "expired"} and not runner.closed_emitted + try: + if kind == "expired": + await runner.expire(reason, emit_expired=not emit_after_cleanup) + else: + await runner.close(reason, emit_closed=not emit_after_cleanup) + except Exception: + logger.exception("duplex runner %s failed for session %s", kind, session.session_id) + session.close() + try: + await self._finalize_pending_cleanup(pending, session) + finally: + # ``emit_after_cleanup`` was decided before the teardown awaits; a + # runtime close racing them may have emitted the terminal already. + if emit_after_cleanup and not runner.closed_emitted: + runner.mark_closed_emitted() + terminal: DuplexEvent = ( + SessionExpired(reason=reason) if kind == "expired" else SessionClosed(reason=reason) + ) + self.emit(session, terminal) + + async def _finalize_pending_cleanup( + self, + pending: _PendingSessionCleanup, + session: DuplexEngineSession | None, + ) -> None: + if pending.submitted_request_ids: + await self.stage_port.cleanup(list(pending.submitted_request_ids), abort=True) + if pending.reserved_request_ids: + await self.stage_port.cleanup(list(pending.reserved_request_ids)) + if session is not None: + session.release_all_requests() + self._unregister_session_requests(pending.session_id) + if self._closing.get(pending.session_id) is pending: + self._closing.pop(pending.session_id, None) + self._session_snapshots.pop(pending.session_id, None) + + async def resume(self, message: ResumeDuplexSessionMessage) -> None: + session: DuplexEngineSession | None = None + try: + runner = self._require_runner(message.session_id) + session = runner.session + try: + session.resume_lease(expected_lease_generation=message.expected_lease_generation) + except ValueError as exc: + raise DuplexSessionError(str(exc), code="session_resume_conflict") from exc + await self._put_result(message, operation="resume", ok=True, session=session) + except Exception as exc: + logger.exception("resume_duplex_session failed: %s", exc) + await self._put_result(message, operation="resume", ok=False, session=session, error=exc) + + async def touch(self, message: TouchDuplexSessionMessage) -> None: + session: DuplexEngineSession | None = None + try: + runner = self._require_runner(message.session_id) + session = runner.session + activity = DuplexLeaseActivity(message.activity) + if activity is DuplexLeaseActivity.DETACH: + session.detach_lease() + else: + session.touch_lease(activity) + await self._put_result(message, operation="touch", ok=True, session=session) + except DuplexSessionError as exc: + # A lease touch racing the session's own close is routine: the + # caller learns the session is gone from the result, not a fault. + logger.debug("touch_duplex_session skipped: %s", exc) + await self._put_result(message, operation="touch", ok=False, session=session, error=exc) + except Exception as exc: + logger.exception("touch_duplex_session failed: %s", exc) + await self._put_result(message, operation="touch", ok=False, session=session, error=exc) + + # ------------------------------------------------------------------ # + # Lease expiry / reaper # + # ------------------------------------------------------------------ # + + async def reaper_loop(self, shutdown_event: asyncio.Event) -> None: + interval = float(self.runtime_config.reaper_interval_s) + while not shutdown_event.is_set(): + try: + await asyncio.wait_for(shutdown_event.wait(), timeout=interval) + except TimeoutError: + try: + await self.reap_expired() + except Exception: + logger.exception("[DuplexSessionManager] expiry cleanup failed; retrying on next tick") + + async def reap_expired(self, now: float | None = None) -> int: + completed = 0 + effective_now = self._clock() if now is None else now + # Retry request cleanups recorded by the orchestrator error paths. + for key, request_cleanup in list(self._pending_request_cleanups.items()): + if key in self._request_cleanups_in_progress: + continue + try: + await self._complete_request_cleanup(key, request_cleanup) + except Exception as exc: + logger.warning( + "duplex request cleanup remains pending for session %s: %s", + request_cleanup.session_id, + exc, + ) + continue + completed += 1 + # Retry closes/expiries whose stage cleanup failed. + for session_id, session_cleanup in list(self._closing.items()): + if session_id in self.runners: + continue + if session_cleanup.kind == "request_cleanup" and any( + key[0] == session_id for key in self._pending_request_cleanups + ): + # The stage cleanup is owned by the request-cleanup path (orchestrator + # or the retry above); ``finalize_closed_sessions`` releases the slot. + continue + try: + await self._finalize_pending_cleanup(session_cleanup, self._session_snapshots.get(session_id)) + except Exception as exc: + logger.warning( + "duplex %s cleanup remains pending for session %s: %s", + session_cleanup.kind, + session_id, + exc, + ) + continue + completed += 1 + # Expire leases. + for session_id, runner in list(self.runners.items()): + lease = runner.session.lease + if lease.disconnect_grace_expired(effective_now): + reason = "disconnect_grace_expired" + elif lease.idle_expired(effective_now): + reason = "idle_ttl_expired" + else: + continue + self._session_snapshots[session_id] = runner.session + try: + await self._close_runner(runner, kind="expired", reason=reason) + except Exception as exc: + logger.warning("duplex expiry cleanup remains pending for session %s: %s", session_id, exc) + continue + completed += 1 + return completed + + # ------------------------------------------------------------------ # + # Request-triggered cleanup (orchestrator error paths) # + # ------------------------------------------------------------------ # + + def close_sessions_for_request_ids( + self, + request_ids: list[str], + *, + abort: bool = False, + cleanup_in_progress: bool = False, + ) -> dict[str, list[str]]: + request_id_set = set(request_ids) + closed: dict[str, list[str]] = {} + for session_id, runner in list(self.runners.items()): + session = runner.session + stale = session.resource_request_ids() + if request_id_set.isdisjoint(stale): + continue + if not session.begin_close(reason="request_cleanup"): + continue + closed[session_id] = stale + key = (session_id, session.lease_generation) + existing = self._pending_request_cleanups.get(key) + merged = tuple(dict.fromkeys([*(existing.request_ids if existing is not None else ()), *stale])) + self._pending_request_cleanups[key] = _PendingRequestCleanup( + session_id=session_id, + lease_generation=session.lease_generation, + request_ids=merged, + abort=abort or (existing.abort if existing is not None else False), + ) + if cleanup_in_progress: + self._request_cleanups_in_progress.add(key) + # Tear the runner down in the background; the orchestrator owns the stage cleanup. + self._session_snapshots[session_id] = session + self._closing[session_id] = _PendingSessionCleanup( + kind="request_cleanup", + session_id=session_id, + lease_generation=session.lease_generation, + reason="request_cleanup", + submitted_request_ids=(), + reserved_request_ids=(), + ) + self.runners.pop(session_id, None) + task = asyncio.create_task(runner.expire("request_cleanup"), name=f"duplex-request-cleanup-{session_id}") + self._track_task(task) + return closed + + def defer_request_cleanups(self, session_ids: Iterable[str]) -> None: + session_id_set = set(session_ids) + active_keys = {key for key in self._request_cleanups_in_progress if key[0] in session_id_set} + self._request_cleanups_in_progress.difference_update(active_keys) + + def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: + session_id_set = set(session_ids) + for key in list(self._pending_request_cleanups): + if key[0] not in session_id_set: + continue + self._pending_request_cleanups.pop(key, None) + self._request_cleanups_in_progress.discard(key) + for session_id in session_id_set: + session = self._session_snapshots.pop(session_id, None) + if session is not None and session.lease.terminal_reason is not None: + session.release_all_requests() + session.close() + self._closing.pop(session_id, None) + self._unregister_session_requests(session_id) + + async def _complete_request_cleanup( + self, + key: tuple[str, int], + pending: _PendingRequestCleanup, + ) -> None: + task = self._request_cleanup_tasks.get(key) + if task is not None and task.done(): + self._request_cleanup_tasks.pop(key, None) + task = None + if task is None: + task = asyncio.create_task( + self._run_request_cleanup(key, pending), + name=f"duplex-request-cleanup-{pending.session_id}", + ) + self._request_cleanup_tasks[key] = task + + def discard(completed: asyncio.Task[None]) -> None: + if self._request_cleanup_tasks.get(key) is completed: + self._request_cleanup_tasks.pop(key, None) + + task.add_done_callback(discard) + await asyncio.shield(task) + + async def _run_request_cleanup(self, key: tuple[str, int], pending: _PendingRequestCleanup) -> None: + await self.stage_port.cleanup(list(pending.request_ids), abort=pending.abort) + if self._pending_request_cleanups.get(key) is pending: + self._pending_request_cleanups.pop(key, None) + self._request_cleanups_in_progress.discard(key) + self.finalize_closed_sessions([pending.session_id]) + + # ------------------------------------------------------------------ # + # Shutdown # + # ------------------------------------------------------------------ # + + async def shutdown(self) -> None: + runners = list(self.runners.values()) + self.runners.clear() + for runner in runners: + try: + await runner.shutdown() + except Exception: + logger.exception("duplex runner shutdown failed for %s", runner.session.session_id) + tasks = tuple(self._dispatched_control_tasks | set(self._request_cleanup_tasks.values())) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._session_control_tails.clear() + self._closing.clear() + self._session_snapshots.clear() + self._request_index.clear() + if self._owns_executor: + self.executor.shutdown(wait=False, cancel_futures=True) + + +__all__ = ["DuplexSessionManager"] diff --git a/vllm_omni/engine/duplex/session/model_channel.py b/vllm_omni/engine/duplex/session/model_channel.py new file mode 100644 index 000000000..322db7ede --- /dev/null +++ b/vllm_omni/engine/duplex/session/model_channel.py @@ -0,0 +1,1100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The model side of one duplex session. + +Everything between the session and the model runtime: planning an append with +the plugin and submitting it to the stage port, turning the stage output that +comes back into session events, and -- when the model chooses to keep listening +-- offering it another unit of silence so the turn can continue. + +Those three read as separate concerns but they call each other in a cycle. An +append submits and immediately projects its own first outputs; a projected +output can schedule a continuation; a continuation is another append. Splitting +them into three modules would mean three mutual imports and no fewer edges, so +they are one component whose seam with the runner is narrow instead. + +That seam is three callbacks, for the three things this component triggers but +does not own: closing the session, scheduling the continuation append as a +tracked task, and aborting a stage request in the background. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import replace +from typing import TYPE_CHECKING, Protocol + +import pybase64 as base64 +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DuplexSessionState +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexOutputContext, + DuplexOutputDecision, + DuplexStageSubmission, + duplex_data_plane_request_info, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexRuntimeConfigError, + coerce_int, + payload_turn_id, +) +from vllm_omni.engine.duplex.session.context import DuplexSessionContext, StageOutput +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.duplex import attach_duplex_output_decision + +if TYPE_CHECKING: + from vllm.outputs import RequestOutput + +logger = init_logger(__name__) + + +class SilenceContinuationScheduler(Protocol): + """Runs one continuation append as a tracked task (see ``_schedule_silence_continuation``).""" + + async def __call__( + self, + payload: dict[str, object], + *, + request_id: str, + owner_id: str, + response_id: str | None, + response_owned: bool, + expected_epoch: int | None, + expected_model_turn_id: int | None, + ) -> bool: ... + + +class ModelChannel: + """Appends out to the model, events back from it, for one session.""" + + # One MiniCPM model unit (1 s at 16 kHz) is the compatibility default. + _SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii") + _RESPONSE_MAX_CONTINUATION_UNITS = 8 + _AUTO_RESPONSE_MAX_CONTINUATION_UNITS = 64 + + def __init__( + self, + ctx: DuplexSessionContext, + out: SessionEmitter, + *, + close_from_runtime: Callable[[str], Awaitable[None]], + schedule_silence_continuation: SilenceContinuationScheduler, + abort_request: Callable[..., Awaitable[None]], + ) -> None: + self._ctx = ctx + self._out = out + self._close_from_runtime = close_from_runtime + self._schedule_silence_continuation = schedule_silence_continuation + self._abort_request = abort_request + + @staticmethod + def should_commit_response_to_history(session: DuplexEngineSession, response_id: str | None) -> bool: + if response_id is not None and response_id != session.active_response_id: + return True + mode = session.response_config.extra_body.get("realtime_response_conversation") + return not isinstance(mode, str) or mode.strip().lower() != "none" + + def response_created_payload( + self, + response_id: str, + *, + epoch: int, + request_id: str | None = None, + ) -> dict[str, object]: + session = self._ctx.session + response_config = session.response_config + payload: dict[str, object] = { + "type": "response.created", + "session_id": session.session_id, + "response_id": response_id, + "epoch": epoch, + "modalities": list(response_config.modalities), + } + if request_id is not None: + payload["request_id"] = request_id + metadata = response_config.extra_body.get("realtime_response_metadata") + if not isinstance(metadata, dict): + metadata = response_config.extra_body.get("realtime_metadata") + if isinstance(metadata, dict): + payload["metadata"] = dict(metadata) + conversation = response_config.extra_body.get("realtime_response_conversation") + if isinstance(conversation, str): + payload["conversation"] = conversation + prompt = response_config.extra_body.get("realtime_response_prompt") + if isinstance(prompt, dict): + payload["prompt"] = dict(prompt) + return payload + + # ------------------------------------------------------------------ # + # Submitting an append # + # ------------------------------------------------------------------ # + + async def append_runtime_input( + self, + payload: object, + *, + operation_id: str | None = None, + final: bool, + expected_epoch: int | None = None, + ) -> tuple[bool, bool]: + session = self._ctx.session + if not session.capabilities.supports_input_append: + return True, False + if expected_epoch is not None and session.epoch != expected_epoch: + return True, False + try: + result = await self._append_via_data_plane( + payload, + final=final, + operation_id=operation_id, + expected_epoch=expected_epoch, + ) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to append duplex runtime input: %s", exc) + self.send_runtime_error("runtime_append_failed", exc) + return False, False + if result is None: + return True, False + if expected_epoch is not None and session.epoch != expected_epoch: + return True, False + request_id, _ = duplex_data_plane_request_info(result) + if request_id is not None: + self._ctx.plugin.data_plane.begin_request(request_id) + close_reason, emitted_response = await self._send_model_output_events(result, expected_epoch=expected_epoch) + emitted_response = emitted_response or request_id is not None + if close_reason is None and self._start_data_plane_stream(result): + emitted_response = True + if close_reason is not None: + await self._close_from_runtime(close_reason) + return False, emitted_response + return True, emitted_response + + async def _append_via_data_plane( + self, + payload: object, + *, + final: bool, + operation_id: str | None, + expected_epoch: int | None, + ) -> dict[str, object] | None: + """Submit one planned append to the stage port and bind the resulting stage request.""" + session = self._ctx.session + if self._ctx.stage_port.stage_count == 0: + raise RuntimeError("duplex_data_plane_has_no_stage") + payload_turn = payload_turn_id(payload) + fence = DuplexFence( + session.session_id, + epoch=session.epoch, + turn_id=( + payload_turn + if payload_turn is not None + else ( + session.active_response_turn_id if session.active_response_turn_id is not None else session.turn_id + ) + ), + ) + lease_operation_id = f"append:{operation_id or uuid.uuid4().hex}" + operation_started = False + stage_id = 0 + request_id = self._ctx.manager.stage_request_id(fence, stage_id=stage_id) + try: + session.begin_lease_operation(fence, lease_operation_id) + operation_started = True + reservation = session.prepare_append(fence) + already_submitted = session.stage_request_submitted(stage_id, request_id) + request_context = self._ctx.manager.ensure_stage_request(session, stage_id=stage_id, fence=fence) + if request_context is None: + raise RuntimeError("duplex_data_plane_has_no_stage") + append_plan = self._ctx.plugin.plan_append( + request_id=request_id, + fence=fence, + session_config=dict(request_context.session_config), + runtime_config=dict(request_context.runtime_config), + seq=reservation.update.seq, + turn_seq=reservation.update.turn_seq, + payload=payload, + final=final, + sampling_params=request_context.stage_sampling_params, + ) + if not isinstance(append_plan, DuplexAppendPlan): + raise TypeError("duplex plugin plan_append() must return DuplexAppendPlan") + submission = DuplexStageSubmission( + context=request_context, + prompt=append_plan.prompt, + already_submitted=already_submitted, + ) + submission_result = await self._ctx.stage_port.submit(submission) + try: + if submission_result.request_id != request_id or submission_result.stage_id != stage_id: + raise RuntimeError("duplex stage adapter returned a mismatched submission result") + if expected_epoch is not None and session.epoch != expected_epoch: + raise DuplexFenceMismatchError(session.fence, fence) + update = session.commit_append(reservation) + session.bind_stage_request(stage_id, request_id, fence=fence) + except BaseException: + try: + await self._ctx.stage_port.cleanup([request_id]) + except Exception as cleanup_exc: + logger.warning( + "duplex append compensation remains pending for session %s: %s", + session.session_id, + cleanup_exc, + ) + raise + session.touch_lease(DuplexLeaseActivity.APPEND) + return { + "ok": True, + "operation": "append", + "session_id": session.session_id, + "stage_results": [ + { + "stage_id": stage_id, + "replica_id": submission_result.replica_id, + "result": { + "supported": True, + "data_plane_append": True, + "request_id": request_id, + "response_stage_id": request_context.final_stage_id, + "seq": update.seq, + "turn_id": update.turn_id, + "turn_seq": update.turn_seq, + "resumable": True, + }, + } + ], + } + finally: + if operation_started and lease_operation_id in session.lease.active_operations: + try: + session.end_lease_operation(lease_operation_id) + except Exception: + session.lease.active_operations.discard(lease_operation_id) + + def _start_data_plane_stream(self, result: object) -> bool: + """Bind the resumable stage request as the session's active data-plane stream.""" + request_id, _ = duplex_data_plane_request_info(result) if isinstance(result, dict) else (None, None) + if request_id is None or self._data_plane_outputs_finished(result): + return False + session = self._ctx.session + session.bind_request(request_id) + if self._ctx.run.stream_request_id == request_id: + return False + self._ctx.run.stream_request_id = request_id + return True + + def cancel_data_plane_stream(self) -> bool: + had_stream = self._ctx.run.stream_request_id is not None + self._ctx.run.stream_request_id = None + return had_stream + + async def signal_cancel_fence(self, cancelled_fence: DuplexFence) -> bool: + """In-process equivalent of the old ``barge_in`` control signal.""" + session = self._ctx.session + try: + next_fence = session.sync_fence() + if next_fence.epoch > cancelled_fence.epoch: + stale_request_ids = session.cancel_fence(cancelled_fence, next_fence) + else: + stale_request_ids = [] + if stale_request_ids: + await self._ctx.stage_port.cleanup(list(stale_request_ids), abort=True) + session.touch_lease(DuplexLeaseActivity.SIGNAL) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to signal duplex runtime session: %s", exc) + self.send_runtime_error("runtime_signal_failed", exc) + return False + return True + + def send_runtime_error(self, code: str, exc: BaseException) -> None: + if isinstance(exc, DuplexRuntimeConfigError): + code = exc.code + self._out.emit_error(code, str(exc), retryable=False) + + @staticmethod + def _data_plane_outputs_finished(result: object) -> bool: + if not isinstance(result, dict): + return False + outputs = result.get("data_plane_outputs") + if not isinstance(outputs, list) or not outputs: + return False + return bool(getattr(outputs[-1], "finished", False)) + + def _runtime_data_plane_context(self) -> object: + session = self._ctx.session + response_config = session.response_config + return self._ctx.plugin.data_plane_context( + epoch=session.epoch, + turn_id=session.turn_id, + active_response_turn_id=session.active_response_turn_id, + active_response_id=session.active_response_id, + auto_responds=self._out.auto_responds(), + response_format=response_config.response_format, + speed=response_config.speed, + modalities=tuple(response_config.modalities), + ) + + # ------------------------------------------------------------------ # + # Stage outputs # + # ------------------------------------------------------------------ # + + def decide_output( + self, stage_id: int, output: RequestOutput, context: DuplexOutputContext + ) -> DuplexOutputDecision | None: + """Pure plugin decision (no session mutation: this runs inline on the orchestrator loop).""" + decision = self._ctx.plugin.decide_output( + stage_id=stage_id, + final_stage_id=context.final_stage_id, + segment_finished=context.segment_finished, + segment_token_ids=context.segment_token_ids, + segment_output_metadata=dict(context.segment_output_metadata), + output=output, + ) + if decision is not None and not isinstance(decision, DuplexOutputDecision): + raise TypeError("duplex plugin decide_output() must return DuplexOutputDecision or None") + return decision + + @staticmethod + def stage_metrics_snapshot(stage_id: int, metrics: object, output: object) -> dict[str, dict[str, object]] | None: + if not isinstance(metrics, StageRequestStats): + return None + event = metrics + if event.stage_id is None: + event = replace(event, stage_id=stage_id) + if event.final_output_type is None: + final_output_type = getattr(output, "final_output_type", None) + if isinstance(final_output_type, str): + event = replace(event, final_output_type=final_output_type) + try: + merged = OrchestratorAggregator._merge_stage_metric_event(None, event) + except Exception: + return None + return {str(stage_id): merged} + + def _build_stage_output(self, item: StageOutput) -> OmniRequestOutput: + output = item.output + finished = bool(getattr(output, "finished", False)) or item.context.segment_finished + if item.decision is not None: + engine_output = OmniRequestOutput.from_stage_output( + output, + request_id=item.request_id, + stage_id=item.stage_id, + final_output_type=item.decision.final_output_type, + ) + # Set after construction, not as a keyword: ``from_stage_output`` + # copies ``finished`` from the source last, and a resumable stage + # request never finishes -- only its segment does. The decision + # closes that segment, and the projector recognises a listen only + # on a finished output. The orchestrator queue used to re-assert + # ``finished`` on every direct-response message; the mailbox hands + # the raw stage output over, so it has to happen here. + engine_output.finished = True + engine_output = attach_duplex_output_decision(engine_output, item.decision) + elif isinstance(output, OmniRequestOutput): + engine_output = output + if finished: + engine_output.finished = True + else: + final_output_type = getattr(output, "final_output_type", None) + engine_output = OmniRequestOutput.from_stage_output( + output, + request_id=item.request_id, + finished=finished, + stage_id=item.stage_id, + final_output_type=final_output_type if isinstance(final_output_type, str) else "audio", + ) + snapshot = self.stage_metrics_snapshot(item.stage_id, item.metrics, output) + if snapshot is not None: + existing = engine_output.metrics if isinstance(engine_output.metrics, dict) else {} + existing_stage_metrics = existing.get("stage_metrics") + merged_stage_metrics = ( + {**existing_stage_metrics, **snapshot} if isinstance(existing_stage_metrics, dict) else snapshot + ) + engine_output.metrics = {**existing, "stage_metrics": merged_stage_metrics} + return engine_output + + async def on_stage_output_item(self, item: StageOutput) -> None: + session = self._ctx.session + if session.state == DuplexSessionState.CLOSED or self._ctx.run.closing: + return + expected_epoch = item.context.identity.fence.epoch + if session.epoch != expected_epoch: + return + if session.lease.terminal_reason is not None: + # Lease already closed/expired: the output belongs to a dead session. + return + session.touch_lease(DuplexLeaseActivity.MODEL_OUTPUT) + if self._ctx.plugin.data_plane.is_terminal(item.request_id): + return + if self._out.auto_responds(): + active_request_id = session.active_request_id + if active_request_id is not None and active_request_id != item.request_id: + return + engine_output = self._build_stage_output(item) + drain_result = {"data_plane_outputs": [engine_output]} + close_reason, emitted_response = await self._send_model_output_events( + drain_result, expected_epoch=expected_epoch + ) + if close_reason is not None: + await self._close_from_runtime(close_reason) + return + finished = self._data_plane_outputs_finished(drain_result) + if finished and emitted_response and not self._out.auto_responds(): + # A finished, emitted response releases the per-request projector + # cursor on its way out and offers the model another + # silence unit. + self._ctx.plugin.data_plane.close_stream(item.request_id) + if self._ctx.run.stream_request_id == item.request_id: + self._ctx.run.stream_request_id = None + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), + name="duplex-continue", + ) + + async def _send_model_output_events( + self, + result: object, + *, + expected_epoch: int | None = None, + ) -> tuple[str | None, bool]: + session = self._ctx.session + if expected_epoch is not None and session.epoch != expected_epoch: + return None, False + close_reason: str | None = None + emitted_response = False + request_id, _ = duplex_data_plane_request_info(result) if isinstance(result, dict) else (None, None) + if self._ctx.plugin.data_plane.is_terminal(request_id): + return None, False + if request_id is not None and session.active_request_id is None: + session.bind_request(request_id) + context = self._runtime_data_plane_context() + for model_result in self._ctx.plugin.data_plane.project(result, context=context): + close_reason_for_result, did_emit = await self._send_one_model_output_event( + model_result, + expected_epoch=expected_epoch, + ) + emitted_response = emitted_response or did_emit + close_reason = close_reason or close_reason_for_result + if expected_epoch is not None and session.epoch != expected_epoch: + return None, emitted_response + return close_reason, emitted_response + + def _fail_response_from_model_error(self, model_result: dict[str, object]) -> None: + """Report a data-plane error and fail the response it interrupted.""" + session = self._ctx.session + response_id = session.active_response_id + self._out.emit_error( + str(model_result.get("error_code")), + str(model_result.get("error") or "Duplex native data-plane error"), + ) + if response_id is None: + return + session.end_response(commit_text=False) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": model_result.get("error_code")}, + "playback": session.playback.as_dict(), + } + ) + + def _complete_model_turn_without_output( + self, + model_result: dict[str, object], + *, + model_turn_id: int | None, + data_plane_request_id: object, + ) -> bool: + """The model ended its turn with nothing to say. + + Returns whether an event was emitted: an auto-response session reports + the decision as a listen, a turn-based one simply releases the request. + """ + session = self._ctx.session + data_plane = self._ctx.plugin.data_plane + auto_response = self._out.auto_responds() + if isinstance(data_plane_request_id, str): + if not auto_response and data_plane_request_id == session.active_request_id: + session.clear_request() + if not auto_response: + data_plane.mark_terminal(data_plane_request_id) + if model_turn_id is not None: + session.complete_model_turn(model_turn_id) + if not auto_response: + return False + self._ctx.model_state.clear_continuation() + payload = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": "model_turn_completed_without_output", + "model_listen": True, + } + self._attach_runtime_metadata(payload, model_result) + self._out.emit(payload) + return True + + async def _on_model_listen( + self, + model_result: dict[str, object], + *, + model_turn_id: int | None, + data_plane_request_id: object, + expected_epoch: int | None, + ) -> tuple[str | None, bool]: + """The model chose to keep listening rather than speak. + + Either it continues the current response with another unit, or the + response ends here: an auto-response session that has continuations + left schedules one, and anything else closes the turn out. + """ + session = self._ctx.session + model_state = self._ctx.model_state + data_plane = self._ctx.plugin.data_plane + auto_response = self._out.auto_responds() + close_reason: str | None = None + emitted_response = False + self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) + if ( + session.active_response_id is not None + and model_turn_id is not None + and not session.active_response_accepts_model_turn(model_turn_id) + ): + return close_reason, emitted_response + active_response_id = session.active_response_id + auto_continuations_remaining = active_response_id is None or self.response_continuations_remaining( + active_response_id + ) + non_terminal_auto_listen = ( + auto_response + and active_response_id is not None + and session.active_request_id is not None + and model_result.get("end_of_turn") is not True + and auto_continuations_remaining + ) + if non_terminal_auto_listen: + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + ) + return close_reason, emitted_response + if not auto_response and data_plane_request_id == session.active_request_id: + session.clear_request() + model_listen = model_result.get("model_listen") + if not isinstance(model_listen, bool): + model_listen = model_result.get("reason") in {None, "", "model_listen"} + response_id = session.active_response_id + if isinstance(data_plane_request_id, str) and not auto_response: + data_plane.mark_terminal(data_plane_request_id) + emitted_response = True + payload = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": model_result.get("reason") or "model_listen", + "model_listen": model_listen, + } + if response_id is not None: + payload["response_id"] = response_id + self._attach_runtime_metadata(payload, model_result) + self._out.emit(payload) + if model_result.get("abort_data_plane_request") is True and isinstance(data_plane_request_id, str): + await self._abort_request(data_plane_request_id, notify=False) + if response_id is not None: + if not auto_response and self.response_continuations_remaining(response_id): + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + ) + return close_reason, emitted_response + if auto_response: + model_state.clear_continuation() + if not auto_continuations_remaining: + completed_turn_id = model_turn_id + if completed_turn_id is None: + completed_turn_id = session.active_response_turn_id + if completed_turn_id is not None: + session.complete_model_turn(completed_turn_id) + session.end_response(commit_text=False, preserve_request=auto_response) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "playback": session.playback.as_dict(), + } + ) + return close_reason, emitted_response + + async def _send_one_model_output_event( + self, + model_result: dict[str, object], + *, + expected_epoch: int | None = None, + ) -> tuple[str | None, bool]: + session = self._ctx.session + data_plane = self._ctx.plugin.data_plane + close_reason: str | None = None + emitted_response = False + if expected_epoch is not None and session.epoch != expected_epoch: + return close_reason, emitted_response + data_plane_request_id = model_result.get("data_plane_request_id") + if isinstance(data_plane_request_id, str) and data_plane.is_terminal(data_plane_request_id): + return close_reason, emitted_response + auto_response = self._out.auto_responds() + active_request_matches = session.active_request_id == data_plane_request_id or ( + auto_response and session.active_request_id is None + ) + if isinstance(data_plane_request_id, str) and not active_request_matches: + return close_reason, emitted_response + if isinstance(model_result.get("error_code"), str): + self._fail_response_from_model_error(model_result) + return close_reason, True + if model_result.get("function_call") is True: + self._out.emit( + { + "type": "function_call.done", + "session_id": session.session_id, + "epoch": session.epoch, + "call_id": model_result.get("call_id"), + "name": model_result.get("name"), + "arguments": model_result.get("arguments", ""), + } + ) + return close_reason, True + if model_result.get("requires_stage_handoff") is True or model_result.get("requires_tts_stage") is True: + # Reserve the protocol response only when Stage1 emits text/audio. + return close_reason, False + is_listen = model_result.get("is_listen") + model_turn_id = coerce_int(model_result.get("model_turn_id")) + if model_result.get("is_buffering") is True or model_result.get("prefill_success") is False: + if model_result.get("data_plane_request_id") == session.active_request_id: + session.clear_request() + payload = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": model_result.get("reason") or "buffering", + "model_listen": False, + "buffering": True, + } + self._attach_runtime_metadata(payload, model_result) + self._out.emit(payload) + return close_reason, emitted_response + if is_listen is True: + return await self._on_model_listen( + model_result, + model_turn_id=model_turn_id, + data_plane_request_id=data_plane_request_id, + expected_epoch=expected_epoch, + ) + + text = model_result.get("text") + audio = model_result.get("audio_data", model_result.get("audio")) + end_of_turn = bool(model_result.get("end_of_turn", False)) + has_text = isinstance(text, str) and bool(text) + has_audio = isinstance(audio, str) and bool(audio) + if not has_text and not has_audio and not end_of_turn: + tts_segment_ended = ( + model_result.get("stage_role") == "tts" and model_result.get("abort_data_plane_request") is True + ) + if ( + tts_segment_ended + and auto_response + and (model_turn_id is not None or session.active_response_id is not None) + ): + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch, expected_model_turn_id=model_turn_id), + name="duplex-continue", + ) + return close_reason, emitted_response + if end_of_turn and not has_text and not has_audio and session.active_response_id is None: + emitted_response = self._complete_model_turn_without_output( + model_result, + model_turn_id=model_turn_id, + data_plane_request_id=data_plane_request_id, + ) + return close_reason, emitted_response + if session.active_response_id is None and model_turn_id is not None and model_turn_id < session.turn_id: + # Late audio of a completed model turn must not reserve a second response. + return close_reason, emitted_response + self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) + if ( + session.active_response_id is not None + and model_turn_id is not None + and not session.active_response_accepts_model_turn(model_turn_id) + ): + return close_reason, emitted_response + emitted_response = True + response_created = False + response_id = session.active_response_id + if response_id is None: + response_id = session.begin_response(turn_id=model_turn_id) + response_created = True + self._out.emit(self.response_created_payload(response_id, epoch=session.epoch)) + response_stage_metrics = session.accumulate_response_stage_metrics( + model_result.get("stage_metrics") if isinstance(model_result.get("stage_metrics"), Mapping) else None + ) + if response_created: + speak_payload = { + "type": "response.speak", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "text": text if isinstance(text, str) else "", + "end_of_turn": end_of_turn, + "model_speak": True, + } + self._attach_runtime_metadata(speak_payload, model_result, stage_metrics=response_stage_metrics) + self._out.emit(speak_payload) + previous_sent_ms = session.playback.sent_ms + text_chars_before_append = len("".join(session.assistant_text_buffer)) + if isinstance(text, str): + session.append_assistant_text(text) + duration_ms = model_result.get("audio_duration_ms") + text_chars = len("".join(session.assistant_text_buffer)) + mark_duration_ms = None + mark_text_chars: int | None = text_chars + if model_result.get("audio_text_mark") is False: + mark_text_chars = None + if isinstance(duration_ms, int | float): + mark_duration_ms = int(duration_ms) + if model_result.get("audio_duration_is_cumulative") is not True: + mark_duration_ms += session.playback.sent_ms + audio_text_marks = model_result.get("audio_text_marks") + audio_text_marks = self._normalize_audio_text_marks( + audio_text_marks if isinstance(audio_text_marks, list) else None, + audio_offset_ms=( + 0 + if model_result.get("audio_text_marks_are_cumulative") is True + or model_result.get("audio_duration_is_cumulative") is True + else previous_sent_ms + ), + text_offset_chars=( + 0 if model_result.get("audio_text_marks_are_cumulative") is True else text_chars_before_append + ), + ) + session.mark_audio_sent( + mark_duration_ms, + text_chars=mark_text_chars if mark_duration_ms is not None else None, + audio_text_marks=audio_text_marks, + ) + payload = { + "type": "response.output_audio.delta", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "text": text if isinstance(text, str) else "", + "audio": audio if isinstance(audio, str) else "", + "format": ( + model_result.get("audio_format") + if isinstance(model_result.get("audio_format"), str) + else session.response_config.response_format + ), + "end_of_turn": end_of_turn, + "model_speak": True, + } + if mark_duration_ms is not None: + payload["audio_duration_ms"] = mark_duration_ms + if audio_text_marks: + payload["audio_text_marks"] = audio_text_marks + elif mark_duration_ms is not None and mark_text_chars is not None: + payload["audio_text_marks"] = [ + {"text_chars": max(0, int(mark_text_chars)), "audio_end_ms": max(0, int(mark_duration_ms))} + ] + payload["playback"] = session.playback.as_dict() + sample_rate_hz = model_result.get("sample_rate_hz") or model_result.get("audio_sample_rate_hz") + if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) > 0: + payload["sample_rate_hz"] = int(sample_rate_hz) + self._attach_runtime_metadata(payload, model_result, stage_metrics=response_stage_metrics) + self._out.emit(payload) + if ( + not end_of_turn + and model_result.get("stage_role") == "tts" + and model_result.get("abort_data_plane_request") is True + and auto_response + ): + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + ) + if end_of_turn: + data_plane_request_id = model_result.get("data_plane_request_id") + if isinstance(data_plane_request_id, str) and not auto_response: + data_plane.close_stream(data_plane_request_id) + if isinstance(data_plane_request_id, str): + if not auto_response and data_plane_request_id == session.active_request_id: + session.clear_request() + if not auto_response: + data_plane.mark_terminal(data_plane_request_id) + should_commit = self.should_commit_response_to_history(session, response_id) + committed_message = session.end_response(commit_text=should_commit, preserve_request=auto_response) + model_turn_id = coerce_int(model_result.get("model_turn_id")) + if model_turn_id is not None: + session.complete_model_turn(model_turn_id) + if should_commit: + session.register_history_item(f"item_{response_id}", committed_message) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": committed_message is not None, + "playback": session.playback.as_dict(), + } + ) + return close_reason, emitted_response + + def _end_active_response_before_future_model_turn(self, *, model_turn_id: int | None) -> None: + session = self._ctx.session + if not self._out.auto_responds(): + return + response_id = session.active_response_id + active_turn_id = session.active_response_turn_id + if response_id is None or model_turn_id is None or active_turn_id is None: + return + if int(model_turn_id) <= int(active_turn_id): + return + session.complete_model_turn(int(model_turn_id) - 1) + should_commit = self.should_commit_response_to_history(session, response_id) + committed_message = session.end_response(commit_text=should_commit, preserve_request=True) + if should_commit: + session.register_history_item(f"item_{response_id}", committed_message) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": committed_message is not None, + "playback": session.playback.as_dict(), + } + ) + + @staticmethod + def _normalize_audio_text_marks( + audio_text_marks: list[object] | None, + *, + audio_offset_ms: int, + text_offset_chars: int, + ) -> list[dict[str, int]] | None: + if not audio_text_marks: + return None + normalized: list[dict[str, int]] = [] + for raw_mark in audio_text_marks: + if not isinstance(raw_mark, dict): + continue + raw_text_chars = raw_mark.get("text_chars") + raw_audio_end_ms = raw_mark.get("audio_end_ms", raw_mark.get("audio_ms")) + if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): + continue + normalized.append( + { + "text_chars": max(0, int(raw_text_chars) + int(text_offset_chars)), + "audio_end_ms": max(0, int(raw_audio_end_ms) + int(audio_offset_ms)), + } + ) + return normalized or None + + @staticmethod + def _attach_runtime_metadata( + payload: dict[str, object], + model_result: dict[str, object], + *, + stage_metrics: Mapping[str, object] | None = None, + ) -> None: + metadata: dict[str, object] = {} + runtime_impl = model_result.get("runtime_impl") + if isinstance(runtime_impl, str) and runtime_impl: + metadata["runtime_impl"] = runtime_impl + owned_runtime = model_result.get("owned_runtime") + if isinstance(owned_runtime, bool): + metadata["owned_runtime"] = owned_runtime + model_turn_id = coerce_int(model_result.get("model_turn_id")) + if model_turn_id is not None: + metadata["model_turn_id"] = model_turn_id + for name in ("uses_model_runner_scheduler", "runner_kv_backed"): + value = model_result.get(name) + if isinstance(value, bool): + metadata[name] = value + effective_stage_metrics = stage_metrics if stage_metrics is not None else model_result.get("stage_metrics") + if isinstance(effective_stage_metrics, Mapping): + metadata["stage_metrics"] = { + str(stage_id): dict(values) + for stage_id, values in effective_stage_metrics.items() + if isinstance(values, Mapping) + } + if metadata: + payload["vllm_omni"] = metadata + + # ------------------------------------------------------------------ # + # Silence continuation # + # ------------------------------------------------------------------ # + + def silence_unit_payload(self) -> dict[str, object]: + samples = int(self._ctx.plugin.silence_continuation_samples) + audio = ( + self._SILENCE_UNIT_PAYLOAD_AUDIO + if samples == 16000 + else base64.b64encode(bytes(samples * 4)).decode("ascii") + ) + return {"type": "audio", "audio": audio, "format": "pcm_f32le", "sample_rate_hz": 16000} + + def response_continuations_remaining(self, response_id: str) -> bool: + model_state = self._ctx.model_state + owner_id = f"response:{response_id}" + count = model_state.continuation_units if model_state.continuation_owner_id == owner_id else 0 + limit = ( + self._AUTO_RESPONSE_MAX_CONTINUATION_UNITS + if self._out.auto_responds() + else self._RESPONSE_MAX_CONTINUATION_UNITS + ) + return count < limit + + async def _finish_bounded_auto_response( + self, + *, + expected_epoch: int | None, + model_turn_id: int | None = None, + ) -> None: + session = self._ctx.session + model_state = self._ctx.model_state + response_id = session.active_response_id + response_epoch = session.epoch + response_turn_id = model_turn_id if model_turn_id is not None else session.active_response_turn_id + if expected_epoch is not None and response_epoch != expected_epoch: + return + model_state.clear_continuation() + payload: dict[str, object] = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": response_epoch, + "reason": "continuation_limit", + "model_listen": False, + } + if response_id is not None: + payload["response_id"] = response_id + self._out.emit(payload) + if response_id is None: + if session.epoch == response_epoch and response_turn_id is not None: + session.complete_model_turn(response_turn_id) + return + if session.epoch != response_epoch or session.active_response_id != response_id: + return + if response_turn_id is not None: + session.complete_model_turn(response_turn_id) + session.end_response(commit_text=False, preserve_request=True) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": response_epoch, + "committed": False, + "playback": session.playback.as_dict(), + } + ) + + def silence_continuation_is_stale( + self, + *, + request_id: str, + response_id: str | None, + response_owned: bool, + expected_epoch: int | None, + expected_model_turn_id: int | None, + ) -> bool: + session = self._ctx.session + stale_common_owner = ( + session.state == DuplexSessionState.CLOSED + or session.active_request_id != request_id + or (expected_epoch is not None and session.epoch != expected_epoch) + ) + stale_response_owner = response_owned and session.active_response_id != response_id + stale_model_turn_owner = not response_owned and ( + session.active_response_id is not None or session.turn_id != expected_model_turn_id + ) + return stale_common_owner or stale_response_owner or stale_model_turn_owner + + async def maybe_continue_response( + self, + *, + expected_epoch: int | None, + expected_model_turn_id: int | None = None, + ) -> None: + session = self._ctx.session + model_state = self._ctx.model_state + response_id = session.active_response_id + if session.state == DuplexSessionState.CLOSED or self._ctx.run.closing: + model_state.clear_continuation() + return + request_id = session.active_request_id + if request_id is None: + model_state.clear_continuation() + return + if expected_epoch is not None and session.epoch != expected_epoch: + return + auto_response = self._out.auto_responds() + response_owned = response_id is not None + if response_owned: + owner_id = f"response:{response_id}" + payload_turn_id_value = ( + session.active_response_turn_id if session.active_response_turn_id is not None else session.turn_id + ) + else: + if not auto_response or expected_model_turn_id is None or session.turn_id != expected_model_turn_id: + model_state.clear_continuation() + return + owner_id = f"model-turn:{expected_model_turn_id}" + payload_turn_id_value = expected_model_turn_id + count = model_state.continuation_units if model_state.continuation_owner_id == owner_id else 0 + continuation_limit = ( + self._AUTO_RESPONSE_MAX_CONTINUATION_UNITS if auto_response else self._RESPONSE_MAX_CONTINUATION_UNITS + ) + if count >= continuation_limit: + if auto_response: + await self._finish_bounded_auto_response( + expected_epoch=expected_epoch, + model_turn_id=payload_turn_id_value, + ) + return + payload = self.silence_unit_payload() + if auto_response and count + 1 == continuation_limit: + payload["force_listen"] = True + payload["duplex_turn_id"] = payload_turn_id_value + try: + scheduled = await self._schedule_silence_continuation( + payload, + request_id=request_id, + owner_id=owner_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to schedule duplex native response continuation: %s", exc) + scheduled = False + if scheduled: + model_state.continuation_owner_id = owner_id + model_state.continuation_units = count + 1 diff --git a/vllm_omni/engine/duplex/session/overlap_policy.py b/vllm_omni/engine/duplex/session/overlap_policy.py new file mode 100644 index 000000000..0aad04368 --- /dev/null +++ b/vllm_omni/engine/duplex/session/overlap_policy.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""How one duplex session reacts to input that arrives while it is speaking. + +Short acknowledgement, barge-in, or keep listening: the rules read the session +config and the incoming append, and nothing else. Extracted from +``DuplexSessionRunner`` because they were a decision function wearing a +method's clothes -- the only runner state they ever touched was ``session`` -- +and because reaching them previously meant driving a whole append through the +runner. +""" + +from __future__ import annotations + +import binascii +from collections.abc import Mapping + +import numpy as np +import pybase64 as base64 + +from vllm_omni.engine.duplex.config import DuplexOverlapPolicy +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + + +def decide( + session: DuplexEngineSession, + event: dict[str, object], + payload: dict[str, object], + *, + auto_responds: bool, +) -> dict[str, object]: + duration_ms = input_audio_duration_ms(event, payload) + is_speech = input_looks_like_speech(session, event, payload) + if not session.capabilities.supports_barge_in and event_requests_barge_in(event): + return defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=is_speech) + explicit = event.get("overlap_action") or event.get("overlap") + if isinstance(explicit, str): + normalized = explicit.strip().lower() + if normalized in {"barge_in", "interrupt", "cancel"}: + return { + "action": "barge_in", + "reason": "client_overlap_action", + "duration_ms": duration_ms, + "buffer_audio": True, + } + if normalized in {"listen", "continue", "continue_output", "ack"}: + session.reset_overlap_speech() + return { + "action": "listen", + "reason": "client_overlap_action", + "duration_ms": duration_ms, + "buffer_audio": ( + normalized == "listen" and is_speech and duration_ms > session.config.overlap_short_ack_ms + ), + "defer_runtime_append": True, + } + if normalized in {"drop", "ignore", "silence"}: + session.reset_overlap_speech() + return { + "action": "drop", + "reason": "client_overlap_action", + "duration_ms": duration_ms, + "buffer_audio": False, + } + + if bool(event.get("force_barge_in", False)): + return { + "action": "barge_in", + "reason": "client_force_barge_in", + "duration_ms": duration_ms, + "buffer_audio": True, + } + if auto_responds: + if is_speech: + session.accumulate_overlap_speech(duration_ms) + speech_started = vad_speech_started(event, payload) + if ( + session.capabilities.supports_barge_in + and is_speech + and session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + and speech_started is not False + ): + return { + "action": "barge_in", + "reason": ("server_vad_speech_started" if speech_started is True else "barge_in_on_speech"), + "cancel_reason": "turn_detected" if speech_started is True else "barge_in", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } + return { + "action": "listen", + "reason": "auto_response_continuous", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": False, + "force_listen": event.get("force_listen") is True or payload.get("force_listen") is True, + "preserve_realtime_input": True, + } + if bool(event.get("force_listen", False)): + session.reset_overlap_speech() + return { + "action": "listen", + "reason": "client_force_listen", + "duration_ms": duration_ms, + "buffer_audio": is_speech, + "defer_runtime_append": True, + } + + policy = session.config.overlap_policy + if not is_speech: + if session.overlap_speech_ms <= 0: + session.reset_overlap_speech() + return { + "action": "drop", + "reason": "silence_or_noise", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": False, + } + + if is_short_ack_transcript_hint(event, payload): + session.reset_overlap_speech() + return { + "action": "listen", + "reason": "short_ack_transcript", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": False, + "defer_runtime_append": True, + } + + if policy == DuplexOverlapPolicy.LISTEN_ONLY.value: + session.accumulate_overlap_speech(duration_ms) + return { + "action": "listen", + "reason": "policy_listen_only", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + + if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value and not session.capabilities.supports_barge_in: + return defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=True) + + session.accumulate_overlap_speech(duration_ms) + if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: + speech_started = vad_speech_started(event, payload) + if speech_started is False: + return { + "action": "listen", + "reason": "server_vad_utterance_active", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": False, + "force_listen": True, + "preserve_realtime_input": True, + } + return { + "action": "barge_in", + "reason": ("server_vad_speech_started" if speech_started is True else "policy_barge_in_on_speech"), + "cancel_reason": "turn_detected" if speech_started is True else "barge_in", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } + + if ( + duration_ms <= session.config.overlap_short_ack_ms + and session.overlap_speech_ms <= session.config.overlap_short_ack_ms + ): + return { + "action": "listen", + "reason": "short_ack", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + if session.overlap_speech_ms >= session.config.overlap_barge_in_ms: + if not session.capabilities.supports_barge_in: + return { + "action": "listen", + "reason": "barge_in_unsupported", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + return { + "action": "barge_in", + "reason": "long_overlap_speech", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } + return { + "action": "listen", + "reason": "accumulating_overlap_speech", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + + +def vad_speech_started(event: Mapping[str, object], payload: Mapping[str, object]) -> bool | None: + for source in (event, payload): + vad = source.get("vad") + if isinstance(vad, Mapping) and isinstance(vad.get("speech_started"), bool): + return bool(vad["speech_started"]) + return None + + +def event_requests_barge_in(event: Mapping[str, object]) -> bool: + if event.get("force_barge_in") is True: + return True + explicit = event.get("overlap_action") or event.get("overlap") + return isinstance(explicit, str) and explicit.strip().lower() in {"barge_in", "interrupt", "cancel"} + + +def defer_unsupported_barge_in( + session: DuplexEngineSession, + *, + duration_ms: int, + is_speech: bool, +) -> dict[str, object]: + if is_speech: + session.accumulate_overlap_speech(duration_ms) + return { + "action": "listen", + "reason": "barge_in_unsupported", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": is_speech, + "defer_runtime_append": True, + } + + +def is_short_ack_transcript_hint(event: dict[str, object], payload: dict[str, object]) -> bool: + raw_text = event.get("transcript") or event.get("text") or payload.get("transcript") or payload.get("text") + if not isinstance(raw_text, str): + return False + normalized = raw_text.strip().lower() + if not normalized: + return False + compact = "".join(ch for ch in normalized if ch.isalnum() or "一" <= ch <= "鿿") + if compact in { + "嗯", + "嗯嗯", + "对", + "对的", + "好", + "好的", + "继续", + "继续说", + "可以", + "是的", + "yes", + "yeah", + "yep", + "ok", + "okay", + "continue", + "goon", + "right", + }: + return True + return normalized in {"go on", "keep going", "please continue"} + + +def input_audio_duration_ms(event: dict[str, object], payload: dict[str, object]) -> int: + for key in ("duration_ms", "audio_duration_ms"): + value = event.get(key) + if isinstance(value, int | float): + return max(0, int(value)) + fmt = payload.get("format") + sample_rate_hz = payload.get("sample_rate_hz") + audio = payload.get("audio") + if fmt == "pcm_f32le" and isinstance(sample_rate_hz, int) and sample_rate_hz > 0 and isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return 0 + return int((len(raw) // 4) * 1000 / sample_rate_hz) + return 0 + + +def merge_audio_payloads(first: dict[str, object], second: dict[str, object]) -> dict[str, object]: + if first.get("format") != "pcm_f32le" or second.get("format") != "pcm_f32le": + return second + first_rate = first.get("sample_rate_hz") + second_rate = second.get("sample_rate_hz") + if not isinstance(first_rate, int) or not isinstance(second_rate, int) or first_rate != second_rate: + return second + first_audio = first.get("audio") + second_audio = second.get("audio") + if not isinstance(first_audio, str) or not isinstance(second_audio, str): + return second + try: + first_raw = base64.b64decode(first_audio, validate=True) + second_raw = base64.b64decode(second_audio, validate=True) + except (binascii.Error, ValueError): + return second + merged = dict(second) + merged["audio"] = base64.b64encode(first_raw + second_raw).decode("ascii") + merged["sample_rate_hz"] = first_rate + merged_frames = [ + frame + for source in (first.get("video_frames"), second.get("video_frames")) + if isinstance(source, list) + for frame in source + if isinstance(frame, str) and frame + ] + if merged_frames: + merged["video_frames"] = merged_frames + else: + merged.pop("video_frames", None) + merged["force_listen"] = bool(first.get("force_listen", False)) or bool(second.get("force_listen", False)) + merged.pop("force_speak", None) + merged["is_speech"] = bool(first.get("is_speech", False)) or bool(second.get("is_speech", False)) + return merged + + +def should_force_listen_for_short_commit( + session: DuplexEngineSession, event: dict[str, object], payload: dict[str, object] +) -> bool: + if event.get("force_listen") is True or payload.get("force_listen") is True: + return True + if event.get("force_barge_in") is True: + return False + if event.get("response_create") is not True: + return False + duration_ms = input_audio_duration_ms(event, payload) + return 0 < duration_ms <= session.config.overlap_short_ack_ms + + +def should_force_listen_for_auto_response_overlap( + event: dict[str, object], payload: dict[str, object], *, auto_responds: bool +) -> bool: + if not auto_responds: + return False + if event.get("force_barge_in") is True: + return False + return event.get("force_listen") is True or payload.get("force_listen") is True + + +def input_looks_like_speech(session: DuplexEngineSession, event: dict[str, object], payload: dict[str, object]) -> bool: + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return value + vad = event.get("vad") + if isinstance(vad, dict): + value = vad.get("is_speech") + if isinstance(value, bool): + return value + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) >= 0.5 + probability = event.get("speech_probability") + if isinstance(probability, int | float): + return float(probability) >= 0.5 + + fmt = payload.get("format") + audio = payload.get("audio") + if fmt in {"pcm_f32le", "pcm16"} and isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return True + if fmt == "pcm_f32le": + if len(raw) < 4 or len(raw) % 4 != 0: + return True + samples = np.frombuffer(raw, dtype=np.float32) + else: + if len(raw) < 2 or len(raw) % 2 != 0: + return True + samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + if samples.size == 0: + return False + rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) + return rms >= session.config.overlap_silence_rms + return True diff --git a/vllm_omni/engine/duplex/session/playback_ledger.py b/vllm_omni/engine/duplex/session/playback_ledger.py new file mode 100644 index 000000000..c6843f64b --- /dev/null +++ b/vllm_omni/engine/duplex/session/playback_ledger.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Applying one ``playback.ack`` to a duplex session. + +The client reports how much of a response it has actually played; the session +uses that to commit or truncate the assistant item in conversation history, and +to release the response's playback state once it is fully played. + +Extracted from ``DuplexSessionRunner`` because the only runner state it touched +was ``session``. It returns the events it produces rather than emitting them, +so the rules can be exercised without a runner. +""" + +from __future__ import annotations + +from vllm_omni.engine.duplex.events import DuplexEvent, PlaybackAcknowledged, error_event +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity + + +def apply_playback_ack(session: DuplexEngineSession, event: dict[str, object]) -> list[DuplexEvent]: + """Apply one ``playback.ack`` and return the events it produces.""" + played_ms = event.get("played_ms", event.get("audio_ms", 0)) + committed_ms = event.get("committed_ms") + if not isinstance(played_ms, int | float): + return [error_event("bad_event", "playback.ack requires played_ms")] + try: + session.touch_lease(DuplexLeaseActivity.PLAYBACK_ACK) + except Exception: + pass + committed_cursor = int(committed_ms) if isinstance(committed_ms, int | float) else int(played_ms) + item_id = event.get("item_id") + response_id = event.get("response_id") + response_id = response_id if isinstance(response_id, str) and response_id else None + if not isinstance(item_id, str) or not item_id: + item_id = f"item_{response_id}" if response_id is not None else None + elif response_id is None and item_id.startswith("item_"): + response_id = item_id.removeprefix("item_") + if response_id is None and item_id is None and len(session.pending_history_item_ids) == 1: + item_id = next(iter(session.pending_history_item_ids)) + if item_id.startswith("item_"): + response_id = item_id.removeprefix("item_") + if response_id is None and item_id is None and session.active_response_id is not None: + response_id = session.active_response_id + item_id = f"item_{response_id}" + if response_id is not None: + expected_item_id = f"item_{response_id}" + if item_id is None: + item_id = expected_item_id + elif item_id != expected_item_id: + return [error_event("playback_item_mismatch", "playback.ack item_id must match item_.")] + if not session.has_assistant_response_item(response_id, item_id): + return [ + error_event("playback_item_not_found", f"No assistant response item is registered for {response_id}.") + ] + if session.playback_ack_is_too_late(response_id, item_id): + return [ + error_event("playback_ack_too_late", "playback.ack arrived after a later user input was committed.") + ] + session.reserve_history_item(item_id) + elif item_id is not None: + return [error_event("playback_item_not_found", "playback.ack requires a response-owned assistant item.")] + hard_truncate = event.get("truncate") is True + if hard_truncate: + playback = session.acknowledge_playback(int(played_ms), committed_cursor, response_id=response_id) + playback = session.truncate_playback_commit(committed_cursor, response_id=response_id) + else: + playback = session.acknowledge_playback(int(played_ms), committed_cursor, response_id=response_id) + committed_history = False + if isinstance(item_id, str) and item_id: + expected_item_id = f"item_{response_id}" if response_id is not None else None + if ( + expected_item_id == item_id + and item_id not in session.history_item_ids + and item_id not in session.pending_history_item_ids + ): + session.register_history_item(item_id, None) + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + elif session.pending_history_item_ids: + pending_ids = list(session.pending_history_item_ids) + if len(pending_ids) == 1: + item_id = pending_ids[0] + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + elif session.active_response_id is not None: + item_id = f"item_{session.active_response_id}" + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + elif session.last_assistant_full_message is not None: + if item_id is None and session.history_item_ids: + assistant_item_ids = [ + known_item_id + for known_item_id, message in session.history_item_ids.items() + if message.get("role") == "assistant" + ] + if len(assistant_item_ids) == 1: + item_id = assistant_item_ids[0] + if isinstance(item_id, str) and item_id: + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + events: list[DuplexEvent] = [ + PlaybackAcknowledged( + details={ + "type": "playback.acknowledged", + "session_id": session.session_id, + "epoch": session.epoch, + "item_id": item_id, + "played_ms": int(played_ms), + "committed_ms": committed_cursor, + "truncate": event.get("truncate") is True, + "playback": playback.as_dict(), + "history_committed": committed_history, + } + ) + ] + if committed_history and committed_cursor >= max(playback.sent_ms, playback.generated_ms): + session.release_response_playback(response_id) + session.release_response_history_snapshot(response_id) + return events diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py new file mode 100644 index 000000000..7915500b7 --- /dev/null +++ b/vllm_omni/engine/duplex/session/runner.py @@ -0,0 +1,1772 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Engine-resident session runner: one ordered mailbox and one session state per duplex session. + +``DuplexSessionRunner`` owns the whole lifecycle of one session on the +orchestrator loop (the session is never touched from another thread): + +* every mutation of ``DuplexEngineSession`` happens on this loop, through the + mailbox worker (commands, stage outputs, internal items) or through tracked + tasks that re-validate ``(epoch, turn_id)`` after each ``await``; +* appends are planned with the model plugin and submitted to the stage port + in wire order on the per-session append tail (no RPC hop); +* stage outputs are pushed in by ``DuplexOrchestrator._intercept_stage_output`` + instead of being polled through request queues; +* everything the session says leaves as typed events through + ``DuplexSessionManager.emit`` after terminal-acceptance / stale-epoch + filtering, so a cancelled epoch can never speak again. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, TypeVar + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate +from vllm_omni.engine.duplex.commands import ( + AckPlayback, + AppendAudio, + AppendText, + BargeIn, + CancelInput, + CancelResponse, + ClearInput, + ClearOutputAudio, + CloseSession, + Commit, + CreateItem, + CreateResponse, + DeleteItem, + DuplexCommand, + Heartbeat, + SignalTurn, + TruncateItem, + UpdateSession, +) +from vllm_omni.engine.duplex.config import ( + DuplexConfigError, + DuplexOverlapPolicy, + DuplexPlaybackCommitPolicy, + DuplexSessionState, + ResponseCreateOptions, +) +from vllm_omni.engine.duplex.contracts import ( + DuplexOutputContext, + DuplexOutputDecision, + DuplexStagePort, +) +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + InputCleared, + SessionExpired, + SessionHeartbeatAck, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexModelPlugin, + DuplexModelSessionState, + PcmAppendReservation, + payload_turn_id, +) +from vllm_omni.engine.duplex.realtime_events import ( + RealtimeProjectionState, + discard_pending_input_audio, + note_input_append, + resolve_cancel_response, + resolve_clear_output_audio, + resolve_commit, + resolve_create_item, + resolve_delete_item, + resolve_truncate_item, + retrieve_item_events, +) +from vllm_omni.engine.duplex.session import helpers, overlap_policy, playback_ledger +from vllm_omni.engine.duplex.session.append_task import AppendAttempt +from vllm_omni.engine.duplex.session.commit_policy import CommitAction, CommitSnapshot, decide_commit_action +from vllm_omni.engine.duplex.session.context import ( + DuplexRunState, + DuplexSessionContext, + DuplexSessionTasks, + StageOutput, +) +from vllm_omni.engine.duplex.session.control import SessionControl +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.model_channel import ModelChannel +from vllm_omni.engine.duplex.turn_detection import ( + TurnDetectionResult, +) +from vllm_omni.metrics.stats import StageRequestStats + +if TYPE_CHECKING: + from vllm.config import ModelConfig + from vllm.outputs import RequestOutput + + from vllm_omni.engine.duplex.session.manager import DuplexSessionManager + +logger = init_logger(__name__) + +_OffloadT = TypeVar("_OffloadT") + + +@dataclass(frozen=True, slots=True) +class _Internal: + """Runner-internal mailbox item (its payload runs through the dict-based handlers).""" + + kind: str + payload: dict[str, object] = field(default_factory=dict) + + +_CANCEL_EVENTS = frozenset({"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}) + + +class DuplexSessionRunner: + """Owns one ``DuplexEngineSession`` on the orchestrator loop (see module docstring).""" + + def __init__( + self, + *, + session: DuplexEngineSession, + plugin: DuplexModelPlugin, + stage_port: DuplexStagePort, + manager: DuplexSessionManager, + model_config: ModelConfig | None, + ) -> None: + self.session = session + self.plugin = plugin + self.stage_port = stage_port + self.manager = manager + self.model_config = model_config + if session.model_state is None: + session.model_state = plugin.create_session_state() + self.model_state: DuplexModelSessionState = session.model_state + self.tasks = DuplexSessionTasks() + self._mailbox: asyncio.Queue[DuplexCommand | StageOutput | _Internal] = asyncio.Queue() + self._worker: asyncio.Task[None] | None = None + self._worker_stopped = False + self._loop: asyncio.AbstractEventLoop | None = None + self._background_tasks: set[asyncio.Task[None]] = set() + #: Flags more than one component reads and writes (see session_context). + self.run = DuplexRunState() + self.ctx = DuplexSessionContext( + session=session, + model_state=self.model_state, + plugin=plugin, + stage_port=stage_port, + manager=manager, + tasks=self.tasks, + run=self.run, + services=self, + ) + self.out = SessionEmitter(self.ctx, promote_deferred_overlap=self._promote_deferred_overlap_later) + self.model = ModelChannel( + self.ctx, + self.out, + close_from_runtime=self._close_from_runtime, + schedule_silence_continuation=self._schedule_silence_continuation, + abort_request=self._abort_request_background, + ) + self.control = SessionControl( + self.ctx, + self.out, + self.model, + wait_for_append_tail=self._wait_for_append_tail, + ) + + # ------------------------------------------------------------------ # + # Public interface # + # ------------------------------------------------------------------ # + + def start(self) -> None: + self._loop = asyncio.get_running_loop() + session = self.session + if self.out.projector is None: + default_payload = session.config.extra_body.get("realtime_session_payload") + self.out.projector = RealtimeProjectionState( + session_id=session.session_id, + model=session.config.model, + default_payload=default_payload if isinstance(default_payload, Mapping) else None, + initial_session_update=True, + ) + # Every client speaks the Realtime protocol now; the old runner forced + # the ACK-only playback ledger for that path. + session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value + self.control.init_turn_detection() + self._worker = self._loop.create_task(self._run(), name=f"duplex-session-{session.session_id}") + self.emit({"type": "session.created", "session": session.as_public_dict()}) + + def submit(self, command: DuplexCommand) -> None: + self._mailbox.put_nowait(command) + + def on_stage_output( + self, + stage_id: int, + output: RequestOutput, + metrics: StageRequestStats | None, + *, + request_id: str, + context: DuplexOutputContext, + ) -> bool: + """Accept one stage output (orchestrator loop); return True when it must not be forwarded.""" + decision: DuplexOutputDecision | None = None + if stage_id < context.final_stage_id: + decision = self.model.decide_output(stage_id, output, context) + consume = decision is not None or stage_id >= context.final_stage_id + if not consume: + # Stage0 text without a direct decision feeds the TTS stage as before. + # Its metrics still have to reach the client: before sessions moved + # into the engine the orchestrator published them as a standalone + # ``StageMetricsMessage``, a path session-owned requests no longer + # take. Hand them to the session instead of dropping them, on the + # mailbox so they stay ordered with this session's other work. + snapshot = self.model.stage_metrics_snapshot(stage_id, metrics, output) + if snapshot is not None and not self.run.closing and self.session.state != DuplexSessionState.CLOSED: + self._mailbox.put_nowait(_Internal("stage_metrics", {"stage_metrics": snapshot})) + return False + if self.run.closing or self.session.state == DuplexSessionState.CLOSED: + return True + self._mailbox.put_nowait( + StageOutput( + stage_id=stage_id, + output=output, + metrics=metrics, + request_id=request_id, + context=context, + decision=decision, + ) + ) + return True + + def on_stage_failure(self, stage_id: int, exc: BaseException) -> None: + """A stage rejected this session's request: fail the active response now. + + Runs synchronously on the loop (no mailbox hop): the orchestrator + expires the session right after this call, so a queued item could be + cancelled with the worker and the client would only see + ``session.expired``. Emitting here keeps the order + ``error`` -> failed ``response.done`` -> ``session.expired``. + """ + session = self.session + if session.state == DuplexSessionState.CLOSED: + return + self._emit_error( + "runtime_data_plane_stream_failed", + f"Stage-{stage_id} input processor failed: {type(exc).__name__}: {exc}", + ) + response_id = session.active_response_id + if response_id is not None: + session.end_response(commit_text=False) + self.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": "runtime_data_plane_stream_failed"}, + "playback": session.playback.as_dict(), + } + ) + + @property + def closed_emitted(self) -> bool: + """Whether ``session.closed`` / ``session.expired`` already left this runner.""" + return self.run.closed_emitted + + def mark_closed_emitted(self) -> None: + """Claim the session's one terminal event for the caller. + + The manager emits the deferred terminal itself, after the stage cleanup; + recording it here stops a late runtime close emitting a second one. + """ + self.run.closed_emitted = True + self.run.closed_deferred = False + + @property + def closing(self) -> bool: + """Whether an irreversible close has begun (commands and control ops are refused).""" + return self.run.closing or self.session.state != DuplexSessionState.OPEN + + async def close(self, reason: str, *, emit_closed: bool = True) -> None: + """Graceful close: cancel work, release the data plane, emit ``session.closed``. + + With ``emit_closed=False`` the manager emits ``session.closed`` itself + once the stage resources are released, so the event also means "the + admission slot is free again". + """ + session = self.session + if session.state == DuplexSessionState.CLOSED and (self.run.closed_emitted or self.run.closed_deferred): + await self._stop_worker() + return + self._begin_close(reason) + self.model_state.audio_buffer.clear() + session.release_all_input_bytes() + self.model_state.input_since_commit = False + self.model_state.speech_since_commit = False + self.model_state.clear_committed_audio() + await self.tasks.cancel_append_tasks() + self.model.cancel_data_plane_stream() + await self._cancel_active_response(self.tasks.active_response_task, reason=reason, notify=False) + self.tasks.active_response_task = None + self._cleanup_duplex_session_state() + if not self.run.closed_emitted and not self.run.closed_deferred: + self.run.close_reason = self.run.close_reason or reason + if emit_closed: + self.run.closed_emitted = True + self.emit({"type": "session.closed", "session_id": session.session_id, "reason": reason}) + else: + self.run.closed_deferred = True + session.close() + await self._stop_worker() + + async def expire(self, reason: str, *, emit_expired: bool = True) -> None: + """Lease expiry / runtime cleanup: emit ``session.expired`` and tear down. + + With ``emit_expired=False`` the manager emits the event after the stage + cleanup (see ``close``). + """ + session = self.session + self._begin_close(reason) + if not self.run.closed_emitted: + if emit_expired: + # Also when the event was deferred: the manager only emits a + # deferred terminal for the teardown it drives itself, and an + # expiry that arrives first (a stage failure racing a wire + # ``session.close``) takes the runner away from it. Emitting + # here keeps "every session ends with one terminal event". + self.run.closed_emitted = True + self._emit_events([SessionExpired(reason=reason)]) + else: + self.run.closed_deferred = True + await self.tasks.cancel_append_tasks() + self.model.cancel_data_plane_stream() + active_response_task = self.tasks.active_response_task + if active_response_task is not None and not active_response_task.done(): + active_response_task.cancel() + await asyncio.gather(active_response_task, return_exceptions=True) + self.tasks.active_response_task = None + if session.active_response_id is not None: + session.end_response(commit_text=False) + self._cleanup_duplex_session_state() + session.close() + await self._stop_worker() + + async def shutdown(self) -> None: + self.run.closing = True + self.run.closed_emitted = True + await self.tasks.cancel_append_tasks() + self.model.cancel_data_plane_stream() + active_response_task = self.tasks.active_response_task + if active_response_task is not None and not active_response_task.done(): + active_response_task.cancel() + await asyncio.gather(active_response_task, return_exceptions=True) + self.session.close() + await self._stop_worker() + + # ------------------------------------------------------------------ # + # Worker # + # ------------------------------------------------------------------ # + + async def _run(self) -> None: + while not self._worker_stopped: + item = await self._mailbox.get() + try: + await self._handle_item(item) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Duplex session %s failed handling %r: %s", self.session.session_id, item, exc) + self._emit_error("internal_error", str(exc)) + + async def _stop_worker(self) -> None: + worker = self._worker + self._worker = None + # Set even when the worker itself is stopping (a wire ``session.close`` + # runs inside it): its loop exits after the current item instead of + # parking on the mailbox forever. + self._worker_stopped = True + for task in list(self._background_tasks): + task.cancel() + if self._background_tasks: + await asyncio.gather(*self._background_tasks, return_exceptions=True) + self._background_tasks.clear() + if worker is not None and worker is not asyncio.current_task() and not worker.done(): + worker.cancel() + await asyncio.gather(worker, return_exceptions=True) + + def spawn(self, coro: Awaitable[None], *, name: str) -> asyncio.Task[None]: + task = asyncio.ensure_future(coro) + task.set_name(name) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + return task + + async def offload(self, fn: Callable[..., _OffloadT], *args: object, **kwargs: object) -> _OffloadT: + loop = self._loop or asyncio.get_running_loop() + if kwargs: + return await loop.run_in_executor(self.manager.executor, lambda: fn(*args, **kwargs)) + return await loop.run_in_executor(self.manager.executor, fn, *args) + + async def _handle_item(self, item: DuplexCommand | StageOutput | _Internal) -> None: + session = self.session + if isinstance(item, StageOutput): + await self.model.on_stage_output_item(item) + return + if isinstance(item, _Internal): + await self._on_internal(item) + return + if self.run.closing or session.state == DuplexSessionState.CLOSED: + if isinstance(item, Commit): + session.release_pending_turn() + elif isinstance(item, AppendAudio): + session.release_input_bytes(len(item.audio)) + return + await self._on_command(item) + + async def _on_internal(self, item: _Internal) -> None: + if item.kind == "stage_metrics": + stage_metrics = item.payload.get("stage_metrics") + if isinstance(stage_metrics, Mapping): + self.session.stash_stage_metrics(stage_metrics) + return + if item.kind == "promote_deferred_overlap": + if self.run.closing or self.session.state != DuplexSessionState.OPEN: + return + payload = item.payload.get("payload") + if not isinstance(payload, dict): + return + await self._start_append( + payload, + final=True, + precreate_response=bool(item.payload.get("precreate_response", False)), + operation_id=self.model_state.committed_audio_operation_id, + retained_committed_payload=(payload if self.model_state.committed_audio_payload is payload else None), + ) + return + if item.kind == "commit": + await self._on_commit(dict(item.payload)) + return + if item.kind == "run_payload": + await self._run_internal_payload(dict(item.payload)) + return + logger.warning("Unknown duplex runner internal item: %s", item.kind) + + async def _run_internal_payload(self, payload: dict[str, object]) -> None: + """Run one internal event dictionary through the matching handler.""" + event_type = payload.get("type") + if event_type == "input_audio_buffer.append": + await self._on_append_audio(payload) + elif event_type in {"input_audio_buffer.commit", "input.commit", "response.create"}: + await self._on_commit(payload) + elif event_type == "playback.ack": + self._emit_events(playback_ledger.apply_playback_ack(self.session, payload)) + elif event_type in _CANCEL_EVENTS: + await self._on_cancel(payload) + elif event_type == "turn.signal": + await self.control.on_turn_signal(payload) + else: + self._emit_error("unknown_event", f"Unknown duplex event: {event_type}") + + # ------------------------------------------------------------------ # + # Commands # + # ------------------------------------------------------------------ # + + async def _on_command(self, command: DuplexCommand) -> None: + session = self.session + projector = self._require_projector() + if isinstance(command, AppendAudio): + # The manager reserved the wire size at admission; the handler + # re-reserves the decoded size around its PCM reservation. + session.release_input_bytes(len(command.audio)) + await self._on_append_audio(command.payload()) + elif isinstance(command, AppendText): + session.mark_user_input_activity() + self._emit_error( + "native_text_append_unsupported", + "The selected native duplex runtime accepts audio append only", + event_id=command.event_id, + ) + elif isinstance(command, Commit): + session.release_pending_turn() + resolved = resolve_commit(projector, command) + self._emit_events(resolved.events) + if resolved.reset_vad: + self.control.reset_vad() + if resolved.payload is not None: + await self._on_commit(resolved.payload) + elif isinstance(command, CreateResponse): + await self._on_commit(command.payload()) + elif isinstance(command, ClearInput): + self._on_clear_input() + elif isinstance(command, CancelResponse): + resolved = resolve_cancel_response(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._on_cancel(payload) + elif isinstance(command, ClearOutputAudio): + resolved = resolve_clear_output_audio(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._on_cancel(payload) + elif isinstance(command, CancelInput | BargeIn): + await self._on_cancel(command.payload()) + elif isinstance(command, SignalTurn): + if command.event == "conversation.item.retrieve": + self._emit_events( + retrieve_item_events( + projector, + { + **dict(command.signal_payload), + **({"event_id": command.event_id} if command.event_id else {}), + }, + ) + ) + return + payload = command.payload() + if command.event in _CANCEL_EVENTS: + normalized = dict(payload.get("payload") or {}) + normalized.update(payload) + normalized["type"] = command.event + await self._on_cancel(normalized) + else: + await self.control.on_turn_signal(payload) + elif isinstance(command, UpdateSession): + await self.control.on_session_update(dict(command.patch), realtime_event_id=command.event_id) + elif isinstance(command, AckPlayback): + self._emit_events(playback_ledger.apply_playback_ack(self.session, command.payload())) + elif isinstance(command, Heartbeat): + self._on_heartbeat(command) + elif isinstance(command, CreateItem): + resolved = resolve_create_item(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._run_internal_payload(payload) + elif isinstance(command, DeleteItem): + resolved = resolve_delete_item(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self.control.on_turn_signal(payload) + elif isinstance(command, TruncateItem): + resolved = resolve_truncate_item(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._run_internal_payload(payload) + elif isinstance(command, CloseSession): + await self._on_close_command(command.reason) + else: + self._emit_error("unknown_event", f"Unknown duplex command: {type(command).__name__}") + + async def _on_close_command(self, reason: str) -> None: + # A close requested through the ordered command stream (the client's + # ``session.close``) tears the runner down here; the manager then + # aborts the stage requests, frees the admission slot and emits + # ``session.closed`` exactly like a close through the control RPC. + # Mark the session closing before the first await so a concurrent + # resume / touch / command is refused instead of racing the teardown. + self._begin_close(reason) + await self.close(reason, emit_closed=False) + self.manager.close_from_runner(self, reason) + + def _on_heartbeat(self, command: Heartbeat) -> None: + try: + self.session.touch_lease(DuplexLeaseActivity.HEARTBEAT) + except Exception as exc: + self._emit_error("runtime_touch_failed", str(exc), event_id=command.event_id) + return + self._emit_events([SessionHeartbeatAck()]) + + def _on_clear_input(self) -> None: + session = self.session + model_state = self.model_state + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.clear_committed_audio() + session.cancel_pending_input() + projector = self.out.projector + if projector is not None: + from vllm_omni.engine.duplex.realtime_events import clear_input_buffer + + clear_input_buffer(projector) + self._emit_events([InputCleared()]) + + # ------------------------------------------------------------------ # + # Emission (delegated to SessionEmitter) # + # ------------------------------------------------------------------ # + + def emit(self, payload: dict[str, object]) -> None: + """Apply an internal event to the session, project it, and send it.""" + self.out.emit(payload) + + def _emit_events(self, events: list[DuplexEvent]) -> None: + self.out.emit_events(events) + + def _emit_error( + self, + code: str, + message: str, + *, + event_id: object | None = None, + retryable: bool | None = None, + ) -> None: + self.out.emit_error(code, message, event_id=event_id, retryable=retryable) + + def _require_projector(self) -> RealtimeProjectionState: + return self.out.require_projector() + + def _is_stale_model_output(self, payload: dict[str, object]) -> bool: + return self.out.is_stale_model_output(payload) + + def _promote_deferred_overlap_later(self, payload: dict[str, object], precreate_response: bool) -> None: + """Re-enter the mailbox so the promoted turn is handled in command order.""" + self._mailbox.put_nowait( + _Internal( + "promote_deferred_overlap", + {"payload": payload, "precreate_response": precreate_response}, + ) + ) + + # ------------------------------------------------------------------ # + # Session lifecycle helpers # + # ------------------------------------------------------------------ # + + def _begin_close(self, reason: str) -> None: + self.run.closing = True + self.run.close_reason = self.run.close_reason or reason + self.session.mark_closing() + + def _session_auto_responds(self) -> bool: + return self.out.auto_responds() + + def _cleanup_duplex_session_state(self) -> None: + session = self.session + self.plugin.data_plane.close_session(session.session_id, active_request_id=session.active_request_id) + self.run.stream_request_id = None + + # ------------------------------------------------------------------ # + # Append path (was the audio-append branch + start_native_append) # + # ------------------------------------------------------------------ # + + async def _barge_in_for_overlap(self, event: dict[str, object], decision: dict[str, object]) -> bool: + """Interrupt the model so the user's overlapping speech becomes the turn. + + Cancels the active response and the data-plane stream, advances the + barge-in epoch so stale output is dropped, and reports what was cut. + Returns whether the append should continue: losing the fence race + means this append belongs to a superseded epoch and is abandoned. + """ + session = self.session + model_state = self.model_state + event["force_barge_in"] = True + cancelled_fence = session.fence + playback_was_active = helpers.assistant_playback_active(self.session) + model_state.audio_buffer.clear_force_listen() + session.reset_overlap_speech() + model_state.input_since_commit = False + model_state.speech_since_commit = False + await self.tasks.cancel_append_tasks() + had_stream = self.run.stream_request_id is not None + cancel_reason = str(decision.get("cancel_reason") or "barge_in") + cancelled = await self._cancel_active_response( + self.tasks.active_response_task, + reason=cancel_reason, + ) + had_stream = self.model.cancel_data_plane_stream() or had_stream + if not cancelled and had_stream: + old_epoch = session.epoch + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, old_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and playback_was_active: + old_epoch = session.epoch + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, session.last_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": session.last_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if session.epoch > cancelled_fence.epoch: + if not await self.model.signal_cancel_fence(cancelled_fence): + return False + self.tasks.active_response_task = None + return True + + async def _on_append_audio(self, event: dict[str, object]) -> None: + session = self.session + model_state = self.model_state + session.mark_user_input_activity() + audio = event.get("audio") or event.get("data") + if not isinstance(audio, str): + self._emit_error("bad_event", "input_audio_buffer.append requires audio") + return + if not session.capabilities.supports_barge_in and overlap_policy.event_requests_barge_in(event): + self._emit_events([helpers.barge_in_unsupported_error()]) + event = dict(event) + event.pop("force_barge_in", None) + for key in ("overlap_action", "overlap"): + value = event.get(key) + if isinstance(value, str) and value.strip().lower() in {"barge_in", "interrupt", "cancel"}: + event.pop(key, None) + fmt = event.get("format") if isinstance(event.get("format"), str) else "pcm16" + sr_raw = event.get("sample_rate_hz") or event.get("sample_rate") + sample_rate_hz = sr_raw if isinstance(sr_raw, int | float) else 16000 + try: + audio, fmt, sample_rate_hz = await self.offload( + convert_input_audio_with_rate, + audio, + fmt, + sample_rate_hz=sample_rate_hz, + ) + except ValueError as exc: + self._emit_error("bad_event", str(exc)) + return + if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le"}: + self._emit_error("bad_audio", "input_audio_buffer.append pcm16 audio could not be decoded") + return + event["audio"] = audio + event["format"] = fmt + event["sample_rate_hz"] = sample_rate_hz + vad_result = await self.control.run_turn_detection(event) + projector = self._require_projector() + self._emit_events(note_input_append(projector, event, vad_result=vad_result)) + if self.run.closing or session.state != DuplexSessionState.OPEN: + return + + force_listen = bool(event.get("force_listen", False)) + payload: dict[str, object] = { + "type": "audio", + "audio": audio, + "format": fmt, + "sample_rate_hz": sample_rate_hz, + "force_listen": force_listen, + } + video_frames = event.get("video_frames") + if isinstance(video_frames, list): + frames = [frame for frame in video_frames if isinstance(frame, str) and frame] + if frames: + payload["video_frames"] = frames + payload["is_speech"] = overlap_policy.input_looks_like_speech(self.session, event, payload) + auto_responds = self._session_auto_responds() + defer_append = False + buffer_overlap_audio = True + self._mark_pending_silence_superseded() + overlap_active = helpers.response_in_progress(self.session, self.tasks) and ( + not auto_responds + or ( + session.capabilities.supports_barge_in + and ( + session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + or overlap_policy.event_requests_barge_in(event) + ) + ) + ) + if overlap_active: + decision = overlap_policy.decide(self.session, event, payload, auto_responds=auto_responds) + self._emit_events([helpers.overlap_decision_event(self.session, decision)]) + action = decision.get("action") + if action == "drop": + duration_ms = overlap_policy.input_audio_duration_ms(event, payload) + self._emit_events(discard_pending_input_audio(projector, duration_ms)) + self._maybe_schedule_vad_commit(vad_result) + return + if action == "listen": + buffer_overlap_audio = bool(decision.get("buffer_audio", True)) + defer_append = bool(decision.get("defer_runtime_append", True)) + if not buffer_overlap_audio and decision.get("preserve_realtime_input") is not True: + self._emit_events( + discard_pending_input_audio(projector, overlap_policy.input_audio_duration_ms(event, payload)) + ) + if decision.get("force_listen", True) is True: + payload["force_listen"] = True + else: + if not await self._barge_in_for_overlap(event, decision): + return + buffer_overlap_audio = True + defer_append = False + elif not auto_responds and not overlap_policy.input_looks_like_speech(self.session, event, payload): + # Turn-mode only: skip silent chunks so they don't open a response. + self.emit( + { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": "silence_or_noise", + } + ) + self._maybe_schedule_vad_commit(vad_result) + return + if overlap_policy.should_force_listen_for_auto_response_overlap(event, payload, auto_responds=auto_responds): + payload["force_listen"] = True + if not buffer_overlap_audio: + self._maybe_schedule_vad_commit(vad_result) + return + session.mark_user_input_activity() + model_state.input_since_commit = True + model_state.speech_since_commit = model_state.speech_since_commit or overlap_policy.input_looks_like_speech( + self.session, event, payload + ) + raw_audio_bytes = helpers.audio_payload_size_bytes(payload) + try: + if not session.reserve_input_bytes( + raw_audio_bytes, + limit=int(self.manager.runtime_config.max_pending_input_bytes_per_session), + ): + self._emit_error("input_backpressure", "Duplex session pending input exceeds server limit") + return + # Full-duplex: emit each ~chunk_period of audio so the model runs + # per-chunk generation without an explicit response.create. + allow_emit = not defer_append and self._session_auto_responds() + pcm_reservation = model_state.audio_buffer.prepare_append( + payload, + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + allow_emit=allow_emit, + ) + except ValueError as exc: + session.release_input_bytes(raw_audio_bytes) + self._emit_error("bad_event", str(exc)) + return + if pcm_reservation is None: + self._maybe_schedule_vad_commit(vad_result) + return + if pcm_reservation.byte_count == 0: + session.release_input_bytes(raw_audio_bytes) + payload = pcm_reservation.payload + await self._start_append(payload, final=False, pcm_reservation=pcm_reservation) + self._maybe_schedule_vad_commit(vad_result) + + def _maybe_schedule_vad_commit(self, vad_result: TurnDetectionResult | None) -> None: + """Server VAD ended the user turn: run the same commit the old translator synthesized.""" + if vad_result is None or not vad_result.should_commit: + return + if self._session_auto_responds(): + # A model-native session decides its own turns; server VAD is there + # to hear the user (speech_started / speech_stopped, barge-in), not + # to end the turn. Committing at the detector's stop cuts the + # utterance short: the model listens on that early final unit, the + # trailing silence lands in a second, near-empty turn, and the + # client's own commit then has nothing left to answer. The old + # translator synthesized this commit for turn-based server VAD only. + return + command = Commit(final=True, create_response=vad_result.create_response) + resolved = resolve_commit(self._require_projector(), command) + self._emit_events(resolved.events) + if resolved.reset_vad: + self.control.reset_vad() + if resolved.payload is not None: + self._mailbox.put_nowait(_Internal("commit", resolved.payload)) + + def _clear_completed_pending_silence(self) -> None: + task = self.model_state.pending_silence_task + if task is not None and task.done(): + self.model_state.pending_silence_task = None + self.model_state.pending_silence_owner_id = None + + def _mark_pending_silence_superseded(self) -> None: + task = self.model_state.pending_silence_task + if task is None: + return + if task.done(): + self.model_state.pending_silence_task = None + # Do not cancel the task here: a silence append may already have + # reached the stage; before_append skips silence that has not started. + self.model_state.pending_silence_owner_id = None + + def _real_input_waiting(self) -> bool: + self._clear_completed_pending_silence() + return self.model_state.audio_buffer.has_pending() or self.model_state.audio_buffer.has_reserved() + + async def _start_append( + self, + payload: dict[str, object], + *, + final: bool, + precreate_response: bool = False, + pcm_reservation: PcmAppendReservation | None = None, + operation_id: str | None = None, + retained_committed_payload: dict[str, object] | None = None, + silence_continuation: bool = False, + before_append: Callable[[], bool] | None = None, + ) -> asyncio.Task[bool]: + session = self.session + model_state = self.model_state + if not silence_continuation: + self._mark_pending_silence_superseded() + append_epoch = session.epoch + append_turn_id = payload_turn_id(payload) + if append_turn_id is None: + append_turn_id = session.turn_id + request_id = helpers.stage0_request_id(self.session, append_epoch) + if final or precreate_response: + session.bind_request(request_id) + if precreate_response: + session.bind_response_turn(append_turn_id) + if precreate_response and session.active_response_id is None: + response_id = session.begin_response(turn_id=append_turn_id) + self.emit(self.model.response_created_payload(response_id, epoch=append_epoch)) + attempt = AppendAttempt( + ctx=self.ctx, + out=self.out, + model=self.model, + fail_session=self._fail_session_from_append, + payload=payload, + epoch=append_epoch, + request_id=request_id, + final=final, + pcm_reservation=pcm_reservation, + operation_id=operation_id, + retained_committed_payload=retained_committed_payload, + precreated_response_id=session.active_response_id if precreate_response else None, + before_append=before_append, + ) + + predecessor = self.tasks.append_tail + if predecessor is not None and predecessor.done(): + try: + predecessor_ok = predecessor.result() + except (asyncio.CancelledError, Exception): + predecessor_ok = False + if not predecessor_ok: + # Appends queued behind a failed predecessor stop; a later + # command is an explicit retry and starts a new chain. + predecessor = None + task = asyncio.create_task(attempt.run_in_wire_order(predecessor)) + task.add_done_callback(attempt.release_on_failure) + self.tasks.append_tail = task + self.tasks.track_append_task( + task, + epoch=append_epoch, + final=final, + response_bound=final or precreate_response, + ) + if silence_continuation: + model_state.pending_silence_task = task + task.add_done_callback(attempt.clear_pending_silence) + # Let this wire-order effect start before the next mailbox item can cancel it. + await asyncio.sleep(0) + return task + + def _fail_session_from_append(self, reason: str) -> None: + """An append task died: mark the session closing now, close it out next.""" + self._begin_close(reason) + self.spawn(self._close_from_runtime(reason), name="duplex-runtime-close") + + async def _wait_for_append_tail(self) -> bool: + predecessor = self.tasks.append_tail + if predecessor is None: + return True + try: + return await predecessor + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + return False + except Exception: + return False + + async def _schedule_silence_continuation( + self, + payload: object, + *, + request_id: str, + owner_id: str, + response_id: str | None, + response_owned: bool, + expected_epoch: int | None, + expected_model_turn_id: int | None, + ) -> bool: + session = self.session + model_state = self.model_state + self._clear_completed_pending_silence() + pending_silence = model_state.pending_silence_task + if pending_silence is not None and not pending_silence.done(): + if pending_silence is asyncio.current_task(): + return False + try: + if not await pending_silence: + return False + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + return False + except Exception: + return False + self._clear_completed_pending_silence() + pending_silence = model_state.pending_silence_task + if pending_silence is not None and not pending_silence.done(): + return False + append_tail = self.tasks.append_tail + if (append_tail is None or append_tail.done()) and self._real_input_waiting(): + return False + continuation_delay_s = max(0.0, float(session.capabilities.chunk_period_ms or 1000) / 1000.0) + if continuation_delay_s > 0: + await asyncio.sleep(continuation_delay_s) + if ( + self.tasks.append_tail is not append_tail + or ((append_tail is None or append_tail.done()) and self._real_input_waiting()) + or self.model.silence_continuation_is_stale( + request_id=request_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + ): + return False + + def _still_valid() -> bool: + return not self._real_input_waiting() and not self.model.silence_continuation_is_stale( + request_id=request_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + + model_state.pending_silence_owner_id = owner_id + task = await self._start_append( + dict(payload) if isinstance(payload, dict) else {}, + final=False, + silence_continuation=True, + before_append=_still_valid, + ) + return task is not None + + # ------------------------------------------------------------------ # + # Runtime-initiated close # + # ------------------------------------------------------------------ # + + async def _close_from_runtime(self, reason: str) -> None: + """A runtime failure closed the session from inside the engine.""" + session = self.session + if session.state == DuplexSessionState.CLOSED: + return + self.run.runtime_closed = True + self._begin_close(reason) + self._cleanup_duplex_session_state() + # ``closed_deferred`` means a manager-driven close already promised the + # terminal after its stage cleanup. Emitting here too would give the + # session two terminal events: this path can run concurrently, when an + # append task abandoned by that close fails and calls back in. + if not self.run.closed_emitted and not self.run.closed_deferred: + self.run.closed_emitted = True + self.emit({"type": "session.closed", "session_id": session.session_id, "reason": reason}) + session.close() + # The manager aborts the stage requests (if any) and frees the + # admission slot; a session without stage requests must not stay + # registered until its idle TTL. + self.manager.close_from_runner(self, reason) + + # ------------------------------------------------------------------ # + # Cancel / barge-in # + # ------------------------------------------------------------------ # + + async def _on_cancel(self, event: dict[str, object]) -> None: + session = self.session + model_state = self.model_state + event_type = str(event.get("type")) + if event_type == "barge_in" and not session.capabilities.supports_barge_in: + self._emit_events([helpers.barge_in_unsupported_error()]) + return + cancel_reason = ( + "output_audio_buffer_clear" + if event_type == "output_audio_buffer.clear" + else "client_cancelled" + if event_type == "response.cancel" + else "barge_in" + ) + cancelled_fence = session.fence + if event_type == "response.cancel": + requested_response_id = event.get("response_id") + has_active_response_work = helpers.response_in_progress(self.session, self.tasks) + if ( + isinstance(requested_response_id, str) + and session.active_response_id is not None + and requested_response_id != session.active_response_id + ): + self._emit_error( + "response_not_active", + f"Response is not active: {requested_response_id}", + event_id=event.get("realtime_event_id"), + ) + return + if not has_active_response_work: + if isinstance(requested_response_id, str): + return + self._emit_error( + "response_not_active", + "response.cancel requires an active response", + event_id=event.get("realtime_event_id"), + ) + return + had_unbuffered_append = model_state.input_since_commit and not model_state.audio_buffer.has_pending() + playback_was_active = helpers.assistant_playback_active(self.session) + if event_type in {"input.cancel", "barge_in"}: + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.clear_committed_audio() + had_append = await self.tasks.cancel_append_tasks( + response_bound_only=event_type in {"response.cancel", "output_audio_buffer.clear"}, + ) + if event_type == "response.cancel": + session.release_input_bytes(model_state.clear_committed_audio()) + had_stream = self.run.stream_request_id is not None + cancelled = await self._cancel_active_response(self.tasks.active_response_task, reason=cancel_reason) + had_stream = self.model.cancel_data_plane_stream() or had_stream + if not cancelled and (had_stream or had_append or had_unbuffered_append): + old_epoch = session.epoch + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, old_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and playback_was_active: + old_epoch = session.epoch + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, session.last_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": session.last_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and event_type == "response.cancel": + old_epoch = session.epoch + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, old_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and event_type == "output_audio_buffer.clear": + old_playback = session.playback.as_dict() + committed_ms = session.playback.committed_ms + session.clear_playback_cursor() + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": session.active_response_id, + "reason": cancel_reason, + "cancelled_epoch": session.epoch, + "epoch": session.epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + self.tasks.active_response_task = None + return + if not cancelled: + self._cancel_pending_input(reason="barge_in") + if not await self.model.signal_cancel_fence(cancelled_fence): + return + self.tasks.active_response_task = None + + async def _cancel_active_response( + self, + active_task: asyncio.Task[None] | None, + *, + reason: str, + notify: bool = True, + ) -> bool: + session = self.session + has_running_task = active_task is not None and not active_task.done() + if not has_running_task and session.active_request_id is None and session.active_response_id is None: + return False + + old_epoch = session.epoch + old_request_id = session.active_request_id + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + committed_message = session.end_response( + commit_text=self.model.should_commit_response_to_history(session, old_response_id), + playback_commit_policy=DuplexPlaybackCommitPolicy.ACK_ONLY.value, + ) + if old_response_id is not None: + item_id = f"item_{old_response_id}" + if committed_message is not None: + session.register_history_item(item_id, committed_message) + elif committed_ms > 0 and not session.playback_ack_is_too_late(old_response_id, item_id): + session.truncate_history_item(item_id, audio_end_ms=committed_ms) + # The epoch bump is the atomic part: from here on every model output + # and append of the old epoch is dropped by the stale-epoch filter in + # ``emit`` / the append tail, whatever the awaits below interleave with. + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + if old_request_id is not None: + # Release projector/parser cursors so cancelled epochs do not + # accumulate until the whole session closes. + self.plugin.data_plane.close_stream(old_request_id) + await self._abort_request_background(old_request_id, notify=notify) + if has_running_task and active_task is not None: + active_task.cancel() + try: + await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25) + except TimeoutError: + pass + if notify: + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + return True + + async def _abort_request_background(self, request_id: str, *, notify: bool) -> None: + try: + await self.stage_port.abort_requests([request_id]) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to abort duplex request %s: %s", request_id, exc) + if notify and self.session.state != DuplexSessionState.CLOSED: + self.model.send_runtime_error("runtime_abort_failed", exc) + + def _cancel_pending_input(self, *, reason: str) -> None: + session = self.session + cancelled = session.cancel_pending_input() + helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "input.cancelled", + "session_id": session.session_id, + "reason": reason, + "epoch": session.epoch, + "cancelled": cancelled, + } + ) + + # ------------------------------------------------------------------ # + # Commit / response.create # + # ------------------------------------------------------------------ # + + def _apply_response_create_options(self, response_payload: dict[str, object]) -> str | None: + session = self.session + try: + options = ResponseCreateOptions.from_realtime( + response_payload, + private_runtime_config_keys=self.plugin.private_runtime_config_keys, + ) + except DuplexConfigError as exc: + return exc.code + try: + session.reserve_response_options(options) + except Exception: + return "response_already_active" + return None + + async def _flush_and_submit_committed_turn( + self, + event: dict[str, object], + *, + event_type: str, + realtime_item_id: object, + should_create_response: bool, + ) -> bool: + """Flush the buffered turn and submit it, with no response yet running. + + Returns whether the turn was flushed. An empty buffer falls through to + the bare acknowledgement at the end of ``_on_commit``. + """ + session = self.session + model_state = self.model_state + commit_reservation = ( + model_state.audio_buffer.prepare_commit( + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + ) + if event_type in {"input.commit", "input_audio_buffer.commit"} + else None + ) + flushed_buffer_reserved_bytes = model_state.audio_buffer.pending_byte_count if commit_reservation is None else 0 + flushed = ( + commit_reservation.payload + if commit_reservation is not None + else model_state.audio_buffer.flush(chunk_period_ms=session.capabilities.chunk_period_ms or 1000) + ) + if model_state.committed_audio_payload is not None: + if flushed is not None: + flushed = overlap_policy.merge_audio_payloads(model_state.committed_audio_payload, flushed) + else: + flushed = model_state.committed_audio_payload + if commit_reservation is not None: + commit_reservation.commit() + if flushed is None: + return False + if overlap_policy.should_force_listen_for_short_commit(self.session, event, flushed): + flushed = dict(flushed) + flushed["force_listen"] = True + model_state.input_since_commit = False + committed = helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + self.emit( + helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + ) + operation_id = commit_reservation.operation_id if commit_reservation is not None else uuid.uuid4().hex + reserved_bytes = ( + commit_reservation.byte_count if commit_reservation is not None else flushed_buffer_reserved_bytes + ) + model_state.retain_committed_audio(flushed, operation_id=operation_id, reserved_bytes=reserved_bytes) + if should_create_response: + await self._start_append( + flushed, + final=True, + precreate_response=True, + operation_id=model_state.committed_audio_operation_id, + retained_committed_payload=flushed, + ) + else: + model_state.deferred_response_create = False + return True + + async def _start_response_from_committed_audio(self) -> None: + """Answer a client ``response.create``. + + A turn starts from committed input. Audio is the usual kind, but the + Realtime protocol also lets a client build the turn out of + ``conversation.item.create`` and ask for a response with no audio at + all; those items are input too. A model-native session still generates + per audio unit, so an item-only turn is primed with the model's own + silence unit -- the one continuations already use -- rather than with + anything this layer invents. + + Refused only when nothing is waiting, or a response is already running, + rather than silently producing an empty turn. + """ + session = self.session + model_state = self.model_state + if ( + helpers.response_in_progress(self.session, self.tasks) + or self.tasks.append_tasks + or self.run.stream_request_id is not None + ): + if session.active_response_id is None and ( + session.active_request_id is not None + or self.tasks.append_tasks + or self.run.stream_request_id is not None + ): + return + self._emit_error( + "response_already_active", "response.create cannot start while another response is active." + ) + session.discard_response_options() + return + if model_state.committed_audio_payload is not None: + committed_payload = model_state.committed_audio_payload + operation_id = model_state.committed_audio_operation_id + if operation_id is None: + operation_id = uuid.uuid4().hex + model_state.committed_audio_operation_id = operation_id + session.reset_unanswered_user_items() + await self._start_append( + committed_payload, + final=True, + precreate_response=True, + operation_id=operation_id, + retained_committed_payload=committed_payload, + ) + return + if session.unanswered_user_items(): + # Conversation items are context, not a turn. A model-native model + # decides to speak from the audio it hears, and there is no audio + # here: opening a response anyway produces one the model never + # fills, which the caller only discovers when the session idles + # out. Text reaches such a model through the seeded opening turn + # (``DuplexSessionConfig.initial_user_text``) instead. + self._emit_error( + "text_only_turn_unsupported", + "This duplex model answers speech input only. Seed the session with " + "initial_user_text to ask it in text." + if session.capabilities.supports_chat_completions + else "This duplex model answers speech input only.", + ) + session.reset_unanswered_user_items() + session.discard_response_options() + return + self._emit_error( + "response_create_without_input", + "Duplex response.create requires committed audio to answer.", + ) + session.discard_response_options() + + def _discard_short_overlap_ack(self) -> None: + """Drop a sub-threshold interjection made while the model is speaking. + + Too short to be a turn, so it neither interrupts the response nor + becomes one: the buffered audio is discarded and the turn stays with + the model. + """ + session = self.session + model_state = self.model_state + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.clear_committed_audio() + self._emit_events(discard_pending_input_audio(self._require_projector(), session.overlap_speech_ms)) + self.emit( + { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": session.turn_id, + "epoch": session.epoch, + "empty": True, + "is_speech": False, + "overlap_ack": True, + "no_response": True, + } + ) + session.reset_overlap_speech() + session.discard_response_options() + + def _defer_commit_behind_active_response( + self, + event: dict[str, object], + *, + realtime_item_id: object, + should_create_response: bool, + precreate_response_requested: bool, + ) -> bool: + """Retain a commit that arrived while a response is still running. + + The audio is kept so the turn is not lost, and the response it should + start is remembered; ``_maybe_promote_deferred_overlap`` replays it once + the active response ends. Returns whether the commit was deferred -- + an empty buffer has nothing to retain and falls through. + """ + session = self.session + model_state = self.model_state + commit_reservation = model_state.audio_buffer.prepare_commit( + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + ) + deferred_payload = commit_reservation.payload + if deferred_payload is None: + commit_reservation.commit() + return False + if model_state.committed_audio_payload is not None: + deferred_payload = overlap_policy.merge_audio_payloads( + model_state.committed_audio_payload, + deferred_payload, + ) + model_state.retain_committed_audio( + deferred_payload, + operation_id=commit_reservation.operation_id, + reserved_bytes=commit_reservation.byte_count, + ) + commit_reservation.commit() + model_state.deferred_response_create = should_create_response + model_state.deferred_precreate_response = precreate_response_requested + model_state.input_since_commit = False + model_state.speech_since_commit = False + committed = helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + committed_payload = helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + committed_payload["overlap_deferred"] = True + committed_payload["response_create_deferred"] = should_create_response + self.emit(committed_payload) + return True + + async def _commit_and_start_auto_response(self, event: dict[str, object], *, realtime_item_id: object) -> None: + """Commit the buffered turn and submit it, letting the model answer it.""" + session = self.session + model_state = self.model_state + commit_reservation = model_state.audio_buffer.prepare_commit( + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + ) + final_payload = commit_reservation.payload + if model_state.committed_audio_payload is not None: + if final_payload is not None: + final_payload = overlap_policy.merge_audio_payloads(model_state.committed_audio_payload, final_payload) + else: + final_payload = model_state.committed_audio_payload + commit_reservation.commit() + if final_payload is not None: + model_state.retain_committed_audio( + final_payload, + operation_id=commit_reservation.operation_id, + reserved_bytes=commit_reservation.byte_count, + ) + model_state.deferred_response_create = False + model_state.input_since_commit = False + model_state.speech_since_commit = False + data_plane_turn_id = session.turn_id + committed = helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + turn_id=data_plane_turn_id, + ) + self.emit( + helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + ) + if final_payload is not None: + await self._start_append( + {**final_payload, "duplex_turn_id": data_plane_turn_id}, + final=True, + precreate_response=False, + operation_id=commit_reservation.operation_id, + retained_committed_payload=final_payload, + ) + + def _commit_silent_input(self) -> None: + """Commit a turn the client itself marked as silence: drop it and keep listening.""" + session = self.session + model_state = self.model_state + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.clear_committed_audio() + self.emit( + { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": session.turn_id, + "epoch": session.epoch, + "empty": True, + "is_speech": False, + "no_response": True, + } + ) + self.emit( + { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": "silence_or_noise", + } + ) + + async def _on_commit(self, event: dict[str, object]) -> None: + session = self.session + model_state = self.model_state + event_type = str(event.get("type")) + realtime_item_id = event.get("realtime_item_id") + realtime_validated_audio_commit = ( + event_type == "input_audio_buffer.commit" and isinstance(realtime_item_id, str) and bool(realtime_item_id) + ) + if event_type in {"input.commit", "input_audio_buffer.commit"} and not await self._wait_for_append_tail(): + # resolve_commit has already applied this commit's side effects + # (speech_stopped emitted, input flags cleared), so returning + # silently leaves the client with those effects, no + # input_audio_buffer.committed and no error. The session.update + # path says so explicitly; this one now does too. + self._emit_error( + "commit_aborted", + "the commit was not applied because the preceding append failed", + event_id=event.get("realtime_event_id"), + ) + return + if event_type == "input_audio_buffer.commit" and event.get("is_speech") is False: + self._commit_silent_input() + return + should_create_response = ( + event_type == "response.create" + or bool(event.get("response_create", event_type == "input.commit")) + or (event_type == "input_audio_buffer.commit" and self._session_auto_responds()) + ) + precreate_response_requested = event_type == "response.create" or bool( + event.get("response_create", event_type == "input.commit") + ) + if event_type == "response.create": + response_payload = event.get("response") + if isinstance(response_payload, dict): + response_options_error = self._apply_response_create_options(response_payload) + if response_options_error is not None: + error_message = ( + "The selected native duplex runtime does not support generation " + "overrides for instructions, voice, temperature, max tokens, tools, or " + "tool_choice." + if response_options_error == "unsupported_native_response_options" + else "response.create cannot reserve options while another response is active." + ) + self._emit_error(response_options_error, error_message, event_id=event.get("realtime_event_id")) + return + if event_type == "input_audio_buffer.commit": + has_pending_audio = ( + model_state.input_since_commit + or model_state.audio_buffer.has_pending() + or model_state.committed_audio_payload is not None + or realtime_validated_audio_commit + ) + if not has_pending_audio and not self.tasks.append_tasks and self.run.stream_request_id is None: + self._emit_error( + "input_audio_buffer_empty", "input_audio_buffer.commit requires a non-empty input audio buffer." + ) + return + commit_action = decide_commit_action( + CommitSnapshot( + auto_responds=self._session_auto_responds(), + speech_since_commit=model_state.speech_since_commit, + active_response_id=session.active_response_id, + overlap_speech_ms=session.overlap_speech_ms, + response_in_progress=helpers.response_in_progress(self.session, self.tasks), + playback_active=helpers.assistant_playback_active(self.session), + ) + ) + if commit_action is CommitAction.DEFER_ACTIVE_RESPONSE: + if session.overlap_speech_ms <= session.config.overlap_short_ack_ms: + self._discard_short_overlap_ack() + return + + if self._defer_commit_behind_active_response( + event, + realtime_item_id=realtime_item_id, + should_create_response=should_create_response, + precreate_response_requested=precreate_response_requested, + ): + return + if commit_action is CommitAction.START_AUTO_RESPONSE: + await self._commit_and_start_auto_response(event, realtime_item_id=realtime_item_id) + return + if event_type == "response.create": + await self._start_response_from_committed_audio() + return + if not helpers.response_in_progress(self.session, self.tasks) and await self._flush_and_submit_committed_turn( + event, + event_type=event_type, + realtime_item_id=realtime_item_id, + should_create_response=should_create_response, + ): + return + # Nothing flushed (or a response is still in progress): acknowledge the + # commit without starting a new response. + had_uncommitted_audio = ( + model_state.input_since_commit + or model_state.audio_buffer.has_pending() + or model_state.committed_audio_payload is not None + or realtime_validated_audio_commit + ) + committed = None + if event_type in {"input_audio_buffer.commit", "input.commit"}: + model_state.input_since_commit = False + model_state.speech_since_commit = False + if event_type != "response.create": + committed = ( + helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + if had_uncommitted_audio + else None + ) + self.emit( + helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + ) + return + if committed is not None: + if isinstance(realtime_item_id, str): + session.register_history_item(realtime_item_id, committed.message) + self.emit(helpers.input_committed_payload(session, committed, realtime_item_id=realtime_item_id)) + # NOTE(refactor): the old generic chat-completion fallback response + # (`_run_response`) is gone; a native session with a response already in + # progress only acknowledges the commit here. + + # ------------------------------------------------------------------ # + # playback.ack (moved from OmniDuplexSessionHandler) # + # ------------------------------------------------------------------ # + + +__all__ = ["DuplexSessionRunner"] diff --git a/vllm_omni/engine/duplex/turn_detection.py b/vllm_omni/engine/duplex/turn_detection.py new file mode 100644 index 000000000..6c5c84f8a --- /dev/null +++ b/vllm_omni/engine/duplex/turn_detection.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Server-side turn detection (``turn_detection: server_vad``) for the session runner. + +Moved from the Realtime input translator: the runner owns one +``ServerTurnDetector`` per session, feeds each appended chunk to it and emits +``input_audio_buffer.speech_started`` / ``speech_stopped`` from the result. +The two-phase ``session.update`` commit (prepare / commit / reject) that used +to live on the translator is expressed by :class:`PendingTurnDetectionUpdate`. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import numpy as np + +from vllm_omni.engine.duplex.vad import ( + SILERO_VAD_MIN_THRESHOLD, + ServerVADUnavailableError, + SileroStreamingVAD, + SileroVADBackendProvider, + SileroVADConfig, + StreamingVADResult, +) + +__all__ = [ + "PendingTurnDetectionUpdate", + "SileroVADBackendProvider", + "ServerTurnDetector", + "ServerVADUnavailableError", + "TurnDetectionConfig", + "TurnDetectionResult", + "apply_turn_detection_result", + "configured_realtime_turn_detection", + "normalize_turn_detection_session_payload", + "validate_realtime_turn_detection", +] + + +def configured_realtime_turn_detection(session_payload: Mapping[str, object]) -> tuple[str | None, object]: + """Return ``(field_path, value)`` for the turn detection setting present in the payload.""" + field = "turn_detection" if "turn_detection" in session_payload else None + selected = session_payload.get("turn_detection") + audio_config = session_payload.get("audio") + if isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict) and "turn_detection" in audio_input: + field, selected = "audio.input.turn_detection", audio_input["turn_detection"] + return field, selected + + +def _valid_vad_number(value: object, *, minimum: float, maximum: float | None = None, strict: bool = False) -> bool: + return ( + isinstance(value, int | float) + and not isinstance(value, bool) + and np.isfinite(float(value)) + and (float(value) > minimum if strict else float(value) >= minimum) + and (maximum is None or float(value) <= maximum) + ) + + +#: Every field a ``server_vad`` object may carry. Unknown keys are refused +#: rather than ignored: a misspelled tuning knob that silently does nothing is +#: worse than a rejected session, because the endpointing still *looks* applied. +_SERVER_VAD_FIELDS = frozenset( + { + "type", + "threshold", + "prefix_padding_ms", + "silence_duration_ms", + "create_response", + "interrupt_response", + "min_speech_duration_ms", + } +) + + +def validate_realtime_turn_detection(session_payload: Mapping[str, object]) -> str | None: + """Validate the ``turn_detection`` object of a Realtime session payload (None when valid).""" + field, turn_detection = configured_realtime_turn_detection(session_payload) + if field is None: + if session_payload.get("overlap_policy") == "barge_in_on_speech": + return ( + "overlap_policy='barge_in_on_speech' requires turn_detection.type='server_vad' on the Realtime endpoint" + ) + return None + if turn_detection is not None and not isinstance(turn_detection, dict): + return f"{field} must be null or an object with type='server_vad'" + if isinstance(turn_detection, dict): + if turn_detection.get("type") != "server_vad": + return f"{field}.type must be 'server_vad'" + unknown = sorted(set(turn_detection) - _SERVER_VAD_FIELDS) + if unknown: + return f"Unknown {field} field(s): {', '.join(unknown)}" + threshold = turn_detection.get("threshold", 0.5) + if not _valid_vad_number(threshold, minimum=SILERO_VAD_MIN_THRESHOLD, strict=True, maximum=1): + return f"{field}.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD} and at most 1" + for name in ("prefix_padding_ms", "silence_duration_ms", "min_speech_duration_ms"): + value = turn_detection.get(name) + if value is not None and not _valid_vad_number(value, minimum=0): + return f"{field}.{name} must be a non-negative number" + if turn_detection.get("interrupt_response", True) is not True: + return ( + f"{field}.interrupt_response=false is unsupported; use turn_detection=null for model-owned listen/speak" + ) + desired_policy = "barge_in_on_speech" if turn_detection is not None else "listen_only" + overlap_policy = session_payload.get("overlap_policy") + if isinstance(overlap_policy, str) and overlap_policy != desired_policy: + return f"overlap_policy={overlap_policy!r} conflicts with turn_detection; expected {desired_policy!r}" + return None + + +@dataclass(frozen=True, slots=True) +class TurnDetectionConfig: + """Normalized ``server_vad`` configuration (all fields filled with defaults).""" + + threshold: float = 0.5 + prefix_padding_ms: int = 300 + silence_duration_ms: int = 500 + min_speech_duration_ms: int = 96 + create_response: bool = True + interrupt_response: bool = True + + @classmethod + def from_realtime(cls, turn_detection: Mapping[str, object]) -> TurnDetectionConfig: + return cls( + threshold=float(turn_detection.get("threshold", 0.5)), + prefix_padding_ms=int(turn_detection.get("prefix_padding_ms", 300)), + silence_duration_ms=int(turn_detection.get("silence_duration_ms", 500)), + min_speech_duration_ms=int(turn_detection.get("min_speech_duration_ms", 96)), + create_response=bool(turn_detection.get("create_response", True)), + interrupt_response=bool(turn_detection.get("interrupt_response", True)), + ) + + @property + def overlap_policy(self) -> str: + return "barge_in_on_speech" + + def build_detector(self, backend_provider: SileroVADBackendProvider | None = None) -> ServerTurnDetector: + return ServerTurnDetector(self, backend_provider=backend_provider) + + +def normalize_turn_detection_session_payload( + session_payload: dict[str, object], +) -> tuple[bool, TurnDetectionConfig | None]: + """Fill ``turn_detection`` defaults and derive ``overlap_policy`` in place. + + Returns ``(configured, config)`` where ``configured`` is False when the + payload does not mention turn detection at all (nothing changes), and + ``config`` is ``None`` for ``turn_detection: null`` (model-owned turns, + ``overlap_policy`` forced to ``listen_only``). + """ + field, configured = configured_realtime_turn_detection(session_payload) + if field is None: + return False, None + if configured is None: + session_payload["overlap_policy"] = "listen_only" + return True, None + assert isinstance(configured, dict) + merged = { + "type": "server_vad", + "interrupt_response": True, + "threshold": 0.5, + "prefix_padding_ms": 300, + "silence_duration_ms": 500, + "min_speech_duration_ms": 96, + **configured, + } + config = TurnDetectionConfig.from_realtime(merged) + session_payload["turn_detection"] = dict(merged) + session_payload["overlap_policy"] = config.overlap_policy + return True, config + + +@dataclass(frozen=True, slots=True) +class TurnDetectionResult: + is_speech: bool + speech_active: bool + speech_started: bool = False + speech_stopped: bool = False + speech_probability: float = 0.0 + audio_start_ms: int | None = None + audio_end_ms: int | None = None + #: A ``speech_stopped`` under ``server_vad`` ends the user turn. + should_commit: bool = False + #: Whether the auto-commit should also request a response (``create_response``). + create_response: bool = True + + def as_payload(self) -> dict[str, object]: + return { + "backend": "silero", + "is_speech": self.is_speech, + "speech_active": self.speech_active, + "speech_started": self.speech_started, + "speech_stopped": self.speech_stopped, + "speech_probability": self.speech_probability, + } + + +class ServerTurnDetector: + """One Silero-backed streaming detector per session.""" + + def __init__( + self, + config: TurnDetectionConfig, + *, + backend_provider: SileroVADBackendProvider | None = None, + ) -> None: + self.config = config + self._vad = SileroStreamingVAD( + SileroVADConfig( + threshold=float(config.threshold), + prefix_padding_ms=int(config.prefix_padding_ms), + silence_duration_ms=int(config.silence_duration_ms), + min_speech_duration_ms=max(32, int(config.min_speech_duration_ms)), + ), + backend_provider=backend_provider, + ) + + @property + def vad(self) -> SileroStreamingVAD: + return self._vad + + def reset(self) -> None: + self._vad.reset() + + def process( + self, + base64_audio: str, + *, + fmt: str, + sample_rate_hz: int | None, + audio_end_ms: int | None = None, + ) -> TurnDetectionResult: + """Score one appended chunk. Raises ``ServerVADUnavailableError`` / ``ValueError`` like the VAD.""" + result: StreamingVADResult = self._vad.process_base64(base64_audio, fmt=fmt, sample_rate_hz=sample_rate_hz) + stop_ms = result.speech_end_ms if result.speech_end_ms is not None else audio_end_ms + return TurnDetectionResult( + is_speech=result.is_speech, + speech_active=result.speech_active, + speech_started=result.speech_started, + speech_stopped=result.speech_stopped, + speech_probability=result.speech_probability, + audio_start_ms=result.speech_start_ms, + audio_end_ms=stop_ms, + should_commit=bool(result.speech_stopped and not result.speech_active), + create_response=self.config.create_response, + ) + + +def apply_turn_detection_result(payload: dict[str, object], result: TurnDetectionResult) -> None: + """Merge a detector result onto an internal ``input_audio_buffer.append`` payload. + + Mirrors what the old translator attached: ``is_speech`` from the detector, + the ``vad`` hint block, and ``force_listen`` while speech is active. + """ + payload["is_speech"] = result.is_speech + payload["vad"] = result.as_payload() + if result.speech_active: + payload["force_listen"] = True + if result.audio_start_ms is not None and result.speech_started: + payload["audio_start_ms"] = result.audio_start_ms + if result.audio_end_ms is not None and result.speech_stopped: + payload["audio_end_ms"] = result.audio_end_ms + + +@dataclass +class PendingTurnDetectionUpdate: + """A ``session.update`` turn-detection change awaiting the engine ACK (two-phase).""" + + config: TurnDetectionConfig | None + detector: ServerTurnDetector | None + + @classmethod + def prepare( + cls, + session_payload: dict[str, object], + *, + backend_provider: SileroVADBackendProvider | None = None, + ) -> PendingTurnDetectionUpdate | None: + """Normalize the payload in place and stage the new detector; None when not configured.""" + configured, config = normalize_turn_detection_session_payload(session_payload) + if not configured: + return None + detector = config.build_detector(backend_provider) if config is not None else None + return cls(config=config, detector=detector) + + def commit( + self, current: ServerTurnDetector | None + ) -> tuple[TurnDetectionConfig | None, ServerTurnDetector | None]: + """Adopt the staged detector; the replaced one is reset. Returns the new (config, detector).""" + if current is not None: + current.reset() + return self.config, self.detector + + def reject(self) -> None: + if self.detector is not None: + self.detector.reset() diff --git a/vllm_omni/engine/duplex/vad.py b/vllm_omni/engine/duplex/vad.py new file mode 100644 index 000000000..ad3c26820 --- /dev/null +++ b/vllm_omni/engine/duplex/vad.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Silero server VAD for the engine-resident duplex session. + +The detector keeps the split upstream's serving-side ``server_vad`` module +used, because the split is what makes one model serve many sessions: a +**backend** scores one 512-sample frame and is otherwise stateless, and the +per-stream state -- the partial frame, the model state, the endpoint counters -- +belongs to :class:`SileroStreamingVAD`, one per session. + +Two backends. :class:`SileroVADBackend` runs the pinned Silero v6.2 ONNX graph on +CPU and is shared process-wide, which is why its state is passed in and out +rather than held on the session object. :class:`TorchSileroBackend` loads the +same model through the ``silero-vad`` package for environments that have no +local ONNX artifact; torch keeps its state inside the module, so that one is per +session. + +Endpointing follows Silero v6.2's streaming hysteresis: activation uses the +configured threshold, while a turn can only *end* on frames below +``max(threshold - 0.15, 0.01)``. Once a silence candidate exists, louder frames +keep the elapsed-silence clock running but cannot themselves close the turn. +""" + +from __future__ import annotations + +import binascii +import hashlib +import threading +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +import numpy as np +import pybase64 as base64 +from vllm.logger import init_logger + +logger = init_logger(__name__) + +SILERO_VAD_REPO_ID = "istupakov/silero-vad-onnx" +SILERO_VAD_REVISION = "8b14476858ef240c50b3884bb38cc67290c1cc70" +SILERO_VAD_FILENAME = "silero_vad.onnx" +SILERO_VAD_SHA256 = "1a153a22f4509e292a94e67d6f9b85e8deb25b4988682b7e174c65279d8788e3" +SILERO_VAD_MIN_THRESHOLD = 0.15 + +_SAMPLE_RATE_HZ = 16_000 +_FRAME_SAMPLES = 512 + + +class ServerVADUnavailableError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class SileroVADConfig: + threshold: float = 0.5 + prefix_padding_ms: int = 300 + silence_duration_ms: int = 500 + min_speech_duration_ms: int = 96 + + +@dataclass(frozen=True, slots=True) +class StreamingVADResult: + is_speech: bool + speech_active: bool + speech_started: bool = False + speech_stopped: bool = False + speech_probability: float = 0.0 + speech_start_ms: int | None = None + speech_end_ms: int | None = None + + +# --------------------------------------------------------------------------- # +# Backends # +# --------------------------------------------------------------------------- # + + +class SpeechDetectorBackend(Protocol): + """Scores one frame. Per-stream state is created here but owned by the caller.""" + + frame_samples: int + + def new_state(self) -> object: ... + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: ... + + +@dataclass(frozen=True, slots=True) +class _SileroVADState: + model_state: np.ndarray + context: np.ndarray + + +class SileroVADBackend: + """Shared ONNX Runtime Silero v6.2 detector running on CPU.""" + + sample_rate_hz = _SAMPLE_RATE_HZ + frame_samples = _FRAME_SAMPLES + context_samples = 64 + model_state_shape = (2, 1, 128) + + def __init__(self, model_path: str | Path) -> None: + try: + import onnxruntime as ort + except ImportError as exc: # pragma: no cover - platform packaging supplies ORT. + raise ServerVADUnavailableError("server_vad requires ONNX Runtime") from exc + + session_options = ort.SessionOptions() + session_options.inter_op_num_threads = 1 + session_options.intra_op_num_threads = 1 + + self.model_path = Path(model_path) + self._session = ort.InferenceSession( + str(self.model_path), + providers=["CPUExecutionProvider"], + sess_options=session_options, + ) + input_names = {item.name for item in self._session.get_inputs()} + if not {"input", "sr", "state"} <= input_names: + raise ServerVADUnavailableError(f"Unsupported Silero ONNX input contract: {sorted(input_names)}") + self._warm_up() + + def _warm_up(self) -> None: + self.infer(np.zeros(self.frame_samples, dtype=np.float32), self.new_state()) + + def new_state(self) -> _SileroVADState: + return _SileroVADState( + model_state=np.zeros(self.model_state_shape, dtype=np.float32), + context=np.zeros((1, self.context_samples), dtype=np.float32), + ) + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + if not isinstance(state, _SileroVADState): + raise TypeError("Silero detector state must be created by SileroVADBackend.new_state()") + model_state = np.ascontiguousarray(state.model_state, dtype=np.float32) + context = np.ascontiguousarray(state.context, dtype=np.float32) + if model_state.shape != self.model_state_shape: + raise ValueError(f"Silero model state must have shape {self.model_state_shape}, got {model_state.shape}") + expected_context_shape = (1, self.context_samples) + if context.shape != expected_context_shape: + raise ValueError(f"Silero context must have shape {expected_context_shape}, got {context.shape}") + + audio = np.ascontiguousarray(frame, dtype=np.float32).reshape(1, -1) + if audio.shape[1] != self.frame_samples: + raise ValueError( + f"Silero detector frame must contain exactly {self.frame_samples} samples, got {audio.shape[1]}" + ) + model_input = np.concatenate((context, audio), axis=1) + # ONNX Runtime permits concurrent Run calls on one CPU session. Stream + # state is explicit, so independent sessions can safely share this one. + output = self._session.run( + None, + { + "input": model_input, + "state": model_state, + "sr": np.asarray(self.sample_rate_hz, dtype=np.int64), + }, + ) + if len(output) < 2: + raise ServerVADUnavailableError("Silero ONNX model did not return probability and model state") + probability = float(np.asarray(output[0]).reshape(-1)[0]) + next_state = _SileroVADState( + model_state=np.ascontiguousarray(output[1], dtype=np.float32), + context=np.ascontiguousarray(model_input[:, -self.context_samples :], dtype=np.float32), + ) + return probability, next_state + + +class TorchSileroBackend: + """The same model through the ``silero-vad`` package, for hosts with no ONNX artifact. + + Torch keeps the recurrent state inside the module, so unlike the ONNX + backend this one cannot be shared: each session gets its own instance and + ``new_state`` only resets it. + """ + + sample_rate_hz = _SAMPLE_RATE_HZ + frame_samples = _FRAME_SAMPLES + + def __init__(self) -> None: + try: + import torch + from silero_vad import load_silero_vad + except ImportError as exc: + raise ServerVADUnavailableError( + "server_vad needs either a local Silero ONNX artifact (see " + "duplex_session.server_vad_model_path) or the optional 'silero-vad' " + "package; install vllm-omni[server-vad]" + ) from exc + self._torch = torch + self._model = load_silero_vad(onnx=False) + + def new_state(self) -> object: + reset_states = getattr(self._model, "reset_states", None) + if callable(reset_states): + reset_states() + return None + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + tensor = self._torch.from_numpy(np.ascontiguousarray(frame, dtype=np.float32)) + return float(self._model(tensor, self.sample_rate_hz).item()), state + + +class SileroVADBackendProvider: + """Resolve, verify and load one detector backend per engine process. + + ONNX is preferred and shared; the torch package is the fallback and is built + per call because its state is internal. + """ + + def __init__(self, *, model_path: str | None = None) -> None: + self.model_path = model_path + self._backend: SileroVADBackend | None = None + self._lock = threading.Lock() + + def get(self) -> SpeechDetectorBackend: + """The shared ONNX backend, or a per-session torch one when it cannot be built. + + An explicitly configured ``server_vad_model_path`` never falls back: the + operator named that artifact, so a missing file or a missing ONNX Runtime + is their error, not something to paper over with a different model. + """ + if self._backend is not None: + return self._backend + with self._lock: + if self._backend is not None: + return self._backend + path = self._resolve_local_artifact() + if path is None: + # No pinned artifact: the torch package keeps state internally, + # so this instance belongs to one session only. + return TorchSileroBackend() + self._verify_checksum(path) + try: + self._backend = SileroVADBackend(path) + except ServerVADUnavailableError: + if self.model_path: + raise + logger.warning( + "Silero ONNX artifact %s found but ONNX Runtime is unavailable; " + "falling back to the torch 'silero-vad' package (one model per session)", + path, + ) + return TorchSileroBackend() + return self._backend + + def _resolve_local_artifact(self) -> Path | None: + if self.model_path: + path = Path(self.model_path).expanduser() + if path.is_file(): + return path + # Explicitly configured: a missing file is an error, not a fallback. + raise ServerVADUnavailableError(f"Configured Silero VAD model does not exist: {path}") + try: + from vllm.transformers_utils.repo_utils import try_get_local_file + + cached = try_get_local_file( + model=SILERO_VAD_REPO_ID, + file_name=SILERO_VAD_FILENAME, + revision=SILERO_VAD_REVISION, + ) + except Exception: + return None + return cached if isinstance(cached, Path) and cached.is_file() else None + + def _verify_checksum(self, path: Path) -> None: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != SILERO_VAD_SHA256: + raise ServerVADUnavailableError( + f"Silero VAD checksum mismatch for {path}: expected {SILERO_VAD_SHA256}, got {digest}" + ) + + +# --------------------------------------------------------------------------- # +# Per-session streaming detector # +# --------------------------------------------------------------------------- # + + +class SileroStreamingVAD: + _SAMPLE_RATE_HZ = _SAMPLE_RATE_HZ + _WINDOW_SAMPLES = _FRAME_SAMPLES + + def __init__( + self, + config: SileroVADConfig, + *, + backend: SpeechDetectorBackend | None = None, + backend_provider: SileroVADBackendProvider | None = None, + frame_scorer: Callable[[np.ndarray], float] | None = None, + ) -> None: + self.config = config + self._explicit_backend = backend + self._provider = backend_provider + self._frame_scorer = frame_scorer + self._backend: SpeechDetectorBackend | None = backend + self._detector_state: object = None + self._clear_stream_state() + + def _clear_stream_state(self) -> None: + self._pending = np.empty(0, dtype=np.float32) + self._speech_active = False + self._candidate_samples = 0 + self._candidate_start_sample: int | None = None + self._silence_samples = 0 + self._processed_samples = 0 + self._stream_start_ms = 0 + self._resample_rate_hz: int | None = None + self._resample_remainder = 0 + + @property + def scratch_bytes(self) -> int: + """Bytes this detector holds between chunks (the partial frame only).""" + return int(self._pending.nbytes) + + @property + def speech_active(self) -> bool: + return self._speech_active + + def reset(self) -> None: + """Drop stream state while preserving the session-relative audio clock. + + The clock survives so ``speech_start_ms`` after a barge-in still refers + to the session's timeline; ``_stream_start_ms`` then stops the prefix + padding reaching back into audio that was discarded. + """ + stream_start_ms = round(self._processed_samples * 1000 / self._SAMPLE_RATE_HZ) + processed = self._processed_samples + self._clear_stream_state() + self._processed_samples = processed + self._stream_start_ms = stream_start_ms + if self._backend is not None: + self._detector_state = self._backend.new_state() + + # ------------------------------------------------------------------ # + # Input # + # ------------------------------------------------------------------ # + + def process_base64( + self, + audio: object, + *, + fmt: object, + sample_rate_hz: object, + ) -> StreamingVADResult: + if not isinstance(audio, str): + raise ValueError("Silero server VAD requires base64 audio") + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("Silero server VAD received invalid base64 audio") from exc + rate = int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else self._SAMPLE_RATE_HZ + if rate <= 0: + raise ValueError("Silero server VAD requires a positive sample rate") + + if fmt == "pcm_f32le": + if len(raw) % 4: + raise ValueError("Silero server VAD received an incomplete pcm_f32le frame") + samples = np.frombuffer(raw, dtype=" np.ndarray: + """Resample to 16 kHz, carrying the fractional sample-count remainder. + + A fallback path: the session runner already normalises input to 16 kHz + through ``convert_input_audio_with_rate`` before the detector sees it, + so this runs only for callers that feed the VAD directly. + + Carrying the remainder keeps the *number* of output samples exact, so + frame boundaries cannot drift over a long stream. The interpolation + itself is per chunk and has no filter state, so sample values near a + chunk edge differ slightly from resampling the same audio in one go -- + enough for the frame grid to stay put, not enough to call it a + streaming resampler. + """ + if rate == self._SAMPLE_RATE_HZ: + self._resample_rate_hz = None + self._resample_remainder = 0 + return samples + if self._resample_rate_hz is not None and rate != self._resample_rate_hz: + raise ValueError("server_vad input sample rate cannot change within a continuous audio stream") + self._resample_rate_hz = rate + numerator = samples.size * self._SAMPLE_RATE_HZ + self._resample_remainder + target_size, self._resample_remainder = divmod(numerator, rate) + if target_size == 0: + return np.empty(0, dtype=np.float32) + if samples.size == 1: + return np.full(target_size, samples[0], dtype=np.float32) + source_x = np.linspace(0.0, 1.0, num=samples.size, endpoint=True) + target_x = np.linspace(0.0, 1.0, num=target_size, endpoint=True) + return np.interp(target_x, source_x, samples).astype(np.float32) + + # ------------------------------------------------------------------ # + # Endpointing # + # ------------------------------------------------------------------ # + + def process(self, samples: np.ndarray) -> StreamingVADResult: + samples = np.asarray(samples, dtype=np.float32).reshape(-1) + if samples.size: + self._pending = np.concatenate((self._pending, samples)) + + started = False + stopped = False + start_ms: int | None = None + end_ms: int | None = None + max_probability = 0.0 + contained_speech = self._speech_active + min_speech_samples = max(1, round(self.config.min_speech_duration_ms * self._SAMPLE_RATE_HZ / 1000)) + min_silence_samples = max(1, round(self.config.silence_duration_ms * self._SAMPLE_RATE_HZ / 1000)) + negative_threshold = max(self.config.threshold - SILERO_VAD_MIN_THRESHOLD, 0.01) + + pending = self._pending + offset = 0 + while pending.size - offset >= self._WINDOW_SAMPLES: + frame = np.ascontiguousarray(pending[offset : offset + self._WINDOW_SAMPLES], dtype=np.float32) + offset += self._WINDOW_SAMPLES + frame_start = self._processed_samples + self._processed_samples += self._WINDOW_SAMPLES + probability = min(1.0, max(0.0, self._score_frame(frame))) + max_probability = max(max_probability, probability) + + if self._speech_active: + contained_speech = True + below_negative_threshold = probability < negative_threshold + if not below_negative_threshold and self._silence_samples == 0: + continue + # A silence candidate is running: louder frames keep its clock + # moving but cannot themselves close the turn. + self._silence_samples += self._WINDOW_SAMPLES + if not below_negative_threshold or self._silence_samples < min_silence_samples: + continue + self._speech_active = False + self._silence_samples = 0 + stopped = True + # OpenAI defines audio_end_ms as the end of the audio sent to + # the model, trailing silence included. + end_ms = max(0, round(self._processed_samples * 1000 / self._SAMPLE_RATE_HZ)) + continue + + if probability >= self.config.threshold: + if self._candidate_samples == 0: + self._candidate_start_sample = frame_start + self._candidate_samples += self._WINDOW_SAMPLES + if self._candidate_samples >= min_speech_samples: + candidate_start = self._candidate_start_sample or 0 + detected_start_ms = round(candidate_start * 1000 / self._SAMPLE_RATE_HZ) + self._speech_active = True + self._candidate_samples = self._silence_samples = 0 + self._candidate_start_sample = None + contained_speech = True + started = True + # Prefix audio from before the last reset is gone. + start_ms = max( + self._stream_start_ms, + max(0, detected_start_ms - self.config.prefix_padding_ms), + ) + else: + self._candidate_samples = 0 + self._candidate_start_sample = None + + # Copy the residual once, after the frames are scored, so a large input + # chunk is neither pinned by a view nor re-copied for every frame. + self._pending = pending[offset:].copy() + + return StreamingVADResult( + contained_speech, self._speech_active, started, stopped, max_probability, start_ms, end_ms + ) + + def _score_frame(self, frame: np.ndarray) -> float: + if self._frame_scorer is not None: + return float(self._frame_scorer(frame)) + backend = self._require_backend() + probability, self._detector_state = backend.infer(frame, self._detector_state) + return float(probability) + + def _require_backend(self) -> SpeechDetectorBackend: + if self._backend is None: + provider = self._provider or SileroVADBackendProvider() + self._backend = provider.get() + self._detector_state = self._backend.new_state() + return self._backend diff --git a/vllm_omni/engine/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py new file mode 100644 index 000000000..56ca3e352 --- /dev/null +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""DuplexOmniEngine: the duplex engine sibling of ``AsyncOmniEngine``. + +It is deliberately small. The engine base owns the stage processes, the +orchestrator thread and the correlated RPC transport; this class only knows +how duplex session messages enter the engine (open/close/resume/touch through +correlated RPC, session commands one-way) and which orchestrator to build +(``DuplexOrchestrator`` with the model plugin loaded). +""" + +from __future__ import annotations + +import asyncio +import queue +import uuid +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionError, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, EncodeAudio, load_duplex_plugin +from vllm_omni.engine.messages import EngineQueueMessage +from vllm_omni.engine.orchestrator import OrchestratorBase + +if TYPE_CHECKING: + from vllm_omni.engine.duplex.commands import DuplexCommand + +logger = init_logger(__name__) + +_DEFAULT_CONTROL_TIMEOUT_S = 10.0 +_COMMAND_PUT_TIMEOUT_S = 30.0 + + +class DuplexOmniEngine(AsyncOmniEngine): + """Engine for full-duplex models; sessions run inside ``DuplexOrchestrator``.""" + + plugin: DuplexModelPlugin | None = None + #: Set by ``_validate_deployment`` from the deploy config before any stage starts. + duplex_session_config: DuplexSessionRuntimeConfig + + # Any: **kwargs is forwarded verbatim to ``OmniEngineBase.__init__``. + def __init__(self, model: str, *, duplex_audio_encoder: EncodeAudio, **kwargs: Any) -> None: + # Injected by ``DuplexOmni`` (entrypoint layer): the engine package + # knows how to hand model audio to the plugin, not how to encode it. + self._audio_encoder = duplex_audio_encoder + super().__init__(model, **kwargs) + + # ---- orchestrator construction (orchestrator thread) ---- + + def _validate_deployment(self) -> None: + """Load the plugin and check the scheduler contract before any stage process starts.""" + pipeline_config = self.pipeline_config + plugin_path = getattr(pipeline_config, "duplex_plugin", None) if pipeline_config is not None else None + if pipeline_config is None: + raise ValueError( + f"{self.model!r}: the model pipeline could not be resolved (models with custom HF code " + "such as MiniCPM-o need trust_remote_code=True)" + ) + if not plugin_path: + raise ValueError(f"{self.model!r} is not a duplex model: the pipeline declares no duplex_plugin") + deploy_config = self.deploy_config + if deploy_config is None: + raise ValueError("DuplexOmniEngine requires a deploy config with session_mode: duplex (none resolved)") + if getattr(deploy_config, "session_mode", "turn") != "duplex": + raise ValueError( + "DuplexOmniEngine requires a deploy config with session_mode: duplex " + f"(got {getattr(deploy_config, 'session_mode', None)!r})" + ) + self.duplex_session_config = deploy_config.duplex_session + self.plugin = load_duplex_plugin(plugin_path, self._audio_encoder) + + # Any: signature of the ``OmniEngineBase._create_orchestrator`` seam it overrides. + def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: + from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator + + assert self.plugin is not None, "_validate_deployment() must run before the orchestrator is created" + stage0_vllm_config = self.stage_pools[0].stage_vllm_config if self.stage_pools else None + model_config = getattr(stage0_vllm_config, "model_config", None) + return DuplexOrchestrator( + plugin=self.plugin, + duplex_session_config=self.duplex_session_config, + model_config=model_config, + **orchestrator_kwargs, + ) + + # ---- deployment facts ---- + + @property + def duplex_capabilities(self) -> DuplexCapabilities: + if self.plugin is None: + raise RuntimeError("duplex plugin is not loaded") + return self.plugin.capabilities(max_sessions=self.duplex_session_config.max_sessions) + + # ---- session message surface ---- + + def _execute_control( + self, + message: EngineQueueMessage, + *, + control_id: str, + operation: str, + session_id: str, + timeout: float | None, + ) -> DuplexControlResultMessage: + try: + result = self.rpc_client.execute( + ("duplex", control_id), + message, + timeout=timeout, + timeout_message=f"duplex {operation} timed out for session {session_id}", + # Block rather than surface a raw queue.Full: the handlers below + # promise a typed DuplexSessionError, and a momentarily full + # request queue is backpressure, not a failed control op. + block_on_submit=True, + ) + except TimeoutError as exc: + raise DuplexSessionError(str(exc), code="timeout", retryable=True, session_id=session_id) from exc + except RuntimeError as exc: + # The RPC router is closed or the orchestrator reported a terminal error. + raise DuplexSessionError(str(exc), code="engine_error", session_id=session_id) from exc + if not isinstance(result, DuplexControlResultMessage): + raise DuplexSessionError( + f"unexpected duplex control result: {type(result).__name__}", + code="internal_error", + session_id=session_id, + ) + if not result.ok: + raise DuplexSessionError( + result.error_message or f"duplex {operation} failed", + code=result.error_code or "internal_error", + retryable=result.error_retryable, + session_id=session_id, + ) + return result + + def _open_session( + self, + session_id: str, + session_config: DuplexSessionConfig, + *, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + OpenDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + session_config=session_config, + ), + control_id=control_id, + operation="open", + session_id=session_id, + timeout=timeout, + ) + + async def open_session_async( + self, + session_id: str, + session_config: DuplexSessionConfig, + *, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, lambda: self._open_session(session_id, session_config, timeout=timeout)) + + def _close_session( + self, + session_id: str, + *, + reason: str = "client_close", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + CloseDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + reason=reason, + ), + control_id=control_id, + operation="close", + session_id=session_id, + timeout=timeout, + ) + + async def close_session_async( + self, + session_id: str, + *, + reason: str = "client_close", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self._close_session(session_id, reason=reason, timeout=timeout), + ) + + def _resume_session( + self, + session_id: str, + *, + expected_lease_generation: int, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + ResumeDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + expected_lease_generation=expected_lease_generation, + ), + control_id=control_id, + operation="resume", + session_id=session_id, + timeout=timeout, + ) + + async def resume_session_async( + self, + session_id: str, + *, + expected_lease_generation: int, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self._resume_session( + session_id, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ), + ) + + def _touch_session( + self, + session_id: str, + *, + activity: str, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + TouchDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + activity=activity, + ), + control_id=control_id, + operation="touch", + session_id=session_id, + timeout=timeout, + ) + + async def touch_session_async( + self, + session_id: str, + *, + activity: str, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self._touch_session(session_id, activity=activity, timeout=timeout), + ) + + def _submit_command(self, session_id: str, command: DuplexCommand) -> None: + """One-way: enqueue a session command in caller order (blocks only on queue backpressure).""" + if not self.is_alive(): + raise DuplexSessionError("engine is not alive", code="engine_dead", session_id=session_id) + message = DuplexSessionCommandMessage(session_id=session_id, command=command) + try: + self.request_queue.sync_q.put(message, timeout=_COMMAND_PUT_TIMEOUT_S) + except queue.Full as exc: + raise DuplexSessionError( + "engine request queue is full", code="engine_backpressure", retryable=True, session_id=session_id + ) from exc + + async def submit_command_async(self, session_id: str, command: DuplexCommand) -> None: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, lambda: self._submit_command(session_id, command)) + + +__all__ = ["DuplexOmniEngine"] diff --git a/vllm_omni/engine/duplex_orchestrator.py b/vllm_omni/engine/duplex_orchestrator.py new file mode 100644 index 000000000..09a767fb1 --- /dev/null +++ b/vllm_omni/engine/duplex_orchestrator.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""DuplexOrchestrator: the duplex sibling of ``Orchestrator``. + +It hosts the ``DuplexSessionManager`` (admission, leases, session runners) on +the orchestrator loop and implements the ``DuplexStagePort`` the runners use +to submit resumable Stage0 requests. Generic stage forwarding, prewarm and +cleanup stay in ``OrchestratorBase``; this class only fills the template seams +and applies the session-owned policy. +""" + +from __future__ import annotations + +import time as _time +from collections.abc import Coroutine +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.contracts import ( + DuplexFence, + DuplexOutputContext, + DuplexRequestIdentity, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, +) +from vllm_omni.engine.duplex.plugin import DuplexModelPlugin +from vllm_omni.engine.duplex.session.engine_session import DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.messages import EngineQueueMessage, OutputMessage +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + build_engine_core_request_from_tokens, +) + +if TYPE_CHECKING: + from vllm.config import ModelConfig + from vllm.outputs import RequestOutput + + from vllm_omni.metrics.stats import StageRequestStats + +logger = init_logger(__name__) + + +@dataclass +class DuplexOrchestratorRequestState(OrchestratorRequestState): + """Request bookkeeping for a session-owned (resumable duplex) stage request.""" + + session_id: str = "" + fence: DuplexFence | None = None + stage_fences: dict[int, DuplexFence] = field(default_factory=dict) + config_generation: int = -1 + + +class DuplexOrchestrator(Orchestrator, DuplexStagePort): + """Stage management for a duplex deployment; owns one ``DuplexSessionManager``. + + Extends the turn-based orchestrator rather than sitting beside it, so one + engine serves both a duplex session and an ordinary request. The + dependency runs duplex -> turn-based and never the reverse, which is what + keeps ``Orchestrator`` free of duplex code. + """ + + def __init__( + self, + # Any: ``*args`` / ``**kwargs`` are forwarded verbatim to ``OrchestratorBase.__init__``. + *args: Any, + plugin: DuplexModelPlugin, + duplex_session_config: DuplexSessionRuntimeConfig, + model_config: ModelConfig | None = None, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + self.plugin = plugin + self.duplex_session_config = duplex_session_config + self.session_manager = DuplexSessionManager( + plugin=plugin, + stage_port=self, + output_sink=self.output_async_queue, + result_sink=self.rpc_async_queue, + runtime_config=self.duplex_session_config, + model_config=model_config, + ) + + # ------------------------------------------------------------------ # + # Template seams # + # ------------------------------------------------------------------ # + + async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: + if self.session_manager.accepts(msg): + self.session_manager.dispatch(msg) + return True + # Not a session message: it is an ordinary turn-based request. + return await super()._dispatch_message(msg) + + # Any: ``Coroutine``'s send/yield parameters, as in the ``OrchestratorBase`` seam. + def _background_tasks(self) -> list[Coroutine[Any, Any, None]]: + return [self.session_manager.reaper_loop(self._shutdown_event)] + + async def _shutdown_extensions(self) -> None: + await self.session_manager.shutdown() + + def _on_stage_submitted( + self, + stage_id: int, + request_id: str, + replica_id: int, + req_state: OrchestratorRequestState, + ) -> None: + del replica_id + if not isinstance(req_state, DuplexOrchestratorRequestState) or req_state.fence is None: + return + runner = self.session_manager.runner_for_request_id(request_id) + if runner is None: + return + fence = req_state.fence + req_state.stage_fences[stage_id] = fence + try: + runner.session.bind_stage_request(stage_id, request_id, fence=fence) + except DuplexFenceMismatchError: + # The session already advanced past this request's epoch (cancel + # raced the submit); the cancel path aborts the stale request id. + logger.debug("[DuplexOrchestrator] stale stage binding ignored for %s stage-%s", request_id, stage_id) + return + req_state.stage_submit_ts[stage_id] = _time.time() + self._register_running_request(req_state) + self.session_manager.register_request(request_id, req_state.session_id) + + async def _intercept_stage_output( + self, + stage_id: int, + replica_id: int, + output: RequestOutput, + req_state: OrchestratorRequestState, + stage_metrics: StageRequestStats | None, + submit_ts: float | None, + ) -> bool: + del replica_id, submit_ts + if not isinstance(req_state, DuplexOrchestratorRequestState) or req_state.fence is None: + return False + request_id = output.request_id + runner = self.session_manager.runner_for_request_id(request_id) + if runner is None: + # Session gone: nothing may forward or reach a client. + return True + segment = req_state.streaming.segment(stage_id) + context = DuplexOutputContext( + identity=DuplexRequestIdentity( + session_id=req_state.session_id, + fence=req_state.stage_fences.get(stage_id, req_state.fence), + ), + final_stage_id=req_state.final_stage_id, + segment_finished=req_state.streaming.enabled and segment.finished, + segment_token_ids=tuple(segment.token_ids), + segment_output_metadata=segment.output_metadata, + ) + return runner.on_stage_output( + stage_id, + output, + stage_metrics, + request_id=request_id, + context=context, + ) + + async def _handle_forward_failure( + self, + req_id: str, + next_stage_id: int, + req_state: OrchestratorRequestState, + exc: BaseException, + ) -> bool: + if not req_state.session_owned: + return False + runner = self.session_manager.runner_for_request_id(req_id) + if runner is not None: + runner.on_stage_failure(next_stage_id, exc) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + release_owners=True, + ) + return True + + async def _cleanup_request_ids( + self, + request_ids: list[str], + *, + abort: bool = False, + release_owners: bool = False, + ) -> list[OutputMessage]: + if not request_ids: + return [] + cleanup_ids = list(dict.fromkeys(request_ids)) + closing_session_ids: list[str] = [] + if release_owners: + closed_sessions = self.session_manager.close_sessions_for_request_ids( + cleanup_ids, + abort=abort, + cleanup_in_progress=True, + ) + closing_session_ids.extend(closed_sessions) + for session_id, stale_request_ids in closed_sessions.items(): + logger.info( + "[DuplexOrchestrator] closed duplex session %s while cleaning failed request ids %s", + session_id, + stale_request_ids, + ) + cleanup_ids.extend(stale_request_ids) + cleanup_ids = list(dict.fromkeys(cleanup_ids)) + try: + outputs = await super()._cleanup_request_ids(cleanup_ids, abort=abort) + except BaseException: + if closing_session_ids: + self.session_manager.defer_request_cleanups(closing_session_ids) + raise + if closing_session_ids: + self.session_manager.finalize_closed_sessions(closing_session_ids) + for request_id in cleanup_ids: + self.session_manager.unregister_request(request_id) + return outputs + + # ------------------------------------------------------------------ # + # DuplexStagePort # + # ------------------------------------------------------------------ # + + @property + def stage_count(self) -> int: + return len(self.stage_pools) + + def sampling_defaults(self) -> tuple[object, ...]: + defaults = [] + for pool in self.stage_pools: + client = pool.stage_client + if client is None: + # Every replica of this stage was evicted (engine core died); a + # session cannot be admitted until the deployment is restarted. + raise RuntimeError(f"stage {pool.stage_id} has no live replica") + defaults.append(client.default_sampling_params) + return tuple(defaults) + + @staticmethod + def _sync_bridge_state( + request_state: OrchestratorRequestState, + context: DuplexStageRequestContext, + ) -> None: + duplex_state = request_state.streaming.bridge_states.setdefault("duplex", {}) + if not isinstance(duplex_state, dict): + duplex_state = {} + request_state.streaming.bridge_states["duplex"] = duplex_state + previous_epoch = duplex_state.get("epoch") + current_model_turn_id = duplex_state.get("model_turn_id") + if ( + not isinstance(current_model_turn_id, int) + or previous_epoch != context.fence.epoch + or current_model_turn_id < context.fence.turn_id + ): + # A safety boundary can close a response before the model emits + # its normal turn_eos. Catch up the engine-owned identity when the + # next fenced append starts. + duplex_state["model_turn_id"] = context.fence.turn_id + duplex_state.update( + { + "session_id": context.session_id, + "fence": context.fence, + "epoch": context.fence.epoch, + "turn_id": context.fence.turn_id, + "session_config": dict(context.session_config), + "runtime_config": dict(context.runtime_config), + } + ) + + def ensure_request(self, context: DuplexStageRequestContext) -> None: + request_state = self.request_states.get(context.request_id) + if request_state is None: + request_state = DuplexOrchestratorRequestState( + request_id=context.request_id, + prompt=None, + sampling_params_list=list(context.sampling_params), + final_stage_id=context.final_stage_id, + session_owned=True, + session_id=context.session_id, + fence=context.fence, + config_generation=context.config_generation, + ) + request_state.streaming.enabled = True + self.request_states[context.request_id] = request_state + elif isinstance(request_state, DuplexOrchestratorRequestState): + if request_state.config_generation != context.config_generation: + request_state.sampling_params_list = list(context.sampling_params) + request_state.config_generation = context.config_generation + request_state.session_id = context.session_id + request_state.fence = context.fence + else: + raise RuntimeError(f"request {context.request_id} is not session-owned") + self._sync_bridge_state(request_state, context) + self.session_manager.register_request(context.request_id, context.session_id) + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: + context = submission.context + request_state = self.request_states.get(context.request_id) + if not isinstance(request_state, DuplexOrchestratorRequestState): + raise RuntimeError(f"duplex request was not preregistered: {context.request_id}") + request = build_engine_core_request_from_tokens( + request_id=context.request_id, + prompt=dict(submission.prompt), + params=context.stage_sampling_params, + model_config=self.stage_pools[context.stage_id].stage_vllm_config.model_config, + resumable=True, + ) + request.external_req_id = request.request_id + pool = self.stage_pools[context.stage_id] + if submission.already_submitted: + replica_id = await pool.submit_update(context.request_id, request_state, request) + else: + replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None) + if self.async_chunk and context.stage_id == 0: + prewarmed = await self._prewarm_async_chunk_stages( + context.request_id, + request, + request_state, + ) + if not prewarmed: + # The prewarm already failed the request, aborted it and + # popped its state; the runner turns this into an append + # error instead of writing onto an orphaned object. + raise RuntimeError( + f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted" + ) + request_state.stage_fences[context.stage_id] = context.fence + request_state.stage_submit_ts[context.stage_id] = _time.time() + self._register_running_request(request_state) + return DuplexStageSubmissionResult( + request_id=context.request_id, + stage_id=context.stage_id, + replica_id=replica_id, + ) + + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + await self._cleanup_request_ids(request_ids, abort=abort) + + async def abort_requests(self, request_ids: list[str]) -> None: + if request_ids: + await self._abort_request_ids(list(dict.fromkeys(request_ids))) + + +__all__ = ["DuplexOrchestrator", "DuplexOrchestratorRequestState"] diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py new file mode 100644 index 000000000..87ef8d5e6 --- /dev/null +++ b/vllm_omni/engine/omni_engine_base.py @@ -0,0 +1,1229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniEngineBase: process/thread ownership and transport shared by every Omni engine.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import json +import queue +import threading +import time +import uuid +import weakref +from collections.abc import Mapping +from dataclasses import asdict +from pathlib import Path +from typing import Any, cast + +import janus +import torch +from vllm import envs as vllm_envs +from vllm.logger import init_logger +from vllm.v1.engine.input_processor import InputProcessor + +from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config +from vllm_omni.config.stage_config import ( + _DEPLOY_DIR, + DuplexSessionRuntimeConfig, + PipelineConfig, + load_deploy_config, +) +from vllm_omni.diffusion.data import ( + DiffusionParallelConfig, + parse_attention_config, +) +from vllm_omni.diffusion.io_support import get_diffusion_output_type +from vllm_omni.engine.async_engine_utils import ( + SHUTDOWN_ENQUEUE_TIMEOUT_S, + SHUTDOWN_JOIN_TIMEOUT_S, + enqueue_orchestrator_shutdown, + is_abort_transport_shutdown, + is_janus_sync_queue_shutdown, + shutdown_runtime_after_orchestrator, + weak_shutdown_async_omni_engine, +) +from vllm_omni.engine.messages import ( + AbortRequestMessage, + AbortResultMessage, + CollectiveRPCRequestMessage, + CollectiveRPCResultMessage, + EngineQueueMessage, + ErrorMessage, + OutputMessage, +) +from vllm_omni.engine.orchestrator import OrchestratorBase +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.engine.stage_client import StageClient +from vllm_omni.engine.stage_init_utils import build_stage0_input_processor +from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.engine.stage_runtime import ( + StageRuntimeInfo, + create_stage_runtime, +) +from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin +from vllm_omni.entrypoints.utils import parse_stage_overrides +from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.metrics.prometheus import OmniRequestCounter + +logger = init_logger(__name__) + +_STARTUP_POLL_INTERVAL_S = 1.0 +_REQUEST_QUEUE_MAXSIZE = 256 +_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None] + + +def load_and_resolve_stage_configs( + model: str, + kwargs: dict[str, Any], + *, + trust_remote_code: bool | None, + deploy_config_path: str | None, + stage_overrides: Mapping[str, Mapping[str, Any]] | None, + strategy_config_path: str | None, +) -> OmniConfigResolution: + """Compatibility seam delegating to the single config resolver.""" + return resolve_omni_config( + model, + trust_remote_code=trust_remote_code, + cli_overrides=kwargs, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ) + + +class OmniEngineBase: + """Generic engine: launches an orchestrator in a background thread. + + Shared by ``AsyncOmniEngine`` (turn-based requests) and ``DuplexOmniEngine`` + (duplex sessions). Subclasses construct their orchestrator in + ``_create_orchestrator``. + + All stage clients, input/output processors, and stage-to-stage transfer + logic live inside the Orchestrator coroutine (running in its own thread + with a dedicated asyncio event loop). This class communicates with it + via janus queues (sync side for callers, async side for orchestrator). + + Args: + model: Model name or path + init_timeout: Total timeout waiting for orchestrator startup (seconds). + stage_init_timeout: Timeout for stage initialization (seconds) + **kwargs: Additional arguments + """ + + # Class-level defaults so tests that bypass __init__ via object.__new__ + # don't AttributeError when stage-init / forward paths touch these attrs. + _log_stats: bool = False + _coordinator_runtime: Any = None + _transfer_emitter: Any = None + _prom_metrics: Any = None + _enable_orch_monitor: bool = False + # Lazily created by get_output_blocking_async(). + _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None + + def __init__( + self, + model: str, + stage_init_timeout: int = 300, + init_timeout: int = 600, + single_stage_mode: bool = False, + transfer_emitter: Any = None, + prom_metrics: Any = None, + log_stats: bool = False, + tokenizer: str | None = None, + trust_remote_code: bool | None = None, + **kwargs: Any, + ) -> None: + self.model = model + self.tokenizer = tokenizer + # Cached by get_diffusion_od_config(). + self._diffusion_od_config_view: Any = None + startup_timeout = int(init_timeout) + # Forwarded into Orchestrator so its _forward_to_next_stage path can + # emit per-edge transfer_tx_s / transfer_size_bytes histograms. + # Optional: when None, Orchestrator silently skips TX emit (existing + # RX path still works via OrchestratorAggregator). + self._transfer_emitter = transfer_emitter + self._prom_metrics = prom_metrics + # Drives upstream EngineCore + scheduler stats production. When False + # the engine skips SchedulerStats / IterationStats; the per-(stage, + # replica) vllm:* wrap stays registered but reads zero. Respects the + # --log-stats CLI flag set by the user via OmniBase. + self._log_stats = log_stats + self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False)) + + logger.info(f"[OmniEngine] Initializing with model {model}") + + # ------------------------------------------------------------------ # + # Single-stage mode detection # + # ------------------------------------------------------------------ # + # Single-stage mode is enabled when the caller explicitly passes # + # single_stage_mode=True, or when a stage_id is provided in the args. # + _stage_id_kwarg = kwargs.get("stage_id") + if isinstance(_stage_id_kwarg, int) and not single_stage_mode: + single_stage_mode = True + + self.single_stage_mode: bool = single_stage_mode + self._single_stage_id_filter: int | None = ( + int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None + ) + self._omni_master_address: str | None = kwargs.get("omni_master_address") + self._omni_master_port: int | None = kwargs.get("omni_master_port") + + # New omni-coordinator flags. Consumed only in single_stage_mode. + # ``omni_dp_size_local`` is process-local: each invocation (head and + # every headless) launches that many replicas for its own stage. + self._omni_dp_size_local: int = int(kwargs.get("omni_dp_size_local") or 1) + if self._omni_dp_size_local < 1: + raise ValueError(f"--omni-dp-size-local must be >= 1, got {self._omni_dp_size_local}") + self._omni_lb_policy: str = str(kwargs.get("omni_lb_policy") or "random") + self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) + if self._omni_heartbeat_timeout <= 0: + raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") + # Concurrent same-device stage init (admission + SH/EX phase locks). + # Sourced from the parallel_stage_init orchestrator/CLI arg (config, + # not an env var); default False preserves serial init. + self._parallel_stage_init: bool = bool(kwargs.get("parallel_stage_init") or False) + + if single_stage_mode: + logger.info( + "[OmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", + self._single_stage_id_filter, + self._omni_master_address, + self._omni_master_port, + ) + + # Keep the historical tuple return from _resolve_stage_configs while + # retaining the richer resolver result for pipeline-wide settings. + # Overrides of that private seam fall back to the factory below. + deploy_config_path = kwargs.get("deploy_config") + # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not + # specified" so stage-config resolution can defer to the deploy yaml's + # per-stage value (see ``with_trust_remote_code_override``). The + # restriction path below loads the top-level HF config via vLLM's + # ``get_config``, which needs a real bool, so collapse ``None`` to the + # default ``False`` here (#5495). + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=bool(trust_remote_code), + deploy_config_path=deploy_config_path, + ) + self.pipeline_config = pipeline_config + self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () + # The resolved deploy profile (explicit --deploy-config or the pipeline + # default). Duplex engines read ``deploy_config.duplex_session`` and + # ``session_mode`` from it; the turn-based engine only keeps it for + # introspection. + self.deploy_config = None + #: Path ``self.deploy_config`` was parsed from, so the duplex runtime + #: config below can reuse it instead of re-reading the same yaml. + self._deploy_config_source: str | None = None + deploy_config_source: str | Path | None = deploy_config_path + if deploy_config_source is None and pipeline_config is not None: + default_name = getattr(pipeline_config, "default_deploy_config_name", None) + if default_name: + deploy_config_source = _DEPLOY_DIR / default_name + elif isinstance(deploy_config_source, str) and not Path(deploy_config_source).exists(): + candidate = _DEPLOY_DIR / deploy_config_source + if candidate.exists(): + deploy_config_source = candidate + if deploy_config_source is not None: + try: + self.deploy_config = load_deploy_config(deploy_config_source) + self._deploy_config_source = str(deploy_config_source) + except Exception: + if deploy_config_path is not None: + # The caller named this file: failing to load it is their error. + raise + logger.warning( + "[OmniEngine] Could not load the pipeline's default deploy config %s", + deploy_config_source, + exc_info=True, + ) + + # Subclasses check deployment invariants here, before any stage + # process is launched (a failure is cheap at this point). + self._validate_deployment() + + # Tri-state: None means "not specified" — the deploy yaml's per-stage + # trust_remote_code stays in effect. An explicit True/False here is a + # global override (precedence: caller > deploy yaml > default False); + # the merge rule lives in with_trust_remote_code_override. + kwargs = with_trust_remote_code_override(kwargs, trust_remote_code) + self._config_resolution: OmniConfigResolution | None = None + self.config_path, self.stage_configs = self._resolve_stage_configs( + model, + kwargs, + trust_remote_code=trust_remote_code, + ) + if self._config_resolution is None: + # Same model, trust_remote_code and deploy path as the resolution + # above, so reuse it rather than paying for the HF config again. + self._set_pipeline_runtime_config(self.pipeline_config, self.config_path) + else: + self._set_pipeline_runtime_config( + self._config_resolution.pipeline_config, + self._config_resolution.config_path, + ) + + self.num_stages = len(self.stage_configs) + stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None + self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) + self.stage_pools: list[StagePool] = [] + self.stage_clients: list[StageClient] = [] # logical-stage view for external readers + self.input_processor: InputProcessor | None = None + self.prompt_transform_func: Any | None = None + self.prompt_expand_func: Any | None = None + self.supported_tasks: tuple[str, ...] = ("generate",) + self.default_sampling_params_list: list[OmniSamplingParams] = [] + self.stage_metadata: list[StageRuntimeInfo] = [] + # Janus queues are constructed eagerly here (not deferred to the + # orchestrator thread) so the master server's ROUTER thread always + # sees a non-None ``self.request_queue`` when on_register fires. + # ``async_q`` lazily binds to whatever event loop first awaits on + # it (the orchestrator loop), so cross-thread use stays correct. + self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) + self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self._shutdown_called = False + self._weak_finalizer: weakref.finalize | None = None + self._correlated_rpc_client: CorrelatedRpcClient | None = None + self._running_counter = OmniRequestCounter() + self._engines_waiting_counter = OmniRequestCounter() + + logger.info(f"[OmniEngine] Launching Orchestrator thread with {self.num_stages} stages") + + # Launch orchestrator background thread + startup_future: concurrent.futures.Future = concurrent.futures.Future() + + self.orchestrator_thread = threading.Thread( + target=self._bootstrap_orchestrator, + args=( + stage_init_timeout, + startup_future, + ), + daemon=True, + name="orchestrator", + ) + self.orchestrator_thread.start() + self._wait_for_orchestrator_init(startup_future, startup_timeout) + self._correlated_rpc_client = CorrelatedRpcClient( + self.request_queue.sync_q, + self.rpc_output_queue.sync_q, + ) + + # Stage runtime fields are assigned directly on self by the bootstrap thread. + self._weak_finalizer = weakref.finalize( + self, + weak_shutdown_async_omni_engine, + self.orchestrator_thread, + self.request_queue, + self.output_queue, + self.rpc_output_queue, + self._correlated_rpc_client, + ) + + logger.info(f"[OmniEngine] Orchestrator ready with {self.num_stages} stages") + + def get_diffusion_od_config(self) -> Any: + """Expose the diffusion ``model_class_name`` to client-side model-extras. + + The worker holds the full config; here we just resolve the pipeline class + name from the model config (cached). ``model_class_name`` may be ``None``. + """ + if self._diffusion_od_config_view is None: + from types import SimpleNamespace + + from vllm_omni.diffusion.data import resolve_model_class_name + from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata + + model_class_name = resolve_model_class_name(self.model) + metadata = get_diffusion_model_metadata(model_class_name) + self._diffusion_od_config_view = SimpleNamespace( + model_class_name=model_class_name, + supports_multimodal_inputs=metadata.supports_multimodal_inputs, + max_multimodal_image_inputs=metadata.max_multimodal_image_inputs, + supports_mixed_reference_inputs=metadata.supports_mixed_reference_inputs, + ) + return self._diffusion_od_config_view + + def _initialize_stages(self, stage_init_timeout: int) -> None: + """Initialize stage clients/processors via StageRuntime and assign to self.""" + self._runtime = create_stage_runtime( + stage_configs=self.stage_configs, + model=self.model, + config_path=self.config_path, + single_stage_mode=self.single_stage_mode, + stage_init_timeout=stage_init_timeout, + async_chunk=self.async_chunk, + tokenizer=self.tokenizer, + parallel_stage_init=self._parallel_stage_init, + single_stage_id_filter=self._single_stage_id_filter, + omni_master_address=self._omni_master_address, + omni_master_port=self._omni_master_port, + omni_dp_size_local=self._omni_dp_size_local, + omni_heartbeat_timeout=self._omni_heartbeat_timeout, + omni_lb_policy=self._omni_lb_policy, + request_queue=self.request_queue, + log_stats=self._log_stats, + ) + self._runtime.initialize() + + self.num_stages = len(self.stage_configs) + self.stage_pools = self._runtime.stage_pools + self.stage_clients = [ + cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None + ] + self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] + self.output_processors = [pool.output_processor for pool in self.stage_pools] + self.input_processor = ( + build_stage0_input_processor(self.stage_vllm_configs[0]) + if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None + else None + ) + self.prompt_transform_func = ( + getattr(self.stage_clients[0], "prompt_transform_func", None) if self.stage_clients else None + ) + self.prompt_expand_func = next( + ( + getattr(client, "prompt_expand_func", None) + for client in self.stage_clients + if getattr(client, "prompt_expand_func", None) is not None + ), + None, + ) + self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] + self.stage_metadata = [ + StageRuntimeInfo( + final_output=client.final_output, + final_output_type=client.final_output_type, + stage_type=client.stage_type, + model_stage=getattr(client, "model_stage", None), + ) + for client in self.stage_clients + ] + supported_tasks: set[str] = set() + if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): + supported_tasks.add("generate") + if any(meta.final_output_type == "audio" for meta in self.stage_metadata): + supported_tasks.add("speech") + self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) + + def _bootstrap_orchestrator( + self, + stage_init_timeout: int, + startup_future: concurrent.futures.Future, + ) -> None: + """Create loop, initialize stages, then run Orchestrator.""" + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def _run_orchestrator() -> None: + self._initialize_stages(stage_init_timeout) + + pd_config = self._detect_pd_config() + + membership_controller = self._runtime.create_membership_controller() + + orchestrator = self._create_orchestrator( + request_async_queue=self.request_queue.async_q, + output_async_queue=self.output_queue.async_q, + rpc_async_queue=self.rpc_output_queue.async_q, + stage_pools=self.stage_pools, + async_chunk=self.async_chunk, + pd_config=pd_config, + membership_controller=membership_controller, + running_counter=self._running_counter, + engines_waiting_counter=self._engines_waiting_counter, + transfer_emitter=self._transfer_emitter, + prom_metrics=self._prom_metrics, + log_stats=self._log_stats, + enable_orch_monitor=self._enable_orch_monitor, + ) + if not startup_future.done(): + startup_future.set_result(asyncio.get_running_loop()) + await orchestrator.run() + + try: + loop.run_until_complete(_run_orchestrator()) + except Exception as e: + if not startup_future.done(): + wrapped = RuntimeError(f"Orchestrator initialization failed: {e}") + wrapped.__cause__ = e + startup_future.set_exception(wrapped) + logger.exception("[OmniEngine] Orchestrator thread crashed") + error_text = str(e) or "Orchestrator thread crashed" + try: + error_msg = ErrorMessage(error=error_text, fatal=True) + if self.output_queue is not None: + self.output_queue.sync_q.put_nowait(error_msg) + if self.rpc_output_queue is not None: + self.rpc_output_queue.sync_q.put_nowait(error_msg) + except Exception: + pass + raise + finally: + try: + pending = [task for task in asyncio.all_tasks(loop) if not task.done()] + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + if hasattr(loop, "shutdown_default_executor"): + loop.run_until_complete(loop.shutdown_default_executor()) + except Exception: + logger.exception("[OmniEngine] Failed during orchestrator loop cleanup") + finally: + asyncio.set_event_loop(None) + loop.close() + + def _validate_deployment(self) -> None: + """Seam: check ``pipeline_config`` / ``deploy_config`` before stages start (default: nothing).""" + + def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: + """Construct this engine's orchestrator (runs on the orchestrator thread after stage init).""" + raise NotImplementedError + + def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None: + """ + Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process. + """ + deadline = time.monotonic() + startup_timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + logger.warning( + "[OmniEngine] Orchestrator startup timed out after %ss. " + "Multi-stage deployments that initialize stages sequentially on one device " + "or load checkpoints from slow storage may need larger --init-timeout and " + "--stage-init-timeout values.", + startup_timeout, + ) + self._try_shutdown("[OmniEngine] Failed to cleanup after orchestrator startup timeout") + raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") + try: + startup_future.result( + timeout=min(remaining, _STARTUP_POLL_INTERVAL_S), + ) + break + except concurrent.futures.TimeoutError: + if not self.orchestrator_thread.is_alive(): + self._try_shutdown("[OmniEngine] Failed to cleanup after orchestrator startup failure") + if startup_future.done(): + startup_future.result() # re-raises the real exception + raise RuntimeError("Orchestrator thread died during startup") + except Exception: + self._try_shutdown("[OmniEngine] Failed to cleanup after orchestrator startup failure") + raise + + @staticmethod + def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None: + if cache_backend == "cache_dit": + return { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": False, + "taylorseer_order": 1, + "scm_steps_mask_policy": None, + "scm_steps_policy": "dynamic", + } + if cache_backend == "tea_cache": + return { + "rel_l1_thresh": 0.2, + } + if cache_backend == "mag_cache": + return { + "mag_threshold": 0.24, + "mag_max_skip_steps": 5, + "mag_retention_ratio": 0.1, + } + if cache_backend == "step_cache": + return { + "step_cache_dit_enabled": True, + "velocity_sim_thresholds": [0.95, 0.93], + "velocity_skip_countdowns": [4, 2], + "step_cache_dit_min_history": 2, + "step_cache_dit_max_history": 2, + } + return None + + @staticmethod + def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None: + if isinstance(cache_config, str): + try: + cache_config = json.loads(cache_config) + except json.JSONDecodeError: + logger.warning("Invalid cache_config JSON, using defaults.") + cache_config = None + if cache_config is None and cache_backend not in (None, "", "none"): + cache_config = OmniEngineBase._get_default_cache_config(cache_backend) + return cache_config + + def _detect_pd_config(self) -> dict[str, Any] | None: + """Detect PD (Prefill-Decode) disaggregation config from stage_configs. + Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. + """ + pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs) + if pd_pair is None: + return None + prefill_idx, decode_idx = pd_pair + + # Extract bootstrap address from prefill stage engine_args + bootstrap_addr: str | None = None + try: + prefill_cfg = self.stage_configs[prefill_idx] + ea = getattr(prefill_cfg, "engine_args", None) + kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None + if kv_cfg is not None: + port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT + kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1" + bootstrap_addr = f"http://{kv_ip}:{port}" + except Exception as exc: + logger.warning("[OmniEngine] Could not extract PD bootstrap address: %s", exc) + + logger.info( + "[OmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s", + prefill_idx, + decode_idx, + bootstrap_addr, + ) + prefill_engine_id: str | None = None + try: + prefill_client = self.stage_clients[prefill_idx] + kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None) + prefill_engine_id = getattr(kv_cfg, "engine_id", None) + except Exception as exc: + logger.warning("[OmniEngine] Could not extract prefill engine_id: %s", exc) + + return { + "pd_pair": (prefill_idx, decode_idx), + "bootstrap_addr": bootstrap_addr, + "prefill_engine_id": prefill_engine_id, + } + + @staticmethod + def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: + """Create a default single-stage diffusion config from kwargs.""" + # We temporally create a default config for diffusion stage. + # In the future, we should merge the default config with the user-provided config. + normalized_kwargs = dict(kwargs) + default_sampling_params = normalized_kwargs.get("default_sampling_params") + if isinstance(default_sampling_params, str): + try: + default_sampling_params = json.loads(default_sampling_params) + except json.JSONDecodeError: + logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") + default_sampling_params = None + if not isinstance(default_sampling_params, dict): + default_sampling_params = None + stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} + if normalized_kwargs.get("dtype") is None: + normalized_kwargs["dtype"] = "auto" + + # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. + if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): + if not isinstance(normalized_kwargs["dtype"], torch.dtype): + raise TypeError( + f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}" + ) + normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.") + + cache_backend = normalized_kwargs.get("cache_backend", "none") + cache_config = OmniEngineBase._normalize_cache_config( + cache_backend, + normalized_kwargs.get("cache_config", None), + ) + + parallel_config = normalized_kwargs.get("parallel_config") + if isinstance(parallel_config, dict): + parallel_config = DiffusionParallelConfig.from_dict(parallel_config) + if parallel_config is None: + ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 + ring_degree = normalized_kwargs.get("ring_degree") or 1 + allgather_degree = normalized_kwargs.get("allgather_degree") or 1 + ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" + ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False)) + sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") + pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 + data_parallel_size = normalized_kwargs.get("data_parallel_size") + tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 + cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 + pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 + vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 + vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile" + text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1 + enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False + use_hsdp = normalized_kwargs.get("use_hsdp", False) + hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) + hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) + if sequence_parallel_size is None: + sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=pipeline_parallel_size, + data_parallel_size=data_parallel_size, + tensor_parallel_size=tensor_parallel_size, + enable_expert_parallel=enable_expert_parallel, + sequence_parallel_size=sequence_parallel_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + allgather_degree=allgather_degree, + ulysses_mode=ulysses_mode, + ulysses_a2a_permute=ulysses_a2a_permute, + cfg_parallel_size=cfg_parallel_size, + vae_patch_parallel_size=vae_patch_parallel_size, + vae_parallel_mode=vae_parallel_mode, + text_encoder_tp_size=text_encoder_tp_size, + use_hsdp=use_hsdp, + hsdp_shard_size=hsdp_shard_size, + hsdp_replicate_size=hsdp_replicate_size, + ) + + num_gpus = normalized_kwargs.get("num_gpus") + if num_gpus is not None: + num_gpus = int(num_gpus) + parallel_config.resolve_data_parallel_size(num_gpus) + + num_devices = max(1, int(parallel_config.world_size)) + devices = ",".join(str(i) for i in range(num_devices)) + model_class_name = kwargs.get("model_class_name", None) + final_output_type = get_diffusion_output_type(model_class_name) + + attention_config = None + if ( + kwargs.get("diffusion_attention_config") is not None + or kwargs.get("diffusion_attention_backend") is not None + or kwargs.get("fastvideo_vsa_topk") is not None + ): + attention_config = parse_attention_config( + kwargs.get("diffusion_attention_config"), + attention_backend=kwargs.get("diffusion_attention_backend"), + fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), + ) + + extras = dict(kwargs.get("extras") or {}) + for key, default in ( + ("auxiliary_text_encoder", None), + ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), + ): + top_level_value = kwargs.get(key) + if top_level_value is not None: + extras[key] = top_level_value + else: + extras.setdefault(key, default) + + stage_engine_args = { + "max_num_seqs": kwargs.get("max_num_seqs") or 1, + "parallel_config": parallel_config, + # Default-stage construction bypasses the structured projection. + # Runner selection remains owned by the selected engine/platform. + "engine_backend": kwargs.get("engine_backend", "default"), + "model_class_name": kwargs.get("model_class_name", None), + "task_type": kwargs.get("task_type", None), + "model_config": kwargs.get("model_config", None), + "additional_config": kwargs.get("additional_config", None), + "step_execution": kwargs.get("step_execution", False), + "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), + "vae_use_slicing": kwargs.get("vae_use_slicing", False), + "vae_use_tiling": kwargs.get("vae_use_tiling", False), + "cache_backend": cache_backend, + "cache_config": cache_config, + "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), + "diffusion_offload_config": kwargs.get("diffusion_offload_config", None), + "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), + "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), + "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), + "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), + "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), + "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), + "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), + "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), + "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), + "diffusion_compile_granularity": ( + "regional" + if kwargs.get("diffusion_compile_granularity") is None + else kwargs["diffusion_compile_granularity"] + ), + "diffusion_compile_dynamic": ( + True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"] + ), + "fa_deterministic": bool(kwargs.get("fa_deterministic", False)), + "boundary_ratio": kwargs.get("boundary_ratio", None), + "flow_shift": kwargs.get("flow_shift", None), + "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), + "lora_path": kwargs.get("lora_path", None), + "lora_scale": kwargs.get("lora_scale", 1.0), + "lora_backend": kwargs.get("lora_backend", "peft"), + "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), + "worker_extension_cls": kwargs.get("worker_extension_cls", None), + "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), + "distributed_executor_backend": kwargs.get("distributed_executor_backend"), + "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), + "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), + "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), + "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), + "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), + "quantization": kwargs.get("quantization", None), + "quantization_config": kwargs.get("quantization_config", None), + "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), + "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), + "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), + **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), + "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), + "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), + "streaming_output": kwargs.get("diffusion_streaming_output", False), + "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), + "extras": extras, + **( + { + "profiler_config": asdict(kwargs["profiler_config"]) + if hasattr(kwargs["profiler_config"], "__dataclass_fields__") + else kwargs["profiler_config"] + } + if kwargs.get("profiler_config") is not None + else {} + ), + } + if num_gpus is not None: + stage_engine_args["num_gpus"] = num_gpus + # Only set dtype if it was already explicitly passed and normalized + if "dtype" in normalized_kwargs: + stage_engine_args["dtype"] = normalized_kwargs["dtype"] + + # New split fields for diffusers adapter kwargs. + if kwargs.get("diffusers_load_kwargs") is not None: + stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] + if kwargs.get("diffusers_call_kwargs") is not None: + stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] + + default_stage_cfg = [ + { + "stage_id": 0, + "stage_type": "diffusion", + "runtime": { + "process": True, + "devices": devices, + }, + "engine_args": stage_engine_args, + "engine_input_source": [], + "default_sampling_params": stage_default_sampling_params, + "final_output": True, + "final_output_type": final_output_type, + } + ] + default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" + return default_stage_cfg + + def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: + """Apply a strategy-derived ``omni_lb_policy`` to the engine. + + Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is + the engine default and is treated as "unset" (indistinguishable from no + flag), so a strategy value overrides it. If the user explicitly passed a + non-default policy that conflicts with the strategy-derived one, raise so + the mismatch is not silently ignored. + """ + if not derived: + return + explicit = kwargs.get("omni_lb_policy") + user_set = explicit is not None and str(explicit) != "random" + if user_set: + if str(explicit) != str(derived): + raise ValueError( + f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given " + f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. " + "Drop --omni-lb-policy to use the strategy value, or make them match." + ) + return + if self._omni_lb_policy != str(derived): + logger.info( + "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).", + derived, + self._omni_lb_policy, + ) + self._omni_lb_policy = str(derived) + + @staticmethod + def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Compatibility seam for the factory-owned diffusion fallback.""" + return StageConfigFactory.create_default_diffusion(kwargs) + + def _set_pipeline_runtime_config( + self, + pipeline_config: PipelineConfig | None, + config_path: str | None, + ) -> None: + """Initialize engine-wide settings resolved from pipeline metadata.""" + self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () + # No duplex_runtime_extension / duplex_serving_adapter / duplex_control + # here: the pre-framework wiring they named is gone, and design rule 3 + # in docs/design/fullduplex.md keeps duplex vocabulary out of the base. + self.duplex_session_config = DuplexSessionRuntimeConfig() + if config_path is not None: + if self.deploy_config is not None and self._deploy_config_source == str(config_path): + # Already parsed in __init__; the same yaml twice is pure cost. + self.duplex_session_config = self.deploy_config.duplex_session + else: + self.duplex_session_config = load_deploy_config(config_path).duplex_session + + def _resolve_stage_configs( + self, + model: str, + kwargs: dict[str, Any], + *, + trust_remote_code: bool | None, + ) -> tuple[str, list[Any]]: + """Resolve stage configs and inject defaults shared by orchestrator/headless.""" + + for legacy_arg in ("stage_configs_path", "stage_configs"): + if legacy_arg in kwargs: + raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") + + # log_stats is captured by __init__; its CLI-only negative alias must + # not cross into per-stage structured config ownership validation. + kwargs.pop("disable_log_stats", None) + deploy_config_path = kwargs.pop("deploy_config", None) + strategy_config_path = kwargs.pop("strategy_config", None) + # CLI callers arrive pre-parsed; offline Python callers may use the + # JSON-string form documented in recipes. + stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None)) + + # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; + # stage configs know the field as ``streaming_output``. Mirror only a + # truthy value so the CLI's False default does not override deploy YAML. + if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: + kwargs["streaming_output"] = True + + resolution = cast( + _ConfigResolutionResult, + load_and_resolve_stage_configs( + model, + kwargs, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ), + ) + if isinstance(resolution, OmniConfigResolution): + self._config_resolution = resolution + config_path = resolution.config_path + stage_configs = list(resolution.stage_configs) + strategy_lb_policy = resolution.omni_lb_policy + else: + # Compatibility for overrides of the historical tuple-returning + # seam. Production always receives OmniConfigResolution above. + config_path, stage_configs, strategy_lb_policy = resolution + + # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is + # an orchestrator-level knob (read once at construction), so apply it here + # rather than as a per-stage config field. + self._apply_strategy_lb_policy(strategy_lb_policy, kwargs) + + return cast(str, config_path), stage_configs + + # ==================== Public API ==================== + + def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: + """Read one output message from the Orchestrator output queue.""" + try: + return self.output_queue.sync_q.get(timeout=timeout) + except queue.Empty: + if not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return None + + async def try_get_output_async(self) -> EngineQueueMessage | None: + """Async read from the Orchestrator output queue.""" + try: + return self.output_queue.sync_q.get_nowait() + except queue.Empty: + if not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return None + + async def get_output_blocking_async(self, timeout: float = 1.0) -> EngineQueueMessage | None: + """Blocking-wait read from the Orchestrator output queue. + + Waits up to ``timeout`` seconds in a dedicated drain thread for the + next message (condition-variable wakeup instead of a poll cadence); + returns ``None`` on timeout so the caller keeps its liveness check, + mirroring ``try_get_output_async``'s contract. Used by the serving + final-output drain when ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` is on. + """ + executor = self._output_drain_executor + if executor is None: + executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="omni-output-drain", + ) + self._output_drain_executor = executor + + sync_q = self.output_queue.sync_q + + def _drain_get() -> EngineQueueMessage | None: + # Exceptions are swallowed to a None sentinel: the queue may be + # closed mid-shutdown, and an exception left on an executor future + # after task cancellation would warn as never-retrieved. + try: + return sync_q.get(timeout=timeout) + except queue.Empty: + return None + except Exception: + return None + + loop = asyncio.get_running_loop() + msg = await loop.run_in_executor(executor, _drain_get) + if msg is None and not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return msg + + def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: + """Get cached metadata for a stage.""" + return self.stage_metadata[stage_id] + + def abort(self, request_ids: list[str]) -> None: + """Fire-and-forget abort: enqueue and return without waiting. + + Prefer :meth:`abort_async` when the caller needs acknowledgment that + stage aborts, binding release, and orchestrator request cleanup finished. + """ + if not request_ids or getattr(self, "_shutdown_called", False): + return + if self.request_queue is None: + raise RuntimeError("request_queue is not initialized") + try: + self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_janus_sync_queue_shutdown(exc): + return + raise + + async def abort_async( + self, + request_ids: list[str], + timeout: float | None = None, + ) -> list[OutputMessage]: + """Abort requests and wait for orchestrator acknowledgment. + + Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the + :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and + raises if the orchestrator reports failure or times out. + + Returns: + Final-stage AR abort ``OutputMessage`` list carrying partial + tokens generated before abort (empty for diffusion / no OP state). + """ + if not request_ids or getattr(self, "_shutdown_called", False): + return [] + if self.request_queue is None: + raise RuntimeError("request_queue is not initialized") + transport = self._correlated_rpc_client + if transport is None: + raise RuntimeError("correlated RPC client is not initialized") + + rpc_id = uuid.uuid4().hex + msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id) + + def _wait() -> AbortResultMessage: + result_msg = transport.execute( + ("abort", rpc_id), + msg, + timeout=timeout, + timeout_message=f"abort timed out after {timeout} seconds", + block_on_submit=True, + ) + if not isinstance(result_msg, AbortResultMessage): + raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}") + return result_msg + + loop = asyncio.get_running_loop() + try: + result_msg = await loop.run_in_executor(None, _wait) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): + return [] + raise + if not result_msg.success: + raise RuntimeError(result_msg.error or "abort failed") + return list(result_msg.abort_outputs or []) + + def collective_rpc( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + stage_ids: list[int] | None = None, + ) -> list[Any]: + """Send a control RPC to the Orchestrator and wait for aggregated results. + + This uses a dedicated RPC output queue so control-plane messages do not + race with the normal request output polling loop. + """ + rpc_id = uuid.uuid4().hex + msg = CollectiveRPCRequestMessage( + rpc_id=rpc_id, + method=method, + timeout=timeout, + args=tuple(args), + kwargs=kwargs or {}, + stage_ids=stage_ids, + ) + + transport = self._correlated_rpc_client + if transport is None: + raise RuntimeError("correlated RPC client is not initialized") + result_msg = transport.execute( + ("collective", rpc_id), + msg, + timeout=timeout, + timeout_message=f"collective_rpc timed out after {timeout} seconds", + block_on_submit=True, + ) + if not isinstance(result_msg, CollectiveRPCResultMessage): + raise RuntimeError(f"unexpected collective RPC result type: {type(result_msg).__name__}") + return list(result_msg.results) + + async def collective_rpc_async( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + stage_ids: list[int] | None = None, + ) -> list[Any]: + """Async wrapper around collective_rpc().""" + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self.collective_rpc( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + stage_ids=stage_ids, + ), + ) + + @property + def rpc_client(self) -> CorrelatedRpcClient: + """The correlated RPC transport (available once the orchestrator is up).""" + if self._correlated_rpc_client is None: + raise RuntimeError("engine RPC transport is not initialized") + return self._correlated_rpc_client + + def is_alive(self) -> bool: + """Whether the orchestrator thread is alive.""" + return bool(self.orchestrator_thread.is_alive()) + + def shutdown(self) -> None: + """Send shutdown message and wait for the Orchestrator thread to exit.""" + if getattr(self, "_shutdown_called", False): + return + self._shutdown_called = True + finalizer = getattr(self, "_weak_finalizer", None) + if finalizer is not None and finalizer.alive: + finalizer.detach() + + logger.info("[OmniEngine] Shutting down Orchestrator") + request_queue_closed = False + shutdown_enqueued = enqueue_orchestrator_shutdown( + self.request_queue, + timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S, + ) + if self.request_queue is not None and not shutdown_enqueued: + logger.error( + "[OmniEngine] Failed to enqueue orchestrator shutdown; " + "closing the request queue to wake the request handler" + ) + try: + self.request_queue.close() + request_queue_closed = True + except Exception: + logger.exception("[OmniEngine] Failed to close the request queue") + + if self._correlated_rpc_client is not None: + try: + self._correlated_rpc_client.close() + except Exception: + logger.exception("[OmniEngine] Failed to close correlated RPC client") + + orchestrator_stopped = False + try: + if self.is_alive(): + self.orchestrator_thread.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S) + orchestrator_stopped = not self.is_alive() + if not orchestrator_stopped: + logger.error( + "[OmniEngine] Orchestrator did not stop within %.1f seconds; continuing cleanup", + SHUTDOWN_JOIN_TIMEOUT_S, + ) + except Exception: + logger.exception("[OmniEngine] Failed to join Orchestrator thread") + + for q in (self.request_queue, self.output_queue, self.rpc_output_queue): + try: + if not (q is self.request_queue and request_queue_closed): + q.close() + except Exception: + pass + + if self._output_drain_executor is not None: + # Any in-flight blocking get bails out within its ≤1 s timeout + # (or immediately via the queue close above), so don't wait. + self._output_drain_executor.shutdown(wait=False) + self._output_drain_executor = None + + if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped: + try: + self._runtime.shutdown() + except Exception: + logger.exception("[OmniEngine] Failed to shutdown StageRuntime") + elif hasattr(self, "_runtime") and self._runtime is not None: + logger.warning("[OmniEngine] Deferring StageRuntime shutdown until the Orchestrator exits") + threading.Thread( + target=shutdown_runtime_after_orchestrator, + args=(self.orchestrator_thread, self._runtime), + daemon=True, + name="omni-stage-runtime-shutdown", + ).start() + + # ── Release CuMem allocator memory pool ────────────────────────────── + # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory + # Management) allocator holds model weights in a singleton memory pool + # that lives in the parent process. Killing the engine-core subprocess + # does NOT release this pool — the weights stay resident on the GPU + # and can cause CUDA OOM for subsequent engine instances (especially + # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). + # + # CuMemAllocator.sleep() is NOT idempotent — calling it on already- + # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator + # cuMemRelease (double-free of the memory handle). Use release_pools() + # instead, which is the designed cleanup path: it drops MemPool refs + # and lets the destructor/free path handle asleep entries correctly + # (returns a null handle so the C extension skips unmap/release). + try: + from vllm.device_allocator.cumem import CuMemAllocator, cumem_available + + if cumem_available: + allocator = CuMemAllocator.get_instance() + allocator.release_pools() + logger.debug("[OmniEngine] Released CuMem memory pool during shutdown") + except Exception: + pass + + def _try_shutdown(self, *args, **kwargs) -> None: + try: + self.shutdown() + except Exception: + logger.exception(*args, **kwargs) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 75620bec9..aced2c1b3 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -18,7 +18,7 @@ import asyncio import os import shutil import time as _time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from dataclasses import dataclass, field from http import HTTPStatus from typing import Any @@ -34,28 +34,10 @@ from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker -from vllm_omni.engine.duplex.contracts import ( - DuplexControlPlanePort, - DuplexOutputContext, - DuplexOutputDecision, - DuplexRequestIdentity, - DuplexRuntimeExtension, - DuplexStageRequestContext, - DuplexStageSubmission, - DuplexStageSubmissionResult, -) -from vllm_omni.engine.duplex.control_plane import DuplexControlPlane -from vllm_omni.engine.duplex.lease import DuplexLeaseConfig -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.session import ( - DuplexSessionRuntimeManager, - DuplexSessionRuntimeState, -) from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -82,7 +64,6 @@ from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import attach_duplex_output_decision logger = init_logger(__name__) @@ -224,9 +205,10 @@ class OrchestratorRequestState: # Per-request pipeline timing accumulator (milliseconds) pipeline_timings: dict[str, float] = field(default_factory=dict) - duplex_identity: DuplexRequestIdentity | None = None - duplex_stage_fences: dict[int, DuplexFence] = field(default_factory=dict) - duplex_config_generation: int = -1 + # A session-owned request's lifetime belongs to a session (duplex), not to + # its own ``finished`` flag: no synthetic terminal, no client emission of + # stage-0 segment ends, no terminal re-forward, no cleanup on finish. + session_owned: bool = False running_counter_registered: bool = False request_artifact_dirs: set[str] = field(default_factory=set) @@ -279,143 +261,14 @@ class StreamingInputState: return segment if segment is not None else StreamingSegmentState() -class _OrchestratorDuplexStagePort: - """Adapts generic stage pools to the model-neutral duplex control plane.""" +class OrchestratorBase: + """Stage-management loop shared by the turn-based and duplex orchestrators. - def __init__( - self, - *, - stage_pools: list[StagePool], - request_states: dict[str, OrchestratorRequestState], - running_counter: OmniRequestCounter | None, - cleanup_request_ids: Callable[..., Any], - async_chunk: bool, - prewarm_async_chunk_stages: Callable[ - [str, Any, OrchestratorRequestState], - Awaitable[bool], - ], - ) -> None: - self._stage_pools = stage_pools - self._request_states = request_states - self._running_counter = running_counter - self._cleanup_request_ids = cleanup_request_ids - self._async_chunk = async_chunk - self._prewarm_async_chunk_stages = prewarm_async_chunk_stages - - @property - def stage_count(self) -> int: - return len(self._stage_pools) - - def sampling_defaults(self) -> tuple[object, ...]: - return tuple(pool.stage_client.default_sampling_params for pool in self._stage_pools) - - @staticmethod - def _sync_bridge_state( - request_state: OrchestratorRequestState, - context: DuplexStageRequestContext, - ) -> None: - duplex_state = request_state.streaming.bridge_states.setdefault("duplex", {}) - if not isinstance(duplex_state, dict): - duplex_state = {} - request_state.streaming.bridge_states["duplex"] = duplex_state - previous_epoch = duplex_state.get("epoch") - current_model_turn_id = duplex_state.get("model_turn_id") - if ( - not isinstance(current_model_turn_id, int) - or previous_epoch != context.fence.epoch - or current_model_turn_id < context.fence.turn_id - ): - # A serving-side safety boundary can close a Realtime response - # before the model emits its normal turn_eos. Catch up the - # engine-owned identity when the next fenced append starts. - duplex_state["model_turn_id"] = context.fence.turn_id - duplex_state.update( - { - "session_id": context.session_id, - "fence": context.fence, - "incarnation": context.fence.incarnation, - "epoch": context.fence.epoch, - "turn_id": context.fence.turn_id, - "response_seq": context.fence.response_seq, - "session_config": dict(context.session_config), - "runtime_config": dict(context.runtime_config), - } - ) - - def ensure_request(self, context: DuplexStageRequestContext) -> None: - request_state = self._request_states.get(context.request_id) - if request_state is None: - request_state = OrchestratorRequestState( - request_id=context.request_id, - prompt=None, - sampling_params_list=list(context.sampling_params), - final_stage_id=context.final_stage_id, - duplex_config_generation=context.config_generation, - ) - request_state.streaming.enabled = True - self._request_states[context.request_id] = request_state - elif request_state.duplex_config_generation != context.config_generation: - request_state.sampling_params_list = list(context.sampling_params) - request_state.duplex_config_generation = context.config_generation - request_state.duplex_identity = DuplexRequestIdentity( - session_id=context.session_id, - fence=context.fence, - ) - self._sync_bridge_state(request_state, context) - - async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: - context = submission.context - request_state = self._request_states.get(context.request_id) - if request_state is None: - raise RuntimeError(f"duplex request was not preregistered: {context.request_id}") - request = build_engine_core_request_from_tokens( - request_id=context.request_id, - prompt=dict(submission.prompt), - params=context.stage_sampling_params, - model_config=self._stage_pools[context.stage_id].stage_vllm_config.model_config, - resumable=True, - ) - request.external_req_id = request.request_id - pool = self._stage_pools[context.stage_id] - if submission.already_submitted: - replica_id = await pool.submit_update(context.request_id, request_state, request) - else: - replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None) - if self._async_chunk and context.stage_id == 0: - prewarmed = await self._prewarm_async_chunk_stages( - context.request_id, - request, - request_state, - ) - if not prewarmed: - # The prewarm already failed the request, aborted it, and - # popped its state. Falling through would write fences and - # a submit timestamp onto that orphaned object and - # re-increment the running counter the cleanup just - # released, then hand the control plane a success result - # for a request that no longer exists. Raise instead: - # ``DuplexControlPlane.handle_append`` turns this into an - # error result for the append. - raise RuntimeError( - f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted" - ) - request_state.duplex_stage_fences[context.stage_id] = context.fence - request_state.stage_submit_ts[context.stage_id] = _time.time() - if not request_state.running_counter_registered and self._running_counter is not None: - self._running_counter.increment() - request_state.running_counter_registered = True - return DuplexStageSubmissionResult( - request_id=context.request_id, - stage_id=context.stage_id, - replica_id=replica_id, - ) - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - await self._cleanup_request_ids(request_ids, abort=abort) - - -class Orchestrator: - """Runs inside a background thread's asyncio event loop.""" + Runs inside a background thread's asyncio event loop. Subclasses fill the + template seams (``_dispatch_message``, ``_background_tasks``, + ``_shutdown_extensions``, ``_on_stage_submitted``, ``_intercept_stage_output``, + ``_handle_forward_failure``, ``_cleanup_request_ids``). + """ # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when transfer / counter emit paths access them. @@ -424,7 +277,6 @@ class Orchestrator: _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None - duplex_control_plane: DuplexControlPlanePort | None = None def __init__( self, @@ -442,9 +294,6 @@ class Orchestrator: prom_metrics: Any = None, log_stats: bool = False, enable_orch_monitor: bool = False, - duplex_runtime_extension: DuplexRuntimeExtension | None = None, - enable_duplex_control: bool = False, - duplex_session_config: DuplexSessionRuntimeConfig | None = None, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -485,32 +334,6 @@ class Orchestrator: self._cfg_tracker = CfgCompanionTracker() self._stage_input_processors: dict[int, Any] = {} - self.duplex_control_plane: DuplexControlPlanePort | None = None - self._duplex_reaper_interval_s = 1.0 - if enable_duplex_control: - runtime_session_config = duplex_session_config or DuplexSessionRuntimeConfig() - self._duplex_reaper_interval_s = runtime_session_config.reaper_interval_s - - self.duplex_control_plane = DuplexControlPlane( - extension=duplex_runtime_extension, - stage_port=_OrchestratorDuplexStagePort( - stage_pools=self.stage_pools, - request_states=self.request_states, - running_counter=self._running_counter, - cleanup_request_ids=self._cleanup_request_ids, - async_chunk=self.async_chunk, - prewarm_async_chunk_stages=self._prewarm_async_chunk_stages, - ), - result_sink=self.rpc_async_queue, - lifecycle_sink=self.output_async_queue, - lease_config=DuplexLeaseConfig( - idle_ttl_s=runtime_session_config.idle_ttl_s, - disconnect_grace_s=runtime_session_config.disconnect_grace_s, - ), - max_sessions=runtime_session_config.max_sessions, - completed_append_limit=runtime_session_config.completed_append_cache_size, - ) - self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._event_driven_orch = _event_driven_orch_enabled() @@ -588,17 +411,6 @@ class Orchestrator: logger.exception("[Orchestrator] OmniPrometheusStatLogger init failed; metrics wrap disabled") self._stat_logger = None - @property - def duplex_sessions(self) -> DuplexSessionRuntimeManager: - if self.duplex_control_plane is None: - raise RuntimeError("duplex control plane is disabled") - return self.duplex_control_plane.sessions - - def _require_duplex_control_plane(self) -> DuplexControlPlanePort: - if self.duplex_control_plane is None: - raise RuntimeError("duplex control plane is disabled") - return self.duplex_control_plane - async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" logger.info("[Orchestrator] Starting event loop") @@ -620,8 +432,8 @@ class Orchestrator: membership_watcher = self._membership.start() tasks = [request_task, output_task] - if self.duplex_control_plane is not None: - tasks.append(asyncio.create_task(self._duplex_reaper_loop(), name="orchestrator-duplex-reaper")) + for index, coro in enumerate(self._background_tasks()): + tasks.append(asyncio.create_task(coro, name=f"orchestrator-background-{index}")) if membership_watcher is not None: tasks.append(membership_watcher) @@ -655,8 +467,7 @@ class Orchestrator: await asyncio.gather(*tasks, return_exceptions=True) except Exception: pass - if self.duplex_control_plane is not None: - await self.duplex_control_plane.shutdown() + await self._shutdown_extensions() if self._membership is not None: await self._membership.drain_tasks(timeout=10.0) @@ -680,18 +491,10 @@ class Orchestrator: msg = await self.request_async_queue.get() msg_type = msg.type - if msg_type == "add_request": - await self._handle_add_request(msg) - elif msg_type == "streaming_update": - await self._handle_streaming_update(msg) - elif msg_type == "add_companion_request": - await self._handle_add_companion(msg) - elif self.duplex_control_plane is not None and self.duplex_control_plane.accepts(msg): - self.duplex_control_plane.dispatch(msg) - elif msg_type == "abort": + if await self._dispatch_message(msg): + continue + if msg_type == "abort": await self._handle_abort(msg) - elif msg_type == "interaction": - await self._handle_interaction(msg) elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) elif isinstance(msg, RegisterRemoteReplicaMessage): @@ -717,183 +520,58 @@ class Orchestrator: else: logger.warning("[Orchestrator] Unknown message type: %s", msg_type) - async def _duplex_reaper_loop(self) -> None: - while not self._shutdown_event.is_set(): - try: - await asyncio.wait_for( - self._shutdown_event.wait(), - timeout=self._duplex_reaper_interval_s, - ) - except TimeoutError: - plane = self.duplex_control_plane - if plane is not None: - try: - await plane.reap_expired() - except Exception: - logger.exception("[Orchestrator] Duplex expiry cleanup failed; retrying on next tick") + # ---- Template seams filled by the turn-based / duplex subclasses ---- - async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: - """Handle an add_request message from the main thread.""" - stage_id = 0 - request_id = msg.request_id - prompt = msg.prompt - original_prompt = msg.original_prompt - sampling_params_list = msg.sampling_params_list - if not sampling_params_list: - raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") - final_stage_id = msg.final_stage_id - final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) + async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: + """Handle a subclass-specific request-queue message; return True when consumed.""" + del msg + return False - if not self.stage_pools[stage_id].live_replica_ids(): - # Stage 0 lost all replicas between the HTTP-layer errored check and - # dispatch. Runs before request state / running counter registration, - # so the helper's cleanup is a no-op here. - await self._fail_request_dead_stage(request_id, stage_id) - return + def _background_tasks(self) -> list[Coroutine[Any, Any, None]]: + """Extra coroutines gathered with the request/output handlers in ``run()``.""" + return [] - logger.debug( - "[Orchestrator] _handle_add_request: stage=%s req=%s " - "prompt_type=%s original_prompt_type=%s final_stage=%s " - "num_sampling_params=%d", - stage_id, - request_id, - type(prompt).__name__, - type(original_prompt).__name__, - final_stage_id, - len(sampling_params_list), - ) + async def _shutdown_extensions(self) -> None: + """Runs in ``run()``'s ``finally`` after task cancellation, before membership drain.""" - req_state = OrchestratorRequestState( - request_id=request_id, - prompt=original_prompt, - sampling_params_list=sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - request_timestamp=float(msg.request_timestamp or _time.time()), - mm_features=getattr(prompt, "mm_features", None), - request_artifact_dirs=set(msg.request_artifact_dirs or ()), - ) - self.request_states[request_id] = req_state - self._register_running_request(req_state) - req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) - req_state.stage_submit_ts[stage_id] = _time.time() - enqueue_ts = msg.enqueue_ts - if enqueue_ts > 0: - req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0 - preprocess_ms = msg.preprocess_ms - if preprocess_ms > 0: - req_state.pipeline_timings["preprocess_ms"] = preprocess_ms - if not await self._dispatch_or_fail_request( - lambda: self.stage_pools[stage_id].submit_initial( - request_id, - req_state, - prompt, - prompt_text=msg.output_prompt_text, - ), - req_id=request_id, - stage_id=stage_id, - operation="add_request", - ): - return + def _on_stage_submitted( + self, + stage_id: int, + request_id: str, + replica_id: int, + req_state: OrchestratorRequestState, + ) -> None: + """Called after every successful ``submit_initial`` / ``submit_update``.""" + del stage_id, request_id, replica_id, req_state - if self.async_chunk and stage_id == 0 and final_stage_id > 0: - await self._prewarm_async_chunk_stages(request_id, prompt, req_state) + async def _intercept_stage_output( + self, + stage_id: int, + replica_id: int, + output: Any, + req_state: OrchestratorRequestState, + stage_metrics: Any, + submit_ts: float | None, + ) -> bool: + """Offer a processed output to the subclass before forwarding. - async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: - """Handle a streaming_update message for an existing request.""" - stage_id = 0 - request_id = msg.request_id - request = msg.prompt - final_stage_id = msg.final_stage_id - req_state = self.request_states.get(request_id) - if req_state is None: - # Streaming updates always follow the first-chunk add_request - # through the same ordered submission queue, so an unknown id - # here can only mean the request already finished or was - # aborted (e.g. the client disconnected). Re-adding it would - # resurrect a headless session that keeps cycling through the - # stages with nobody consuming its outputs (issue #4271). - logger.warning( - "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)", - request_id, - ) - return + Return True when the output must not be forwarded to the next stage + (the subclass consumed it). Client emission for session-owned requests + is already suppressed by ``OrchestratorRequestState.session_owned``. + """ + del stage_id, replica_id, output, req_state, stage_metrics, submit_ts + return False - if msg.sampling_params_list: - req_state.sampling_params_list = msg.sampling_params_list - - req_state.streaming.enabled = True - req_state.stage_submit_ts[stage_id] = _time.time() - if not await self._dispatch_or_fail_request( - lambda: self.stage_pools[stage_id].submit_update( - request_id, - req_state, - request, - prompt_text=msg.output_prompt_text, - ), - req_id=request_id, - stage_id=stage_id, - operation="streaming_update", - ): - return - - if self.async_chunk and stage_id == 0 and final_stage_id > 0: - await self._prewarm_async_chunk_stages(request_id, request, req_state) - - async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: - """Handle an add_companion_request message: submit companion to stage 0.""" - companion_id = msg.companion_id - parent_id = msg.parent_id - role = msg.role - companion_prompt = msg.prompt - sampling_params_list = msg.sampling_params_list - - parent_state = self.request_states.get(parent_id) - if parent_state is None: - logger.info( - "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", - companion_id, - role, - parent_id, - ) - return - - self._cfg_tracker.register_companion(parent_id, role, companion_id) - - companion_state = OrchestratorRequestState( - request_id=companion_id, - prompt=companion_prompt, - sampling_params_list=sampling_params_list, - final_stage_id=0, - final_output_stage_ids={0}, - request_timestamp=parent_state.request_timestamp, - ) - self.request_states[companion_id] = companion_state - companion_state.stage_submit_ts[0] = _time.time() - # A dead-stage failure is attributed to the parent request; its cleanup - # also releases this companion. - if not await self._dispatch_or_fail_request( - lambda: self.stage_pools[0].submit_initial( - companion_id, - companion_state, - companion_prompt, - prompt_text=msg.companion_prompt_text, - affinity_request_id=parent_id, - ), - req_id=parent_id, - stage_id=0, - operation=f"companion {companion_id}", - dispatch_req_id=companion_id, - ): - return - - logger.info( - "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", - companion_id, - role, - parent_id, - self.stage_pools[0].get_bound_replica_id(companion_id), - ) + async def _handle_forward_failure( + self, + req_id: str, + next_stage_id: int, + req_state: OrchestratorRequestState, + exc: BaseException, + ) -> bool: + """Return True when a next-stage input-processing failure was absorbed (else it re-raises).""" + del req_id, next_stage_id, req_state, exc + return False async def _handle_abort(self, msg: AbortRequestMessage) -> None: """Handle an abort message from the main thread. @@ -935,44 +613,6 @@ class Orchestrator: for output_msg in abort_outputs: await self.output_async_queue.put(output_msg) - async def _handle_interaction(self, msg: InteractionMessage) -> None: - """Handle a midway interaction for an active streaming diffusion request.""" - stage_id = 0 - request_id = msg.request_id - event_id = msg.interaction.get("event_id") - req_state = self.request_states.get(request_id) - if req_state is None: - logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) - await self.output_async_queue.put( - ErrorMessage( - error=f"No active request for interaction: {request_id}", - fatal=False, - request_id=request_id, - event_id=event_id, - stage_id=stage_id, - ) - ) - return - - try: - await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) - except Exception as exc: - logger.info( - "[Orchestrator] Failed interaction for req %s: %s", - request_id, - exc, - exc_info=True, - ) - await self.output_async_queue.put( - ErrorMessage( - error=f"Failed interaction for request {request_id}: {exc}", - fatal=False, - request_id=request_id, - event_id=event_id, - stage_id=stage_id, - ) - ) - async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]: """Forward abort requests to all stage pools. @@ -1539,7 +1179,7 @@ class Orchestrator: await self._cleanup_request_ids( [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], abort=True, - close_duplex_sessions=True, + release_owners=True, ) def _update_stage_replica_waiting(self, stage_id: int, replica_id: int, n_waiting: int) -> None: @@ -1615,6 +1255,13 @@ class Orchestrator: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, + # A session-owned request has an owner that outlives the stage + # request: without this the runner and its admission slot + # survive the replica that was serving them, and the client is + # told nothing -- the request-scoped ErrorMessage has no + # frontend request_states entry to land in. No-op for + # turn-based requests, which have no owner. + release_owners=True, ) async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: @@ -1642,6 +1289,7 @@ class Orchestrator: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, + release_owners=True, ) async def _fail_request_client_error( @@ -1652,7 +1300,7 @@ class Orchestrator: *, status_code: int = HTTPStatus.BAD_REQUEST.value, error_type: str = DEFAULT_CLIENT_ERROR_TYPE, - close_duplex_sessions: bool = False, + release_owners: bool = False, ) -> None: """Fail one request with a non-fatal client error (400 by default). @@ -1672,7 +1320,7 @@ class Orchestrator: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, - close_duplex_sessions=close_duplex_sessions, + release_owners=release_owners, ) async def _dispatch_or_fail_request( @@ -1767,7 +1415,7 @@ class Orchestrator: request_ids: list[str], *, abort: bool = False, - close_duplex_sessions: bool = False, + release_owners: bool = False, ) -> list[OutputMessage]: """Release pool bindings and logical request state for the given ids. @@ -1778,10 +1426,14 @@ class Orchestrator: loss, membership unregister) funnels through here, so tracker state cannot outlive its requests. + ``release_owners`` asks the subclass to also release any session that + owns these requests (error paths pass True); the base ignores it. + Returns: Final-stage AR abort ``OutputMessage`` list when ``abort=True``; otherwise an empty list. """ + del release_owners if not request_ids: return [] @@ -1808,42 +1460,18 @@ class Orchestrator: if cid not in batch: batch.add(cid) cleanup_ids.append(cid) - closing_session_ids: list[str] = [] - if close_duplex_sessions and self.duplex_control_plane is not None: - closed_sessions = self.duplex_control_plane.close_sessions_for_request_ids( - cleanup_ids, - abort=abort, - cleanup_in_progress=True, - ) - closing_session_ids.extend(closed_sessions) - for session_id, stale_request_ids in closed_sessions.items(): - logger.info( - "[Orchestrator] closed duplex session %s while cleaning failed request ids %s", - session_id, - stale_request_ids, - ) - cleanup_ids.extend(stale_request_ids) - cleanup_ids = list(dict.fromkeys(cleanup_ids)) - abort_outputs: list[OutputMessage] = [] - try: - if abort: - abort_outputs = await self._abort_request_ids(cleanup_ids) - self._release_request_bindings(cleanup_ids) - for request_id in cleanup_ids: - self._pd_kv_params.pop(request_id, None) - req_state = self.request_states.pop(request_id, None) - if req_state is not None: - cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) - if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: - self._running_counter.decrement() - req_state.running_counter_registered = False - except BaseException: - if closing_session_ids and self.duplex_control_plane is not None: - self.duplex_control_plane.defer_request_cleanups(closing_session_ids) - raise - if closing_session_ids and self.duplex_control_plane is not None: - self.duplex_control_plane.finalize_closed_sessions(closing_session_ids) + if abort: + abort_outputs = await self._abort_request_ids(cleanup_ids) + self._release_request_bindings(cleanup_ids) + for request_id in cleanup_ids: + self._pd_kv_params.pop(request_id, None) + req_state = self.request_states.pop(request_id, None) + if req_state is not None: + cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) + if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: + self._running_counter.decrement() + req_state.running_counter_registered = False return abort_outputs async def _apply_raw_terminal_stage_finish( @@ -1888,7 +1516,7 @@ class Orchestrator: for request_id in request_ids: req_state = self.request_states.get(request_id) - if req_state is None or self._is_duplex_session_request(req_state): + if req_state is None or req_state.session_owned: continue final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} @@ -1937,28 +1565,6 @@ class Orchestrator: params.cfg_kv_request_ids = companion_request_ids return params - def _duplex_session_for_req_state(self, req_state: OrchestratorRequestState) -> DuplexSessionRuntimeState | None: - if self.duplex_control_plane is None: - return None - return self.duplex_control_plane.session_for_identity(req_state.duplex_identity) - - def _record_duplex_stage_submission( - self, - stage_id: int, - request_id: str, - replica_id: int, - req_state: OrchestratorRequestState, - ) -> None: - del replica_id - identity = req_state.duplex_identity - session = self._duplex_session_for_req_state(req_state) - if identity is None or session is None: - return - req_state.duplex_stage_fences[stage_id] = identity.fence - session.bind_stage_request(stage_id, request_id, fence=identity.fence) - req_state.stage_submit_ts[stage_id] = _time.time() - self._register_running_request(req_state) - def _register_running_request(self, req_state: OrchestratorRequestState) -> None: if req_state.running_counter_registered or self._running_counter is None: return @@ -1991,17 +1597,11 @@ class Orchestrator: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) - # Duplex stage-0 segment boundaries are not client-visible outputs: - # direct decisions are emitted by the model runtime extension below, - # while spoken content flows through the next stage. Forwarding - # the raw stage-0 output as well injects one cumulative-text, - # no-audio message per unit that every downstream consumer must - # filter out again (the official implementation returns exactly one - # result per audio chunk). - is_duplex_stage0_segment = ( - stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished - ) - if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment: + if req_state.session_owned: + # Session-owned outputs are delivered to their session through + # ``_intercept_stage_output`` below, never to a request queue. + pass + elif self.stage_pools[stage_id].final_output: await self.output_async_queue.put( OutputMessage( request_id=req_id, @@ -2009,13 +1609,7 @@ class Orchestrator: replica_id=replica_id, engine_outputs=output, metrics=stage_metrics, - finished=( - request_finished - or ( - self._is_duplex_session_request(req_state) - and req_state.streaming.segment(stage_id).finished - ) - ), + finished=request_finished, stage_submit_ts=submit_ts, ) ) @@ -2036,16 +1630,7 @@ class Orchestrator: self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params) req_state.pd_prefill_multimodal_output = getattr(output, "multimodal_output", None) - duplex_output_decision = self._duplex_output_decision(stage_id, output, req_state) - if duplex_output_decision is not None: - await self._emit_duplex_direct_output( - stage_id, - req_id, - output, - duplex_output_decision, - stage_metrics, - submit_ts, - ) + if await self._intercept_stage_output(stage_id, replica_id, output, req_state, stage_metrics, submit_ts): return if ( @@ -2082,7 +1667,7 @@ class Orchestrator: and self.request_states.get(req_id) is req_state and finished and not final_only_finished - and not self._is_duplex_session_request(req_state) + and not req_state.session_owned ): # For streaming sessions, send the terminal (resumable=False) update only on a finish await self._forward_to_next_stage( @@ -2095,7 +1680,7 @@ class Orchestrator: is_final_update=True, ) - if request_finished and not self._is_duplex_session_request(req_state): + if request_finished and not req_state.session_owned: await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: @@ -2189,92 +1774,6 @@ class Orchestrator: request.payload_sender_info = payload_sender_info return request - @staticmethod - def _duplex_output_context( - req_state: OrchestratorRequestState, - *, - # Required: segment state is stage-keyed, so a caller that omitted this - # would silently read "no boundary" rather than the intended stage's. - stage_id: int | None, - ) -> DuplexOutputContext | None: - identity = req_state.duplex_identity - if identity is None: - return None - fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence - segment = req_state.streaming.segment(stage_id) - return DuplexOutputContext( - identity=DuplexRequestIdentity( - session_id=identity.session_id, - fence=fence, - ), - final_stage_id=req_state.final_stage_id, - segment_finished=req_state.streaming.enabled and segment.finished, - segment_token_ids=tuple(segment.token_ids), - segment_output_metadata=segment.output_metadata, - ) - - @staticmethod - def _is_duplex_session_request(req_state: OrchestratorRequestState) -> bool: - return req_state.duplex_identity is not None - - @staticmethod - def _duplex_fence_for_req_state( - req_state: OrchestratorRequestState, - *, - stage_id: int | None = None, - ) -> DuplexFence | None: - context = Orchestrator._duplex_output_context(req_state, stage_id=stage_id) - return context.identity.fence if context is not None else None - - def _duplex_output_decision( - self, - stage_id: int, - output: Any, - req_state: OrchestratorRequestState, - ) -> DuplexOutputDecision | None: - if self.duplex_control_plane is None: - return None - context = self._duplex_output_context(req_state, stage_id=stage_id) - decision = self.duplex_control_plane.decide_output( - stage_id, - output, - context, - ) - return decision - - async def _emit_duplex_direct_output( - self, - stage_id: int, - req_id: str, - output: Any, - decision: DuplexOutputDecision, - stage_metrics: Any, - submit_ts: float | None, - ) -> None: - action = getattr(decision.action, "value", decision.action) - if action != "direct_response": - raise ValueError(f"Unsupported duplex output action: {action}") - engine_output = attach_duplex_output_decision( - OmniRequestOutput.from_stage_output( - output, - request_id=req_id, - finished=True, - stage_id=stage_id, - final_output_type=decision.final_output_type, - ), - decision, - ) - await self.output_async_queue.put( - OutputMessage( - request_id=req_id, - stage_id=stage_id, - engine_outputs=engine_output, - metrics=stage_metrics, - finished=True, - stage_submit_ts=submit_ts, - ) - ) - @staticmethod def _completion_multimodal_output(output: Any, completion: Any) -> dict[str, Any]: mm_output = getattr(output, "multimodal_output", None) @@ -2653,7 +2152,7 @@ class Orchestrator: }, params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_logical, req_id, replica_id, @@ -2706,7 +2205,7 @@ class Orchestrator: replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_logical, req_id, replica_id, @@ -2751,7 +2250,7 @@ class Orchestrator: str(exc), status_code=exc.status_code, error_type=exc.error_type, - close_duplex_sessions=self._is_duplex_session_request(req_state), + release_owners=req_state.session_owned, ) return except Exception as exc: @@ -2760,21 +2259,8 @@ class Orchestrator: req_id, next_logical, ) - if not self._is_duplex_session_request(req_state): + if not await self._handle_forward_failure(req_id, next_logical, req_state, exc): raise - await self.output_async_queue.put( - ErrorMessage( - request_id=req_id, - stage_id=next_logical, - error=f"Stage-{next_logical} input processor failed: {type(exc).__name__}: {exc}", - error_type=type(exc).__name__, - ) - ) - await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], - abort=True, - close_duplex_sessions=True, - ) return finally: req_state.streaming.source_token_decoder = previous_decoder @@ -2842,7 +2328,7 @@ class Orchestrator: replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_logical, req_id, replica_id, @@ -2900,7 +2386,7 @@ class Orchestrator: 0, "async_chunk requires prompt_token_ids on the stage-0 request; " "an embeds-only prompt cannot prewarm downstream stages", - close_duplex_sessions=True, + release_owners=True, ) return False @@ -2996,7 +2482,7 @@ class Orchestrator: # The guard swallows submit_initial's return value; the pool binding # it recorded carries the same replica. bound_replica_id = next_pool.get_bound_replica_id(request_id) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_stage_id, request_id, bound_replica_id if bound_replica_id is not None else 0, @@ -3081,3 +2567,222 @@ class Orchestrator: for pool in self.stage_pools: for replica_id in pool.live_replica_ids(): pool.shutdown_replica(replica_id) + + +class Orchestrator(OrchestratorBase): + """Turn-based orchestrator: admits ``add_request`` / streaming / companion / interaction messages.""" + + async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: + msg_type = msg.type + if msg_type == "add_request": + await self._handle_add_request(msg) + elif msg_type == "streaming_update": + await self._handle_streaming_update(msg) + elif msg_type == "add_companion_request": + await self._handle_add_companion(msg) + elif msg_type == "interaction": + await self._handle_interaction(msg) + else: + return False + return True + + async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: + """Handle an add_request message from the main thread.""" + stage_id = 0 + request_id = msg.request_id + prompt = msg.prompt + original_prompt = msg.original_prompt + sampling_params_list = msg.sampling_params_list + if not sampling_params_list: + raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") + final_stage_id = msg.final_stage_id + final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) + + if not self.stage_pools[stage_id].live_replica_ids(): + # Stage 0 lost all replicas between the HTTP-layer errored check and + # dispatch. Runs before request state / running counter registration, + # so the helper's cleanup is a no-op here. + await self._fail_request_dead_stage(request_id, stage_id) + return + + logger.debug( + "[Orchestrator] _handle_add_request: stage=%s req=%s " + "prompt_type=%s original_prompt_type=%s final_stage=%s " + "num_sampling_params=%d", + stage_id, + request_id, + type(prompt).__name__, + type(original_prompt).__name__, + final_stage_id, + len(sampling_params_list), + ) + + req_state = OrchestratorRequestState( + request_id=request_id, + prompt=original_prompt, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + request_timestamp=float(msg.request_timestamp or _time.time()), + mm_features=getattr(prompt, "mm_features", None), + request_artifact_dirs=set(msg.request_artifact_dirs or ()), + ) + self.request_states[request_id] = req_state + self._register_running_request(req_state) + req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) + req_state.stage_submit_ts[stage_id] = _time.time() + enqueue_ts = msg.enqueue_ts + if enqueue_ts > 0: + req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0 + preprocess_ms = msg.preprocess_ms + if preprocess_ms > 0: + req_state.pipeline_timings["preprocess_ms"] = preprocess_ms + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[stage_id].submit_initial( + request_id, + req_state, + prompt, + prompt_text=msg.output_prompt_text, + ), + req_id=request_id, + stage_id=stage_id, + operation="add_request", + ): + return + + if self.async_chunk and stage_id == 0 and final_stage_id > 0: + await self._prewarm_async_chunk_stages(request_id, prompt, req_state) + + async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: + """Handle a streaming_update message for an existing request.""" + stage_id = 0 + request_id = msg.request_id + request = msg.prompt + final_stage_id = msg.final_stage_id + req_state = self.request_states.get(request_id) + if req_state is None: + # Streaming updates always follow the first-chunk add_request + # through the same ordered submission queue, so an unknown id + # here can only mean the request already finished or was + # aborted (e.g. the client disconnected). Re-adding it would + # resurrect a headless session that keeps cycling through the + # stages with nobody consuming its outputs (issue #4271). + logger.warning( + "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)", + request_id, + ) + return + + if msg.sampling_params_list: + req_state.sampling_params_list = msg.sampling_params_list + + req_state.streaming.enabled = True + req_state.stage_submit_ts[stage_id] = _time.time() + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[stage_id].submit_update( + request_id, + req_state, + request, + prompt_text=msg.output_prompt_text, + ), + req_id=request_id, + stage_id=stage_id, + operation="streaming_update", + ): + return + + if self.async_chunk and stage_id == 0 and final_stage_id > 0: + await self._prewarm_async_chunk_stages(request_id, request, req_state) + + async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: + """Handle an add_companion_request message: submit companion to stage 0.""" + companion_id = msg.companion_id + parent_id = msg.parent_id + role = msg.role + companion_prompt = msg.prompt + sampling_params_list = msg.sampling_params_list + + parent_state = self.request_states.get(parent_id) + if parent_state is None: + logger.info( + "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", + companion_id, + role, + parent_id, + ) + return + + self._cfg_tracker.register_companion(parent_id, role, companion_id) + + companion_state = OrchestratorRequestState( + request_id=companion_id, + prompt=companion_prompt, + sampling_params_list=sampling_params_list, + final_stage_id=0, + final_output_stage_ids={0}, + request_timestamp=parent_state.request_timestamp, + ) + self.request_states[companion_id] = companion_state + companion_state.stage_submit_ts[0] = _time.time() + # A dead-stage failure is attributed to the parent request; its cleanup + # also releases this companion. + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[0].submit_initial( + companion_id, + companion_state, + companion_prompt, + prompt_text=msg.companion_prompt_text, + affinity_request_id=parent_id, + ), + req_id=parent_id, + stage_id=0, + operation=f"companion {companion_id}", + dispatch_req_id=companion_id, + ): + return + + logger.info( + "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", + companion_id, + role, + parent_id, + self.stage_pools[0].get_bound_replica_id(companion_id), + ) + + async def _handle_interaction(self, msg: InteractionMessage) -> None: + """Handle a midway interaction for an active streaming diffusion request.""" + stage_id = 0 + request_id = msg.request_id + event_id = msg.interaction.get("event_id") + req_state = self.request_states.get(request_id) + if req_state is None: + logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) + await self.output_async_queue.put( + ErrorMessage( + error=f"No active request for interaction: {request_id}", + fatal=False, + request_id=request_id, + event_id=event_id, + stage_id=stage_id, + ) + ) + return + + try: + await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) + except Exception as exc: + logger.info( + "[Orchestrator] Failed interaction for req %s: %s", + request_id, + exc, + exc_info=True, + ) + await self.output_async_queue.put( + ErrorMessage( + error=f"Failed interaction for request {request_id}: {exc}", + fatal=False, + request_id=request_id, + event_id=event_id, + stage_id=stage_id, + ) + ) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 993f0ca1e..6114bbea2 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -11,7 +11,6 @@ StageEngineCoreClient instances) instead of OmniStage with worker processes. from __future__ import annotations import asyncio -import os import time import uuid from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence @@ -21,31 +20,17 @@ from vllm import TokensPrompt from vllm.engine.protocol import EngineClient, StreamingInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest -from vllm.outputs import CompletionOutput, PoolingRequestOutput -from vllm.plugins.io_processors import get_io_processor +from vllm.outputs import PoolingRequestOutput from vllm.pooling_params import PoolingParams from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tasks import SupportedTask -from vllm.utils import random_uuid -from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import ( - DuplexFence, - DuplexSessionLifecycleMessage, -) -from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.messages import ErrorMessage +from vllm_omni.entrypoints.async_omni_base import ABORT_TIMEOUT_S, AsyncOmniBase from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.duplex_request_client import ( - DuplexRequestClient, - DuplexRequestOutputPort, -) -from vllm_omni.entrypoints.omni_base import ( - OmniBase, - OmniEngineDeadError, -) from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics @@ -59,74 +44,9 @@ if TYPE_CHECKING: from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType logger = init_logger(__name__) -_FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 -# Blocking-wait interval for the event-driven final-output drain -# (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): a message wakes the drain immediately via -# the janus queue's condition variable; this timeout only bounds how often the -# orchestrator liveness check runs while the pipeline is idle. -_FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 -# Shared DELETE / generate() cleanup abort bound. Env is the documented knob. -ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) -class AsyncEventResolver: - """ - A generic signal aggregator designed for synchronized handshakes in - distributed or multi-stage environments. Supports waiting for a specified - number (expected_count) of worker signals in both inline and multiprocess modes. - """ - - def __init__(self, orchestrator=None): - self._pending_tasks: dict[str, dict] = {} - self.orchestrator = orchestrator - self._lock = asyncio.Lock() - - def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: - loop = asyncio.get_running_loop() - fut = loop.create_future() - self._pending_tasks[task_id] = { - "future": fut, - "expected_count": expected_count, - "received": [], - "start_time": time.time(), - } - return fut - - async def resolve(self, ack: OmniACK): - tid = getattr(ack, "task_id", None) - - if tid is None and isinstance(ack, dict): - tid = ack.get("task_id") - - async with self._lock: - task_info = self._pending_tasks.get(tid) - if task_info is None: - logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.") - return - - task_info["received"].append(ack) - current_count = len(task_info["received"]) - expected = task_info["expected_count"] - - orchestrator = self.orchestrator - if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics: - freed = getattr(ack, "freed_bytes", 0) - if freed == 0 and isinstance(ack, dict): - freed = ack.get("freed_bytes", 0) - orchestrator.metrics.record_vram_reclaimed(freed) - - logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.") - - if current_count >= expected: - self._pending_tasks.pop(tid) - fut = task_info["future"] - if not fut.done(): - elapsed = time.time() - task_info["start_time"] - logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.") - fut.set_result(task_info["received"]) - - -class AsyncOmni(EngineClient, OmniBase): +class AsyncOmni(AsyncOmniBase, EngineClient): """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine. This is the refactored version that uses AsyncOmniEngine instead of @@ -155,8 +75,13 @@ class AsyncOmni(EngineClient, OmniBase): ... print(output) """ - def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: - OmniBase.__init__(self, model=model, **kwargs) + def _create_engine(self, **engine_kwargs: Any) -> AsyncOmniEngine: + return AsyncOmniEngine(**engine_kwargs) + + def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None: + tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) + AsyncOmniBase.__init__(self, *args, model=model, **kwargs) + self.tts_max_instructions_length = tts_max_instructions_length self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False # In-flight EngineCore submits (non-streaming add_request, or each @@ -173,294 +98,6 @@ class AsyncOmni(EngineClient, OmniBase): self._sleeping_tags: set[str] = set() self._stage_sleeping_tags: dict[int, set[str]] = {} self._level2_sleeping: bool = False - self._duplex_request_client: DuplexRequestClient | None = None - self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue() - self.final_output_task: asyncio.Task | None = None - self.event_resolver = AsyncEventResolver(orchestrator=self) - self.config_path = self.engine.config_path - self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) - self.input_processor = self.engine.input_processor - self.endpoint_restrictions = self.engine.endpoint_restrictions - self.duplex_session_config = self.engine.duplex_session_config - self.duplex_serving_adapter_path = self.engine.duplex_serving_adapter_path - - stage_index = self._get_comprehension_stage_index() - if stage_index is None: - self.io_processor = None - else: - vllm_config = self.engine.stage_vllm_configs[stage_index] - io_processor_plugin = vllm_config.model_config.io_processor_plugin - renderer = self.renderer - if renderer is None: - from vllm.renderers import renderer_from_config - - renderer = renderer_from_config(vllm_config) - self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) - - def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None: - """Look up the sticky-routed replica for (stage_id, request_id). - - Used as the ``replica_resolver`` callback by ``OrchestratorAggregator`` - to label transfer_* metrics without plumbing replica ids through - ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters. - Returns None when stage_id is out of range or the request hasn't been - bound to a replica yet — the metric emit then defensive-skips. - """ - pools = getattr(self.engine, "stage_pools", None) - if pools is None or not (0 <= stage_id < len(pools)): - return None - return pools[stage_id].get_bound_replica_id(request_id) - - def _get_comprehension_stage_index(self) -> int | None: - fallback_idx: int | None = None - for idx, stage_client in enumerate(self.engine.stage_clients): - stage_vllm_config = self.engine.stage_vllm_configs[idx] - if stage_vllm_config is None: - continue - if fallback_idx is None: - fallback_idx = idx - if stage_client.is_comprehension: - return idx - return fallback_idx - - @property - def renderer(self): - """Return the renderer from the engine input processor when available.""" - if self.input_processor is None: - return None - return self.input_processor.renderer - - @property - def vllm_config(self): - """Return the vLLM config for the comprehension stage when present.""" - stage_index = self._get_comprehension_stage_index() - if stage_index is None: - return None - return self.engine.stage_vllm_configs[stage_index] - - async def get_vllm_config(self) -> Any: - """Compatibility helper for call sites expecting async vllm config access.""" - return self.vllm_config - - def get_diffusion_od_config(self) -> Any | None: - """Return the diffusion-stage config when the pipeline has one.""" - saw_diffusion_stage = False - for stage_client in self.engine.stage_clients: - if getattr(stage_client, "stage_type", None) != "diffusion": - continue - - saw_diffusion_stage = True - - od_config = getattr(stage_client, "od_config", None) - if od_config is not None: - return od_config - - inner_engine = getattr(stage_client, "_engine", None) - od_config = getattr(inner_engine, "od_config", None) - if od_config is not None: - return od_config - - # Out-of-process diffusion clients don't carry od_config (it lives in the - # worker); fall back to the engine's model_class_name resolution. - if saw_diffusion_stage: - return self.engine.get_diffusion_od_config() - - return None - - @property - def model_config(self): - """Return the model config for the comprehension stage when present.""" - vllm_config = self.vllm_config - if vllm_config is None: - return None - return vllm_config.model_config - - @staticmethod - def _get_unique_request_id(external_request_id: str): - """Get a random new request ID for this request; at the server level, - this is usually set by the calling entrypoint, but in direct calls, we - need to set it explicitly since we do not allow empty IDs. - - NOTE: in the upstream vLLM, this is done in the InputProcessor's - `assign_request_id`. - """ - uuid = random_uuid() - prefix = "" if not external_request_id else f"{external_request_id}-" - return f"{prefix}{uuid:.8}" - - async def open_duplex_session_async( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Open an engine-level duplex session when the backend supports it.""" - return await self._get_duplex_request_client().open( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - runtime_config=runtime_config, - fence=fence, - timeout=timeout, - ) - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - expected_epoch: int | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - collect_outputs: bool = True, - ) -> dict[str, object]: - """Append input to an engine-level duplex session.""" - return await self._get_duplex_request_client().append( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - expected_epoch=expected_epoch, - fence=fence, - timeout=timeout, - collect_outputs=collect_outputs, - ) - - async def collect_duplex_data_plane_outputs_async( - self, - request_id: str, - *, - response_stage_id: int | None = None, - timeout: float | None = 10.0, - ) -> list[OmniRequestOutput]: - """Collect the next duplex data-plane output batch for a live request.""" - return await self._get_duplex_request_client().collect_registered_outputs( - request_id, - response_stage_id=response_stage_id, - timeout=timeout, - ) - - async def signal_duplex_turn_async( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Send a turn/control signal to an engine-level duplex session.""" - return await self._get_duplex_request_client().signal( - session_id, - event=event, - fence=fence, - next_fence=next_fence, - session_config=session_config, - runtime_config=runtime_config, - timeout=timeout, - ) - - async def close_duplex_session_async( - self, - session_id: str, - *, - reason: str = "client_close", - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Close an engine-level duplex session.""" - return await self._get_duplex_request_client().close( - session_id, - reason=reason, - fence=fence, - timeout=timeout, - ) - - async def touch_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return await self._get_duplex_request_client().touch( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ) - - async def resume_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return await self._get_duplex_request_client().resume( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ) - - def _get_duplex_request_client(self) -> DuplexRequestClient: - client = getattr(self, "_duplex_request_client", None) - if client is None: - engine = getattr(self, "engine", None) - client = DuplexRequestClient( - engine, - DuplexRequestOutputPort( - request_states=getattr(self, "request_states", {}), - num_stages=getattr(engine, "num_stages", 1), - log_stats=getattr(self, "log_stats", False), - start_output_handler=self._final_output_handler, - process_single_result=self._process_single_result, - ), - ) - self._duplex_request_client = client - return client - - @staticmethod - def _duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]: - return DuplexRequestClient.request_info(result) - - async def _collect_duplex_data_plane_outputs( - self, - request_id: str, - req_state: ClientRequestState, - *, - response_stage_id: int | None, - timeout: float | None, - ) -> list[OmniRequestOutput]: - return await self._get_duplex_request_client().collect_outputs( - request_id, - req_state, - response_stage_id=response_stage_id, - timeout=timeout, - ) - - @classmethod - def _is_direct_duplex_data_plane_response(cls, output: object) -> bool: - return DuplexRequestClient.is_direct_response(output) - - @classmethod - def _duplex_multimodal_output(cls, output: object) -> dict[str, object]: - return DuplexRequestClient.multimodal_output(output) # ==================== Generate Method ==================== @@ -857,198 +494,6 @@ class AsyncOmni(EngineClient, OmniBase): """ raise NotImplementedError("AsyncOmni.encode is not implemented.") - # ==================== Processing Methods ==================== - - async def _process_orchestrator_results( - self, - request_id: str, - metrics: OrchestratorMetrics, - final_stage_id_for_e2e: int, - req_start_ts: dict[str, float], - wall_start_ts: float, - ) -> AsyncGenerator[OmniRequestOutput, None]: - """Read results from the Orchestrator (via the request's asyncio.Queue) - and yield OmniRequestOutput objects. - - The Orchestrator handles all stage-to-stage transfers. This method - only processes final outputs that arrive on the per-request queue. - """ - req_state = self.request_states.get(request_id) - if req_state is None: - return - - while True: - result = await req_state.queue.get() - - if isinstance(result, ErrorMessage): - logger.error( - "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s", - request_id, - result.stage_id, - result.error, - ) - if result.fatal: - raise OmniEngineDeadError( - result.error, - error_stage_id=result.stage_id, - ) - self._raise_nonfatal_error_message(result) - - if not isinstance(result, OutputMessage): - logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) - continue - - stage_id = result.stage_id - - self._check_engine_output_error(result, request_id, stage_id) - - # Process the result (constructs OmniRequestOutput) - output_to_yield = self._process_single_result( - result, - stage_id, - metrics, - req_start_ts, - wall_start_ts, - final_stage_id_for_e2e, - ) - - if output_to_yield: - # Set the external request ID back to the user yielded input - output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id - logger.debug( - "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s", - request_id, - stage_id, - getattr(output_to_yield, "final_output_type", None), - ) - yield output_to_yield - - # The Orchestrator sets "finished" when the final stage is done - if result.finished: - break - - # ==================== Output Handler ==================== - - def _final_output_handler(self) -> None: - """Start the final output handler if not already running. - - This handler reads messages from the Orchestrator output queue and - routes them to per-request asyncio.Queues. - """ - if self.final_output_task is not None: - return - - engine = self.engine - - # Event-driven drain (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): block on the - # queue's condition variable in a dedicated thread instead of the - # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side - # event-driven loop (vllm_omni/engine/orchestrator.py). - from vllm_omni.engine.orchestrator import _event_driven_orch_enabled - - event_driven_drain = _event_driven_orch_enabled() and hasattr(engine, "get_output_blocking_async") - - async def _final_output_loop(): - """Background coroutine that dispatches final outputs to request queues.""" - try: - while True: - if event_driven_drain: - msg = await engine.get_output_blocking_async(timeout=_FINAL_OUTPUT_BLOCKING_WAIT_S) - if msg is None: - # Timed out with the orchestrator alive; loop for - # the periodic liveness check. - continue - else: - msg = await engine.try_get_output_async() - if msg is None: - await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) - continue - - if isinstance(msg, dict) and msg.get("type") == "ack": - ack_data = msg.get("ack") - tid = getattr(ack_data, "task_id", "unknown") - logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}") - await self.event_resolver.resolve(ack_data) - continue - if isinstance(msg, OmniACK): - logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}") - await self.event_resolver.resolve(msg) - continue - if hasattr(msg, "task_id"): - tid = getattr(msg, "task_id") - logger.info(f"[{self._name}] Intercepted task-ID object: {tid}") - await self.event_resolver.resolve(msg) - continue - - if getattr(msg, "type", None) == "duplex_session_lifecycle": - await self.duplex_lifecycle_events.put(msg) - continue - - if isinstance(msg, ErrorMessage): - # Route request-scoped errors to that request's queue and - # keep the loop alive. A request whose stage replica died - # and was evicted gets a fatal error delivered here; only - # that request fails (its consumer raises), the server - # stays up for other stages/requests (#4285). A fatal - # error without a request_id is a genuine engine-wide - # death and falls through to the except handler below. - if msg.request_id is not None: - req_state = self.request_states.get(msg.request_id) - if req_state is not None: - await req_state.queue.put(msg) - else: - logger.warning( - "[%s] dropping error for unknown req %s", - self._name, - msg.request_id, - ) - continue - if not msg.fatal: - continue - - should_continue, _, stage_id, req_state = self._handle_output_message(msg) - if should_continue: - continue - - req_state.stage_id = stage_id - - # Route to the per-request queue - await req_state.queue.put(msg) - - except asyncio.CancelledError: - raise - except OmniEngineDeadError as e: - logger.error("[AsyncOmni] Engine dead: %s", e) - for req_state in list(self.request_states.values()): - error_msg = ErrorMessage( - error=str(e), - fatal=True, - request_id=req_state.request_id, - stage_id=e.error_stage_id, - ) - await req_state.queue.put(error_msg) - except EngineDeadError as e: - logger.error("[AsyncOmni] Engine dead: %s", e) - for req_state in list(self.request_states.values()): - error_msg = ErrorMessage( - error=str(e), - fatal=True, - request_id=req_state.request_id, - ) - await req_state.queue.put(error_msg) - except Exception as e: - logger.exception("[AsyncOmni] final_output_loop failed.") - for req_state in list(self.request_states.values()): - error_msg = ErrorMessage( - request_id=req_state.request_id, - error=str(e), - ) - await req_state.queue.put(error_msg) - self.final_output_task = None - - self.final_output_task = asyncio.create_task(_final_output_loop()) - logger.debug("[AsyncOmni] Final output handler started") - # ==================== Control Methods ==================== async def collective_rpc( @@ -1172,102 +617,6 @@ class AsyncOmni(EngineClient, OmniBase): if self.log_stats: logger.info("[AsyncOmni] Queued interaction for request %s", request_id) - async def _abort_internal_requests( - self, - request_id: str | Iterable[str], - *, - timeout: float = ABORT_TIMEOUT_S, - ): - """Abort request(s) via the Orchestrator given internal request IDs, - which take the format -. - """ - request_ids = [request_id] if isinstance(request_id, str) else list(request_id) - # Request IDs are already internal, so we just need to get the matching states. - internal_req_ids = [rid for rid in request_ids if rid in self.request_states] - try: - # Unbind generate() if abort_async blocks in the executor. - await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) - except TimeoutError: - logger.warning( - "[AsyncOmni] Timed out aborting %s after %.1fs; " - "engine abort is best-effort until the current batch drains", - ",".join(internal_req_ids), - timeout, - ) - except Exception: - logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) - - async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: - """Abort request IDs via the engine and enqueue terminal abort outputs. - - Waits for orchestrator abort acknowledgment, enqueues any AR terminal - abort outputs (partial tokens) into each request's asyncio queue, then - cancels the input pump. Frontend ``request_states`` stay registered so - ``generate()`` can consume the terminal message in - ``_process_orchestrator_results`` and run normal cleanup. - - When ``abort_async`` returns no output for an active request (OP not - registered yet, unbound replica, or orchestrator id drop), enqueue a - synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. - """ - abort_outputs = await self.engine.abort_async(request_ids, timeout=timeout) or [] - delivered: set[str] = set() - for output_msg in abort_outputs: - req_id = getattr(output_msg, "request_id", None) - if req_id is None: - continue - state = self.request_states.get(req_id) - if state is None: - logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) - continue - await state.queue.put(output_msg) - delivered.add(req_id) - for rid in request_ids: - state = self.request_states.get(rid) - if state is not None and rid not in delivered: - queue = getattr(state, "queue", None) - if queue is not None: - await state.queue.put(self._synthetic_abort_output_message(rid)) - delivered.add(rid) - for rid in request_ids: - self._record_request_failure_once(rid, reason="client_abort") - state = self.request_states.get(rid) - input_stream_task = getattr(state, "input_stream_task", None) - if input_stream_task is not None and not input_stream_task.done(): - input_stream_task.cancel() - if self.log_stats: - logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) - - @staticmethod - def _synthetic_abort_output_message(request_id: str) -> OutputMessage: - """Terminal abort OutputMessage used when the engine returned none.""" - engine_output = OmniRequestOutput( - request_id=request_id, - finished=True, - stage_id=0, - final_output_type="text", - outputs=[ - CompletionOutput( - index=0, - text="", - token_ids=[], - cumulative_logprob=None, - logprobs=None, - finish_reason="abort", - stop_reason=None, - ) - ], - ) - return OutputMessage( - request_id=request_id, - stage_id=0, - replica_id=None, - engine_outputs=engine_output, - metrics=None, - finished=True, - stage_submit_ts=None, - ) - def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" n_stages = len(self.engine.stage_clients) @@ -1697,42 +1046,6 @@ class AsyncOmni(EngineClient, OmniBase): # ==================== Properties ==================== - @property - def is_running(self) -> bool: - """Check if the engine is running.""" - orchestrator_alive = self.engine.is_alive() - task_alive = self.final_output_task is not None and not self.final_output_task.done() - return orchestrator_alive and task_alive - - @property - def errored(self) -> bool: - """Whether the engine is in a process-fatal error state. - - Delegates to ``OmniBase.errored``, which is true only when the - orchestrator thread is dead; per-stage liveness is reported via - ``check_health`` instead. Redeclared here to satisfy the - ``EngineClient`` abstract-property requirement (Python's ABC - mechanism does not resolve abstract methods from sibling MRO - entries). - """ - return OmniBase.errored.fget(self) # type: ignore[union-attr] - - @property - def _name(self) -> str: - return "AsyncOrchestrator" - - @property - def is_stopped(self) -> bool: - """EngineClient abstract property implementation.""" - return self.errored - - @property - def dead_error(self) -> BaseException: - """EngineClient abstract property implementation.""" - return OmniEngineDeadError() - - # ==================== EngineClient Interface ==================== - async def get_input_preprocessor(self) -> InputProcessor: """Get input preprocessor.""" return self.input_processor @@ -1793,16 +1106,3 @@ class AsyncOmni(EngineClient, OmniBase): async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: """Return the task set exposed by the orchestrator-backed engine.""" return tuple(self.engine.supported_tasks) - - async def check_health(self) -> None: - """Check engine health by verifying the Orchestrator process is alive.""" - OmniBase.check_health(self) - - # ==================== Shutdown ==================== - - def shutdown(self, timeout: float | None = None) -> None: - """Shutdown the engine.""" - if self.final_output_task is not None: - self.final_output_task.cancel() - self.final_output_task = None - OmniBase.shutdown(self) diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py new file mode 100644 index 000000000..e7c86de97 --- /dev/null +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -0,0 +1,557 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""AsyncOmniBase: the asyncio-side foundation shared by ``AsyncOmni`` and ``DuplexOmni``. + +Owns the engine output pump (one consumer of the engine output queue that +routes messages to per-request queues, ACK resolvers, or a subclass hook), +engine-dead fan-out, health/shutdown, and the config accessors. Turn-based +request handling lives in ``AsyncOmni``; duplex sessions in ``DuplexOmni``. +""" + +from __future__ import annotations + +import asyncio +import os +import time +from collections.abc import AsyncGenerator, Iterable +from typing import Any + +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput +from vllm.plugins.io_processors import get_io_processor +from vllm.utils import random_uuid +from vllm.v1.engine.exceptions import EngineDeadError + +from vllm_omni.diffusion.data import OmniACK +from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.entrypoints.omni_base import ( + OmniBase, + OmniEngineDeadError, +) +from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) +_FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 +# Blocking-wait interval for the event-driven final-output drain +# (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): a message wakes the drain immediately via +# the janus queue's condition variable; this timeout only bounds how often the +# orchestrator liveness check runs while the pipeline is idle. +_FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 +# Shared DELETE / generate() cleanup abort bound. Env is the documented knob. +ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) + + +class AsyncEventResolver: + """ + A generic signal aggregator designed for synchronized handshakes in + distributed or multi-stage environments. Supports waiting for a specified + number (expected_count) of worker signals in both inline and multiprocess modes. + """ + + def __init__(self, orchestrator=None): + self._pending_tasks: dict[str, dict] = {} + self.orchestrator = orchestrator + self._lock = asyncio.Lock() + + def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: + loop = asyncio.get_running_loop() + fut = loop.create_future() + self._pending_tasks[task_id] = { + "future": fut, + "expected_count": expected_count, + "received": [], + "start_time": time.time(), + } + return fut + + async def resolve(self, ack: OmniACK): + tid = getattr(ack, "task_id", None) + + if tid is None and isinstance(ack, dict): + tid = ack.get("task_id") + + async with self._lock: + task_info = self._pending_tasks.get(tid) + if task_info is None: + logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.") + return + + task_info["received"].append(ack) + current_count = len(task_info["received"]) + expected = task_info["expected_count"] + + orchestrator = self.orchestrator + if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics: + freed = getattr(ack, "freed_bytes", 0) + if freed == 0 and isinstance(ack, dict): + freed = ack.get("freed_bytes", 0) + orchestrator.metrics.record_vram_reclaimed(freed) + + logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.") + + if current_count >= expected: + self._pending_tasks.pop(tid) + fut = task_info["future"] + if not fut.done(): + elapsed = time.time() - task_info["start_time"] + logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.") + fut.set_result(task_info["received"]) + + +class AsyncOmniBase(OmniBase): + """Shared asyncio foundation of the async entrypoints (see module docstring).""" + + def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: + OmniBase.__init__(self, model=model, **kwargs) + self.final_output_task: asyncio.Task | None = None + self.event_resolver = AsyncEventResolver(orchestrator=self) + self.config_path = self.engine.config_path + self.input_processor = self.engine.input_processor + self.endpoint_restrictions = self.engine.endpoint_restrictions + + stage_index = self._get_comprehension_stage_index() + if stage_index is None: + self.io_processor = None + else: + vllm_config = self.engine.stage_vllm_configs[stage_index] + io_processor_plugin = vllm_config.model_config.io_processor_plugin + renderer = self.renderer + if renderer is None: + from vllm.renderers import renderer_from_config + + renderer = renderer_from_config(vllm_config) + self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) + + # ==================== Subclass seam ==================== + + def _route_engine_message(self, msg: Any) -> bool: + """Consume an engine output message the subclass owns (return True when consumed).""" + del msg + return False + + def _on_engine_dead(self, error: str) -> None: + """Seam: the output pump is ending because the engine died (default: nothing). + + ``DuplexOmni`` ends every session handle here; turn-based requests are + already failed through their per-request queues. + """ + del error + + def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None: + """Look up the sticky-routed replica for (stage_id, request_id). + + Used as the ``replica_resolver`` callback by ``OrchestratorAggregator`` + to label transfer_* metrics without plumbing replica ids through + ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters. + Returns None when stage_id is out of range or the request hasn't been + bound to a replica yet — the metric emit then defensive-skips. + """ + pools = getattr(self.engine, "stage_pools", None) + if pools is None or not (0 <= stage_id < len(pools)): + return None + return pools[stage_id].get_bound_replica_id(request_id) + + def _get_comprehension_stage_index(self) -> int | None: + fallback_idx: int | None = None + for idx, stage_client in enumerate(self.engine.stage_clients): + stage_vllm_config = self.engine.stage_vllm_configs[idx] + if stage_vllm_config is None: + continue + if fallback_idx is None: + fallback_idx = idx + if stage_client.is_comprehension: + return idx + return fallback_idx + + @property + def renderer(self): + """Return the renderer from the engine input processor when available.""" + if self.input_processor is None: + return None + return self.input_processor.renderer + + @property + def vllm_config(self): + """Return the vLLM config for the comprehension stage when present.""" + stage_index = self._get_comprehension_stage_index() + if stage_index is None: + return None + return self.engine.stage_vllm_configs[stage_index] + + async def get_vllm_config(self) -> Any: + """Compatibility helper for call sites expecting async vllm config access.""" + return self.vllm_config + + def get_diffusion_od_config(self) -> Any | None: + """Return the diffusion-stage config when the pipeline has one.""" + saw_diffusion_stage = False + for stage_client in self.engine.stage_clients: + if getattr(stage_client, "stage_type", None) != "diffusion": + continue + + saw_diffusion_stage = True + + od_config = getattr(stage_client, "od_config", None) + if od_config is not None: + return od_config + + inner_engine = getattr(stage_client, "_engine", None) + od_config = getattr(inner_engine, "od_config", None) + if od_config is not None: + return od_config + + # Out-of-process diffusion clients don't carry od_config (it lives in the + # worker); fall back to the engine's model_class_name resolution. + if saw_diffusion_stage: + return self.engine.get_diffusion_od_config() + + return None + + @property + def model_config(self): + """Return the model config for the comprehension stage when present.""" + vllm_config = self.vllm_config + if vllm_config is None: + return None + return vllm_config.model_config + + @staticmethod + def _get_unique_request_id(external_request_id: str): + """Get a random new request ID for this request; at the server level, + this is usually set by the calling entrypoint, but in direct calls, we + need to set it explicitly since we do not allow empty IDs. + + NOTE: in the upstream vLLM, this is done in the InputProcessor's + `assign_request_id`. + """ + uuid = random_uuid() + prefix = "" if not external_request_id else f"{external_request_id}-" + return f"{prefix}{uuid:.8}" + + # ==================== Processing Methods ==================== + + async def _process_orchestrator_results( + self, + request_id: str, + metrics: OrchestratorMetrics, + final_stage_id_for_e2e: int, + req_start_ts: dict[str, float], + wall_start_ts: float, + ) -> AsyncGenerator[OmniRequestOutput, None]: + """Read results from the Orchestrator (via the request's asyncio.Queue) + and yield OmniRequestOutput objects. + + The Orchestrator handles all stage-to-stage transfers. This method + only processes final outputs that arrive on the per-request queue. + """ + req_state = self.request_states.get(request_id) + if req_state is None: + return + + while True: + result = await req_state.queue.get() + + if isinstance(result, ErrorMessage): + logger.error( + "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s", + request_id, + result.stage_id, + result.error, + ) + if result.fatal: + raise OmniEngineDeadError( + result.error, + error_stage_id=result.stage_id, + ) + self._raise_nonfatal_error_message(result) + + if not isinstance(result, OutputMessage): + logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) + continue + + stage_id = result.stage_id + + self._check_engine_output_error(result, request_id, stage_id) + + # Process the result (constructs OmniRequestOutput) + output_to_yield = self._process_single_result( + result, + stage_id, + metrics, + req_start_ts, + wall_start_ts, + final_stage_id_for_e2e, + ) + + if output_to_yield: + # Set the external request ID back to the user yielded input + output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id + logger.debug( + "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s", + request_id, + stage_id, + getattr(output_to_yield, "final_output_type", None), + ) + yield output_to_yield + + # The Orchestrator sets "finished" when the final stage is done + if result.finished: + break + + # ==================== Output Handler ==================== + + def _final_output_handler(self) -> None: + """Start the final output handler if not already running. + + This handler reads messages from the Orchestrator output queue and + routes them to per-request asyncio.Queues. + """ + if self.final_output_task is not None: + return + + engine = self.engine + + # Event-driven drain (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): block on the + # queue's condition variable in a dedicated thread instead of the + # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side + # event-driven loop (vllm_omni/engine/orchestrator.py). + from vllm_omni.engine.orchestrator import _event_driven_orch_enabled + + event_driven_drain = _event_driven_orch_enabled() and hasattr(engine, "get_output_blocking_async") + + async def _final_output_loop(): + """Background coroutine that dispatches final outputs to request queues.""" + try: + while True: + if event_driven_drain: + msg = await engine.get_output_blocking_async(timeout=_FINAL_OUTPUT_BLOCKING_WAIT_S) + if msg is None: + # Timed out with the orchestrator alive; loop for + # the periodic liveness check. + continue + else: + msg = await engine.try_get_output_async() + if msg is None: + await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) + continue + + if self._route_engine_message(msg): + continue + + if isinstance(msg, dict) and msg.get("type") == "ack": + ack_data = msg.get("ack") + tid = getattr(ack_data, "task_id", "unknown") + logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}") + await self.event_resolver.resolve(ack_data) + continue + if isinstance(msg, OmniACK): + logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}") + await self.event_resolver.resolve(msg) + continue + if hasattr(msg, "task_id"): + tid = getattr(msg, "task_id") + logger.info(f"[{self._name}] Intercepted task-ID object: {tid}") + await self.event_resolver.resolve(msg) + continue + + if isinstance(msg, ErrorMessage): + # Route request-scoped errors to that request's queue and + # keep the loop alive. A request whose stage replica died + # and was evicted gets a fatal error delivered here; only + # that request fails (its consumer raises), the server + # stays up for other stages/requests (#4285). A fatal + # error without a request_id is a genuine engine-wide + # death and falls through to the except handler below. + if msg.request_id is not None: + req_state = self.request_states.get(msg.request_id) + if req_state is not None: + await req_state.queue.put(msg) + else: + logger.warning( + "[%s] dropping error for unknown req %s", + self._name, + msg.request_id, + ) + continue + if not msg.fatal: + continue + + should_continue, _, stage_id, req_state = self._handle_output_message(msg) + if should_continue: + continue + + req_state.stage_id = stage_id + + # Route to the per-request queue + await req_state.queue.put(msg) + + except asyncio.CancelledError: + raise + except OmniEngineDeadError as e: + logger.error("[%s] Engine dead: %s", self._name, e) + for req_state in list(self.request_states.values()): + error_msg = ErrorMessage( + error=str(e), + fatal=True, + request_id=req_state.request_id, + stage_id=e.error_stage_id, + ) + await req_state.queue.put(error_msg) + self._on_engine_dead(str(e)) + except EngineDeadError as e: + logger.error("[%s] Engine dead: %s", self._name, e) + for req_state in list(self.request_states.values()): + error_msg = ErrorMessage( + error=str(e), + fatal=True, + request_id=req_state.request_id, + ) + await req_state.queue.put(error_msg) + self._on_engine_dead(str(e)) + except Exception as e: + logger.exception("[%s] final_output_loop failed.", self._name) + for req_state in list(self.request_states.values()): + error_msg = ErrorMessage( + request_id=req_state.request_id, + error=str(e), + ) + await req_state.queue.put(error_msg) + self.final_output_task = None + self._on_engine_dead(str(e)) + + self.final_output_task = asyncio.create_task(_final_output_loop()) + logger.debug("[AsyncOmni] Final output handler started") + + async def _abort_internal_requests( + self, + request_id: str | Iterable[str], + *, + timeout: float = ABORT_TIMEOUT_S, + ): + """Abort request(s) via the Orchestrator given internal request IDs, + which take the format -. + """ + request_ids = [request_id] if isinstance(request_id, str) else list(request_id) + # Request IDs are already internal, so we just need to get the matching states. + internal_req_ids = [rid for rid in request_ids if rid in self.request_states] + try: + # Unbind generate() if abort_async blocks in the executor. + await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) + except TimeoutError: + logger.warning( + "[AsyncOmni] Timed out aborting %s after %.1fs; " + "engine abort is best-effort until the current batch drains", + ",".join(internal_req_ids), + timeout, + ) + except Exception: + logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) + + async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: + """Abort request IDs via the engine and enqueue terminal abort outputs. + + Waits for orchestrator abort acknowledgment, enqueues any AR terminal + abort outputs (partial tokens) into each request's asyncio queue, then + cancels the input pump. Frontend ``request_states`` stay registered so + ``generate()`` can consume the terminal message in + ``_process_orchestrator_results`` and run normal cleanup. + + When ``abort_async`` returns no output for an active request (OP not + registered yet, unbound replica, or orchestrator id drop), enqueue a + synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. + """ + abort_outputs = await self.engine.abort_async(request_ids, timeout=timeout) or [] + delivered: set[str] = set() + for output_msg in abort_outputs: + req_id = getattr(output_msg, "request_id", None) + if req_id is None: + continue + state = self.request_states.get(req_id) + if state is None: + logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) + continue + await state.queue.put(output_msg) + delivered.add(req_id) + for rid in request_ids: + state = self.request_states.get(rid) + if state is not None and rid not in delivered: + queue = getattr(state, "queue", None) + if queue is not None: + await state.queue.put(self._synthetic_abort_output_message(rid)) + delivered.add(rid) + for rid in request_ids: + self._record_request_failure_once(rid, reason="client_abort") + state = self.request_states.get(rid) + input_stream_task = getattr(state, "input_stream_task", None) + if input_stream_task is not None and not input_stream_task.done(): + input_stream_task.cancel() + if self.log_stats: + logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) + + @staticmethod + def _synthetic_abort_output_message(request_id: str) -> OutputMessage: + """Terminal abort OutputMessage used when the engine returned none.""" + engine_output = OmniRequestOutput( + request_id=request_id, + finished=True, + stage_id=0, + final_output_type="text", + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="abort", + stop_reason=None, + ) + ], + ) + return OutputMessage( + request_id=request_id, + stage_id=0, + replica_id=None, + engine_outputs=engine_output, + metrics=None, + finished=True, + stage_submit_ts=None, + ) + + @property + def is_running(self) -> bool: + """Check if the engine is running.""" + orchestrator_alive = self.engine.is_alive() + task_alive = self.final_output_task is not None and not self.final_output_task.done() + return orchestrator_alive and task_alive + + @property + def _name(self) -> str: + return type(self).__name__ + + @property + def is_stopped(self) -> bool: + """EngineClient abstract property implementation.""" + return self.errored + + @property + def dead_error(self) -> BaseException: + """EngineClient abstract property implementation.""" + return OmniEngineDeadError() + + # ==================== EngineClient Interface ==================== + + async def check_health(self) -> None: + """Check engine health by verifying the Orchestrator process is alive.""" + OmniBase.check_health(self) + + # ==================== Shutdown ==================== + + def shutdown(self, timeout: float | None = None) -> None: + """Shutdown the engine.""" + if self.final_output_task is not None: + self.final_output_task.cancel() + self.final_output_task = None + OmniBase.shutdown(self) diff --git a/vllm_omni/entrypoints/cli/benchmark/cli_args.py b/vllm_omni/entrypoints/cli/benchmark/cli_args.py index 21f528954..626951298 100644 --- a/vllm_omni/entrypoints/cli/benchmark/cli_args.py +++ b/vllm_omni/entrypoints/cli/benchmark/cli_args.py @@ -206,8 +206,10 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: "--seed-tts-turns-per-session", type=int, default=1, - help="Group this many Seed-TTS target texts into one Realtime session. " - "The first row's reference audio and transcript are reused for every turn.", + help="Group this many Seed-TTS target texts into one benchmark request. " + "The first row's reference audio and transcript are reused for every turn. " + "A model-native duplex session takes its target text once, in the session " + "context, so each turn runs in its own session.", ) group.add_argument( "--seed-tts-root", diff --git a/vllm_omni/entrypoints/duplex/__init__.py b/vllm_omni/entrypoints/duplex/__init__.py index 736e98a6a..327eab07f 100644 --- a/vllm_omni/entrypoints/duplex/__init__.py +++ b/vllm_omni/entrypoints/duplex/__init__.py @@ -1,6 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from .websocket import DuplexWebSocketActor +"""Thin websocket serving for engine-resident duplex sessions. -__all__ = ["DuplexWebSocketActor"] +The session logic lives in ``vllm_omni.engine.duplex`` and is reached through +``vllm_omni.entrypoints.duplex_omni.DuplexOmni``; this package only speaks the +OpenAI Realtime wire protocol over ``/v1/realtime?duplex=1``. +""" + +from .serving import OmniDuplexSessionHandler + +__all__ = ["OmniDuplexSessionHandler"] diff --git a/vllm_omni/entrypoints/duplex/audio_encoding.py b/vllm_omni/entrypoints/duplex/audio_encoding.py new file mode 100644 index 000000000..f790b029b --- /dev/null +++ b/vllm_omni/entrypoints/duplex/audio_encoding.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Audio encoding for duplex model output (entrypoint layer, injected into ``DuplexOmniEngine``).""" + +from __future__ import annotations + +import numpy as np +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def encode_audio( + audio_data: object, + sample_rate_hz: int, + response_format: str, + speed: float | None, +) -> str | None: + """Encode a model audio tensor/array into base64 in ``response_format``. + + Moved from the serving runtime bridge; uses the shared ``AudioMixin`` + encoder directly instead of going through the chat service. + """ + if audio_data is None: + return None + try: + import torch + + from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin + from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio + + if isinstance(audio_data, torch.Tensor): + audio_tensor = audio_data.detach().cpu().float().numpy() + else: + audio_tensor = np.asarray(audio_data, dtype=np.float32) + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.reshape(-1) + audio_response = AudioMixin().create_audio( + CreateAudio( + audio_tensor=audio_tensor, + sample_rate=sample_rate_hz, + response_format=response_format, + speed=float(speed) if isinstance(speed, int | float) and speed > 0 else 1.0, + stream_format="audio", + base64_encode=True, + ) + ) + return str(audio_response.audio_data) + except Exception: + logger.exception("Failed to encode duplex data-plane audio output") + return None + + +__all__ = ["encode_audio"] diff --git a/vllm_omni/entrypoints/duplex/capability.py b/vllm_omni/entrypoints/duplex/capability.py deleted file mode 100644 index 50126886b..000000000 --- a/vllm_omni/entrypoints/duplex/capability.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -def should_enable_duplex_endpoint( - stage_configs: list | None, - *, - config_path: str | None = None, -) -> bool: - """Enable the realtime session handler for explicitly configured deployments.""" - if stage_configs: - for stage in stage_configs: - session_mode = ( - stage.get("session_mode") if isinstance(stage, dict) else getattr(stage, "session_mode", None) - ) - if session_mode == "duplex": - return True - if config_path: - try: - from vllm_omni.config.stage_config import resolve_deploy_yaml - - raw_config = resolve_deploy_yaml(config_path) - if raw_config.get("session_mode") == "duplex" or isinstance(raw_config.get("duplex_session"), dict): - return True - except Exception as exc: - logger.warning("Failed to inspect realtime session configuration from %s: %s", config_path, exc) - return False - - -__all__ = ["should_enable_duplex_endpoint"] diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py deleted file mode 100644 index bd682c4f3..000000000 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ /dev/null @@ -1,302 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import json -from collections.abc import AsyncGenerator -from functools import partial -from typing import Any - -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.serve.engine.protocol import ErrorResponse -from vllm.logger import init_logger - -from vllm_omni.entrypoints.duplex.protocol import DuplexSession - -logger = init_logger(__name__) - - -class ChatFallbackProjectorMixin: - """Project generic chat completion output into duplex response events.""" - - async def _run_response(self, session: DuplexSession, send_json) -> None: - response_id = session.begin_response() - epoch = session.epoch - request_id = f"duplex-{session.session_id}-{epoch}-{session.input_commit_seq}" - session.bind_request(f"chatcmpl-{request_id}") - await send_json( - self._response_created_payload( - session, - response_id, - epoch=epoch, - request_id=session.active_request_id, - ) - ) - - try: - request = self._build_chat_request(session, request_id) - result = await self._chat_service.create_chat_completion(request, raw_request=None) - if isinstance(result, ErrorResponse): - error = result.error - await send_json( - { - "type": "error", - "error": error.message if error else "Chat request failed", - "code": error.type if error else "chat_error", - } - ) - session.end_response(commit_text=False) - return - if hasattr(result, "__aiter__"): - await self._drain_streaming_response(session, result, epoch, response_id, send_json) - else: - await self._emit_full_response(session, result, epoch, response_id, send_json) - if session.epoch == epoch: - final_stage_metrics = session.accumulate_response_stage_metrics(None) - should_commit = self._should_commit_response_to_history(session, response_id) - committed_message = session.end_response(commit_text=should_commit) - if should_commit: - session.register_history_item(f"item_{response_id}", committed_message) - done_payload: dict[str, object] = { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - if final_stage_metrics: - done_payload["vllm_omni"] = {"stage_metrics": final_stage_metrics} - await send_json(done_payload) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.exception("Duplex response failed: %s", exc) - session.end_response(commit_text=False) - await send_json( - { - "type": "error", - "session_id": session.session_id, - "response_id": response_id, - "error": str(exc), - "code": "response_error", - } - ) - - def _build_chat_request(self, session: DuplexSession, request_id: str) -> ChatCompletionRequest: - response_config = session.response_config - messages: list[dict[str, object]] = [] - if response_config.instructions: - messages.append({"role": "system", "content": response_config.instructions}) - messages.extend(session.history) - - kwargs: dict[str, Any] = { - "model": response_config.model or self._chat_service.model_config.model, - "messages": messages, - "stream": True, - } - if response_config.temperature is not None: - kwargs["temperature"] = response_config.temperature - if response_config.max_tokens is not None: - kwargs["max_tokens"] = response_config.max_tokens - model_extra = dict(response_config.extra_body) - model_extra.pop("native_duplex", None) - model_extra.pop("minicpmo45_native_duplex", None) - tools = model_extra.pop("realtime_response_tools", model_extra.pop("realtime_tools", None)) - tool_choice = model_extra.pop( - "realtime_response_tool_choice", - model_extra.pop("realtime_tool_choice", None), - ) - for protocol_key in ( - "realtime_response_conversation", - "realtime_response_metadata", - "realtime_response_prompt", - ): - model_extra.pop(protocol_key, None) - kwargs.update(model_extra) - if "audio" in response_config.modalities: - audio_options = kwargs.get("audio") - audio_options = dict(audio_options) if isinstance(audio_options, dict) else {} - # Chat defaults to WAV. Request the session's internal output format - # explicitly rather than relabeling the returned bytes as PCM. - audio_options["format"] = response_config.response_format - kwargs["audio"] = audio_options - if isinstance(tools, list): - kwargs["tools"] = tools - if isinstance(tool_choice, str | dict): - kwargs["tool_choice"] = tool_choice - - request = ChatCompletionRequest(**kwargs) - object.__setattr__(request, "modalities", response_config.modalities) - object.__setattr__(request, "request_id", request_id) - object.__setattr__( - request, - "chat_template_kwargs", - {"use_tts_template": response_config.use_tts_template}, - ) - return request - - async def _drain_streaming_response( - self, - session: DuplexSession, - result: AsyncGenerator[str, None], - epoch: int, - response_id: str, - send_json, - ) -> None: - async for raw_chunk in result: - if session.epoch != epoch: - return - for payload in self._parse_sse_payloads(raw_chunk): - if payload == "[DONE]": - continue - if isinstance(payload, dict): - await self._emit_chat_payload(session, payload, epoch, response_id, send_json) - - async def _emit_full_response( - self, - session: DuplexSession, - result: Any, - epoch: int, - response_id: str, - send_json, - ) -> None: - if hasattr(result, "model_dump"): - payload = result.model_dump(mode="json", exclude_unset=True) - else: - payload = {"response": str(result)} - await self._emit_chat_payload(session, payload, epoch, response_id, send_json) - - def _parse_sse_payloads(self, raw_chunk: str) -> list[dict[str, object] | str]: - payloads: list[dict[str, object] | str] = [] - for line in raw_chunk.splitlines(): - line = line.strip() - if not line.startswith("data:"): - continue - data = line[5:].strip() - if not data: - continue - if data == "[DONE]": - payloads.append(data) - continue - try: - parsed = json.loads(data) - except json.JSONDecodeError: - logger.debug("Skipping non-JSON duplex stream payload: %s", data) - continue - if isinstance(parsed, dict): - payloads.append(parsed) - return payloads - - async def _emit_chat_payload( - self, - session: DuplexSession, - payload: dict[str, object], - epoch: int, - response_id: str, - send_json, - ) -> None: - if session.epoch != epoch or session.active_response_id != response_id: - return - metrics = payload.get("metrics") - stage_metrics = metrics.get("stage_metrics") if isinstance(metrics, dict) else None - if isinstance(stage_metrics, dict): - session.replace_response_stage_metric_snapshots(stage_metrics) - - modality = payload.get("modality") - if modality not in {None, "text", "audio"}: - await send_json( - { - "type": "error", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "code": "unsupported_response_modality", - "error": f"Unsupported chat response modality: {modality}", - } - ) - return - choices = payload.get("choices") - if not isinstance(choices, list): - await send_json( - { - "type": "response.message", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "payload": payload, - } - ) - return - - for choice in choices: - if not isinstance(choice, dict): - continue - delta = choice.get("delta") - message = choice.get("message") - content = None - if isinstance(delta, dict): - content = delta.get("content") - elif isinstance(message, dict): - content = message.get("content") - - if isinstance(content, str) and content: - if modality != "audio": - session.append_assistant_text(content) - await send_json( - { - "type": "response.text.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "delta": content, - } - ) - - audio_content = content if modality == "audio" else None - if isinstance(message, dict): - message_audio = message.get("audio") - if isinstance(message_audio, dict): - audio_content = message_audio.get("data") - if isinstance(audio_content, str) and audio_content: - # The OpenAI package imports duplex serving during initialization. - # Defer this import until chat fallback has finished loading. - from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata - - audio_metadata = AudioChunkMetadata.model_validate(choice.get("audio_metadata")) - duration_ms = session.record_generated_audio( - response_id, - frame_count=audio_metadata.frame_count, - sample_rate_hz=audio_metadata.sample_rate_hz, - ) - if duration_ms is None: - return - await send_json( - { - "type": "response.output_audio.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "audio": audio_content, - "format": audio_metadata.format, - "sample_rate_hz": audio_metadata.sample_rate_hz, - "channels": audio_metadata.channels, - "audio_duration_ms": duration_ms, - }, - on_accepted=partial(session.mark_audio_sent, duration_ms, response_id=response_id), - ) - - finish_reason = choice.get("finish_reason") - if finish_reason is not None and modality != "audio": - await send_json( - { - "type": "response.output_item.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "finish_reason": finish_reason, - "modality": modality, - } - ) diff --git a/vllm_omni/entrypoints/duplex/realtime_input.py b/vllm_omni/entrypoints/duplex/realtime_input.py index 386b4cfc3..be700ef62 100644 --- a/vllm_omni/entrypoints/duplex/realtime_input.py +++ b/vllm_omni/entrypoints/duplex/realtime_input.py @@ -1,1584 +1,179 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Thin OpenAI Realtime wire envelope for one websocket connection. + +Everything that needed session state now lives engine-side +(``vllm_omni.engine.duplex.realtime_commands`` for the command mapping, +``vllm_omni.engine.duplex.realtime_events`` for the projection state). What is +left here is the per-connection handshake policy (query-param defaults, +autostart / resume-only rules, ``session.resume`` parsing), the wire defaults +used to translate appends, and error rendering through the typed +:class:`~vllm_omni.engine.duplex.events.ErrorEvent`. +""" + from __future__ import annotations -import asyncio -import base64 -import binascii -from collections.abc import Callable -from uuid import uuid4 +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING -import numpy as np - -from vllm_omni.entrypoints.duplex.audio import ( - convert_input_audio_with_rate, - encode_float32_mono_wav_base64, - validate_input_sample_rate_hz, +from vllm_omni.engine.duplex.commands import DuplexCommand, DuplexCommandError +from vllm_omni.engine.duplex.events import error_event +from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + translate_realtime_command, ) -from vllm_omni.entrypoints.duplex.protocol import normalize_native_duplex_key -from vllm_omni.entrypoints.duplex.realtime_state import ( - REALTIME_INPUT_AUDIO_FORMATS, - REALTIME_OUTPUT_AUDIO_FORMATS, - RealtimeStateOwner, - realtime_default_sample_rate_hz, -) -from vllm_omni.entrypoints.duplex.server_vad import ServerVADConfig, parse_session_turn_detection + +if TYPE_CHECKING: + from starlette.websockets import WebSocket + +__all__ = ["RealtimeEnvelope", "RealtimeHandshake", "ResumeRequest", "parse_resume_request"] + +_TRUE_VALUES = {"1", "true", "yes", "on"} +_FALSE_VALUES = {"0", "false", "no", "off"} + +#: Client events the envelope handles itself (never translated into commands). +#: ``conversation.item.retrieve`` is deliberately absent: the conversation items +#: it reads live engine-side, so it travels as a command like any other event. +ENVELOPE_EVENT_TYPES = frozenset({"session.resume", "session.event_ack"}) -class RealtimeInputTranslator(RealtimeStateOwner): - """Translate and validate client Realtime events for the duplex core.""" +@dataclass(frozen=True, slots=True) +class RealtimeHandshake: + """What the first client message asks for.""" - _opened: bool - _initial_session_update: bool - _input_speech_started: bool - _input_audio_buffer_has_audio: bool - _input_audio_buffer_had_non_speech: bool - _input_audio_buffer_transcript_parts: list[str] - _active_input_item_id: str | None - _last_conversation_item_id: str | None - _active_response_id: str | None - _last_response_id: str | None - _input_audio_format: str - _output_audio_format: str - _input_sample_rate_hz: int - _conversation_items: dict[str, dict[str, object]] - _item_truncation_cursors: dict[str, tuple[int, int]] - _pending_outbound: asyncio.Queue[dict[str, object]] + #: ``"open"`` (session object in ``session_payload``) or ``"resume"`` (resume payload). + kind: str + session_payload: dict[str, object] = field(default_factory=dict) + resume_payload: dict[str, object] = field(default_factory=dict) + #: A client event that arrived before the session existed and must run right after opening. + pending_command_payload: dict[str, object] | None = None - async def _send_realtime_payload(self, payload: dict[str, object]) -> None: - raise NotImplementedError - def _realtime_error_payload( - self, +@dataclass(frozen=True, slots=True) +class ResumeRequest: + """A validated ``session.resume`` client event.""" + + session_id: str + resume_token: str + last_received_server_event_seq: int + + +def parse_resume_request(event: Mapping[str, object]) -> ResumeRequest | None: + """Validate the shape of ``session.resume``; ``None`` when any field is missing or malformed.""" + session_id = event.get("session_id") + resume_token = event.get("resume_token") + last_received = event.get("last_received_server_event_seq", 0) + if ( + not isinstance(session_id, str) + or not session_id + or not isinstance(resume_token, str) + or not resume_token + or not isinstance(last_received, int) + or last_received < 0 + ): + return None + return ResumeRequest( + session_id=session_id, + resume_token=resume_token, + last_received_server_event_seq=last_received, + ) + + +@dataclass +class RealtimeEnvelope: + """Per-connection Realtime wire policy (no session state).""" + + default_model: str | None = None + #: ``?resume=1`` / ``?autostart=0``: the first message must be ``session.resume``/``session.update``. + resume_only: bool = False + opened: bool = False + autostarted_default_session: bool = False + defaults: RealtimeInputDefaults = field(default_factory=RealtimeInputDefaults) + + @classmethod + def from_query_params(cls, query_params: Mapping[str, str] | WebSocket) -> RealtimeEnvelope: + if hasattr(query_params, "query_params"): + query_params = query_params.query_params + getter = query_params.get if hasattr(query_params, "get") else None + resume_only = getter("resume") if getter is not None else None + autostart = getter("autostart") if getter is not None else None + model = getter("model") if getter is not None else None + return cls( + default_model=model if isinstance(model, str) and model else None, + resume_only=( + str(resume_only).strip().lower() in _TRUE_VALUES or str(autostart).strip().lower() in _FALSE_VALUES + ), + ) + + # ---- handshake ---- + + def default_session_payload(self) -> dict[str, object]: + return {"model": self.default_model} + + def initial_open_payload(self) -> dict[str, object] | None: + """Session object to open with *before* any client message (``?model=`` autostart), else None.""" + if self.opened or self.resume_only or self.autostarted_default_session or not self.default_model: + return None + self.opened = True + self.autostarted_default_session = True + return self.default_session_payload() + + def first_message(self, payload: Mapping[str, object]) -> RealtimeHandshake: + """Classify the first client message (call only while ``not opened``). + + ``session.update`` opens with its session object; ``session.resume`` + resumes; any other event autostarts the default session and is then + treated as a command (``pending_command_payload``). + """ + event_type = payload.get("type") + self.opened = True + if event_type == "session.resume": + return RealtimeHandshake(kind="resume", resume_payload=dict(payload)) + if event_type == "session.update": + session = payload.get("session") + session_payload = dict(session) if isinstance(session, dict) else dict(payload) + session_payload.pop("type", None) + self.note_session_payload(session_payload) + return RealtimeHandshake(kind="open", session_payload=session_payload) + self.autostarted_default_session = True + return RealtimeHandshake( + kind="open", + session_payload=self.default_session_payload(), + pending_command_payload=dict(payload), + ) + + def note_session_payload(self, session_payload: Mapping[str, object]) -> None: + """Track wire defaults declared by a session object (open or session.update).""" + self.defaults = self.defaults.with_session_payload(session_payload) + + # ---- commands ---- + + def is_envelope_event(self, payload: Mapping[str, object]) -> bool: + return payload.get("type") in ENVELOPE_EVENT_TYPES + + def translate(self, payload: Mapping[str, object]) -> DuplexCommand: + """Wire event -> command (raises :class:`DuplexCommandError`).""" + if payload.get("type") == "session.update": + session = payload.get("session") + if isinstance(session, dict): + self.note_session_payload(session) + return translate_realtime_command(payload, defaults=self.defaults) + + # ---- outbound helpers ---- + + @staticmethod + def error_payload( code: str, message: str, *, event_id: object | None = None, param: object | None = None, ) -> dict[str, object]: - raise NotImplementedError - - _response_is_done: Callable[[object], bool] - - @property - def _turn_based_server_vad(self) -> bool: - return self._turn_detection is not None and not self._native_input_append - - def _reset_realtime_input_buffer_state(self) -> None: - self._input_speech_started = False - self._active_input_item_id = None - self._input_audio_buffer_has_audio = False - self._input_audio_buffer_had_non_speech = False - self._input_audio_buffer_transcript_parts.clear() - - async def discard_pending_input_audio( - self, - *, - audio_end_ms: int | None = None, - ) -> None: - """Drop Realtime input-buffer state that was consumed as overlap. - - The serving overlap policy may classify a short acknowledgement such - as "continue" as a side-channel signal instead of user turn input. The - Realtime protocol has already observed the append and may have emitted - ``speech_started``; reset that transient buffer so a later commit does - not accidentally attach the acknowledgement transcript/audio to the - next real user turn. - """ - if self._input_speech_started and self._active_input_item_id is not None: - await self._send_realtime_payload( - { - "type": "input_audio_buffer.speech_stopped", - "audio_end_ms": max(0, int(audio_end_ms or 0)), - "item_id": self._active_input_item_id, - } - ) - self._reset_realtime_input_buffer_state() - - async def _send_realtime_input_ack(self, event: dict[str, object]) -> None: - if event.get("type") != "conversation.item.create": - return - item = event.get("item") - if not isinstance(item, dict): - return - previous_item_id = event.get("previous_item_id") - if isinstance(previous_item_id, str): - item["_previous_item_id"] = previous_item_id - if item.get("role") != "user": - return - self._conversation_items[str(item["id"])] = item - for payload in self._conversation_item_added_events(item): - await self._send_realtime_payload(payload) - - async def _to_duplex_event(self, event: dict[str, object]) -> dict[str, object] | None: - event_type = event.get("type") - if event_type == "session.update": - session = event.get("session") - session_payload: dict[str, object] = session if isinstance(session, dict) else event - try: - turn_detection_configured, server_vad = parse_session_turn_detection(session_payload) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_turn_detection", - str(exc), - event_id=event.get("event_id"), - param="turn_detection", - ) - ) - return None - normalized_turn_detection = server_vad.as_dict() if server_vad is not None else None - previous_protocol_config = self._realtime_protocol_config() - try: - self._apply_realtime_session_defaults(session_payload) - except ValueError as exc: - self._set_realtime_protocol_config(previous_protocol_config) - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - str(exc), - event_id=event.get("event_id"), - param="audio.input.format", - ) - ) - return None - turn_detection_error = self._validate_realtime_turn_detection( - session_payload, - configured=turn_detection_configured, - normalized=normalized_turn_detection, - ) - if turn_detection_error is not None: - self._set_realtime_protocol_config(previous_protocol_config) - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_turn_detection", - turn_detection_error, - event_id=event.get("event_id"), - param="turn_detection", - ) - ) - return None - turn_detection_candidate = self._prepare_realtime_turn_detection_update( - session_payload, - configured=turn_detection_configured, - normalized=normalized_turn_detection, - ) - try: - session_payload.update(self._realtime_overlap_fields(session_payload)) - if not self._opened: - self._opened = True - self._initial_session_update = True - translated = self._session_create_from_realtime(session_payload) - event_id = event.get("event_id") - if isinstance(event_id, str) and event_id: - translated["_realtime_event_id"] = event_id - else: - translated = { - "type": "turn.signal", - "event": "session.update", - "payload": session_payload, - "realtime_event_id": event.get("event_id"), - } - candidate_protocol_config = self._realtime_protocol_config() - finally: - self._set_realtime_protocol_config(previous_protocol_config) - if turn_detection_candidate[0] or candidate_protocol_config != previous_protocol_config: - self._pending_turn_detection_update = ( - *turn_detection_candidate, - asyncio.Event(), - candidate_protocol_config, - ) - return translated - if event_type == "conversation.item.create": - return await self._conversation_item_to_duplex(event) - if event_type == "conversation.item.delete": - item_id = event.get("item_id") - if not isinstance(item_id, str) or not item_id: - await self._send_realtime_payload( - self._realtime_error_payload( - "missing_item_id", - "conversation.item.delete requires item_id", - event_id=event.get("event_id"), - ) - ) - return None - if item_id not in self._conversation_items: - await self._send_realtime_payload( - self._realtime_error_payload( - "item_not_found", - f"Conversation item not found: {item_id}", - event_id=event.get("event_id"), - ) - ) - return None - self._remove_conversation_item(item_id) - await self._pending_outbound.put( - { - "type": "turn.signal", - "event": "conversation.item.delete", - "payload": {"item_id": item_id}, - } - ) - return None - if event_type == "conversation.item.retrieve": - item_id = event.get("item_id") - if not isinstance(item_id, str) or not item_id: - await self._send_realtime_payload( - self._realtime_error_payload( - "missing_item_id", - "conversation.item.retrieve requires item_id", - event_id=event.get("event_id"), - ) - ) - return None - item = self._conversation_items.get(item_id) - if item is None: - await self._send_realtime_payload( - self._realtime_error_payload( - "item_not_found", - f"Conversation item not found: {item_id}", - event_id=event.get("event_id"), - ) - ) - return None - await self._send_realtime_payload({"type": "conversation.item.retrieved", "item": item}) - return None - if event_type == "conversation.item.truncate": - item_id = event.get("item_id") - audio_end_ms = event.get("audio_end_ms") - content_index = event.get("content_index", 0) - if not isinstance(item_id, str) or not item_id: - await self._send_realtime_payload( - self._realtime_error_payload( - "missing_item_id", - "conversation.item.truncate requires item_id", - event_id=event.get("event_id"), - ) - ) - return None - if not isinstance(audio_end_ms, int | float): - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - "conversation.item.truncate requires numeric audio_end_ms", - event_id=event.get("event_id"), - param="audio_end_ms", - ) - ) - return None - item = self._conversation_items.get(item_id) - if not isinstance(item, dict): - await self._send_realtime_payload( - self._realtime_error_payload( - "item_not_found", - f"Conversation item not found: {item_id}", - event_id=event.get("event_id"), - ) - ) - return None - truncate_error = self._validate_realtime_item_truncate( - item, - content_index=int(content_index) if isinstance(content_index, int | float) else 0, - audio_end_ms=int(audio_end_ms), - ) - if truncate_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - truncate_error, - event_id=event.get("event_id"), - ) - ) - return None - self._item_truncation_cursors[item_id] = ( - int(content_index) if isinstance(content_index, int | float) else 0, - int(audio_end_ms), - ) - self._truncate_realtime_item_content( - item, - content_index=int(content_index) if isinstance(content_index, int | float) else 0, - audio_end_ms=int(audio_end_ms), - ) - await self._pending_outbound.put( - { - "type": "turn.signal", - "event": "conversation.item.truncate", - "payload": { - "item_id": item_id, - "content_index": content_index, - "audio_end_ms": audio_end_ms, - }, - } - ) - await self._pending_outbound.put( - { - "type": "playback.ack", - "item_id": item_id, - "committed_ms": int(audio_end_ms), - "played_ms": int(audio_end_ms), - "truncate": True, - } - ) - return None - if event_type == "input_audio_buffer.append": - audio = event.get("audio") or event.get("delta") - event_format = event.get("format") or event.get("input_audio_format") - native_server_vad = self._turn_detection is not None and self._native_input_append - turn_based_server_vad = self._turn_based_server_vad - try: - fmt, format_rate = self._parse_realtime_audio_format(event_format or self._input_audio_format) - sample_rate_hz = event.get("sample_rate_hz") or event.get("sample_rate") or format_rate - if sample_rate_hz is None: - sample_rate_hz = ( - realtime_default_sample_rate_hz(fmt) - if event_format is not None and fmt != self._input_audio_format - else self._input_sample_rate_hz - ) - if turn_based_server_vad: - sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - str(exc), - event_id=event.get("event_id"), - param="sample_rate_hz", - ) - ) - return None - channels = event.get("channels", 1) - if not self._is_supported_realtime_input_format(fmt): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - f"Unsupported input_audio_format: {fmt}", - event_id=event.get("event_id"), - ) - ) - return None - if turn_based_server_vad and (fmt != "pcm16" or channels != 1 or sample_rate_hz not in {16_000, 24_000}): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - "server_vad requires mono PCM16 input at 16 kHz or 24 kHz", - event_id=event.get("event_id"), - param="audio.input.format", - ) - ) - return None - looks_like_speech: bool | None = None - if not turn_based_server_vad: - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, - ) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - str(exc), - event_id=event.get("event_id"), - param="sample_rate_hz", - ) - ) - return None - if not turn_based_server_vad and not native_server_vad: - looks_like_speech = self._input_looks_like_speech(event, audio=audio, fmt=fmt) - has_audio_payload = isinstance(audio, str) and bool(audio) - self._input_audio_buffer_has_audio = self._input_audio_buffer_has_audio or ( - looks_like_speech and has_audio_payload - ) - self._input_audio_buffer_had_non_speech = self._input_audio_buffer_had_non_speech or ( - not looks_like_speech and has_audio_payload - ) - if looks_like_speech: - await self._emit_input_speech_started(event) - self._remember_input_transcript_hint(event) - payload = { - "type": "input_audio_buffer.append", - "audio": audio, - "format": fmt, - "sample_rate_hz": sample_rate_hz, - } - event_id = event.get("event_id") - if self._turn_detection is not None and isinstance(event_id, str) and event_id: - payload["_realtime_event_id"] = event_id - video_frames = event.get("video_frames") - if video_frames is not None: - frames_error = self._validate_realtime_video_frames(video_frames, event.get("max_slice_nums")) - if frames_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "invalid_video_frames", - frames_error, - event_id=event.get("event_id"), - param="video_frames", - ) - ) - return None - if isinstance(video_frames, list): - payload["video_frames"] = [frame for frame in video_frames if isinstance(frame, str) and frame] - self._copy_realtime_input_hints(event, payload) - if looks_like_speech is not None: - payload["is_speech"] = looks_like_speech - return payload - if event_type == "input_audio_buffer.commit": - if self._turn_based_server_vad: - await self._send_realtime_payload( - self._realtime_error_payload( - "server_vad_manual_commit_unsupported", - ( - "input_audio_buffer.commit is not supported while server_vad is enabled; " - "wait for server endpointing, or start a session with turn_detection=null " - "to use manual commit" - ), - event_id=event.get("event_id"), - ) - ) - return None - if not self._input_audio_buffer_has_audio: - if self._input_audio_buffer_had_non_speech: - self._input_audio_buffer_had_non_speech = False - self._active_input_item_id = None - return { - "type": "input_audio_buffer.commit", - "final": event.get("final", True), - "response_create": False, - "is_speech": False, - } - await self._send_realtime_payload( - self._realtime_error_payload( - "input_audio_buffer_empty", - "input_audio_buffer.commit requires a non-empty input audio buffer", - event_id=event.get("event_id"), - ) - ) - return None - item_id = self._active_input_item_id or f"item_{uuid4().hex}" - await self._emit_input_speech_stopped(event, item_id=item_id) - transcript = self._consume_input_transcript_hint() - self._active_input_item_id = None - self._input_audio_buffer_has_audio = False - self._input_audio_buffer_had_non_speech = False - payload = { - "type": "input_audio_buffer.commit", - "final": event.get("final", True), - "item_id": item_id, - "response_create": bool(event.get("response_create", False)), - } - if transcript: - payload["transcript"] = transcript - return payload - if event_type == "input_audio_buffer.clear": - self._reset_realtime_input_buffer_state() - return {"type": "input_audio_buffer.clear", "reason": event_type} - if event_type == "output_audio_buffer.clear": - payload = {"type": "output_audio_buffer.clear", "reason": event_type} - response_id = event.get("response_id") - if not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id or self._last_response_id - if self._response_is_done(response_id): - await self._send_realtime_payload( - { - "type": "output_audio_buffer.cleared", - "response_id": response_id, - } - ) - return None - if isinstance(response_id, str) and response_id: - payload["response_id"] = response_id - return payload - if event_type == "response.cancel": - payload = {"type": "response.cancel", "reason": event_type} - response_id = event.get("response_id") - if not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id or self._last_response_id - if self._response_is_done(response_id): - await self._send_realtime_payload( - self._realtime_error_payload( - "response_not_active", - f"Response is already complete: {response_id}", - event_id=event.get("event_id"), - ) - ) - return None - if isinstance(response_id, str) and response_id: - payload["response_id"] = response_id - return payload - if event_type == "response.create": - response_payload = event.get("response") - if isinstance(response_payload, dict): - format_error = self._validate_realtime_response_audio_formats(response_payload) - if format_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - format_error, - event_id=event.get("event_id"), - ) - ) - return None - return { - "type": "response.create", - "response": response_payload if isinstance(response_payload, dict) else {}, - } - if event_type in {"session.close", "close"}: - return {"type": "session.close"} - return event + """Wire JSON of a transport-level error (derived from the typed ``ErrorEvent``).""" + return error_event(code, message, event_id=event_id, param=param).to_realtime() @staticmethod - def _normalize_conversation_item(item: dict[str, object]) -> dict[str, object]: - normalized = dict(item) - normalized.setdefault("id", f"item_{uuid4().hex}") - normalized.setdefault("object", "realtime.item") - normalized.setdefault("type", "message") - normalized.setdefault("status", "completed") - if "role" not in normalized and normalized.get("type") == "message": - normalized["role"] = "user" - if not isinstance(normalized.get("content"), list): - normalized["content"] = [] - return normalized - - @staticmethod - def _truncate_realtime_item_content( - item: dict[str, object], - *, - content_index: int, - audio_end_ms: int, - ) -> None: - content = item.get("content") - if not isinstance(content, list) or not content: - return - index = max(0, int(content_index)) - if index >= len(content): - return - part = content[index] - if not isinstance(part, dict): - return - transcript = part.get("transcript") - if not isinstance(transcript, str) or not transcript: - return - marks = part.get("audio_text_marks") - if isinstance(marks, list): - keep_chars = RealtimeInputTranslator._text_chars_for_audio_ms_from_marks( - audio_end_ms, - len(transcript), - marks, - final_ms=part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms"), - ) - part["transcript"] = transcript[:keep_chars].rstrip() - return - duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") - if isinstance(duration_ms, int | float) and duration_ms > 0: - keep_chars = int(len(transcript) * max(0.0, min(1.0, int(audio_end_ms) / float(duration_ms)))) - part["transcript"] = transcript[:keep_chars].rstrip() - elif audio_end_ms <= 0: - part["transcript"] = "" - - @staticmethod - def _validate_realtime_item_truncate( - item: dict[str, object], - *, - content_index: int, - audio_end_ms: int, - ) -> str | None: - if item.get("type") != "message" or item.get("role") != "assistant": - return "conversation.item.truncate only supports assistant message items" - if audio_end_ms < 0: - return "conversation.item.truncate requires non-negative audio_end_ms" - content = item.get("content") - if not isinstance(content, list) or not content: - return None - index = max(0, int(content_index)) - if index >= len(content): - return f"conversation.item.truncate content_index out of range: {content_index}" - part = content[index] - if not isinstance(part, dict): - return "conversation.item.truncate target content part is invalid" - if part.get("type") not in {"audio", "output_audio"}: - return "conversation.item.truncate target content part must be audio" - duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") - if isinstance(duration_ms, int | float) and int(duration_ms) >= 0 and audio_end_ms > int(duration_ms): - return "conversation.item.truncate audio_end_ms exceeds item audio duration" - return None - - @staticmethod - def _text_chars_for_audio_ms_from_marks( - audio_end_ms: int, - text_len: int, - marks: list[object], - *, - final_ms: object | None = None, - ) -> int: - if text_len <= 0: - return 0 - clean_marks: list[tuple[int, int]] = [] - for mark in marks: - if not isinstance(mark, dict): - continue - raw_text_chars = mark.get("text_chars") - raw_audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) - if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): - continue - clean_marks.append((max(0, int(raw_audio_end_ms)), min(text_len, max(0, int(raw_text_chars))))) - if not clean_marks: - return 0 if audio_end_ms <= 0 else text_len - clean_marks.sort(key=lambda item: item[0]) - audio_end_ms = max(0, int(audio_end_ms)) - if audio_end_ms <= 0: - return 0 - previous_ms = 0 - previous_chars = 0 - for mark_ms, mark_chars in clean_marks: - mark_ms = max(previous_ms, mark_ms) - mark_chars = max(previous_chars, min(text_len, mark_chars)) - if audio_end_ms <= mark_ms: - if mark_ms <= previous_ms: - return mark_chars - ratio = (audio_end_ms - previous_ms) / max(1, mark_ms - previous_ms) - return int(previous_chars + (mark_chars - previous_chars) * max(0.0, min(1.0, ratio))) - previous_ms = mark_ms - previous_chars = mark_chars - if isinstance(final_ms, int | float) and int(final_ms) > previous_ms: - if audio_end_ms >= int(final_ms): - return text_len - ratio = (audio_end_ms - previous_ms) / max(1, int(final_ms) - previous_ms) - return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) - return text_len if audio_end_ms >= previous_ms else previous_chars - - def _conversation_item_added_events(self, item: dict[str, object]) -> list[dict[str, object]]: - item_id = item.get("id") - explicit_previous_item_id = item.pop("_previous_item_id", None) - previous_item_id = ( - explicit_previous_item_id if isinstance(explicit_previous_item_id, str) else self._last_conversation_item_id - ) - if isinstance(item_id, str) and item_id: - self._last_conversation_item_id = item_id - return [ - { - "type": "conversation.item.added", - "previous_item_id": previous_item_id, - "item": item, - }, - ] - - def _conversation_item_done_event(self, item: dict[str, object]) -> dict[str, object]: - item_id = item.get("id") - previous_item_id = self._previous_item_id(item_id) if isinstance(item_id, str) else None - if isinstance(item_id, str) and item_id: - already_known = item_id in self._conversation_items - self._conversation_items[item_id] = item - # Completing an item that was already added must not move it to the - # end of the conversation. A later user item may have been added - # while this response was still in progress. - if not already_known: - self._last_conversation_item_id = item_id - return { - "type": "conversation.item.done", - "previous_item_id": previous_item_id, - "item": item, - } - - def _remove_conversation_item(self, item_id: str) -> bool: - removed = self._conversation_items.pop(item_id, None) is not None - if self._last_conversation_item_id == item_id: - remaining_ids = list(self._conversation_items) - self._last_conversation_item_id = remaining_ids[-1] if remaining_ids else None - self._item_truncation_cursors.pop(item_id, None) - return removed - - def _response_output_item_added_events( - self, - *, - response_id: object, - item: dict[str, object], - ) -> list[dict[str, object]]: - return [ - { - "type": "response.output_item.added", - "response_id": response_id, - "output_index": 0, - "item": item, - }, - ] - - @staticmethod - def _response_content_part(*, transcript: str = "") -> dict[str, object]: - return { - "type": "audio", - "transcript": transcript, - } - - @staticmethod - def _response_text_content_part(*, text: str = "") -> dict[str, object]: - return { - "type": "text", - "text": text, - } - - @staticmethod - def _response_item_content_part( - *, - transcript: str = "", - audio_duration_ms: int | None = None, - audio_text_marks: list[dict[str, int]] | None = None, - ) -> dict[str, object]: - part: dict[str, object] = { - "type": "output_audio", - "transcript": transcript, - } - if audio_duration_ms is not None: - part["audio_duration_ms"] = int(audio_duration_ms) - if audio_text_marks: - part["audio_text_marks"] = [dict(mark) for mark in audio_text_marks] - return part - - @staticmethod - def _response_item_text_content_part(*, text: str = "") -> dict[str, object]: - return { - "type": "output_text", - "text": text, - } - - def _previous_item_id(self, item_id: str) -> str | None: - previous: str | None = None - for known_id in self._conversation_items: - if known_id == item_id: - return previous - previous = known_id - if self._last_conversation_item_id == item_id: - return previous - return self._last_conversation_item_id - - def _input_audio_buffer_committed_event( - self, - *, - item_id: str, - ) -> dict[str, object]: - return { - "type": "input_audio_buffer.committed", - "previous_item_id": self._previous_item_id(item_id), - "item_id": item_id, - } - - def _session_create_from_realtime(self, session_payload: dict[str, object]) -> dict[str, object]: - model = session_payload.get("model") - audio_config = session_payload.get("audio") - audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None - audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - extra_body_payload = session_payload.get("extra_body") - extra_body = ( - normalize_native_duplex_key(dict(extra_body_payload)) if isinstance(extra_body_payload, dict) else {} - ) - extra_body["realtime_session_payload"] = self._json_safe_realtime_payload(session_payload) - if isinstance(session_payload.get("tools"), list): - extra_body["realtime_tools"] = session_payload["tools"] - if isinstance(session_payload.get("tool_choice"), str | dict): - extra_body["realtime_tool_choice"] = session_payload["tool_choice"] - metadata = session_payload.get("metadata") - if isinstance(metadata, dict): - extra_body["realtime_metadata"] = dict(metadata) - include = session_payload.get("include") - if isinstance(include, list): - extra_body["realtime_include"] = list(include) - prompt = session_payload.get("prompt") - if isinstance(prompt, dict): - extra_body["realtime_prompt"] = dict(prompt) - input_audio_transcription = self._input_audio_transcription_config(session_payload) - if isinstance(input_audio_transcription, dict): - extra_body["realtime_input_audio_transcription"] = dict(input_audio_transcription) - noise_reduction = session_payload.get("input_audio_noise_reduction") - if isinstance(noise_reduction, dict): - extra_body["realtime_input_audio_noise_reduction"] = dict(noise_reduction) - audio_noise_reduction = audio_input.get("noise_reduction") if isinstance(audio_input, dict) else None - if isinstance(audio_noise_reduction, dict): - extra_body["realtime_input_audio_noise_reduction"] = dict(audio_noise_reduction) - audio = session_payload.get("audio") - if isinstance(audio, dict): - extra_body["realtime_audio"] = dict(audio) - if isinstance(session_payload.get("tracing"), str | dict): - extra_body["realtime_tracing"] = session_payload["tracing"] - response_format = self._duplex_response_format(self._output_audio_format) - extra_body.setdefault("realtime_output_audio_format", self._output_audio_format) - overlap_fields = self._realtime_overlap_fields(session_payload) - voice = session_payload.get("voice") - if not isinstance(voice, str) and isinstance(audio_output, dict): - voice = audio_output.get("voice") - speed = session_payload.get("speed") - if not isinstance(speed, int | float) and isinstance(audio_output, dict): - speed = audio_output.get("speed") - return { - "type": "session.create", - "session_id": session_payload.get("session_id") or session_payload.get("id"), - "session": { - "model": model, - "modalities": ( - session_payload.get("modalities") or session_payload.get("output_modalities") or ["text", "audio"] - ), - "instructions": session_payload.get("instructions"), - "voice": voice, - "ref_audio": session_payload.get("ref_audio"), - "response_format": response_format, - "temperature": session_payload.get("temperature"), - "max_tokens": self.realtime_max_output_tokens( - session_payload.get("max_response_output_tokens") - or session_payload.get("max_output_tokens") - or session_payload.get("max_tokens") - ), - "speed": speed, - "idle_timeout_s": session_payload.get("idle_timeout_s") or 300.0, - **overlap_fields, - "extra_body": extra_body, - }, - } - - def _apply_realtime_session_defaults(self, session_payload: dict[str, object]) -> None: - input_format: object = session_payload.get("input_audio_format") - audio_config = session_payload.get("audio") - if input_format is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - input_format = audio_input.get("format") - input_format_configured = input_format is not None - input_format, input_rate = self._parse_realtime_audio_format(input_format) - if input_format is not None and not ( - isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS - ): - raise ValueError(f"Unsupported input_audio_format: {input_format}") - if isinstance(input_format, str): - self._input_audio_format = input_format - output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") - output_rate_raw: object | None = None - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") - output_format, output_rate = self._parse_realtime_audio_format(output_format) - if output_rate is None and output_rate_raw is not None: - output_rate = validate_input_sample_rate_hz(output_rate_raw) - if output_format is not None and not ( - isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - raise ValueError(f"Unsupported output_audio_format: {output_format}") - if isinstance(output_format, str): - self._output_audio_format = self._realtime_output_format(output_format) - sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") - if sample_rate is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") - if sample_rate is None: - sample_rate = input_rate - if sample_rate is not None: - self._input_sample_rate_hz = validate_input_sample_rate_hz(sample_rate) - elif input_format_configured: - default_input_rate = realtime_default_sample_rate_hz(input_format) - if default_input_rate is not None: - self._input_sample_rate_hz = default_input_rate - if output_rate is not None: - self._output_sample_rate_hz = output_rate - overlap_fields = self._realtime_overlap_fields(session_payload) - overlap_silence_rms = overlap_fields.get("overlap_silence_rms") - if isinstance(overlap_silence_rms, int | float): - self._overlap_silence_rms = max(0.0, float(overlap_silence_rms)) - - @classmethod - def _validate_realtime_response_audio_formats(cls, response_payload: dict[str, object]) -> str | None: - output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") - audio_config = response_payload.get("audio") - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - try: - parsed_output, _ = cls._parse_realtime_audio_format(output_format) - except ValueError as exc: - return str(exc) - if output_format is not None and not ( - isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - return f"Unsupported output_audio_format: {output_format}" - return None - - @staticmethod - def _json_safe_realtime_payload(payload: dict[str, object]) -> dict[str, object]: - clean: dict[str, object] = {} - for key, value in payload.items(): - if key == "extra_body": - continue - if isinstance(value, str | int | float | bool) or value is None: - clean[key] = value - elif isinstance(value, dict): - clean[key] = RealtimeInputTranslator._json_safe_realtime_payload(value) - elif isinstance(value, list): - clean[key] = [ - (RealtimeInputTranslator._json_safe_realtime_payload(item) if isinstance(item, dict) else item) - for item in value - if isinstance(item, str | int | float | bool | dict) or item is None - ] - return clean - - @staticmethod - def _parse_realtime_audio_format(raw_format: object) -> tuple[object, int | None]: - def normalize_format(fmt: str) -> str: - normalized = fmt.lower() - if normalized in {"audio/pcm", "pcm"}: - return "pcm16" - if normalized in {"audio/wav", "wav"}: - return "wav" - if normalized in {"audio/pcm16", "pcm16", "pcm_s16le", "s16le"}: - return "pcm16" - if normalized in {"audio/pcm_f32le", "pcm_f32le", "f32le"}: - return "pcm_f32le" - if normalized in {"audio/g711_ulaw", "g711_ulaw", "g711-ulaw", "ulaw", "mulaw"}: - return "g711_ulaw" - if normalized in {"audio/g711_alaw", "g711_alaw", "g711-alaw", "alaw"}: - return "g711_alaw" - return fmt - - if isinstance(raw_format, str): - return normalize_format(raw_format), None - if not isinstance(raw_format, dict): - return raw_format, None - rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") - sample_rate_hz = validate_input_sample_rate_hz(rate) if rate is not None else None - fmt = raw_format.get("type") or raw_format.get("format") - if not isinstance(fmt, str): - return raw_format, sample_rate_hz - return normalize_format(fmt), sample_rate_hz - - @staticmethod - def _duplex_response_format(realtime_format: str) -> str: - normalized = realtime_format.lower() - if normalized in {"pcm16", "pcm_s16le", "s16le"}: - return "pcm" - if normalized in {"g711_ulaw", "g711_alaw"}: - return "pcm" - if normalized in {"wav", "pcm"}: - return normalized - return "wav" - - @staticmethod - def _realtime_output_format(duplex_format: object) -> str: - if isinstance(duplex_format, str) and duplex_format.lower() in {"g711_ulaw", "g711_alaw"}: - return duplex_format.lower() - if isinstance(duplex_format, str) and duplex_format.lower() == "pcm": - return "pcm16" - return str(duplex_format or "wav") - - @staticmethod - def _is_supported_realtime_input_format(fmt: object) -> bool: - return isinstance(fmt, str) and fmt.lower() in REALTIME_INPUT_AUDIO_FORMATS - - @staticmethod - def _input_explicitly_non_speech(event: dict[str, object]) -> bool: - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return not value - vad = event.get("vad") - if isinstance(vad, dict): - value = vad.get("is_speech") - if isinstance(value, bool): - return not value - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) < 0.5 - probability = event.get("speech_probability") - return isinstance(probability, int | float) and float(probability) < 0.5 - - def _input_looks_like_speech(self, event: dict[str, object], *, audio: object, fmt: object) -> bool: - if RealtimeInputTranslator._input_explicitly_non_speech(event): - return False - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return value - vad = event.get("vad") - if isinstance(vad, dict): - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) >= 0.5 - probability = event.get("speech_probability") - if isinstance(probability, int | float): - return float(probability) >= 0.5 - if fmt != "pcm_f32le" or not isinstance(audio, str): - return True - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return True - if len(raw) < 4 or len(raw) % 4 != 0: - return True - samples = np.frombuffer(raw, dtype=np.float32) - if samples.size == 0: - return False - rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) - threshold = event.get("overlap_silence_rms") - if not isinstance(threshold, int | float): - vad = event.get("vad") - if isinstance(vad, dict): - threshold = vad.get("silence_rms") - silence_rms = float(threshold) if isinstance(threshold, int | float) else self._overlap_silence_rms - return rms >= max(0.0, silence_rms) - - @staticmethod - def _copy_realtime_input_hints(source: dict[str, object], target: dict[str, object]) -> None: - for key in ( - "duration_ms", - "audio_duration_ms", - "audio_start_ms", - "audio_end_ms", - "is_speech", - "speech", - "speech_probability", - "vad", - "overlap_action", - "overlap", - "force_barge_in", - "force_listen", - "text", - "transcript", - ): - if key in source: - target[key] = source[key] - - @classmethod - def _realtime_overlap_fields(cls, session_payload: dict[str, object]) -> dict[str, object]: - fields: dict[str, object] = {} - if isinstance(session_payload.get("overlap_policy"), str): - fields["overlap_policy"] = session_payload["overlap_policy"] - for key in ("overlap_short_ack_ms", "overlap_barge_in_ms", "overlap_silence_rms"): - value = session_payload.get(key) - if isinstance(value, int | float): - fields[key] = value - - if isinstance(session_payload.get("playback_commit_policy"), str): - fields["playback_commit_policy"] = session_payload["playback_commit_policy"] - return fields - - @staticmethod - def _validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: - """Validate omni-duplex camera frames on input_audio_buffer.append. - - Wire contract matches the official MiniCPM-o duplex loop: one base - base64 JPEG per ~1 s audio chunk, optionally followed by that unit's - stacked composite tiling the sub-frames captured inside it (at most 2 - images either way). A caller-supplied ``max_slice_nums`` is rejected - rather than silently ignored: slicing is Stage 0's decision here, and - Stage 0 already applies the official HD suggestion for a stacked unit - (``max_slice_nums=[2, 1]``). The wire simply does not let the client - choose it. - """ - if max_slice_nums not in (None, 1): - return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" - if not isinstance(video_frames, list): - return "video_frames must be a list of base64-encoded images" - frames = [frame for frame in video_frames if frame is not None] - if len(frames) > 2: - return "video_frames carries more than 2 frames for one append; send ~1 frame per 1 s chunk" - for frame in frames: - if not isinstance(frame, str) or not frame: - return "video_frames entries must be non-empty base64 strings" - if len(frame) > 4_000_000: - return "video_frames entry exceeds 4MB base64; reduce capture resolution or JPEG quality" - try: - header = base64.b64decode(frame[:64] + "=" * (-len(frame[:64]) % 4)) - except (binascii.Error, ValueError): - return "video_frames entries must be valid base64" - if not (header.startswith(b"\xff\xd8") or header.startswith(b"\x89PNG")): - return "video_frames entries must be JPEG or PNG images" - return None - - def _validate_realtime_turn_detection( - self, - session_payload: dict[str, object], - *, - configured: bool, - normalized: dict[str, object] | None, - ) -> str | None: - effective_turn_detection = normalized if configured else self._turn_detection - if effective_turn_detection is None: - if session_payload.get("overlap_policy") == "barge_in_on_speech": - return "overlap_policy='barge_in_on_speech' requires turn_detection.type='server_vad'" - return None - - interrupt_response = effective_turn_detection.get("interrupt_response") - if interrupt_response is None: - # Serving resolves the default and validates overlap policy once - # the session's runtime capabilities are known. - return None - desired_policy = "barge_in_on_speech" if interrupt_response else "listen_only" - overlap_policy = session_payload.get("overlap_policy") - if isinstance(overlap_policy, str) and overlap_policy != desired_policy: - return f"overlap_policy={overlap_policy!r} conflicts with turn_detection; expected {desired_policy!r}" - return None - - def validate_server_vad_input_audio(self, session_payload: dict[str, object]) -> str | None: - """Validate turn-based VAD audio after serving resolves runtime capabilities.""" - candidate = self._pending_turn_detection_update - input_format, input_rate, *_ = candidate[3] if candidate is not None else self._realtime_protocol_config() - audio_config = session_payload.get("audio") - audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None - channels = audio_input.get("channels") if isinstance(audio_input, dict) else None - raw_format = session_payload.get("input_audio_format") - if raw_format is None and isinstance(audio_input, dict): - raw_format = audio_input.get("format") - if channels is None and isinstance(raw_format, dict): - channels = raw_format.get("channels") - if channels is None: - channels = 1 - if input_format != "pcm16" or channels != 1 or input_rate not in {16_000, 24_000}: - return "server_vad requires mono PCM16 input at 16 kHz or 24 kHz" - return None - - def _prepare_realtime_turn_detection_update( - self, - session_payload: dict[str, object], - *, - configured: bool, - normalized: dict[str, object] | None, - ) -> tuple[bool, dict[str, object] | None]: - if self._pending_turn_detection_update is not None: - raise RuntimeError("A Realtime session update is already pending") - - turn_detection = dict(normalized) if configured and normalized is not None else None - if configured: - if turn_detection is None: - session_payload.setdefault("overlap_policy", "listen_only") - elif "interrupt_response" in turn_detection: - session_payload.setdefault( - "overlap_policy", "barge_in_on_speech" if turn_detection["interrupt_response"] else "listen_only" - ) - session_payload["turn_detection"] = turn_detection - return configured, turn_detection - - def _realtime_protocol_config(self) -> tuple[str, int, str, int | None, float]: - return ( - self._input_audio_format, - self._input_sample_rate_hz, - self._output_audio_format, - self._output_sample_rate_hz, - self._overlap_silence_rms, - ) - - def _set_realtime_protocol_config( - self, - config: tuple[str, int, str, int | None, float], - ) -> None: - ( - self._input_audio_format, - self._input_sample_rate_hz, - self._output_audio_format, - self._output_sample_rate_hz, - self._overlap_silence_rms, - ) = config - - def pending_realtime_input_audio_config_changed(self) -> bool: - candidate = self._pending_turn_detection_update - if candidate is None: - return False - candidate_protocol_config = candidate[3] - return candidate_protocol_config[:2] != self._realtime_protocol_config()[:2] - - async def wait_for_realtime_turn_detection_update(self) -> None: - candidate = self._pending_turn_detection_update - if candidate is not None: - await candidate[2].wait() - - def commit_realtime_turn_detection_update(self, server_vad: ServerVADConfig | None = None) -> None: - self._resolve_realtime_turn_detection_update(commit=True, server_vad=server_vad) - - def reject_realtime_turn_detection_update(self) -> None: - self._resolve_realtime_turn_detection_update(commit=False) - - def _resolve_realtime_turn_detection_update( - self, *, commit: bool, server_vad: ServerVADConfig | None = None - ) -> None: - candidate = self._pending_turn_detection_update - if candidate is None: - return - configured, turn_detection, resolved, protocol_config = candidate - if commit: - self._set_realtime_protocol_config(protocol_config) - if configured: - self._turn_detection = server_vad.as_dict() if server_vad is not None else turn_detection - self._turn_detection_configured = True - self._pending_turn_detection_update = None - resolved.set() - - @staticmethod - def _input_audio_transcription_config(session_payload: dict[str, object]) -> dict[str, object] | None: - transcription = session_payload.get("input_audio_transcription") - if isinstance(transcription, dict): - return transcription - audio_config = session_payload.get("audio") - if not isinstance(audio_config, dict): - return None - audio_input = audio_config.get("input") - if not isinstance(audio_input, dict): - return None - transcription = audio_input.get("transcription") - return transcription if isinstance(transcription, dict) else None - - @staticmethod - def realtime_max_output_tokens(value: object) -> int | None: - """Normalize Realtime max output tokens. - - OpenAI Realtime clients commonly send ``"inf"`` for unbounded output. - The duplex core represents that as ``None`` because model-specific - defaults remain in force. - """ - if value is None: - return None - if isinstance(value, str) and value.strip().lower() in {"inf", "infinity", "unlimited"}: - return None - if isinstance(value, int) and value > 0: - return int(value) - return None - - async def _conversation_item_to_duplex(self, event: dict[str, object]) -> dict[str, object] | None: - raw_item = event.get("item") - if not isinstance(raw_item, dict): - return None - public_item = self._normalize_conversation_item(raw_item) - previous_item_id = event.get("previous_item_id") - if isinstance(previous_item_id, str): - public_item["_previous_item_id"] = previous_item_id - item_id = str(public_item["id"]) - item_type = public_item.get("type") - role = public_item.get("role") - if item_type == "function_call_output": - call_id = public_item.get("call_id") - output = public_item.get("output") - matching_call = any( - known.get("type") == "function_call" and known.get("call_id") == call_id - for known in self._conversation_items.values() - ) - duplicate_output = any( - known.get("type") == "function_call_output" and known.get("call_id") == call_id - for known in self._conversation_items.values() - ) - if not isinstance(call_id, str) or not call_id or not matching_call: - await self._send_realtime_payload( - self._realtime_error_payload( - "invalid_function_call_output", - "function_call_output requires the call_id of a completed function call", - event_id=event.get("event_id"), - param="item.call_id", - ) - ) - return None - if not isinstance(output, str) or duplicate_output: - await self._send_realtime_payload( - self._realtime_error_payload( - "invalid_function_call_output", - ( - "function_call_output requires a string output" - if not isinstance(output, str) - else f"function_call_output already exists for call_id {call_id}" - ), - event_id=event.get("event_id"), - param="item.output", - ) - ) - return None - # Do not retain the result until SessionRunner has delivered its - # runtime update. On success, the resulting - # ``conversation.item.created`` event records it in the shared - # input/output protocol state; on failure the client may retry the - # same call_id instead of being rejected by provisional history. - content = public_item.get("content") - if not isinstance(content, list): - return None - parsed_audio_formats: dict[int, tuple[object, int | None]] = {} - for index, part in enumerate(content): - if not isinstance(part, dict) or part.get("type") not in {"input_audio", "audio"}: - continue - raw_format = part.get("format") - try: - parsed_format = self._parse_realtime_audio_format(raw_format or self._input_audio_format) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - str(exc), - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - if not self._is_supported_realtime_input_format(parsed_format[0]): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - f"Unsupported input_audio format in conversation.item.create: {raw_format}", - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - parsed_audio_formats[index] = parsed_format - if item_type != "message" or role in {"assistant", "system"}: - event["item"] = public_item - return { - "type": "turn.signal", - "event": "conversation.item.create", - "payload": {"item": public_item}, - } - server_vad_full_audio_item = self._turn_based_server_vad - text_chunks: list[str] = [] - audio_events: list[dict[str, object]] = [] - normalized_content: list[object] = [] - normalized_full_audio = False - for index, part in enumerate(content): - if not isinstance(part, dict): - normalized_content.append(part) - continue - if part.get("type") in {"input_text", "text"} and isinstance(part.get("text"), str): - text_chunks.append(str(part["text"])) - if part.get("type") in {"input_audio", "audio"}: - audio = part.get("audio") or part.get("data") - if not isinstance(audio, str) or not audio: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - "conversation.item.create input_audio must not be empty", - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - raw_part_format = part.get("format") - fmt, format_rate = parsed_audio_formats[index] - explicit_sample_rate_hz = part.get("sample_rate_hz") - if explicit_sample_rate_hz is None: - explicit_sample_rate_hz = part.get("sample_rate") - if explicit_sample_rate_hz is None and isinstance(raw_part_format, dict): - explicit_sample_rate_hz = ( - raw_part_format.get("rate") - if "rate" in raw_part_format - else raw_part_format.get("sample_rate_hz", raw_part_format.get("sample_rate")) - ) - sample_rate_hz = explicit_sample_rate_hz or format_rate or self._input_sample_rate_hz - if server_vad_full_audio_item and ( - fmt != self._input_audio_format - or (explicit_sample_rate_hz is not None and explicit_sample_rate_hz != self._input_sample_rate_hz) - ): - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - "conversation.item.create input_audio must use the Session input format and sample rate", - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - if server_vad_full_audio_item: - try: - normalized_audio = self._normalize_complete_audio_item_part(audio, self._input_sample_rate_hz) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - str(exc), - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - normalized_part = dict(part) - normalized_part.pop("data", None) - normalized_part.update(audio=normalized_audio, format="wav", sample_rate_hz=16_000) - normalized_content.append(normalized_part) - normalized_full_audio = True - continue - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, - ) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - str(exc), - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - speech_hints = dict(event) - speech_hints.update(part) - if not self._input_looks_like_speech(speech_hints, audio=audio, fmt=fmt): - continue - self._input_speech_started = True - payload = { - "type": "input_audio_buffer.append", - "audio": audio, - "format": fmt, - "sample_rate_hz": sample_rate_hz, - } - self._copy_realtime_input_hints(part, payload) - self._copy_realtime_input_hints(event, payload) - audio_events.append(payload) - continue - normalized_content.append(part) - if server_vad_full_audio_item and normalized_full_audio: - history_item = dict(public_item) - history_item["content"] = normalized_content - history_item["status"] = "completed" - event["item"] = public_item - return { - "type": "turn.signal", - "event": "conversation.item.create", - "payload": { - "item": public_item, - "history_item": history_item, - }, - } - if audio_events: - transcript = self._input_transcript_from_item(public_item) - for extra_event in audio_events[1:]: - self._pending_outbound.put_nowait(extra_event) - commit_payload: dict[str, object] = { - "type": "input_audio_buffer.commit", - "final": True, - "item_id": item_id, - "response_create": False, - } - if transcript: - commit_payload["transcript"] = transcript - self._pending_outbound.put_nowait(commit_payload) - # This complete item is already committed locally; a later ack - # must not be needed to start accepting the next input turn. - self._reset_realtime_input_buffer_state() - event["item"] = public_item - return audio_events[0] - if not text_chunks: - return None - text_item = dict(public_item) - text_item["status"] = "completed" - event["item"] = text_item - return { - "type": "turn.signal", - "event": "conversation.item.create", - "payload": {"item": text_item}, - } - - @staticmethod - def _normalize_complete_audio_item_part( - audio: str, - sample_rate_hz: int, - ) -> str: - """Package a complete Realtime audio item for the chat fallback path.""" - try: - raw = base64.b64decode(audio.strip(), validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError("conversation.item.create input_audio is not valid Base64") from exc - if not raw or len(raw) % np.dtype(" None: - transcript = event.get("transcript") - if not isinstance(transcript, str): - transcript = event.get("text") if isinstance(event.get("text"), str) else None - if not isinstance(transcript, str): - hints = event.get("hints") - if isinstance(hints, dict): - transcript = hints.get("transcript") - if not isinstance(transcript, str): - transcript = hints.get("text") if isinstance(hints.get("text"), str) else None - if isinstance(transcript, str) and transcript: - if ( - self._input_audio_buffer_transcript_parts - and self._input_audio_buffer_transcript_parts[-1] == transcript - ): - return - self._input_audio_buffer_transcript_parts.append(transcript) - - def _consume_input_transcript_hint(self) -> str: - transcript = "".join(self._input_audio_buffer_transcript_parts).strip() - self._input_audio_buffer_transcript_parts.clear() - return transcript - - @staticmethod - def _input_transcript_from_item(item: dict[str, object]) -> str: - content = item.get("content") - if not isinstance(content, list): - return "" - parts: list[str] = [] - for part in content: - if not isinstance(part, dict): - continue - for key in ("transcript", "text"): - value = part.get(key) - if isinstance(value, str) and value: - if part.get("type") in {"input_audio", "audio", "audio_transcript", "transcript"}: - parts.append(value) - break - return "".join(parts).strip() - - @staticmethod - def _input_audio_transcription_completed_event( - item_id: str, - item: dict[str, object], - ) -> dict[str, object] | None: - content = item.get("content") - if not isinstance(content, list): - return None - transcript_parts: list[str] = [] - for index, part in enumerate(content): - if not isinstance(part, dict): - continue - if part.get("type") != "input_audio": - continue - transcript = part.get("transcript") - if isinstance(transcript, str) and transcript: - transcript_parts.append(transcript) - return { - "type": "conversation.item.input_audio_transcription.completed", - "item_id": item_id, - "content_index": index, - "transcript": "".join(transcript_parts), - } - return None - - async def _emit_input_speech_started(self, event: dict[str, object]) -> None: - if self._input_speech_started: - return - self._input_speech_started = True - if self._active_input_item_id is None: - self._active_input_item_id = f"item_{uuid4().hex}" - audio_start_ms = event.get("audio_start_ms", 0) - await self._send_realtime_payload( - { - "type": "input_audio_buffer.speech_started", - "audio_start_ms": int(audio_start_ms) if isinstance(audio_start_ms, int | float) else 0, - "item_id": self._active_input_item_id, - } - ) - - async def _emit_input_speech_stopped(self, event: dict[str, object], *, item_id: str) -> None: - if not self._input_speech_started: - return - self._input_speech_started = False - audio_end_ms = event.get("audio_end_ms", event.get("audio_ms", 0)) - await self._send_realtime_payload( - { - "type": "input_audio_buffer.speech_stopped", - "audio_end_ms": int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, - "item_id": item_id, - } - ) + def command_error_payload(exc: DuplexCommandError) -> dict[str, object]: + return error_event(exc.code, str(exc), event_id=exc.event_id).to_realtime() diff --git a/vllm_omni/entrypoints/duplex/realtime_output.py b/vllm_omni/entrypoints/duplex/realtime_output.py deleted file mode 100644 index fcbdf87d7..000000000 --- a/vllm_omni/entrypoints/duplex/realtime_output.py +++ /dev/null @@ -1,1007 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -from typing import Any -from uuid import uuid4 - -from vllm_omni.entrypoints.duplex.audio import convert_output_audio -from vllm_omni.entrypoints.duplex.realtime_state import ( - _RealtimeResponseState, -) - - -class RealtimeOutputProjector: - """Project internal duplex events onto the OpenAI Realtime schema.""" - - def _from_duplex_event(self, event: dict[str, Any]) -> list[dict[str, object]]: - event_type = event.get("type") - if event_type == "session.created": - session = self._realtime_session_payload(event.get("session")) - created: dict[str, object] = {"type": "session.created", "session": session} - for key in ( - "incarnation", - "attachment_generation", - "resume_token", - ): - if key in event: - created[key] = event[key] - payloads: list[dict[str, object]] = [created] - if self._initial_session_update: - payloads.append({"type": "session.updated", "session": session}) - self._initial_session_update = False - self._hold_realtime_output_until_session_created = False - if self._held_realtime_payloads: - payloads.extend(self._held_realtime_payloads) - self._held_realtime_payloads = [] - return payloads - if event_type == "session.updated": - session = self._realtime_session_payload(event.get("session")) - return [{"type": "session.updated", "session": session}] - if event_type == "input_audio_buffer.speech_started": - item_id = event.get("item_id") - if isinstance(item_id, str): - self._active_input_item_id = item_id - self._input_speech_started = True - self._input_audio_buffer_has_audio = True - return [ - { - "type": "input_audio_buffer.speech_started", - "audio_start_ms": max(0, int(event.get("audio_start_ms", 0) or 0)), - "item_id": item_id, - } - ] - if event_type == "input_audio_buffer.speech_stopped": - self._input_speech_started = False - return [ - { - "type": "input_audio_buffer.speech_stopped", - "audio_end_ms": max(0, int(event.get("audio_end_ms", 0) or 0)), - "item_id": event.get("item_id") or self._active_input_item_id, - } - ] - if event_type in { - "session.resumed", - "session.heartbeat_ack", - "session.replaced", - "session.expired", - "session.resync_required", - }: - if event_type == "session.resumed": - self._hold_realtime_output_until_session_created = False - return [dict(event)] - if event_type == "response.created": - response_id = event.get("response_id") - if isinstance(response_id, str) and response_id: - self._active_response_id = response_id - self._last_response_id = response_id - item_id = self._response_item_id(response_id) - modalities = event.get("modalities") - has_audio_modality = not isinstance(modalities, list) or "audio" in modalities - item = { - "id": item_id, - "object": "realtime.item", - "type": "message", - "role": "assistant", - "status": "in_progress", - "content": [], - } - self._conversation_items[item_id] = item - payloads = [ - self._response_created_event(event), - *self._conversation_item_added_events(item), - *self._response_output_item_added_events(response_id=response_id, item=item), - ] - if has_audio_modality: - payloads.extend(self._ensure_response_audio_part_added(response_id)) - return payloads - if event_type == "response.listen": - listen_response: dict[str, object] = { - "object": "realtime.response", - "status": "listening", - "metadata": event, - } - listen_payload: dict[str, object] = { - "type": "response.listen", - "session_id": event.get("session_id"), - "epoch": event.get("epoch"), - "response": listen_response, - } - listen_response_id = event.get("response_id") - if isinstance(listen_response_id, str) and listen_response_id: - # Keep the response identity visible on the Realtime dialect: - # a listen that terminates a precreated response is matched to - # it by this id (clients demultiplex on response_id). - listen_response["id"] = listen_response_id - listen_payload["response_id"] = listen_response_id - return [listen_payload] - if event_type == "response.speak": - response_id = event.get("response_id") - state = self._response_state(response_id) - if state is not None: - if state.speak_emitted: - return [] - state.speak_emitted = True - return [ - { - "type": "response.speak", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 0, - "metadata": self._response_speak_metadata(event), - } - ] - if event_type == "overlap.decision": - return [ - { - "type": "overlap.decision", - "session_id": event.get("session_id"), - "epoch": event.get("epoch"), - "policy": event.get("policy"), - "action": event.get("action"), - "reason": event.get("reason"), - "metadata": event, - } - ] - if event_type == "response.output_audio.delta": - response_id = event.get("response_id") - audio = event.get("audio", "") - payloads: list[dict[str, object]] = [] - if isinstance(audio, str) and audio: - payloads.extend(self._ensure_response_audio_part_added(response_id)) - payloads.extend(self._realtime_audio_delta_events(event, response_id, audio)) - self._refresh_in_progress_response_item(response_id) - text = event.get("text") - has_text = isinstance(text, str) and bool(text) - if has_text: - self._append_response_transcript(response_id, text) - self._refresh_in_progress_response_item(response_id) - # Keep the audio.delta + transcript.delta pair invariant even for - # text-less units (deduplicated continuations, turn-end flush): - # clients that treat the pair as unit-complete would otherwise - # wait on a transcript that never comes. - emit_transcript = has_text or (isinstance(audio, str) and bool(audio)) - if emit_transcript: - if not has_text: - text = "" - payloads.append( - { - "type": "response.output_audio_transcript.delta", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 0, - "delta": text, - } - ) - if event.get("end_of_turn") is True: - payloads.extend(self._realtime_audio_done_events(event, response_id)) - payloads.extend( - self._realtime_response_terminal_events( - event, - response_id, - status="completed", - status_details={ - "type": "completed", - "reason": event.get("finish_reason") or "stop", - }, - ) - ) - return payloads - if event_type == "response.text.delta": - response_id = event.get("response_id") - text = event.get("delta", "") - state = self._response_state(response_id) - if state is not None and isinstance(text, str) and text: - state.text_parts.append(text) - self._refresh_in_progress_response_item(response_id) - payloads = self._ensure_response_text_part_added(response_id) - payloads.append( - { - "type": "response.output_text.delta", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 1 if state is not None and state.audio_part_added else 0, - "delta": text, - } - ) - return payloads - if event_type == "response.done": - response_id = event.get("response_id") - status = event.get("status") if isinstance(event.get("status"), str) else "completed" - status_details = event.get("status_details") if isinstance(event.get("status_details"), dict) else None - return [ - *self._realtime_audio_done_events(event, response_id), - *self._realtime_response_terminal_events( - event, - response_id, - status=status, - status_details=status_details, - ), - ] - if event_type == "input.committed": - event_item_id = event.get("item_id") - # Acknowledgements carry their own identity; untracked silence - # commits must not consume another turn's item ID. - item_id = event_item_id if isinstance(event_item_id, str) and event_item_id else f"item_{uuid4().hex}" - # The reader may already have accepted audio for the next turn. - if self._active_input_item_id == item_id: - self._reset_realtime_input_buffer_state() - item = self._conversation_items.get(item_id) - payloads = [self._input_audio_buffer_committed_event(item_id=item_id)] - if item is None: - message = event.get("message") - no_response = event.get("no_response") is True - is_speech = event.get("is_speech") - item = { - "id": item_id, - "object": "realtime.item", - "type": "message", - "role": "user", - "status": "completed", - "content": ( - [{"type": "input_audio", "transcript": "", "is_speech": False}] - if no_response and is_speech is False - else self._user_item_content_from_duplex_message(message) - ), - } - self._conversation_items[item_id] = item - payloads.extend(self._conversation_item_added_events(item)) - item["status"] = "completed" - transcription_event = self._input_audio_transcription_completed_event(item_id, item) - if transcription_event is not None: - payloads.append(transcription_event) - payloads.append(self._conversation_item_done_event(item)) - return payloads - if event_type == "input.cancelled": - self._reset_realtime_input_buffer_state() - return [{"type": "input_audio_buffer.cleared"}] - if event_type == "input_audio_buffer.cleared": - self._reset_realtime_input_buffer_state() - return [{"type": "input_audio_buffer.cleared"}] - if event_type == "audio.cancelled": - response_id = event.get("response_id") - payloads: list[dict[str, object]] = [] - if event.get("reason") == "output_audio_buffer_clear": - if not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id or self._last_response_id - payloads.append( - { - "type": "output_audio_buffer.cleared", - "response_id": response_id, - } - ) - if not isinstance(response_id, str) or not response_id: - return payloads - elif not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id - if not isinstance(response_id, str) or not response_id: - return payloads - if self._response_is_done(response_id): - return payloads - committed_ms = event.get("committed_ms") - if isinstance(committed_ms, int | float): - item_id = self._response_item_id(response_id) - committed_audio_ms = max(0, int(committed_ms)) - self._item_truncation_cursors[item_id] = (0, committed_audio_ms) - item = self._conversation_items.get(item_id) - if item is not None: - self._truncate_realtime_item_content( - item, - content_index=0, - audio_end_ms=committed_audio_ms, - ) - # Align server history with acknowledged playback; only explicit - # client truncation emits ``conversation.item.truncated``. - payloads.extend(self._realtime_audio_done_events(event, response_id)) - payloads.extend( - self._realtime_response_terminal_events( - event, - response_id, - status="cancelled", - status_details={ - "type": "cancelled", - "reason": event.get("reason") or "client_cancelled", - }, - ) - ) - if isinstance(response_id, str) and response_id == self._active_response_id: - self._active_response_id = None - return payloads - if event_type == "playback.acknowledged": - return [{"type": "playback.acknowledged", "event": event}] - if event_type == "conversation.item.created": - item = event.get("item") - if isinstance(item, dict) and isinstance(item.get("id"), str): - item_id = str(item["id"]) - already_known = item_id in self._conversation_items - self._conversation_items[item_id] = item - if already_known: - if item.get("status") == "completed": - return [self._conversation_item_done_event(item)] - return [] - payloads = self._conversation_item_added_events(item) - if item.get("status") == "completed": - payloads.append(self._conversation_item_done_event(item)) - return payloads - return [{"type": "conversation.item.created", "item": item, "event": event}] - if event_type == "conversation.item.deleted": - item_id = event.get("item_id") - if isinstance(item_id, str): - self._remove_conversation_item(item_id) - return [ - { - "type": "conversation.item.deleted", - "item_id": item_id, - "event": event, - } - ] - if event_type == "conversation.item.truncated": - item_id = event.get("item_id") - audio_end_ms = event.get("audio_end_ms") - content_index = event.get("content_index", 0) - if isinstance(item_id, str): - item = self._conversation_items.get(item_id) - if item is not None: - self._truncate_realtime_item_content( - item, - content_index=int(content_index) if isinstance(content_index, int | float) else 0, - audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, - ) - return [ - { - "type": "conversation.item.truncated", - "item_id": item_id, - "content_index": content_index, - "audio_end_ms": audio_end_ms, - "event": event, - } - ] - if event_type == "response.output_item.done": - response_id = event.get("response_id") - return self._realtime_response_terminal_events( - event, - response_id, - status="completed", - status_details={ - "type": "completed", - "reason": event.get("finish_reason") or "stop", - }, - ) - if event_type == "function_call.done": - call_id = event.get("call_id") - name = event.get("name") - arguments = event.get("arguments", "") - if not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name: - return [{"type": "duplex.function_call.done", "event": event}] - if not isinstance(arguments, str): - arguments = str(arguments) - response_id = f"resp_{uuid4().hex}" - item = { - "id": f"item_{uuid4().hex}", - "object": "realtime.item", - "type": "function_call", - "status": "completed", - "name": name, - "call_id": call_id, - "arguments": arguments, - } - self._conversation_items[str(item["id"])] = item - response = { - "id": response_id, - "object": "realtime.response", - "status": "completed", - "status_details": {"type": "completed", "reason": "completed"}, - "output": [dict(item)], - "metadata": {}, - } - return [ - {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, - *self._conversation_item_added_events(item), - { - "type": "response.output_item.added", - "response_id": response_id, - "output_index": 0, - "item": dict(item), - }, - { - "type": "response.function_call_arguments.delta", - "response_id": response_id, - "item_id": item["id"], - "output_index": 0, - "call_id": call_id, - "delta": arguments, - }, - { - "type": "response.function_call_arguments.done", - "response_id": response_id, - "item_id": item["id"], - "output_index": 0, - "call_id": call_id, - "arguments": arguments, - }, - { - "type": "response.output_item.done", - "response_id": response_id, - "output_index": 0, - "item": dict(item), - }, - self._conversation_item_done_event(item), - {"type": "response.done", "response": response}, - ] - if event_type == "error": - raw_error = event.get("error") - if isinstance(raw_error, dict): - return [event] - message = str(raw_error or event.get("message") or "Duplex runtime error") - code = str(event.get("code") or "duplex_error") - return [ - self._realtime_error_payload( - code, - message, - event_id=event.get("realtime_event_id") or event.get("event_id"), - param=event.get("param"), - ) - ] - if event_type == "session.closed": - return [{"type": "session.closed", "event": event}] - return [{"type": f"duplex.{event_type}", "event": event}] - - @staticmethod - def _user_item_content_from_duplex_message(message: object) -> list[dict[str, object]]: - if not isinstance(message, dict): - return [{"type": "input_audio", "transcript": ""}] - message_transcript = message.get("transcript") - content = message.get("content") - if isinstance(content, str): - return [{"type": "input_text", "text": content}] - if not isinstance(content, list): - return [ - { - "type": "input_audio", - "transcript": message_transcript if isinstance(message_transcript, str) else "", - } - ] - parts: list[dict[str, object]] = [] - for part in content: - if not isinstance(part, dict): - continue - part_type = part.get("type") - if part_type in {"text", "input_text"} and isinstance(part.get("text"), str): - parts.append({"type": "input_text", "text": part["text"]}) - continue - if part_type == "audio_url": - transcript = ( - part.get("transcript") - if isinstance(part.get("transcript"), str) - else message_transcript - if isinstance(message_transcript, str) - else "" - ) - parts.append({"type": "input_audio", "transcript": transcript}) - continue - if part_type in {"audio", "input_audio"}: - transcript = part.get("transcript") if isinstance(part.get("transcript"), str) else "" - parts.append({"type": "input_audio", "transcript": transcript}) - return parts or [{"type": "input_audio", "transcript": ""}] - - def _realtime_session_payload(self, session: object) -> object: - if not isinstance(session, dict): - return session - payload = dict(session) - payload.setdefault("object", "realtime.session") - payload.setdefault("type", "realtime") - payload.setdefault("id", payload.get("id") or self._default_session_id) - payload.setdefault("model", payload.get("model") or self._default_model) - payload.setdefault("input_audio_format", self._input_audio_format) - payload.setdefault("output_audio_format", self._output_audio_format) - payload.setdefault("modalities", payload.get("modalities") or ["text", "audio"]) - payload.setdefault("output_modalities", payload.get("output_modalities") or payload.get("modalities")) - payload.setdefault("object", "realtime.session") - payload.setdefault( - "audio", - { - "input": { - "format": self._realtime_audio_format_object( - self._input_audio_format, - sample_rate_hz=self._input_sample_rate_hz, - ), - "sample_rate_hz": self._input_sample_rate_hz, - }, - "output": { - "format": self._realtime_audio_format_object( - self._output_audio_format, - sample_rate_hz=self._output_sample_rate_hz, - ), - }, - }, - ) - if self._turn_detection_configured or "turn_detection" in payload: - turn_detection = payload.get("turn_detection", self._turn_detection) - audio = dict(payload["audio"]) if isinstance(payload.get("audio"), dict) else {} - audio_input = dict(audio["input"]) if isinstance(audio.get("input"), dict) else {} - audio_input["turn_detection"] = turn_detection - audio["input"] = audio_input - payload["audio"] = audio - payload["turn_detection"] = turn_detection - payload.setdefault("input_audio_transcription", payload.get("input_audio_transcription")) - payload.setdefault("tracing", payload.get("tracing")) - return payload - - @staticmethod - def _realtime_audio_format_object(fmt: object, *, sample_rate_hz: int | None = None) -> dict[str, object]: - if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le", "pcm"}: - payload: dict[str, object] = {"type": "audio/pcm"} - elif isinstance(fmt, str) and fmt.lower() == "pcm_f32le": - payload = {"type": "audio/pcm_f32le"} - elif isinstance(fmt, str) and fmt.lower() == "g711_ulaw": - payload = {"type": "audio/g711_ulaw"} - elif isinstance(fmt, str) and fmt.lower() == "g711_alaw": - payload = {"type": "audio/g711_alaw"} - else: - payload = {"type": "audio/wav"} - if sample_rate_hz is not None: - payload["rate"] = int(sample_rate_hz) - return payload - - def _response_state( - self, - response_id: object, - *, - event: dict[str, Any] | None = None, - create: bool = True, - ) -> _RealtimeResponseState | None: - if isinstance(response_id, str) and response_id: - key: str | int = response_id - item_id = f"item_{response_id}" - elif event is not None: - key = id(event) - item_id = f"item_{uuid4().hex}" - else: - return None - state = self._response_states.get(key) - if state is None and create: - state = _RealtimeResponseState(item_id=item_id) - self._response_states[key] = state - return state - - def _response_is_done(self, response_id: object) -> bool: - state = self._response_state(response_id, create=False) - return state is not None and state.done_emitted - - def _response_item_id(self, response_id: object) -> str: - state = self._response_state(response_id) - if state is not None: - return state.item_id - return f"item_{uuid4().hex}" - - def _response_done_output_item( - self, - response_id: object, - *, - status: str, - ) -> dict[str, object]: - item_id = self._response_item_id(response_id) - state = self._response_state(response_id) - transcript = state.transcript if state is not None else "" - text = state.text if state is not None else "" - audio_duration_ms = state.audio_duration_ms if state is not None else None - audio_text_marks = state.audio_text_marks if state is not None else None - content: list[dict[str, object]] = [] - if (state is not None and state.audio_part_added) or transcript or audio_duration_ms is not None: - content.append( - self._response_item_content_part( - transcript=transcript, - audio_duration_ms=audio_duration_ms, - audio_text_marks=audio_text_marks, - ) - ) - if text: - content.append(self._response_item_text_content_part(text=text)) - item = { - "id": item_id, - "object": "realtime.item", - "type": "message", - "role": "assistant", - "status": status, - "content": content, - } - self._apply_pending_item_truncation(item) - return item - - def _apply_pending_item_truncation(self, item: dict[str, object]) -> None: - item_id = item.get("id") - if not isinstance(item_id, str) or not item_id: - return - cursor = self._item_truncation_cursors.get(item_id) - if cursor is None: - return - content_index, audio_end_ms = cursor - self._truncate_realtime_item_content( - item, - content_index=content_index, - audio_end_ms=audio_end_ms, - ) - - def _refresh_in_progress_response_item(self, response_id: object) -> None: - if not isinstance(response_id, str) or not response_id: - return - item_id = self._response_item_id(response_id) - item = self._conversation_items.get(item_id) - if not isinstance(item, dict): - return - content = item.get("content") - if not isinstance(content, list): - content = [] - item["content"] = content - - state = self._response_state(response_id) - if state is None: - return - transcript = state.transcript - audio_duration_ms = state.audio_duration_ms - audio_text_marks = state.audio_text_marks - has_audio = state.audio_part_added or bool(transcript) or audio_duration_ms is not None - if has_audio: - audio_part = self._response_item_content_part( - transcript=transcript, - audio_duration_ms=audio_duration_ms, - audio_text_marks=audio_text_marks, - ) - if content and isinstance(content[0], dict) and content[0].get("type") in {"audio", "output_audio"}: - content[0] = audio_part - else: - content.insert(0, audio_part) - - text = state.text - if text: - text_index = ( - 1 - if content - and isinstance(content[0], dict) - and content[0].get("type") - in { - "audio", - "output_audio", - } - else 0 - ) - text_part = self._response_item_text_content_part(text=text) - if ( - len(content) > text_index - and isinstance(content[text_index], dict) - and content[text_index].get("type") in {"text", "output_text"} - ): - content[text_index] = text_part - else: - content.insert(text_index, text_part) - - self._apply_pending_item_truncation(item) - - def _append_response_transcript(self, response_id: object, text: str) -> None: - state = self._response_state(response_id) - if state is None or not text: - return - state.transcript_parts.append(text) - - def _ensure_response_text_part_added(self, response_id: object) -> list[dict[str, object]]: - state = self._response_state(response_id) - if state is None: - return [] - if state.text_part_added: - return [] - state.text_part_added = True - content_index = 1 if state.audio_part_added else 0 - return [ - { - "type": "response.content_part.added", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": content_index, - "part": self._response_text_content_part(), - } - ] - - def _ensure_response_audio_part_added(self, response_id: object) -> list[dict[str, object]]: - state = self._response_state(response_id) - if state is None: - return [] - if state.audio_part_added: - return [] - state.audio_part_added = True - return [ - { - "type": "response.content_part.added", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 0, - "part": self._response_content_part(), - } - ] - - def _remember_response_audio_metadata(self, response_id: object, event: dict[str, Any]) -> None: - state = self._response_state(response_id) - if state is None: - return - duration = event.get("audio_duration_ms") - playback = event.get("playback") - if not isinstance(duration, int | float) and isinstance(playback, dict): - duration = playback.get("sent_ms") or playback.get("generated_ms") - if isinstance(duration, int | float): - state.audio_duration_ms = max(state.audio_duration_ms or 0, int(duration)) - marks = event.get("audio_text_marks") - if not isinstance(marks, list): - return - clean_marks: list[dict[str, int]] = [] - for mark in marks: - if not isinstance(mark, dict): - continue - text_chars = mark.get("text_chars") - audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) - if not isinstance(text_chars, int | float) or not isinstance(audio_end_ms, int | float): - continue - clean_marks.append( - { - "text_chars": max(0, int(text_chars)), - "audio_end_ms": max(0, int(audio_end_ms)), - } - ) - if clean_marks: - merged = list(state.audio_text_marks) - merged.extend(clean_marks) - deduped: dict[tuple[int, int], dict[str, int]] = {} - for mark in merged: - deduped[(int(mark["audio_end_ms"]), int(mark["text_chars"]))] = mark - state.audio_text_marks = sorted( - deduped.values(), - key=lambda mark: (mark["audio_end_ms"], mark["text_chars"]), - ) - - def _response_created_event(self, event: dict[str, Any]) -> dict[str, object]: - response_id = event.get("response_id") - metadata = event.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - metadata = {**metadata, "duplex_event": event} - return { - "type": "response.created", - "response_id": response_id, - "response": { - "id": response_id, - "object": "realtime.response", - "status": "in_progress", - "status_details": None, - "output": [], - "modalities": event.get("modalities") or ["audio", "text"], - "metadata": metadata, - }, - } - - def _realtime_audio_delta_events( - self, - event: dict[str, Any], - response_id: object, - audio: str, - ) -> list[dict[str, object]]: - item_id = self._response_item_id(response_id) - fmt, format_rate = self._parse_realtime_audio_format(event.get("format", "wav")) - source_fmt = self._realtime_output_format(fmt) - source_sample_rate_hz = event.get("sample_rate_hz") or format_rate - target_sample_rate_hz = ( - self._output_sample_rate_hz if self._output_audio_format in {"g711_ulaw", "g711_alaw"} else None - ) - audio, fmt, converted_sample_rate_hz = convert_output_audio( - audio, - source_fmt=source_fmt, - target_fmt=self._output_audio_format, - source_sample_rate_hz=( - int(source_sample_rate_hz) if isinstance(source_sample_rate_hz, int | float) else None - ), - target_sample_rate_hz=target_sample_rate_hz, - ) - sample_rate_hz = ( - converted_sample_rate_hz - if fmt in {"g711_ulaw", "g711_alaw"} - else event.get("sample_rate_hz") or format_rate or self._output_sample_rate_hz - ) - metadata: dict[str, object] = {} - for key in ("session_id", "epoch", "model_speak", "end_of_turn", "playback", "vllm_omni"): - if key in event: - metadata[key] = event[key] - duration_ms = event.get("audio_duration_ms") - if isinstance(duration_ms, int | float): - metadata["audio_duration_ms"] = int(duration_ms) - marks = event.get("audio_text_marks") - if isinstance(marks, list): - metadata["audio_text_marks"] = marks - state = self._response_state(response_id) - if state is not None: - state.audio_delta_emitted = True - self._remember_response_audio_metadata(response_id, event) - payloads: list[dict[str, object]] = [] - if state is not None and not state.speak_emitted and metadata.get("model_speak") is True: - state.speak_emitted = True - payloads.insert( - 0, - { - "type": "response.speak", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "metadata": self._response_speak_metadata(event), - }, - ) - payloads.append( - { - "type": "response.output_audio.delta", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "delta": audio, - "format": fmt, - **({"sample_rate_hz": int(sample_rate_hz)} if isinstance(sample_rate_hz, int | float) else {}), - **({"metadata": metadata} if metadata else {}), - } - ) - return payloads - - @staticmethod - def _response_speak_metadata(event: dict[str, Any]) -> dict[str, object]: - return {key: event[key] for key in ("session_id", "epoch", "model_speak") if key in event} - - def _realtime_audio_done_events( - self, - event: dict[str, Any], - response_id: object, - ) -> list[dict[str, object]]: - item_id = self._response_item_id(response_id) - state = self._response_state(response_id, event=event) - transcript = state.transcript if state is not None else "" - payloads: list[dict[str, object]] = [] - if isinstance(response_id, str) and state is not None and not state.audio_delta_emitted and not transcript: - return payloads - if state is None or not state.audio_done_emitted: - if state is not None: - state.audio_done_emitted = True - payloads.append( - { - "type": "response.output_audio.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - } - ) - if transcript: - payloads.append( - { - "type": "response.output_audio_transcript.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "transcript": transcript, - } - ) - return payloads - - def _realtime_response_terminal_events( - self, - event: dict[str, Any], - response_id: object, - *, - status: str = "completed", - status_details: dict[str, object] | None = None, - ) -> list[dict[str, object]]: - item_id = self._response_item_id(response_id) - state = self._response_state(response_id, event=event) - transcript = state.transcript if state is not None else "" - payloads: list[dict[str, object]] = [] - if state is not None and state.audio_part_added and not state.audio_part_done: - state.audio_part_done = True - payloads.append( - { - "type": "response.content_part.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "part": self._response_content_part(transcript=transcript), - } - ) - if state is not None and state.text_parts and not state.output_text_done: - state.output_text_done = True - content_index = 1 if state.audio_part_added else 0 - payloads.append( - { - "type": "response.output_text.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": content_index, - "text": state.text, - } - ) - if state is not None and state.text_parts and not state.text_part_done: - state.text_part_done = True - content_index = 1 if state.audio_part_added else 0 - payloads.append( - { - "type": "response.content_part.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": content_index, - "part": self._response_text_content_part(text=state.text), - } - ) - if state is None or not state.output_item_done: - if state is not None: - state.output_item_done = True - item = self._response_done_output_item(response_id, status=status) - payloads.append( - { - "type": "response.output_item.done", - "response_id": response_id, - "output_index": 0, - "item": item, - } - ) - if state is None or not state.conversation_item_done: - if state is not None: - state.conversation_item_done = True - payloads.append(self._conversation_item_done_event(item)) - done_event = self._realtime_response_done_event( - {**event, "response_id": response_id}, - state=state, - status=status, - status_details=status_details, - ) - if done_event is not None: - payloads.append(done_event) - payloads.append(self._rate_limits_updated_event()) - if isinstance(response_id, str) and response_id == self._active_response_id: - self._active_response_id = None - return payloads - - def _realtime_response_done_event( - self, - event: dict[str, Any], - *, - state: _RealtimeResponseState | None = None, - status: str = "completed", - status_details: dict[str, object] | None = None, - ) -> dict[str, object] | None: - response_id = event.get("response_id") - if state is None: - state = self._response_state(response_id, event=event) - if state is not None: - if state.done_emitted: - return None - state.done_emitted = True - return { - "type": "response.done", - "response_id": response_id, - "response": { - "id": response_id, - "object": "realtime.response", - "status": status, - "status_details": status_details, - "output": [self._response_done_output_item(response_id, status=status)], - "metadata": event, - }, - } - - @staticmethod - def _rate_limits_updated_event() -> dict[str, object]: - # vLLM-Omni does not yet expose a Realtime-specific quota budget. Emit - # the terminal event with an empty list so clients that sequence on - # rate_limits.updated do not need a private vLLM branch. - return {"type": "rate_limits.updated", "rate_limits": []} diff --git a/vllm_omni/entrypoints/duplex/realtime_session.py b/vllm_omni/entrypoints/duplex/realtime_session.py deleted file mode 100644 index 122a7ab52..000000000 --- a/vllm_omni/entrypoints/duplex/realtime_session.py +++ /dev/null @@ -1,147 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import json -from typing import Any -from uuid import uuid4 - -from fastapi import WebSocket - -from vllm_omni.entrypoints.duplex import realtime_state as _state -from vllm_omni.entrypoints.duplex.realtime_input import ( - RealtimeInputTranslator, -) -from vllm_omni.entrypoints.duplex.realtime_output import ( - RealtimeOutputProjector, -) - -REALTIME_ERROR_TYPES_BY_CODE = _state.REALTIME_ERROR_TYPES_BY_CODE -REALTIME_INPUT_AUDIO_FORMATS = _state.REALTIME_INPUT_AUDIO_FORMATS -REALTIME_OUTPUT_AUDIO_FORMATS = _state.REALTIME_OUTPUT_AUDIO_FORMATS -RealtimeSessionState = _state.RealtimeSessionState -_RealtimeResponseState = _state._RealtimeResponseState - - -class NativeRealtimeSessionProtocol( - RealtimeInputTranslator, - RealtimeOutputProjector, - _state.RealtimeStateOwner, -): - """Compatibility facade for one native Realtime protocol session. - - Input translation, output projection, and mutable state have separate - owners. The facade retains the established import path and private helper - surface used by serving and protocol contract tests. - """ - - def __init__(self, query_params: Any) -> None: - self._state = RealtimeSessionState.from_query_params(query_params) - - def bind_sender(self, send_realtime_json) -> None: - self._send_realtime_json = send_realtime_json - - def bind_native_input_append(self, enabled: bool) -> None: - """Use the serving-selected input path, independent of interruption policy.""" - self._native_input_append = enabled - - def _default_session_payload(self) -> dict[str, object]: - payload: dict[str, object] = { - "model": self._default_model, - "session_id": self._default_session_id, - } - if self._default_extra_body: - payload["extra_body"] = dict(self._default_extra_body) - return payload - - async def receive_internal_event_text(self, websocket: WebSocket) -> str: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - if not self._opened and not self._resume_only and not self._autostarted_default_session and self._default_model: - self._opened = True - self._autostarted_default_session = True - return json.dumps(self._session_create_from_realtime(self._default_session_payload())) - while True: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - await self.wait_for_realtime_turn_detection_update() - raw = await websocket.receive_text() - try: - event = json.loads(raw) - except json.JSONDecodeError: - return raw - if not isinstance(event, dict): - return raw - if not self._opened and event.get("type") == "session.resume": - self._opened = True - translated = await self._to_duplex_event(event) - if translated is None: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - continue - return json.dumps(translated) - if not self._opened and self._resume_only: - translated = await self._to_duplex_event(event) - if translated is None: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - continue - return json.dumps(translated) - if not self._opened and event.get("type") != "session.update": - self._opened = True - await self._pending_outbound.put(self._session_create_from_realtime(self._default_session_payload())) - translated = await self._to_duplex_event(event) - if translated is not None: - await self._send_realtime_input_ack(event) - await self._pending_outbound.put(translated) - return json.dumps(await self._pending_outbound.get()) - translated = await self._to_duplex_event(event) - if translated is None: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - continue - await self._send_realtime_input_ack(event) - return json.dumps(translated) - - def encode_outbound_event(self, data: dict[str, Any]) -> list[dict[str, object]]: - payloads = self._from_duplex_event(data) - for payload in payloads: - self._attach_event_id(payload) - return payloads - - @staticmethod - def _attach_event_id(payload: dict[str, object]) -> None: - payload.setdefault("event_id", f"event_{uuid4().hex}") - - async def _send_realtime_payload(self, payload: dict[str, object]) -> None: - self._attach_event_id(payload) - if ( - self._hold_realtime_output_until_session_created - and payload.get("type") != "error" - and payload.get("type") != "session.created" - ): - self._held_realtime_payloads.append(payload) - return - if self._send_realtime_json is None: - raise RuntimeError("Native Realtime sender is not bound") - await self._send_realtime_json(payload) - - @staticmethod - def _realtime_error_payload( - code: str, - message: str, - *, - event_id: object | None = None, - param: object | None = None, - ) -> dict[str, object]: - error: dict[str, object] = { - "type": REALTIME_ERROR_TYPES_BY_CODE.get(code, "invalid_request_error"), - "code": code, - "message": message, - } - if isinstance(event_id, str) and event_id: - error["event_id"] = event_id - if isinstance(param, str) and param: - error["param"] = param - return {"type": "error", "error": error} diff --git a/vllm_omni/entrypoints/duplex/realtime_state.py b/vllm_omni/entrypoints/duplex/realtime_state.py deleted file mode 100644 index babb470a5..000000000 --- a/vllm_omni/entrypoints/duplex/realtime_state.py +++ /dev/null @@ -1,245 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -from dataclasses import dataclass, field -from typing import Any, Generic, TypeVar - -REALTIME_INPUT_AUDIO_FORMATS = { - "pcm16", - "pcm_s16le", - "s16le", - "pcm_f32le", - "g711_ulaw", - "g711_alaw", -} -REALTIME_OUTPUT_AUDIO_FORMATS = { - "pcm16", - "pcm_s16le", - "s16le", - "wav", - "pcm", - "g711_ulaw", - "g711_alaw", -} -REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ = 24_000 -REALTIME_PCM_F32_DEFAULT_SAMPLE_RATE_HZ = 16_000 -REALTIME_G711_DEFAULT_SAMPLE_RATE_HZ = 8_000 - -_RealtimeProtocolConfig = tuple[str, int, str, int | None, float] -_PendingTurnDetectionUpdate = tuple[ - bool, - dict[str, object] | None, - asyncio.Event, - _RealtimeProtocolConfig, -] - - -def realtime_default_sample_rate_hz(fmt: object) -> int | None: - if not isinstance(fmt, str): - return None - normalized = fmt.lower() - if normalized in {"pcm16", "pcm_s16le", "s16le", "pcm", "audio/pcm"}: - return REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ - if normalized in {"pcm_f32le", "audio/pcm_f32le"}: - return REALTIME_PCM_F32_DEFAULT_SAMPLE_RATE_HZ - if normalized in {"g711_ulaw", "g711_alaw", "audio/g711_ulaw", "audio/g711_alaw"}: - return REALTIME_G711_DEFAULT_SAMPLE_RATE_HZ - return None - - -REALTIME_ERROR_TYPES_BY_CODE = { - "bad_event": "invalid_request_error", - "bad_audio": "invalid_request_error", - "config_timeout": "invalid_request_error", - "invalid_json": "invalid_request_error", - "event_too_large": "invalid_request_error", - "unknown_event": "invalid_request_error", - "internal_error": "server_error", - "runtime_open_failed": "server_error", - "runtime_open_unsupported": "server_error", - "runtime_append_failed": "server_error", - "runtime_append_task_failed": "server_error", - "runtime_signal_failed": "server_error", - "runtime_close_failed": "server_error", - "runtime_abort_failed": "server_error", - "runtime_data_plane_stream_failed": "server_error", - "runtime_data_plane_text_without_audio": "server_error", - "response_error": "server_error", - "chat_error": "server_error", - "server_vad_initialization_failed": "server_error", - "server_vad_inference_failed": "server_error", - "duplex_session_busy": "rate_limit_error", - "resource_exhausted": "rate_limit_error", - "input_backpressure": "rate_limit_error", - "response_already_active": "invalid_request_error", - "response_not_active": "invalid_request_error", - "response_create_without_input": "invalid_request_error", - "input_audio_buffer_empty": "invalid_request_error", - "server_vad_manual_commit_unsupported": "invalid_request_error", - "missing_item_id": "invalid_request_error", - "item_not_found": "invalid_request_error", - "playback_item_mismatch": "invalid_request_error", - "playback_item_not_found": "invalid_request_error", - "playback_ack_too_late": "invalid_request_error", - "unsupported_audio_format": "invalid_request_error", - "unsupported_ref_audio_path": "invalid_request_error", - "ref_audio_required": "invalid_request_error", - "model_update_unsupported": "invalid_request_error", - "voice_update_after_audio_unsupported": "invalid_request_error", - "ref_audio_update_unsupported": "invalid_request_error", - "native_text_append_unsupported": "invalid_request_error", - "runtime_native_stage_role_required": "server_error", - "runtime_native_runner_kv_required": "server_error", -} - - -@dataclass(slots=True) -class _RealtimeResponseState: - item_id: str - transcript_parts: list[str] = field(default_factory=list) - text_parts: list[str] = field(default_factory=list) - audio_duration_ms: int | None = None - audio_text_marks: list[dict[str, int]] = field(default_factory=list) - audio_delta_emitted: bool = False - audio_done_emitted: bool = False - audio_part_added: bool = False - audio_part_done: bool = False - text_part_added: bool = False - text_part_done: bool = False - output_text_done: bool = False - output_item_done: bool = False - conversation_item_done: bool = False - speak_emitted: bool = False - done_emitted: bool = False - - @property - def transcript(self) -> str: - return "".join(self.transcript_parts) - - @property - def text(self) -> str: - return "".join(self.text_parts) - - -@dataclass(slots=True) -class RealtimeSessionState: - """Single mutable authority for one Realtime protocol session.""" - - _opened: bool = False - _autostarted_default_session: bool = False - _resume_only: bool = False - _pending_outbound: asyncio.Queue[dict[str, object]] = field(default_factory=asyncio.Queue) - _held_realtime_payloads: list[dict[str, object]] = field(default_factory=list) - _hold_realtime_output_until_session_created: bool = True - _default_model: object | None = None - _default_session_id: object | None = None - _default_extra_body: dict[str, object] = field(default_factory=dict) - _input_audio_format: str = "pcm16" - _input_sample_rate_hz: int = REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ - _output_audio_format: str = "pcm16" - _overlap_silence_rms: float = 0.003 - _turn_detection: dict[str, object] | None = None - _native_input_append: bool = False - _pending_turn_detection_update: _PendingTurnDetectionUpdate | None = None - _send_realtime_json: Any = None - _initial_session_update: bool = False - _input_speech_started: bool = False - _response_states: dict[str | int, _RealtimeResponseState] = field(default_factory=dict) - _item_truncation_cursors: dict[str, tuple[int, int]] = field(default_factory=dict) - _active_response_id: str | None = None - _last_response_id: str | None = None - _conversation_items: dict[str, dict[str, object]] = field(default_factory=dict) - _last_conversation_item_id: str | None = None - _output_sample_rate_hz: int | None = None - _active_input_item_id: str | None = None - _input_audio_buffer_has_audio: bool = False - _input_audio_buffer_had_non_speech: bool = False - _input_audio_buffer_transcript_parts: list[str] = field(default_factory=list) - _turn_detection_configured: bool = False - - @classmethod - def from_query_params(cls, query_params: Any) -> RealtimeSessionState: - if hasattr(query_params, "query_params"): - query_params = query_params.query_params - getter = query_params.get if hasattr(query_params, "get") else None - # Canonical query param plus the deprecated model-prefixed alias; - # both seed the canonical extra_body key. - native_duplex = None - if getter is not None: - native_duplex = getter("native_duplex") - if native_duplex is None: - native_duplex = getter("minicpmo45_native_duplex") - resume_only = getter("resume") if getter is not None else None - autostart = getter("autostart") if getter is not None else None - default_extra_body = {} - if str(native_duplex).strip().lower() in {"1", "true", "yes", "on"}: - default_extra_body["native_duplex"] = True - return cls( - _resume_only=( - str(resume_only).strip().lower() in {"1", "true", "yes", "on"} - or str(autostart).strip().lower() in {"0", "false", "no", "off"} - ), - _default_model=getter("model") if getter is not None else None, - _default_session_id=(getter("session_id") if getter is not None else None), - _default_extra_body=default_extra_body, - ) - - -_StateValue = TypeVar("_StateValue") - - -class _RealtimeStateField(Generic[_StateValue]): - """Explicit descriptor binding one protocol attribute to session state.""" - - def __set_name__(self, owner: type, name: str) -> None: - del owner - self._name = name - - def __get__(self, instance: Any, owner: type | None = None) -> _StateValue | _RealtimeStateField[_StateValue]: - del owner - if instance is None: - return self - return getattr(instance._state, self._name) - - def __set__(self, instance: Any, value: _StateValue) -> None: - setattr(instance._state, self._name, value) - - -class RealtimeStateOwner: - """Declare the mutable state surface shared by input/output projectors.""" - - _state: RealtimeSessionState - _opened: bool = _RealtimeStateField() - _autostarted_default_session: bool = _RealtimeStateField() - _resume_only: bool = _RealtimeStateField() - _pending_outbound: asyncio.Queue[dict[str, object]] = _RealtimeStateField() - _held_realtime_payloads: list[dict[str, object]] = _RealtimeStateField() - _hold_realtime_output_until_session_created: bool = _RealtimeStateField() - _default_model: object | None = _RealtimeStateField() - _default_session_id: object | None = _RealtimeStateField() - _default_extra_body: dict[str, object] = _RealtimeStateField() - _input_audio_format: str = _RealtimeStateField() - _input_sample_rate_hz: int = _RealtimeStateField() - _output_audio_format: str = _RealtimeStateField() - _overlap_silence_rms: float = _RealtimeStateField() - _turn_detection: dict[str, object] | None = _RealtimeStateField() - _native_input_append: bool = _RealtimeStateField() - _pending_turn_detection_update: _PendingTurnDetectionUpdate | None = _RealtimeStateField() - _send_realtime_json: Any = _RealtimeStateField() - _initial_session_update: bool = _RealtimeStateField() - _input_speech_started: bool = _RealtimeStateField() - _response_states: dict[str | int, _RealtimeResponseState] = _RealtimeStateField() - _item_truncation_cursors: dict[str, tuple[int, int]] = _RealtimeStateField() - _active_response_id: str | None = _RealtimeStateField() - _last_response_id: str | None = _RealtimeStateField() - _conversation_items: dict[str, dict[str, object]] = _RealtimeStateField() - _last_conversation_item_id: str | None = _RealtimeStateField() - _output_sample_rate_hz: int | None = _RealtimeStateField() - _active_input_item_id: str | None = _RealtimeStateField() - _input_audio_buffer_has_audio: bool = _RealtimeStateField() - _input_audio_buffer_had_non_speech: bool = _RealtimeStateField() - _input_audio_buffer_transcript_parts: list[str] = _RealtimeStateField() - _turn_detection_configured: bool = _RealtimeStateField() diff --git a/vllm_omni/entrypoints/duplex/runtime_adapter.py b/vllm_omni/entrypoints/duplex/runtime_adapter.py deleted file mode 100644 index a873f33a8..000000000 --- a/vllm_omni/entrypoints/duplex/runtime_adapter.py +++ /dev/null @@ -1,331 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping -from dataclasses import dataclass, field -from importlib import import_module -from typing import Any, Protocol - - -class ServingRuntimeConfigError(ValueError): - """A model serving plugin rejected client-visible runtime configuration.""" - - def __init__(self, message: str, *, code: str = "invalid_duplex_runtime_config") -> None: - super().__init__(message) - self.code = code - - -def reject_changed_runtime_value( - new_value: object, - current_value: object, - *, - message: str, - code: str, - error_cls: type[ServingRuntimeConfigError] = ServingRuntimeConfigError, -) -> None: - if new_value != current_value: - raise error_cls(message, code=code) - - -class PcmAppendReservation(Protocol): - operation_id: str - payload: dict[str, object] | None - - @property - def active(self) -> bool: ... - - @property - def byte_count(self) -> int: ... - - def commit(self) -> None: ... - - def rollback(self) -> None: ... - - -class PcmAppendBuffer(Protocol): - @property - def pending_byte_count(self) -> int: ... - - def clear(self) -> None: ... - - def clear_force_listen(self) -> None: ... - - def has_pending(self) -> bool: ... - - def has_reserved(self) -> bool: ... - - def prepare_append( - self, - payload: dict[str, object], - *, - operation_id: str, - chunk_period_ms: int, - allow_emit: bool, - ) -> PcmAppendReservation | None: ... - - def prepare_commit( - self, - *, - operation_id: str, - chunk_period_ms: int, - ) -> PcmAppendReservation: ... - - def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: ... - - -class ServingRuntimeSessionState(Protocol): - audio_buffer: PcmAppendBuffer - input_since_commit: bool - speech_since_commit: bool - native_context_locked: bool - committed_audio_payload: dict[str, object] | None - committed_audio_operation_id: str | None - committed_audio_reserved_bytes: int - deferred_response_create: bool - deferred_precreate_response: bool - data_plane_task: asyncio.Task[None] | None - data_plane_restart_requested: bool - continuation_owner_id: str | None - continuation_units: int - pending_silence_task: asyncio.Task[bool] | None - pending_silence_owner_id: str | None - silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None - - def retain_committed_audio( - self, - payload: dict[str, object], - *, - operation_id: str | None, - reserved_bytes: int = 0, - ) -> None: ... - - def clear_committed_audio(self) -> int: ... - - def clear_continuation(self) -> None: ... - - -class _TurnBasedPcmAppendBuffer: - """Empty native-audio buffer for chat-fallback sessions. - - Turn-based sessions keep input audio in ``DuplexSession`` and never use - the model-native append path. The runner still shares lifecycle and - cancellation code with native sessions, so it needs an inert buffer for - those common operations. - """ - - @property - def pending_byte_count(self) -> int: - return 0 - - def clear(self) -> None: - return None - - def clear_force_listen(self) -> None: - return None - - def has_pending(self) -> bool: - return False - - def has_reserved(self) -> bool: - return False - - def prepare_append( - self, - payload: dict[str, object], - *, - operation_id: str, - chunk_period_ms: int, - allow_emit: bool, - ) -> PcmAppendReservation | None: - del payload, operation_id, chunk_period_ms, allow_emit - raise RuntimeError("Model-native audio append is not enabled for this session") - - def prepare_commit( - self, - *, - operation_id: str, - chunk_period_ms: int, - ) -> PcmAppendReservation: - del operation_id, chunk_period_ms - raise RuntimeError("Model-native audio append is not enabled for this session") - - def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: - del chunk_period_ms - return None - - -@dataclass(slots=True) -class TurnBasedServingSessionState: - """Runner-local state for sessions served through chat fallback.""" - - audio_buffer: PcmAppendBuffer = field(default_factory=_TurnBasedPcmAppendBuffer) - input_since_commit: bool = False - speech_since_commit: bool = False - native_context_locked: bool = False - committed_audio_payload: dict[str, object] | None = None - committed_audio_operation_id: str | None = None - committed_audio_reserved_bytes: int = 0 - deferred_response_create: bool = False - deferred_precreate_response: bool = False - data_plane_task: asyncio.Task[None] | None = None - data_plane_restart_requested: bool = False - continuation_owner_id: str | None = None - continuation_units: int = 0 - pending_silence_task: asyncio.Task[bool] | None = None - pending_silence_owner_id: str | None = None - silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None = None - - def retain_committed_audio( - self, - payload: dict[str, object], - *, - operation_id: str | None, - reserved_bytes: int = 0, - ) -> None: - self.committed_audio_payload = payload - self.committed_audio_operation_id = operation_id - self.committed_audio_reserved_bytes += max(0, int(reserved_bytes)) - - def clear_committed_audio(self) -> int: - reserved_bytes = self.committed_audio_reserved_bytes - self.committed_audio_payload = None - self.committed_audio_operation_id = None - self.committed_audio_reserved_bytes = 0 - self.deferred_response_create = False - self.deferred_precreate_response = False - return reserved_bytes - - def clear_continuation(self) -> None: - self.continuation_owner_id = None - self.continuation_units = 0 - self.pending_silence_task = None - self.pending_silence_owner_id = None - - -class RuntimeDataPlane(Protocol): - def begin_request(self, request_id: str) -> None: ... - - def is_terminal(self, request_id: str | None) -> bool: ... - - def mark_terminal(self, request_id: str) -> None: ... - - def close_stream(self, request_id: str) -> None: ... - - def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: ... - - def project(self, result: object, *, context: object | None = None) -> Iterable[dict[str, object]]: ... - - -class ServingRuntimeAdapter(Protocol): - adapter_id: str - session_states: MutableMapping[str, ServingRuntimeSessionState] - data_plane: RuntimeDataPlane - clean_response_done_prefix: str - interrupted_tts_prefix: str - private_runtime_config_keys: frozenset[str] - collect_outputs_on_append: bool - - def create_session_state(self) -> ServingRuntimeSessionState: ... - - def session_state(self, session_id: str) -> ServingRuntimeSessionState: ... - - def remove_session_state(self, session_id: str) -> None: ... - - def is_enabled(self, config: object) -> bool: ... - - def capabilities(self, *, max_sessions: int) -> object: ... - - def validate_client_extra_body(self, extra_body: object) -> None: ... - - async def prepare_runtime_config(self, config: object, *, model_config: Any) -> dict[str, object]: ... - - def runtime_config_for_update( - self, - config: object, - current: Mapping[str, object], - ) -> dict[str, object]: ... - - def data_plane_context( - self, - *, - epoch: int, - turn_id: int, - active_response_turn_id: int | None, - active_response_id: str | None, - auto_responds: bool, - response_format: str, - speed: float | None, - modalities: tuple[str, ...], - ) -> object: ... - - -def load_serving_runtime_adapter( - path: str, - encode_audio, -) -> ServingRuntimeAdapter: - module_name, separator, attribute_name = path.rpartition(".") - if not separator: - raise ValueError(f"Invalid duplex serving runtime adapter path: {path!r}") - adapter_type = getattr(import_module(module_name), attribute_name) - return validate_serving_runtime_adapter(adapter_type(encode_audio)) - - -def validate_serving_runtime_adapter(adapter: object) -> ServingRuntimeAdapter: - required_methods = ( - "create_session_state", - "session_state", - "remove_session_state", - "is_enabled", - "capabilities", - "validate_client_extra_body", - "prepare_runtime_config", - "runtime_config_for_update", - "data_plane_context", - ) - missing = [name for name in required_methods if not callable(getattr(adapter, name, None))] - if missing: - raise TypeError(f"Duplex serving runtime adapter is missing callable method(s): {', '.join(missing)}") - if not isinstance(getattr(adapter, "adapter_id", None), str) or not adapter.adapter_id: - raise TypeError("Duplex serving runtime adapter must declare adapter_id") - if not isinstance(getattr(adapter, "session_states", None), MutableMapping): - raise TypeError("Duplex serving runtime adapter must declare mutable session_states") - for prefix_name in ("clean_response_done_prefix", "interrupted_tts_prefix"): - if not isinstance(getattr(adapter, prefix_name, None), str): - raise TypeError(f"Duplex serving runtime adapter must declare {prefix_name}") - private_keys = getattr(adapter, "private_runtime_config_keys", None) - if not isinstance(private_keys, frozenset) or any(not isinstance(key, str) for key in private_keys): - raise TypeError("Duplex serving runtime adapter must declare private_runtime_config_keys as frozenset[str]") - data_plane = getattr(adapter, "data_plane", None) - if data_plane is None: - raise TypeError("Duplex serving runtime adapter must declare data_plane") - data_plane_methods = ( - "begin_request", - "is_terminal", - "mark_terminal", - "close_stream", - "close_session", - "project", - ) - missing_data_plane_methods = [name for name in data_plane_methods if not callable(getattr(data_plane, name, None))] - if missing_data_plane_methods: - raise TypeError( - "Duplex serving runtime adapter data_plane is missing callable method(s): " - + ", ".join(missing_data_plane_methods) - ) - return adapter # type: ignore[return-value] - - -def payload_turn_id(payload: object) -> int | None: - if not isinstance(payload, Mapping): - return None - return coerce_int(payload.get("duplex_turn_id", payload.get("model_turn_id"))) - - -def coerce_int(value: object) -> int | None: - try: - return int(value) if value is not None else None - except (TypeError, ValueError): - return None diff --git a/vllm_omni/entrypoints/duplex/runtime_bridge.py b/vllm_omni/entrypoints/duplex/runtime_bridge.py deleted file mode 100644 index 96a82c061..000000000 --- a/vllm_omni/entrypoints/duplex/runtime_bridge.py +++ /dev/null @@ -1,1454 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import base64 -import inspect -import time -from collections.abc import Mapping -from dataclasses import asdict, is_dataclass -from enum import Enum - -import numpy as np -from vllm.logger import init_logger - -from vllm_omni.engine.duplex.control_client import DuplexControlRequestError -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import duplex_data_plane_request_info -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexSession, - DuplexSessionState, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - coerce_int, - payload_turn_id, -) - -logger = init_logger(__name__) - - -class NativeRuntimeBridgeMixin: - """Bridge serving sessions to native runtime control and data-plane APIs.""" - - _NATIVE_RUNTIME_FENCE_CONTRACT = { - "open_duplex_session_async": ("fence",), - "append_duplex_input_async": ("fence",), - "signal_duplex_turn_async": ("fence", "next_fence"), - "close_duplex_session_async": ("fence",), - } - - async def _open_runtime_session(self, session: DuplexSession, send_json) -> dict[str, object] | bool: - # Turn-based sessions use the chat-fallback path and do not own a - # model-native duplex runtime. Do not send control messages merely - # because the shared engine client exposes the duplex RPC methods. - if not self._uses_serving_runtime_adapter(session.config): - return True - contract_error = self._native_runtime_contract_error(session) - if contract_error is not None: - await send_json( - { - "type": "error", - "code": "runtime_contract_invalid", - "error": contract_error, - } - ) - return False - open_session = getattr(self._chat_service.engine_client, "open_duplex_session_async", None) - if not callable(open_session): - return True - try: - open_kwargs = { - "session_mode": "duplex", - "capabilities": session.capabilities.as_dict(), - "session_config": session.config.as_dict(), - "timeout": self._runtime_control_timeout_s(session), - } - if self._callable_accepts_keyword(open_session, "runtime_config"): - open_kwargs["runtime_config"] = dict(session.runtime_config) - if self._callable_accepts_keyword(open_session, "fence"): - open_kwargs["fence"] = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - result = await open_session(session.session_id, **open_kwargs) - except Exception as exc: - if isinstance(exc, DuplexControlRequestError) and exc.code == "resource_exhausted": - logger.info("Duplex runtime session admission rejected: %s", exc) - else: - logger.exception("Failed to open duplex runtime session: %s", exc) - await self._send_runtime_error(send_json, "runtime_open_failed", exc, session=session) - return False - if ( - isinstance(result, dict) - and session.capabilities.implementation_level == "model_native_duplex" - and self._runtime_control_failed(result) - ): - await send_json( - { - "type": "error", - "error": "Native duplex runtime is not available for this session", - "code": ("runtime_open_unsupported" if result.get("unsupported_count") else "runtime_open_failed"), - "runtime_control": self._redact_runtime_control_result(result), - } - ) - return False - return result if isinstance(result, dict) else True - - async def _append_runtime_input( - self, - session: DuplexSession, - payload: object, - *, - operation_id: str | None = None, - final: bool, - send_json, - mode: str = "append_tokens", - expected_epoch: int | None = None, - ) -> tuple[bool, bool]: - if not session.capabilities.supports_input_append: - return True, False - append_input = getattr(self._chat_service.engine_client, "append_duplex_input_async", None) - if not callable(append_input): - return True, False - if expected_epoch is not None and session.epoch != expected_epoch: - return True, False - try: - append_kwargs = { - "mode": mode, - "payload": payload, - "final": final, - "timeout": self._runtime_control_timeout_s(session), - } - if operation_id is not None and self._callable_accepts_keyword(append_input, "operation_id"): - append_kwargs["operation_id"] = operation_id - if expected_epoch is not None and self._callable_accepts_keyword(append_input, "expected_epoch"): - append_kwargs["expected_epoch"] = expected_epoch - if self._callable_accepts_keyword(append_input, "fence"): - payload_turn = payload_turn_id(payload) - append_kwargs["fence"] = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=( - payload_turn - if payload_turn is not None - else ( - session.active_response_turn_id - if session.active_response_turn_id is not None - else session.turn_id - ) - ), - incarnation=session.incarnation, - ) - if self._callable_accepts_keyword(append_input, "collect_outputs"): - append_kwargs["collect_outputs"] = bool( - getattr(self._serving_runtime_adapter, "collect_outputs_on_append", False) - ) - result = await append_input(session.session_id, **append_kwargs) - except Exception as exc: - logger.exception("Failed to append duplex runtime input: %s", exc) - await self._send_runtime_error(send_json, "runtime_append_failed", exc, session=session) - return False, False - if isinstance(result, dict) and self._runtime_control_failed(result): - await self._send_runtime_control_error( - send_json, - "runtime_append_failed", - "Duplex runtime append failed", - result, - session=session, - ) - return False, False - if expected_epoch is not None and session.epoch != expected_epoch: - return True, False - await self._send_runtime_control_if_needed(send_json, result, session=session) - request_id, _ = self._data_plane_request_info(result) if isinstance(result, dict) else (None, None) - if request_id is not None: - self._serving_runtime_adapter.data_plane.begin_request(request_id) - close_reason, emitted_response = await self._send_native_duplex_events( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ) - # A resumable data-plane append remains active even when its persistent - # drain task already exists. Treating only a newly-created drain as - # activity clears active_request_id after later appends and prevents a - # terminal TTS segment from scheduling the next model decision. - emitted_response = emitted_response or request_id is not None - if close_reason is None and await self._start_native_data_plane_stream_task( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ): - emitted_response = True - if close_reason is not None: - if not await self._close_runtime_session(session, reason=close_reason, send_json=send_json): - return False, emitted_response - session.close() - await send_json( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": close_reason, - } - ) - return False, emitted_response - return True, emitted_response - - @staticmethod - def _callable_accepts_keyword(fn, name: str) -> bool: - try: - params = inspect.signature(fn).parameters.values() - except (TypeError, ValueError): - return True - return any(param.kind == inspect.Parameter.VAR_KEYWORD or param.name == name for param in params) - - def _native_runtime_contract_error(self, session: DuplexSession) -> str | None: - if session.capabilities.implementation_level != "model_native_duplex": - return None - engine_client = self._chat_service.engine_client - invalid: list[str] = [] - for method_name, required_keywords in self._NATIVE_RUNTIME_FENCE_CONTRACT.items(): - method = getattr(engine_client, method_name, None) - if not callable(method): - invalid.append(f"{method_name} is missing") - continue - try: - parameters = tuple(inspect.signature(method).parameters.values()) - except (TypeError, ValueError): - invalid.append(f"{method_name} signature is not inspectable") - continue - accepts_kwargs = any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters) - declared = {parameter.name for parameter in parameters} - missing = [name for name in required_keywords if not accepts_kwargs and name not in declared] - if missing: - invalid.append(f"{method_name} is missing keyword(s): {', '.join(missing)}") - if not invalid: - return None - return "Native duplex runtime requires the fenced control contract: " + "; ".join(invalid) - - async def _start_native_data_plane_stream_task( - self, - send_json, - result: object, - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> bool: - request_id, _ = self._data_plane_request_info(result) - if request_id is None or self._data_plane_outputs_finished(result): - return False - session.bind_request(request_id) - - native = self._runtime_session_state(session) - old_task = native.data_plane_task - if old_task is not None and not old_task.done(): - if self._session_auto_responds(session): - # One persistent drain per session, like the official worker's - # single synchronous loop: the resumable data-plane request id - # is stable, and cancel/restart on every append orphans any - # decision that lands in the swap window. - native.data_plane_restart_requested = True - return False - old_task.cancel() - try: - await asyncio.wait_for(asyncio.gather(old_task, return_exceptions=True), timeout=0.25) - except asyncio.TimeoutError: - pass - - async def _run() -> None: - close_reason: str | None = None - try: - close_reason = await self._drain_native_data_plane_stream( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ) - if close_reason is not None: - if await self._close_runtime_session(session, reason=close_reason, send_json=send_json): - session.close() - await send_json( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": close_reason, - } - ) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.exception("Failed to drain duplex data-plane stream: %s", exc) - await self._send_runtime_error(send_json, "runtime_data_plane_stream_failed", exc, session=session) - if session.state != DuplexSessionState.CLOSED: - close_reason = "runtime_data_plane_stream_failed" - if await self._close_runtime_session(session, reason=close_reason, send_json=send_json): - session.close() - await send_json( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": close_reason, - } - ) - finally: - if native.data_plane_task is task: - native.data_plane_task = None - restart_requested = ( - native.data_plane_restart_requested - and close_reason is None - and session.state != DuplexSessionState.CLOSED - and session.active_request_id == request_id - and (expected_epoch is None or session.epoch == expected_epoch) - ) - native.data_plane_restart_requested = False - if close_reason is None and not self._session_auto_responds(session): - # Auto-respond sessions keep one resumable stage-1 stream - # whose audio accumulates across speak units; the offset - # must survive drain-task turnover or every unit re-sends - # the reply audio from the start. - self._serving_runtime_adapter.data_plane.close_stream(request_id) - if restart_requested: - await self._start_native_data_plane_stream_task( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ) - elif close_reason is None: - await self._maybe_continue_native_response(send_json, session=session, expected_epoch=expected_epoch) - - task = asyncio.create_task(_run()) - native.data_plane_task = task - return True - - # One MiniCPM model unit (1 s at 16 kHz) is the compatibility default. - # Other native-duplex adapters may use a different frame-locked unit. - _NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii") - _NATIVE_RESPONSE_MAX_CONTINUATION_UNITS = 8 - # Each unit is one second of silence. Auto-response needs substantially - # more room than the eight decision beats above, but must still terminate - # well before the 300-second benchmark playback horizon. The final unit - # asks the model to LISTEN; a protocol fallback closes the response if that - # logit hint is ignored. - _NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS = 64 - - def _native_silence_unit_payload(self) -> dict[str, object]: - samples = int(getattr(self._serving_runtime_adapter, "silence_continuation_samples", 16000)) - audio = ( - self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - if samples == 16000 - else base64.b64encode(bytes(samples * 4)).decode("ascii") - ) - return { - "type": "audio", - "audio": audio, - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - - def _native_response_continuations_remaining(self, session: DuplexSession, response_id: str) -> bool: - native = self._runtime_session_state(session) - owner_id = f"response:{response_id}" - count = native.continuation_units if native.continuation_owner_id == owner_id else 0 - limit = ( - self._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - if self._session_auto_responds(session) - else self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS - ) - return count < limit - - async def _finish_bounded_native_auto_response( - self, - send_json, - *, - session: DuplexSession, - expected_epoch: int | None, - model_turn_id: int | None = None, - ) -> None: - """Close an auto-response whose final silence unit did not end it.""" - native = self._runtime_session_state(session) - response_id = session.active_response_id - response_epoch = session.epoch - response_turn_id = model_turn_id if model_turn_id is not None else session.active_response_turn_id - if expected_epoch is not None and response_epoch != expected_epoch: - return - native.clear_continuation() - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": response_epoch, - "reason": "continuation_limit", - "model_listen": False, - } - if response_id is not None: - # Clients demultiplex a listen against a precreated response by - # its id; every listen that terminates a response must carry it. - payload["response_id"] = response_id - await send_json(payload) - if response_id is None: - if session.epoch == response_epoch and response_turn_id is not None: - session.complete_model_turn(response_turn_id) - return - if session.epoch != response_epoch or session.active_response_id != response_id: - return - if response_turn_id is not None: - session.complete_model_turn(response_turn_id) - session.end_response(commit_text=False, preserve_request=True) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": response_epoch, - "committed": False, - "playback": session.playback.as_dict(), - } - ) - - def _native_silence_continuation_is_stale( - self, - session: DuplexSession, - *, - request_id: str, - response_id: str | None, - response_owned: bool, - expected_epoch: int | None, - expected_incarnation: int, - expected_model_turn_id: int | None, - ) -> bool: - stale_common_owner = ( - session.state == DuplexSessionState.CLOSED - or session.incarnation != expected_incarnation - or session.active_request_id != request_id - or (expected_epoch is not None and session.epoch != expected_epoch) - ) - stale_response_owner = response_owned and session.active_response_id != response_id - stale_model_turn_owner = not response_owned and ( - session.active_response_id is not None or session.turn_id != expected_model_turn_id - ) - return stale_common_owner or stale_response_owner or stale_model_turn_owner - - async def _maybe_continue_native_response( - self, - send_json, - *, - session: DuplexSession, - expected_epoch: int | None, - expected_model_turn_id: int | None = None, - ) -> None: - """Give an active model turn another silence unit. - - A model turn can need another unit before it has emitted visible - text/audio and therefore before a Realtime response exists. Bind that - continuation to the model-turn identity; once a response exists, its - response/turn ownership remains the continuation fence. - """ - response_id = session.active_response_id - native = self._runtime_session_state(session) - if session.state == DuplexSessionState.CLOSED: - native.clear_continuation() - return - request_id = session.active_request_id - if request_id is None: - native.clear_continuation() - return - if expected_epoch is not None and session.epoch != expected_epoch: - return - auto_response = self._session_auto_responds(session) - response_owned = response_id is not None - if response_owned: - owner_id = f"response:{response_id}" - payload_turn_id = ( - session.active_response_turn_id if session.active_response_turn_id is not None else session.turn_id - ) - else: - if not auto_response or expected_model_turn_id is None or session.turn_id != expected_model_turn_id: - native.clear_continuation() - return - owner_id = f"model-turn:{expected_model_turn_id}" - payload_turn_id = expected_model_turn_id - count = native.continuation_units if native.continuation_owner_id == owner_id else 0 - continuation_limit = ( - self._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - if auto_response - else self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS - ) - if count >= continuation_limit: - if auto_response: - await self._finish_bounded_native_auto_response( - send_json, - session=session, - expected_epoch=expected_epoch, - model_turn_id=payload_turn_id, - ) - return - payload = self._native_silence_unit_payload() - if auto_response and count + 1 == continuation_limit: - payload["force_listen"] = True - payload["duplex_turn_id"] = payload_turn_id - - scheduler = native.silence_continuation_scheduler - if scheduler is None: - return - try: - scheduled = await scheduler( - payload, - request_id=request_id, - owner_id=owner_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_incarnation=session.incarnation, - expected_model_turn_id=expected_model_turn_id, - send_json=send_json, - ) - except Exception as exc: - logger.exception("Failed to schedule duplex native response continuation: %s", exc) - scheduled = False - if scheduled: - native.continuation_owner_id = owner_id - native.continuation_units = count + 1 - - async def _cancel_native_data_plane_stream(self, session: DuplexSession) -> bool: - native = self._runtime_session_state(session) - task = native.data_plane_task - native.data_plane_task = None - native.data_plane_restart_requested = False - if task is None or task.done(): - return False - task.cancel() - try: - await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=0.25) - except asyncio.TimeoutError: - # Keep barge-in/cancel responsive. The task still carries the old - # expected_epoch and late model output is filtered by the writer. - pass - return True - - async def _signal_runtime_session( - self, - session: DuplexSession, - event: str, - send_json=None, - *, - fence: DuplexFence | None = None, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - ) -> bool: - if not self._uses_serving_runtime_adapter(session.config): - return True - signal_turn = getattr(self._chat_service.engine_client, "signal_duplex_turn_async", None) - if not callable(signal_turn): - return True - try: - signal_kwargs = { - "event": event, - "timeout": self._runtime_control_timeout_s(session), - } - if self._callable_accepts_keyword(signal_turn, "fence"): - signal_kwargs["fence"] = fence or DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - if next_fence is not None and self._callable_accepts_keyword(signal_turn, "next_fence"): - signal_kwargs["next_fence"] = next_fence - if session_config is not None and self._callable_accepts_keyword(signal_turn, "session_config"): - signal_kwargs["session_config"] = session_config - if runtime_config is not None and self._callable_accepts_keyword(signal_turn, "runtime_config"): - signal_kwargs["runtime_config"] = runtime_config - result = await signal_turn(session.session_id, **signal_kwargs) - except Exception as exc: - logger.exception("Failed to signal duplex runtime session: %s", exc) - if send_json is not None: - await self._send_runtime_error(send_json, "runtime_signal_failed", exc, session=session) - return False - if isinstance(result, dict) and self._runtime_signal_failed(result): - logger.warning( - "Duplex runtime signal failed: session=%s event=%s result=%s", - session.session_id, - event, - self._redact_runtime_control_result(result), - ) - if send_json is not None: - await self._send_runtime_control_error( - send_json, - "runtime_signal_failed", - "Duplex runtime signal failed", - result, - session=session, - ) - return False - await self._send_runtime_control_if_needed(send_json, result, session=session) - return True - - async def _close_runtime_session(self, session: DuplexSession, *, reason: str, send_json=None) -> bool: - if not self._uses_serving_runtime_adapter(session.config): - return True - close_session = getattr(self._chat_service.engine_client, "close_duplex_session_async", None) - if not callable(close_session): - return True - try: - close_kwargs = { - "reason": reason, - "timeout": self._runtime_control_timeout_s(session), - } - if self._callable_accepts_keyword(close_session, "fence"): - close_kwargs["fence"] = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - result = await close_session(session.session_id, **close_kwargs) - except Exception as exc: - logger.exception("Failed to close duplex runtime session: %s", exc) - if send_json is not None: - await self._send_runtime_error(send_json, "runtime_close_failed", exc, session=session) - return False - if isinstance(result, dict) and self._runtime_control_failed(result): - if send_json is not None: - await self._send_runtime_control_error( - send_json, - "runtime_close_failed", - "Duplex runtime close failed", - result, - session=session, - ) - return False - await self._send_runtime_control_if_needed(send_json, result, session=session) - return True - - @staticmethod - def _runtime_control_timeout_s(session: DuplexSession) -> float: - raw = session.config.extra_body.get("duplex_control_timeout_s") or session.config.extra_body.get( - "runtime_control_timeout_s" - ) - if isinstance(raw, int | float) and raw > 0: - return float(raw) - if session.capabilities.implementation_level == "model_native_duplex": - return 60.0 - return 10.0 - - async def _send_runtime_control_if_needed( - self, - send_json, - result: object, - *, - session: DuplexSession, - ) -> None: - if send_json is None: - return - if not isinstance(result, dict): - return - if not result.get("unsupported_count") and not result.get("error_count"): - return - await send_json( - { - "type": "runtime.control", - "session_id": session.session_id, - "epoch": session.epoch, - "result": self._redact_runtime_control_result(result), - } - ) - - @staticmethod - def _runtime_control_failed(result: dict[str, object]) -> bool: - if result.get("ok") is False: - return True - for key in ("unsupported_count", "error_count"): - value = result.get(key) - if isinstance(value, int | float) and value > 0: - return True - return False - - @classmethod - def _runtime_signal_failed(cls, result: dict[str, object]) -> bool: - error_count = result.get("error_count") - if isinstance(error_count, int | float) and error_count > 0: - return True - if result.get("ok") is not False: - return False - return not cls._runtime_signal_has_data_plane_ack(result) - - @classmethod - def _runtime_signal_has_data_plane_ack(cls, value: object) -> bool: - if isinstance(value, dict): - if value.get("data_plane_signal") is True and value.get("supported") is True: - return True - return any(cls._runtime_signal_has_data_plane_ack(child) for child in value.values()) - if isinstance(value, list | tuple): - return any(cls._runtime_signal_has_data_plane_ack(item) for item in value) - return False - - @classmethod - def _redact_runtime_control_result(cls, value: object) -> object: - if isinstance(value, dict): - redacted = { - key: cls._redact_runtime_control_result(child) - for key, child in value.items() - if key - not in { - "stage_handoff", - "tts_handoff", - "omni_payload", - "tts_hidden_states", - "tts_token_ids", - "traceback", - } - } - native_result = redacted.get("native_result") - if isinstance(native_result, dict): - if native_result.get("requires_stage_handoff") is True: - native_result.pop("requires_stage_handoff", None) - if native_result.get("requires_tts_stage") is True: - native_result.pop("requires_tts_stage", None) - return redacted - if isinstance(value, list | tuple): - return [cls._redact_runtime_control_result(item) for item in value] - if is_dataclass(value) and not isinstance(value, type): - # Control results carry typed values straight from the engine: - # `fence` and `accepted_fence` are `DuplexFence` instances added by - # `DuplexControlClient.execute`. This redacted value is emitted on - # the wire (for example `session.created.runtime_control` in - # `session_runner`), where `WebSocket.send_json` would raise - # `TypeError: Object of type DuplexFence is not JSON serializable` - # and kill the session during the handshake. - return {key: cls._redact_runtime_control_result(child) for key, child in asdict(value).items()} - if isinstance(value, Enum): - return value.value - return value - - async def _send_native_duplex_events( - self, - send_json, - result: object, - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> tuple[str | None, bool]: - if send_json is None: - return None, False - if expected_epoch is not None and session.epoch != expected_epoch: - return None, False - close_reason: str | None = None - emitted_response = False - request_id, _ = self._data_plane_request_info(result) - if self._serving_runtime_adapter.data_plane.is_terminal(request_id): - return None, False - if request_id is not None and session.active_request_id is None: - session.bind_request(request_id) - context = self._runtime_data_plane_context(session) - for native_result in self._serving_runtime_adapter.data_plane.project(result, context=context): - close_reason_for_result, did_emit = await self._send_one_native_duplex_event( - send_json, - native_result, - session=session, - expected_epoch=expected_epoch, - ) - emitted_response = emitted_response or did_emit - close_reason = close_reason or close_reason_for_result - if expected_epoch is not None and session.epoch != expected_epoch: - return None, emitted_response - return close_reason, emitted_response - - async def _drain_native_data_plane_stream( - self, - send_json, - result: object, - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> str | None: - request_id, response_stage_id = self._data_plane_request_info(result) - if request_id is None: - return None - if self._data_plane_outputs_finished(result): - return None - collect_outputs = getattr( - self._chat_service.engine_client, - "collect_duplex_data_plane_outputs_async", - None, - ) - if not callable(collect_outputs): - return None - - close_reason: str | None = None - empty_polls = 0 - while close_reason is None: - if self._serving_runtime_adapter.data_plane.is_terminal(request_id): - return None - if self._session_auto_responds(session): - active_request_id = session.active_request_id - if active_request_id is not None and active_request_id != request_id: - return None - if expected_epoch is not None and session.epoch != expected_epoch: - return None - if session.state == DuplexSessionState.CLOSED: - return None - outputs = await collect_outputs( - request_id, - response_stage_id=response_stage_id, - timeout=self._runtime_control_timeout_s(session), - ) - if expected_epoch is not None and session.epoch != expected_epoch: - return None - if not outputs: - # An empty poll means no output arrived within one control - # window, not that the stream is over. Exiting here orphans a - # decision that lands moments later: it would sit queued until - # the NEXT append starts a fresh drain task, adding one full - # chunk of latency to every model decision. - empty_polls += 1 - if empty_polls >= 3: - return None - continue - empty_polls = 0 - drain_result = { - "data_plane_outputs": outputs, - } - close_reason, emitted_response = await self._send_native_duplex_events( - send_json, - drain_result, - session=session, - expected_epoch=expected_epoch, - ) - if ( - self._data_plane_outputs_finished(drain_result) - and emitted_response - and not self._session_auto_responds(session) - ): - # Auto-respond sessions are resumable: every segment ends - # with finished=True but the stream continues with the next - # audio chunk. Exiting here would make each decision wait for - # the next append to start a fresh drain task (one full chunk - # of added latency); keep draining until epoch change, - # session close, or cancellation by a newer drain task. - return close_reason - # A batch without a client-visible event (e.g. the generic - # stage-0 final-output message that precedes the listen/speak - # decision in the same segment) must not terminate the drain; - # the real decision is still in flight. - return close_reason - - @staticmethod - def _data_plane_outputs_finished(result: object) -> bool: - if not isinstance(result, dict): - return False - outputs = result.get("data_plane_outputs") - if not isinstance(outputs, list) or not outputs: - return False - return bool(getattr(outputs[-1], "finished", False)) - - @staticmethod - def _data_plane_request_info(result: object) -> tuple[str | None, int | None]: - if not isinstance(result, dict): - return None, None - return duplex_data_plane_request_info(result) - - def _runtime_data_plane_context(self, session: DuplexSession) -> object: - response_config = session.response_config - return self._serving_runtime_adapter.data_plane_context( - epoch=session.epoch, - turn_id=session.turn_id, - active_response_turn_id=session.active_response_turn_id, - active_response_id=session.active_response_id, - auto_responds=self._session_auto_responds(session), - response_format=response_config.response_format, - speed=response_config.speed, - modalities=tuple(response_config.modalities), - ) - - async def _send_one_native_duplex_event( - self, - send_json, - native_result: dict[str, object], - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> tuple[str | None, bool]: - close_reason: str | None = None - emitted_response = False - if expected_epoch is not None and session.epoch != expected_epoch: - return close_reason, emitted_response - data_plane_request_id = native_result.get("data_plane_request_id") - if isinstance(data_plane_request_id, str) and self._serving_runtime_adapter.data_plane.is_terminal( - data_plane_request_id - ): - return close_reason, emitted_response - active_request_matches = session.active_request_id == data_plane_request_id or ( - self._session_auto_responds(session) and session.active_request_id is None - ) - if isinstance(data_plane_request_id, str) and not active_request_matches: - return close_reason, emitted_response - if isinstance(native_result.get("error_code"), str): - response_id = session.active_response_id - await send_json( - { - "type": "error", - "code": native_result.get("error_code"), - "session_id": session.session_id, - "epoch": session.epoch, - "error": str(native_result.get("error") or "Duplex native data-plane error"), - } - ) - if response_id is not None: - session.end_response(commit_text=False) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": False, - "status": "failed", - "status_details": { - "type": "failed", - "reason": native_result.get("error_code"), - }, - "playback": session.playback.as_dict(), - } - ) - return close_reason, True - if native_result.get("function_call") is True: - await send_json( - { - "type": "function_call.done", - "session_id": session.session_id, - "epoch": session.epoch, - "call_id": native_result.get("call_id"), - "name": native_result.get("name"), - "arguments": native_result.get("arguments", ""), - } - ) - return close_reason, True - if native_result.get("requires_stage_handoff") is True or native_result.get("requires_tts_stage") is True: - # Stage0 announces a talker handoff before Stage1 has produced - # client-visible audio. A single client input stream may contain - # several model turns, including terminal-only turns. Reserve the - # protocol response only when Stage1 emits text/audio so empty - # model turns cannot create empty responses or steal later audio. - return close_reason, False - is_listen = native_result.get("is_listen") - model_turn_id = coerce_int(native_result.get("model_turn_id")) - if native_result.get("is_buffering") is True or native_result.get("prefill_success") is False: - if native_result.get("data_plane_request_id") == session.active_request_id: - session.clear_request() - # No response_id here: a buffering listen is informational, not a - # response terminal — clients must not close a handle over it. - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": native_result.get("reason") or "buffering", - "model_listen": False, - "buffering": True, - } - self._attach_native_runtime_metadata(payload, native_result) - await send_json(payload) - return close_reason, emitted_response - if is_listen is True: - await self._end_active_response_before_future_model_turn( - send_json, - session=session, - model_turn_id=model_turn_id, - ) - if ( - session.active_response_id is not None - and model_turn_id is not None - and not session.active_response_accepts_model_turn(model_turn_id) - ): - return close_reason, emitted_response - auto_response = self._session_auto_responds(session) - active_response_id = session.active_response_id - auto_continuations_remaining = active_response_id is None or self._native_response_continuations_remaining( - session, active_response_id - ) - non_terminal_auto_listen = ( - auto_response - and active_response_id is not None - and session.active_request_id is not None - and native_result.get("end_of_turn") is not True - and auto_continuations_remaining - ) - if non_terminal_auto_listen: - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - ) - return close_reason, emitted_response - if not auto_response and data_plane_request_id == session.active_request_id: - session.clear_request() - model_listen = native_result.get("model_listen") - if not isinstance(model_listen, bool): - model_listen = native_result.get("reason") in {None, "", "model_listen"} - response_id = session.active_response_id - if isinstance(data_plane_request_id, str) and not auto_response: - self._serving_runtime_adapter.data_plane.mark_terminal(data_plane_request_id) - emitted_response = True - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": native_result.get("reason") or "model_listen", - "model_listen": model_listen, - } - if response_id is not None: - # The listen is this response's terminal decision (a - # response.done follows below with the same id); carry the id - # so clients can bind the decision to the precreated response. - payload["response_id"] = response_id - self._attach_native_runtime_metadata(payload, native_result) - await send_json(payload) - if native_result.get("abort_data_plane_request") is True and isinstance(data_plane_request_id, str): - await self._abort_request_background( - session, - data_plane_request_id, - send_json, - notify=False, - ) - if response_id is not None: - if not self._session_auto_responds(session) and self._native_response_continuations_remaining( - session, response_id - ): - # The official model often listens for a silence beat or - # two before answering; keep the response open and give it - # the next silence unit as a decision point. - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - ) - return close_reason, emitted_response - if auto_response: - self._runtime_session_state(session).clear_continuation() - if not auto_continuations_remaining: - completed_turn_id = model_turn_id - if completed_turn_id is None: - completed_turn_id = session.active_response_turn_id - if completed_turn_id is not None: - session.complete_model_turn(completed_turn_id) - session.end_response(commit_text=False, preserve_request=auto_response) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": False, - "playback": session.playback.as_dict(), - } - ) - return close_reason, emitted_response - - text = native_result.get("text") - audio = native_result.get("audio_data", native_result.get("audio")) - end_of_turn = bool(native_result.get("end_of_turn", False)) - has_text = isinstance(text, str) and bool(text) - has_audio = isinstance(audio, str) and bool(audio) - if not has_text and not has_audio and not end_of_turn: - tts_segment_ended = ( - native_result.get("stage_role") == "tts" and native_result.get("abort_data_plane_request") is True - ) - if ( - tts_segment_ended - and self._session_auto_responds(session) - and (model_turn_id is not None or session.active_response_id is not None) - ): - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - expected_model_turn_id=model_turn_id, - ) - return close_reason, emitted_response - if end_of_turn and not has_text and not has_audio and session.active_response_id is None: - if isinstance(data_plane_request_id, str): - if not self._session_auto_responds(session) and data_plane_request_id == session.active_request_id: - session.clear_request() - if not self._session_auto_responds(session): - self._serving_runtime_adapter.data_plane.mark_terminal(data_plane_request_id) - if model_turn_id is not None: - session.complete_model_turn(model_turn_id) - if self._session_auto_responds(session): - self._runtime_session_state(session).clear_continuation() - emitted_response = True - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": "model_turn_completed_without_output", - "model_listen": True, - } - self._attach_native_runtime_metadata(payload, native_result) - await send_json(payload) - return close_reason, emitted_response - if session.active_response_id is None and model_turn_id is not None and model_turn_id < session.turn_id: - # A continuation append can already be in flight when the prior - # response reaches turn EOS. Its late audio still carries the - # completed model turn, so it must not reserve a second Realtime - # response for that turn. - return close_reason, emitted_response - await self._end_active_response_before_future_model_turn( - send_json, - session=session, - model_turn_id=model_turn_id, - ) - if ( - session.active_response_id is not None - and model_turn_id is not None - and not session.active_response_accepts_model_turn(model_turn_id) - ): - return close_reason, emitted_response - emitted_response = True - response_created = False - response_id = session.active_response_id - if response_id is None: - response_id = session.begin_response(turn_id=model_turn_id) - response_created = True - response_request_metrics = session.mark_response_first_outputs( - observed_at_s=time.monotonic(), - has_text=has_text, - has_audio=has_audio, - ) - if response_created: - response_created_payload = self._response_created_payload( - session, - response_id, - epoch=session.epoch, - ) - if response_request_metrics: - response_created_payload["response_request_metrics"] = response_request_metrics - await send_json(response_created_payload) - response_stage_metrics = session.accumulate_response_stage_metrics( - native_result.get("stage_metrics") if isinstance(native_result.get("stage_metrics"), Mapping) else None - ) - if response_created: - speak_payload = { - "type": "response.speak", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "text": text if isinstance(text, str) else "", - "end_of_turn": end_of_turn, - "model_speak": True, - } - self._attach_native_runtime_metadata( - speak_payload, - native_result, - stage_metrics=response_stage_metrics, - response_request_metrics=response_request_metrics, - ) - await send_json(speak_payload) - previous_sent_ms = session.playback.sent_ms - text_chars_before_append = len("".join(session.assistant_text_buffer)) - if isinstance(text, str): - session.append_assistant_text(text) - duration_ms = native_result.get("audio_duration_ms") - text_chars = len("".join(session.assistant_text_buffer)) - mark_duration_ms = None - mark_text_chars: int | None = text_chars - if native_result.get("audio_text_mark") is False: - mark_text_chars = None - if isinstance(duration_ms, int | float): - mark_duration_ms = int(duration_ms) - if native_result.get("audio_duration_is_cumulative") is not True: - mark_duration_ms += session.playback.sent_ms - audio_text_marks = native_result.get("audio_text_marks") - audio_text_marks = self._normalize_native_audio_text_marks( - audio_text_marks if isinstance(audio_text_marks, list) else None, - audio_offset_ms=( - 0 - if native_result.get("audio_text_marks_are_cumulative") is True - or native_result.get("audio_duration_is_cumulative") is True - else previous_sent_ms - ), - text_offset_chars=( - 0 if native_result.get("audio_text_marks_are_cumulative") is True else text_chars_before_append - ), - ) - session.mark_audio_sent( - mark_duration_ms, - text_chars=mark_text_chars if mark_duration_ms is not None else None, - audio_text_marks=audio_text_marks, - ) - payload = { - "type": "response.output_audio.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "text": text if isinstance(text, str) else "", - "audio": audio if isinstance(audio, str) else "", - "format": ( - native_result.get("audio_format") - if isinstance(native_result.get("audio_format"), str) - else session.response_config.response_format - ), - "end_of_turn": end_of_turn, - "model_speak": True, - } - if mark_duration_ms is not None: - payload["audio_duration_ms"] = mark_duration_ms - if audio_text_marks: - payload["audio_text_marks"] = audio_text_marks - elif mark_duration_ms is not None and mark_text_chars is not None: - payload["audio_text_marks"] = [ - { - "text_chars": max(0, int(mark_text_chars)), - "audio_end_ms": max(0, int(mark_duration_ms)), - } - ] - payload["playback"] = session.playback.as_dict() - sample_rate_hz = native_result.get("sample_rate_hz") or native_result.get("audio_sample_rate_hz") - if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) > 0: - payload["sample_rate_hz"] = int(sample_rate_hz) - self._attach_native_runtime_metadata( - payload, - native_result, - stage_metrics=response_stage_metrics, - response_request_metrics=response_request_metrics, - ) - await send_json(payload) - if ( - not end_of_turn - and native_result.get("stage_role") == "tts" - and native_result.get("abort_data_plane_request") is True - and self._session_auto_responds(session) - ): - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - ) - if end_of_turn: - data_plane_request_id = native_result.get("data_plane_request_id") - if isinstance(data_plane_request_id, str) and not self._session_auto_responds(session): - self._serving_runtime_adapter.data_plane.close_stream(data_plane_request_id) - if isinstance(data_plane_request_id, str): - if not self._session_auto_responds(session) and data_plane_request_id == session.active_request_id: - session.clear_request() - if not self._session_auto_responds(session): - self._serving_runtime_adapter.data_plane.mark_terminal(data_plane_request_id) - should_commit = self._should_commit_response_to_history(session, response_id) - committed_message = session.end_response( - commit_text=should_commit, - preserve_request=self._session_auto_responds(session), - ) - model_turn_id = coerce_int(native_result.get("model_turn_id")) - if model_turn_id is not None: - session.complete_model_turn(model_turn_id) - if should_commit: - session.register_history_item(f"item_{response_id}", committed_message) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - ) - return close_reason, emitted_response - - async def _end_active_response_before_future_model_turn( - self, - send_json, - *, - session: DuplexSession, - model_turn_id: int | None, - ) -> None: - if not self._session_auto_responds(session): - return - response_id = session.active_response_id - active_turn_id = session.active_response_turn_id - if response_id is None or model_turn_id is None or active_turn_id is None: - return - if int(model_turn_id) <= int(active_turn_id): - return - session.complete_model_turn(int(model_turn_id) - 1) - should_commit = self._should_commit_response_to_history(session, response_id) - committed_message = session.end_response( - commit_text=should_commit, - preserve_request=True, - ) - if should_commit: - session.register_history_item(f"item_{response_id}", committed_message) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - ) - - @staticmethod - def _normalize_native_audio_text_marks( - audio_text_marks: list[object] | None, - *, - audio_offset_ms: int, - text_offset_chars: int, - ) -> list[dict[str, int]] | None: - if not audio_text_marks: - return None - normalized: list[dict[str, int]] = [] - for raw_mark in audio_text_marks: - if not isinstance(raw_mark, dict): - continue - raw_text_chars = raw_mark.get("text_chars") - raw_audio_end_ms = raw_mark.get("audio_end_ms", raw_mark.get("audio_ms")) - if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): - continue - normalized.append( - { - "text_chars": max(0, int(raw_text_chars) + int(text_offset_chars)), - "audio_end_ms": max(0, int(raw_audio_end_ms) + int(audio_offset_ms)), - } - ) - return normalized or None - - def _cleanup_duplex_session_state(self, session: DuplexSession) -> None: - session_id = session.session_id - pipeline = self._server_vad_pipelines.pop(session_id, None) - if pipeline is not None: - pipeline.reset() - self._realtime_vad_metrics.session_finished() - adapter = self._serving_runtime_adapter - if adapter is None: - self._serving_session_states.pop(session_id, None) - else: - adapter.remove_session_state(session_id) - adapter.data_plane.close_session( - session_id, - active_request_id=session.active_request_id, - ) - - def _encode_native_data_plane_audio( - self, - audio_data: object, - sample_rate_hz: int, - response_format: str, - speed: float | None, - ) -> str | None: - if audio_data is None: - return None - try: - import torch - - from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio - - if isinstance(audio_data, torch.Tensor): - audio_tensor = audio_data.detach().cpu().float().numpy() - else: - audio_tensor = np.asarray(audio_data, dtype=np.float32) - if audio_tensor.ndim > 1: - audio_tensor = audio_tensor.reshape(-1) - audio_response = self._chat_service.create_audio( - CreateAudio( - audio_tensor=audio_tensor, - sample_rate=sample_rate_hz, - response_format=response_format, - speed=float(speed) if isinstance(speed, int | float) and speed > 0 else 1.0, - stream_format="audio", - base64_encode=True, - ) - ) - return str(audio_response.audio_data) - except Exception: - logger.exception("Failed to encode duplex data-plane audio output") - return None - - @staticmethod - def _attach_native_runtime_metadata( - payload: dict[str, object], - native_result: dict[str, object], - *, - stage_metrics: Mapping[str, object] | None = None, - response_request_metrics: Mapping[str, object] | None = None, - ) -> None: - metadata: dict[str, object] = {} - runtime_impl = native_result.get("runtime_impl") - if isinstance(runtime_impl, str) and runtime_impl: - metadata["runtime_impl"] = runtime_impl - owned_runtime = native_result.get("owned_runtime") - if isinstance(owned_runtime, bool): - metadata["owned_runtime"] = owned_runtime - model_turn_id = coerce_int(native_result.get("model_turn_id")) - if model_turn_id is not None: - metadata["model_turn_id"] = model_turn_id - for name in ( - "uses_model_runner_scheduler", - "runner_kv_backed", - ): - value = native_result.get(name) - if isinstance(value, bool): - metadata[name] = value - effective_stage_metrics = stage_metrics if stage_metrics is not None else native_result.get("stage_metrics") - if isinstance(effective_stage_metrics, Mapping): - metadata["stage_metrics"] = { - str(stage_id): dict(values) - for stage_id, values in effective_stage_metrics.items() - if isinstance(values, Mapping) - } - if response_request_metrics: - metadata["response_request_metrics"] = dict(response_request_metrics) - if metadata: - payload["vllm_omni"] = metadata - - async def _send_runtime_error( - self, - send_json, - code: str, - exc: Exception, - *, - session: DuplexSession | None = None, - ) -> None: - retryable = False - message = str(exc) - if isinstance(exc, DuplexControlRequestError): - code = exc.code - retryable = exc.retryable - error_data = exc.result.get("error") - if isinstance(error_data, dict) and isinstance(error_data.get("message"), str): - message = error_data["message"] - payload: dict[str, object] = { - "type": "error", - "code": code, - "error": message, - "retryable": retryable, - } - if session is not None: - payload["session_id"] = session.session_id - payload["epoch"] = session.epoch - await send_json(payload) - - async def _send_runtime_control_error( - self, - send_json, - code: str, - message: str, - result: dict[str, object], - *, - session: DuplexSession, - ) -> None: - await send_json( - { - "type": "error", - "code": code, - "error": message, - "session_id": session.session_id, - "epoch": session.epoch, - "runtime_control": self._redact_runtime_control_result(result), - } - ) diff --git a/vllm_omni/entrypoints/duplex/server_vad.py b/vllm_omni/entrypoints/duplex/server_vad.py deleted file mode 100644 index 610e51c8c..000000000 --- a/vllm_omni/entrypoints/duplex/server_vad.py +++ /dev/null @@ -1,528 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import hashlib -import threading -import time -from collections.abc import Mapping -from dataclasses import dataclass, replace -from pathlib import Path -from typing import TYPE_CHECKING, Literal, Protocol - -import numpy as np -from vllm.transformers_utils.repo_utils import try_get_local_file - -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.audio_utils_mixin import StreamingAudioResampler - -SILERO_VAD_REPO_ID = "istupakov/silero-vad-onnx" -SILERO_VAD_REVISION = "8b14476858ef240c50b3884bb38cc67290c1cc70" -SILERO_VAD_FILENAME = "silero_vad.onnx" -SILERO_VAD_SHA256 = "1a153a22f4509e292a94e67d6f9b85e8deb25b4988682b7e174c65279d8788e3" -SILERO_VAD_MIN_THRESHOLD = 0.15 -SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS = 96 - -_SERVER_VAD_FIELDS = { - "type", - "threshold", - "prefix_padding_ms", - "silence_duration_ms", - "create_response", - "interrupt_response", - "min_speech_duration_ms", -} - - -@dataclass(frozen=True, slots=True) -class ServerVADConfig: - """Validated OpenAI-compatible ``server_vad`` session configuration.""" - - type: Literal["server_vad"] = "server_vad" - threshold: float = 0.5 - prefix_padding_ms: int = 300 - silence_duration_ms: int = 500 - create_response: bool = True - # Omission is resolved from runtime capabilities before the session is opened. - interrupt_response: bool | None = None - min_speech_duration_ms: int | None = None - - @classmethod - def from_value(cls, value: object) -> ServerVADConfig: - if not isinstance(value, dict): - raise ValueError("turn_detection must be null or an object") - unknown = sorted(set(value) - _SERVER_VAD_FIELDS) - if unknown: - raise ValueError(f"Unknown server_vad field(s): {', '.join(unknown)}") - vad_type = value.get("type") - if vad_type != "server_vad": - raise ValueError("turn_detection.type must be 'server_vad'") - - threshold = value.get("threshold", 0.5) - if isinstance(threshold, bool) or not isinstance(threshold, int | float) or not 0 <= threshold <= 1: - raise ValueError("server_vad.threshold must be a number between 0 and 1") - - prefix_padding_ms = value.get("prefix_padding_ms", 300) - if isinstance(prefix_padding_ms, bool) or not isinstance(prefix_padding_ms, int) or prefix_padding_ms < 0: - raise ValueError("server_vad.prefix_padding_ms must be a non-negative integer") - - silence_duration_ms = value.get("silence_duration_ms", 500) - if ( - isinstance(silence_duration_ms, bool) - or not isinstance(silence_duration_ms, int) - or silence_duration_ms <= 0 - ): - raise ValueError("server_vad.silence_duration_ms must be a positive integer") - - create_response = value.get("create_response", True) - if not isinstance(create_response, bool): - raise ValueError("server_vad.create_response must be a boolean") - - interrupt_response = value.get("interrupt_response") - if "interrupt_response" in value and not isinstance(interrupt_response, bool): - raise ValueError("server_vad.interrupt_response must be a boolean") - - min_speech_duration_ms = value.get("min_speech_duration_ms") - if min_speech_duration_ms is not None and ( - isinstance(min_speech_duration_ms, bool) - or not isinstance(min_speech_duration_ms, int | float) - or not np.isfinite(float(min_speech_duration_ms)) - or min_speech_duration_ms < 0 - ): - raise ValueError("server_vad.min_speech_duration_ms must be a non-negative number") - - return cls( - type=vad_type, - threshold=float(threshold), - prefix_padding_ms=prefix_padding_ms, - silence_duration_ms=silence_duration_ms, - create_response=create_response, - interrupt_response=interrupt_response, - min_speech_duration_ms=(int(min_speech_duration_ms) if min_speech_duration_ms is not None else None), - ) - - def as_dict(self) -> dict[str, object]: - result: dict[str, object] = { - "type": self.type, - "threshold": self.threshold, - "prefix_padding_ms": self.prefix_padding_ms, - "silence_duration_ms": self.silence_duration_ms, - "create_response": self.create_response, - } - if self.interrupt_response is not None: - result["interrupt_response"] = self.interrupt_response - if self.min_speech_duration_ms is not None: - result["min_speech_duration_ms"] = self.min_speech_duration_ms - return result - - -def parse_session_turn_detection( - payload: Mapping[str, object], -) -> tuple[bool, ServerVADConfig | None]: - """Extract and validate OpenAI-compatible turn detection aliases.""" - configured_values: list[tuple[str, object]] = [] - if "turn_detection" in payload: - configured_values.append(("turn_detection", payload["turn_detection"])) - audio = payload.get("audio") - if isinstance(audio, Mapping): - audio_input = audio.get("input") - if isinstance(audio_input, Mapping) and "turn_detection" in audio_input: - configured_values.append(("audio.input.turn_detection", audio_input["turn_detection"])) - if not configured_values: - return False, None - - first_path, first_value = configured_values[0] - first_config = None if first_value is None else ServerVADConfig.from_value(first_value) - for field_path, value in configured_values[1:]: - config = None if value is None else ServerVADConfig.from_value(value) - if first_config is not None and config is not None: - # An omitted option must not conflict with an explicit alias value. - if first_config.interrupt_response is None: - first_config = replace(first_config, interrupt_response=config.interrupt_response) - if config.interrupt_response is None: - config = replace(config, interrupt_response=first_config.interrupt_response) - if config != first_config: - raise ValueError(f"{first_path} and {field_path} must not specify conflicting values") - return True, first_config - - -class SpeechDetectorBackend(Protocol): - """Detector contract; per-stream state is owned by the pipeline.""" - - frame_samples: int - - def new_state(self) -> object: ... - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: ... - - -class SpeechDetectorBackendProvider(Protocol): - def get(self) -> SpeechDetectorBackend: ... - - -@dataclass(frozen=True, slots=True) -class SpeechEndpointDecision: - speech_started: bool = False - speech_stopped: bool = False - audio_start_ms: int | None = None - audio_end_ms: int | None = None - endpoint_delay_ms: int | None = None - - -class SpeechEndpointPolicy(Protocol): - @property - def speech_active(self) -> bool: ... - - def update( - self, - probability: float, - *, - frame_start_sample: int, - frame_samples: int, - ) -> SpeechEndpointDecision: ... - - def reset(self) -> None: ... - - -class ThresholdEndpointPolicy: - """Apply threshold and trailing-silence endpoint rules.""" - - def __init__(self, config: ServerVADConfig, *, sample_rate_hz: int) -> None: - self.config = config - self.sample_rate_hz = sample_rate_hz - self._speech_active = False - self._candidate_samples = 0 - self._candidate_start_sample: int | None = None - self._silence_samples = 0 - - @property - def speech_active(self) -> bool: - return self._speech_active - - def update( - self, - probability: float, - *, - frame_start_sample: int, - frame_samples: int, - ) -> SpeechEndpointDecision: - if probability >= self.config.threshold: - self._silence_samples = 0 - if not self._speech_active: - if self._candidate_samples == 0: - self._candidate_start_sample = frame_start_sample - self._candidate_samples += frame_samples - min_speech_duration_ms = self.config.min_speech_duration_ms or 0 - min_speech_samples = max(1, round(min_speech_duration_ms * self.sample_rate_hz / 1000)) - if self._candidate_samples < min_speech_samples: - return SpeechEndpointDecision() - self._speech_active = True - detected_start_sample = ( - self._candidate_start_sample if self._candidate_start_sample is not None else frame_start_sample - ) - detected_start_ms = round(detected_start_sample * 1000 / self.sample_rate_hz) - self._candidate_samples = 0 - self._candidate_start_sample = None - return SpeechEndpointDecision( - speech_started=True, - audio_start_ms=max(0, detected_start_ms - self.config.prefix_padding_ms), - ) - return SpeechEndpointDecision() - - if not self._speech_active: - self._candidate_samples = 0 - self._candidate_start_sample = None - return SpeechEndpointDecision() - - # Match Silero v6.2's streaming VAD hysteresis: activation uses the - # configured threshold, while potential speech end starts only below - # ``max(threshold - 0.15, 0.01)``. Once a silence candidate exists, - # intermediate frames keep elapsed time moving but cannot themselves - # close the turn. - negative_threshold = max(self.config.threshold - SILERO_VAD_MIN_THRESHOLD, 0.01) - below_negative_threshold = probability < negative_threshold - if not below_negative_threshold and self._silence_samples == 0: - return SpeechEndpointDecision() - - self._silence_samples += frame_samples - if not below_negative_threshold: - return SpeechEndpointDecision() - silence_limit = max(1, round(self.config.silence_duration_ms * self.sample_rate_hz / 1000)) - if self._silence_samples < silence_limit: - return SpeechEndpointDecision() - - # OpenAI defines ``audio_end_ms`` as the end of the audio sent to the - # model, including the trailing silence used for endpoint detection. - audio_end_sample = frame_start_sample + frame_samples - endpoint_delay_ms = round(self._silence_samples * 1000 / self.sample_rate_hz) - self.reset() - return SpeechEndpointDecision( - speech_stopped=True, - audio_end_ms=max(0, round(audio_end_sample * 1000 / self.sample_rate_hz)), - endpoint_delay_ms=endpoint_delay_ms, - ) - - def reset(self) -> None: - self._speech_active = False - self._candidate_samples = 0 - self._candidate_start_sample = None - self._silence_samples = 0 - - -@dataclass(frozen=True, slots=True) -class ServerVADFrame: - samples: np.ndarray - decision: SpeechEndpointDecision - - -@dataclass(frozen=True, slots=True) -class ServerVADBatch: - frames: tuple[ServerVADFrame, ...] - inference_ms: float - - -class ServerVADPipeline: - """Frame continuous 16 kHz audio and apply acoustic endpointing. - - The pipeline retains only input-resampling residuals, a partial frame, - detector state, and timing counters. Commit-eligible prefix and utterance - audio remain session-owned. - """ - - sample_rate_hz = 16_000 - - def __init__( - self, - backend: SpeechDetectorBackend, - config: ServerVADConfig, - *, - endpoint_policy: SpeechEndpointPolicy | None = None, - ) -> None: - self.backend = backend - self.config = config - self.endpoint_policy = endpoint_policy or ThresholdEndpointPolicy( - config, - sample_rate_hz=self.sample_rate_hz, - ) - self._input_resampler: StreamingAudioResampler | None = None - self._source_sample_rate_hz: int | None = None - self._scratch = np.empty(0, dtype=np.float32) - self._detector_state = backend.new_state() - self._processed_samples = 0 - self._stream_start_ms = 0 - - @property - def speech_active(self) -> bool: - return self.endpoint_policy.speech_active - - @property - def source_sample_rate_hz(self) -> int | None: - return self._source_sample_rate_hz - - @property - def scratch_bytes(self) -> int: - resampler_bytes = self._input_resampler.scratch_bytes if self._input_resampler is not None else 0 - return int(self._scratch.nbytes) + resampler_bytes - - async def push_pcm16( - self, - samples: bytes, - *, - source_sample_rate_hz: int, - ) -> ServerVADBatch: - """Normalize chunked mono PCM16 input and run endpoint detection.""" - if isinstance(source_sample_rate_hz, bool) or source_sample_rate_hz not in {16_000, 24_000}: - raise ValueError("server_vad PCM16 input sample rate must be 16000 or 24000 Hz") - - if len(samples) % np.dtype(" ServerVADBatch: - normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1) - return await asyncio.to_thread(self._push_sync, normalized) - - def _push_sync(self, samples: np.ndarray) -> ServerVADBatch: - started_at = time.perf_counter() - if self._scratch.size: - samples = np.concatenate((self._scratch, samples)) - frame_samples = int(self.backend.frame_samples) - complete_samples = samples.size - samples.size % frame_samples - # Retained state must not keep the whole input chunk alive through a view. - self._scratch = samples[complete_samples:].copy() - - results: list[ServerVADFrame] = [] - for offset in range(0, complete_samples, frame_samples): - frame = np.ascontiguousarray(samples[offset : offset + frame_samples], dtype=np.float32) - frame_start = self._processed_samples - probability, self._detector_state = self.backend.infer(frame, self._detector_state) - probability = min(1.0, max(0.0, float(probability))) - decision = self.endpoint_policy.update( - probability, - frame_start_sample=frame_start, - frame_samples=frame_samples, - ) - if decision.speech_started and decision.audio_start_ms is not None: - # Prefix audio before the most recent reset is no longer available. - decision = replace(decision, audio_start_ms=max(self._stream_start_ms, decision.audio_start_ms)) - - self._processed_samples += frame_samples - results.append( - ServerVADFrame( - samples=frame, - decision=decision, - ) - ) - return ServerVADBatch( - frames=tuple(results), - inference_ms=(time.perf_counter() - started_at) * 1000, - ) - - def reset(self) -> None: - """Discard stream state while preserving the session-relative audio clock.""" - self._stream_start_ms = round(self._processed_samples * 1000 / self.sample_rate_hz) - self._input_resampler = None - self._source_sample_rate_hz = None - self._scratch = np.empty(0, dtype=np.float32) - self._detector_state = self.backend.new_state() - self.endpoint_policy.reset() - - -@dataclass(frozen=True, slots=True) -class _SileroVADState: - model_state: np.ndarray - context: np.ndarray - - -class SileroVADBackend: - """Shared ONNX Runtime Silero v6.2 detector running on CPU.""" - - sample_rate_hz = 16_000 - frame_samples = 512 - context_samples = 64 - model_state_shape = (2, 1, 128) - - def __init__(self, model_path: str | Path) -> None: - try: - import onnxruntime as ort - except ImportError as exc: # pragma: no cover - platform packaging supplies ORT. - raise RuntimeError("server_vad requires ONNX Runtime") from exc - - session_options = ort.SessionOptions() - session_options.inter_op_num_threads = 1 - session_options.intra_op_num_threads = 1 - - self.model_path = Path(model_path) - self._session = ort.InferenceSession( - str(self.model_path), - providers=["CPUExecutionProvider"], - sess_options=session_options, - ) - input_names = {item.name for item in self._session.get_inputs()} - if "input" not in input_names or "sr" not in input_names or "state" not in input_names: - raise RuntimeError(f"Unsupported Silero ONNX input contract: {sorted(input_names)}") - self._warm_up() - - def _warm_up(self) -> None: - self.infer(np.zeros(self.frame_samples, dtype=np.float32), self.new_state()) - - def new_state(self) -> _SileroVADState: - return _SileroVADState( - model_state=np.zeros(self.model_state_shape, dtype=np.float32), - context=np.zeros((1, self.context_samples), dtype=np.float32), - ) - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - if not isinstance(state, _SileroVADState): - raise TypeError("Silero detector state must be created by SileroVADBackend.new_state()") - model_state = np.ascontiguousarray(state.model_state, dtype=np.float32) - context = np.ascontiguousarray(state.context, dtype=np.float32) - if model_state.shape != self.model_state_shape: - raise ValueError(f"Silero model state must have shape {self.model_state_shape}, got {model_state.shape}") - expected_context_shape = (1, self.context_samples) - if context.shape != expected_context_shape: - raise ValueError(f"Silero context must have shape {expected_context_shape}, got {context.shape}") - - audio = np.ascontiguousarray(frame, dtype=np.float32).reshape(1, -1) - if audio.shape[1] != self.frame_samples: - raise ValueError( - f"Silero detector frame must contain exactly {self.frame_samples} samples, got {audio.shape[1]}" - ) - model_input = np.concatenate((context, audio), axis=1) - inputs = { - "input": model_input, - "state": model_state, - "sr": np.asarray(self.sample_rate_hz, dtype=np.int64), - } - # ONNX Runtime permits concurrent Run calls on one CPU session. Stream - # state is explicit, so independent pipelines can safely share the session. - output = self._session.run(None, inputs) - if len(output) < 2: - raise RuntimeError("Silero ONNX model did not return probability and model state") - probability = float(np.asarray(output[0]).reshape(-1)[0]) - next_state = _SileroVADState( - model_state=np.ascontiguousarray(output[1], dtype=np.float32), - context=np.ascontiguousarray(model_input[:, -self.context_samples :], dtype=np.float32), - ) - return probability, next_state - - -class SileroVADBackendProvider: - """Resolve, verify, and load one Silero model instance per process.""" - - def __init__(self, *, model_path: str | None = None) -> None: - self.model_path = model_path - self._backend: SileroVADBackend | None = None - self._lock = threading.Lock() - - def get(self) -> SileroVADBackend: - if self._backend is not None: - return self._backend - with self._lock: - if self._backend is None: - path = self._resolve_local_artifact() - self._verify_checksum(path) - self._backend = SileroVADBackend(path) - return self._backend - - def _resolve_local_artifact(self) -> Path: - if self.model_path: - path = Path(self.model_path).expanduser() - if path.is_file(): - return path - raise RuntimeError(f"Configured Silero VAD model does not exist: {path}") - cached = try_get_local_file( - model=SILERO_VAD_REPO_ID, - file_name=SILERO_VAD_FILENAME, - revision=SILERO_VAD_REVISION, - ) - if isinstance(cached, Path) and cached.is_file(): - return cached - raise RuntimeError( - "Silero VAD artifact is not available locally. Pre-download the pinned artifact " - "or configure duplex_session.server_vad_model_path." - ) - - def _verify_checksum(self, path: Path) -> None: - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != SILERO_VAD_SHA256: - raise RuntimeError(f"Silero VAD checksum mismatch for {path}: expected {SILERO_VAD_SHA256}, got {digest}") diff --git a/vllm_omni/entrypoints/duplex/serving.py b/vllm_omni/entrypoints/duplex/serving.py index 7bd690841..364f715c0 100644 --- a/vllm_omni/entrypoints/duplex/serving.py +++ b/vllm_omni/entrypoints/duplex/serving.py @@ -1,2130 +1,543 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Thin websocket serving for ``/v1/realtime?duplex=1``. + +Everything about a duplex *session* lives inside the engine (see +``vllm_omni.entrypoints.duplex_omni``). This handler only does transport: + +* accept the socket and run the Realtime handshake (``session.update`` opens a + session through ``DuplexOmni.open_session``; ``session.resume`` re-attaches + to an existing one); +* translate each client JSON event into a ``DuplexCommand`` and submit it in + arrival order; +* pump the session's typed events to the wire, journaling them for resume; +* manage attachments: resume tokens, replay, takeover, detach on disconnect. +""" + from __future__ import annotations import asyncio -import base64 -import binascii -import inspect import json -from collections.abc import Mapping, MutableMapping +from collections.abc import Mapping from contextlib import suppress from dataclasses import dataclass, replace -from typing import TYPE_CHECKING -import numpy as np -from fastapi import WebSocket +from fastapi import WebSocket, WebSocketDisconnect from vllm.logger import init_logger -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig -from vllm_omni.engine.duplex.messages import DuplexFence, DuplexSessionLifecycleMessage -from vllm_omni.engine.duplex.runtime import duplex_resource_request_id -from vllm_omni.entrypoints.duplex.capability import ( - should_enable_duplex_endpoint, -) -from vllm_omni.entrypoints.duplex.chat_fallback import ( - ChatFallbackProjectorMixin, -) -from vllm_omni.entrypoints.duplex.protocol import ( - NATIVE_DUPLEX_KEY, - DuplexCapabilities, - DuplexCommittedInput, - DuplexOverlapPolicy, - DuplexPlaybackCommitPolicy, - DuplexSession, - DuplexSessionConfig, - DuplexSessionRegistry, - DuplexSessionState, - DuplexTurnController, - ResponseCreateOptions, - native_duplex_opt_in, - normalize_native_duplex_key, -) -from vllm_omni.entrypoints.duplex.realtime_session import ( - REALTIME_OUTPUT_AUDIO_FORMATS, - NativeRealtimeSessionProtocol, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - ServingRuntimeAdapter, - ServingRuntimeConfigError, - ServingRuntimeSessionState, - TurnBasedServingSessionState, - load_serving_runtime_adapter, - validate_serving_runtime_adapter, -) -from vllm_omni.entrypoints.duplex.runtime_bridge import ( - NativeRuntimeBridgeMixin, -) -from vllm_omni.entrypoints.duplex.server_vad import ( - SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS, - SILERO_VAD_MIN_THRESHOLD, - ServerVADPipeline, - SileroVADBackendProvider, - SpeechDetectorBackendProvider, - parse_session_turn_detection, +from vllm_omni.engine.duplex.commands import DuplexCommand, DuplexCommandError +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + SessionClosed, + SessionCreated, + SessionReplaced, + SessionResumed, + SessionResyncRequired, ) +from vllm_omni.engine.duplex.messages import DuplexSessionError +from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults +from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request from vllm_omni.entrypoints.duplex.session_attachment import ( DuplexJournalGapError, + DuplexJournalOverflowError, DuplexSessionAttachmentRegistry, InvalidResumeTokenError, -) -from vllm_omni.entrypoints.duplex.session_runner import ( - DuplexSessionRunnerMixin, + ResumeToken, ) from vllm_omni.entrypoints.duplex.websocket import ( - DOMAIN_TERMINAL_EVENTS, - DuplexSessionTasks, - DuplexWebSocketActor, + MAX_EVENT_BYTES, + SendJson, + attachment_callbacks, + receive_text_with_timeout, ) -from vllm_omni.metrics.realtime import RealtimeVADMetrics - -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle logger = init_logger(__name__) -__all__ = ["OmniDuplexSessionHandler", "should_enable_duplex_endpoint"] +__all__ = ["OmniDuplexSessionHandler"] _DEFAULT_CONFIG_TIMEOUT_S = 10.0 _DEFAULT_IDLE_TIMEOUT_S = 300.0 +#: How long the endpoint waits for the session pump to deliver its terminal +#: event once the session is closed. Bounded so a wedged pump cannot pin the +#: connection open; the pump itself closes the socket on its way out. +_PUMP_DRAIN_TIMEOUT_S = 5.0 +#: session.created is journaled by the registry when it hands out the resume credential. +_UNJOURNALED_EVENTS = (SessionCreated,) + + +@dataclass +class _Attachment: + """What one websocket connection knows about the session it is attached to.""" + + handle: DuplexSessionHandle + generation: int @dataclass(frozen=True) -class _DuplexSessionHandshake: - session: DuplexSession - resumed: bool = False - attachment_generation: int | None = None - event_id: str | None = None +class _ResumeCredentials: + """Transport credentials stamped onto ``session.created`` for a resumable session.""" + + attachment_generation: int + resume_token: str -class OmniDuplexSessionHandler( - DuplexSessionRunnerMixin, - NativeRuntimeBridgeMixin, - ChatFallbackProjectorMixin, -): - """WebSocket handler for RFC-style full-duplex session control. - - This owns the serving-side session actor, ordered inbound mailbox, - barge-in epoch, turn controller, and playback commit state. Generic sessions - can still fall back to chat requests, while MiniCPM-o 4.5 native sessions - route audio appends through scheduler data-plane stage requests. It - deliberately does not claim core persistent KV lease support. - """ +class OmniDuplexSessionHandler: + """WebSocket transport for engine-resident duplex sessions.""" def __init__( self, *, - chat_service: OmniOpenAIServingChat, - served_model_name: str | None = None, + duplex_omni: DuplexOmni, config_timeout_s: float = _DEFAULT_CONFIG_TIMEOUT_S, idle_timeout_s: float = _DEFAULT_IDLE_TIMEOUT_S, - log_stats: bool = True, - duplex_session_config: DuplexSessionRuntimeConfig | None = None, - serving_runtime_adapter: ServingRuntimeAdapter | None = None, - serving_runtime_adapter_path: str | None = None, - server_vad_backend_provider: SpeechDetectorBackendProvider | None = None, ) -> None: - self._chat_service = chat_service + self._omni = duplex_omni self._config_timeout_s = config_timeout_s self._idle_timeout_s = idle_timeout_s - self._duplex_session_config = duplex_session_config or DuplexSessionRuntimeConfig() - self._server_vad_backend_provider = server_vad_backend_provider or SileroVADBackendProvider( - model_path=self._duplex_session_config.server_vad_model_path, - ) - self._server_vad_pipelines: dict[str, ServerVADPipeline] = {} - self._realtime_vad_metrics = RealtimeVADMetrics( - served_model_name or chat_service.model_config.model, - log_stats=log_stats, - ) - self._registry = DuplexSessionRegistry( - DuplexCapabilities( - supports_model_native_turn_policy=False, - supports_input_append=False, - supports_replace_latest_chunk=True, - supports_reencode_context=True, - supports_turn_commit_only=True, - supports_kv_lease=False, - ) - ) - self._turn_controller = DuplexTurnController() - adapter_path = serving_runtime_adapter_path or getattr( - chat_service, - "duplex_serving_adapter_path", - None, - ) - self._serving_runtime_adapter: ServingRuntimeAdapter | None - if serving_runtime_adapter is not None: - self._serving_runtime_adapter = validate_serving_runtime_adapter(serving_runtime_adapter) - elif isinstance(adapter_path, str) and adapter_path: - self._serving_runtime_adapter = load_serving_runtime_adapter( - adapter_path, - self._encode_native_data_plane_audio, - ) - else: - self._serving_runtime_adapter = None - self._serving_session_states: MutableMapping[str, ServingRuntimeSessionState] = ( - self._serving_runtime_adapter.session_states if self._serving_runtime_adapter is not None else {} - ) - self._session_tasks: dict[str, DuplexSessionTasks] = {} - self._realtime_protocols: dict[str, NativeRealtimeSessionProtocol] = {} - self._lease_generations: dict[str, int] = {} - self._resync_required_sessions: set[str] = set() - self._lifecycle_queue = getattr(self._chat_service.engine_client, "duplex_lifecycle_events", None) - self._lifecycle_task: asyncio.Task[None] | None = None + runtime_config = duplex_omni.duplex_session_config self._attachment_registry = DuplexSessionAttachmentRegistry( - replay_ttl_s=self._duplex_session_config.resume_replay_ttl_s, - replay_max_bytes_per_session=self._duplex_session_config.resume_replay_max_bytes_per_session, - disconnect_grace_s=self._duplex_session_config.disconnect_grace_s, + replay_ttl_s=runtime_config.resume_replay_ttl_s, + replay_max_bytes_per_session=runtime_config.resume_replay_max_bytes_per_session, ) + self._resync_required_sessions: set[str] = set() + #: Wire defaults (input/output audio format and rate) per session. They + #: are negotiated on the session but live on the per-connection + #: envelope, so a reconnect has to be handed them back. + self._input_defaults: dict[str, RealtimeInputDefaults] = {} + self._pumps: dict[str, asyncio.Task[None]] = {} + + # ------------------------------------------------------------------ # + # Entry point # + # ------------------------------------------------------------------ # async def handle_realtime_session(self, websocket: WebSocket) -> None: - await self.handle_session( - websocket, - realtime_protocol=NativeRealtimeSessionProtocol(websocket.query_params), - ) + await websocket.accept() + envelope = RealtimeEnvelope.from_query_params(websocket.query_params) - def _ensure_lifecycle_listener(self) -> None: - if not isinstance(self._lifecycle_queue, asyncio.Queue): - return - if self._lifecycle_task is None or self._lifecycle_task.done(): - self._lifecycle_task = asyncio.create_task( - self._run_lifecycle_listener(), - name="duplex-serving-lifecycle", - ) + async def send_json(payload: Mapping[str, object]) -> None: + await websocket.send_json(dict(payload)) - async def _run_lifecycle_listener(self) -> None: - queue = self._lifecycle_queue - if not isinstance(queue, asyncio.Queue): - return + attachment: _Attachment | None = None + pending_command: dict[str, object] | None = None try: - while True: - message = await queue.get() - try: - if isinstance(message, DuplexSessionLifecycleMessage): - await self._apply_runtime_lifecycle(message) - finally: - queue.task_done() - if self._registry.active_count() == 0: + open_payload = envelope.initial_open_payload() + if open_payload is None: + first = await self._receive_first_message(websocket, envelope, send_json) + if first is None: return - finally: - if self._lifecycle_task is asyncio.current_task(): - self._lifecycle_task = None - - async def _apply_runtime_lifecycle(self, message: DuplexSessionLifecycleMessage) -> None: - session = self._registry.get(message.session_id) - if session is None: - return - fence = message.fence - if ( - fence.session_id != session.session_id - or fence.incarnation != session.incarnation - or fence.epoch < session.epoch - ): - return - current_generation = self._lease_generations.get(session.session_id, 0) - if message.lease_generation < current_generation: - return - protocol = self._realtime_protocols.get(session.session_id) - expired_payload: dict[str, object] = { - "type": "session.expired", - "session_id": session.session_id, - "incarnation": session.incarnation, - "reason": message.reason, - } - if protocol is not None: - expired_payload = protocol.encode_outbound_event(expired_payload)[0] - with suppress(Exception): - await self._attachment_registry.send_event( - session.session_id, - expired_payload, - journal=False, - ) - - tasks = self._session_tasks.pop(session.session_id, None) - if tasks is not None: - await tasks.cancel_append_tasks() - active_response_task = tasks.active_response_task - if active_response_task is not None and not active_response_task.done(): - active_response_task.cancel() - with suppress(asyncio.CancelledError): - await active_response_task - native = self._serving_session_states.get(session.session_id) - if native is not None and native.data_plane_task is not None: - data_plane_task = native.data_plane_task - native.data_plane_task = None - data_plane_task.cancel() - with suppress(asyncio.CancelledError): - await data_plane_task - if session.active_response_id is not None: - session.end_response(commit_text=False) - self._cleanup_duplex_session_state(session) - self._registry.close(session.session_id) - self._realtime_protocols.pop(session.session_id, None) - self._lease_generations.pop(session.session_id, None) - self._resync_required_sessions.discard(session.session_id) - attachment = await self._attachment_registry.close(session.session_id) - if attachment is not None: + handshake = envelope.first_message(first) + if handshake.kind == "resume": + attachment = await self._resume(websocket, envelope, handshake.resume_payload, send_json) + if attachment is None: + return + else: + open_payload = handshake.session_payload + pending_command = handshake.pending_command_payload + if attachment is None: + assert open_payload is not None + attachment = await self._open(websocket, envelope, open_payload, send_json) + if attachment is None: + return + if pending_command is not None: + await self._submit_wire_event(attachment, envelope, pending_command, send_json) + self._input_defaults[attachment.handle.session_id] = envelope.defaults + await self._read_loop(websocket, envelope, attachment, send_json) + await self._drain_terminal_pump(attachment) + except WebSocketDisconnect: + if attachment is not None: + await self._on_disconnect(attachment) + except Exception as exc: + logger.exception("Duplex websocket session failed: %s", exc) + if attachment is not None and await self._attachment_registry.is_current_attachment( + attachment.handle.session_id, attachment.generation + ): + with suppress(Exception): + await self._attachment_registry.close(attachment.handle.session_id) + with suppress(Exception): + await attachment.handle.close(reason="transport_error") with suppress(Exception): - await attachment.close("session_expired") + await send_json(envelope.error_payload("realtime_input_failed", str(exc))) - def _stop_lifecycle_listener_if_idle(self) -> None: - task = self._lifecycle_task - if self._registry.active_count() != 0 or task is None or task.done(): - return - if task is not asyncio.current_task(): - task.cancel() - self._lifecycle_task = None - - @staticmethod - def _native_audio_payload_size_bytes(payload: Mapping[str, object]) -> int: - audio = payload.get("audio") or payload.get("data") - if not isinstance(audio, str): - return 0 + async def _receive_first_message( + self, websocket: WebSocket, envelope: RealtimeEnvelope, send_json: SendJson + ) -> dict[str, object] | None: + raw = await receive_text_with_timeout(websocket, self._config_timeout_s) + if raw is None: + await send_json(envelope.error_payload("config_timeout", "Timeout waiting for session.update")) + return None try: - return len(base64.b64decode(audio, validate=True)) - except (ValueError, binascii.Error): - return 0 + first = json.loads(raw) + except json.JSONDecodeError: + await send_json(envelope.error_payload("invalid_json", "Invalid JSON in session.update")) + return None + if not isinstance(first, dict): + await send_json(envelope.error_payload("bad_event", "Realtime event must be a JSON object")) + return None + return first - async def _apply_outbound_session_event( - self, - payload: dict[str, object], - *, - session: DuplexSession | None, - actor: DuplexWebSocketActor, - native: ServingRuntimeSessionState, - realtime_protocol: NativeRealtimeSessionProtocol | None, - ) -> tuple[bool, dict[str, object] | None]: - """Apply domain transitions before an event is queued for transport.""" - payload_type = payload.get("type") - is_terminal = payload_type in DOMAIN_TERMINAL_EVENTS - if is_terminal and session is not None: - payload_epoch = payload.get("epoch") - if isinstance(payload_epoch, int) and payload_epoch != session.epoch: - return False, None - if payload_type in {"response.done", "response.listen"} and ( - actor.closing or session.state == DuplexSessionState.CLOSED - ): - return False, None - elif actor._is_stale_model_output(payload): - return False, None + # ------------------------------------------------------------------ # + # Handshake # + # ------------------------------------------------------------------ # - if payload_type == "session.closed": - actor.close_reason = actor.close_reason or str(payload.get("reason") or "closed") - if session is not None: - session.mark_closing() - - if not is_terminal or session is None: - return True, None - - terminal_status = payload.get("status") - terminal_status_details = payload.get("status_details") - if terminal_status is None and isinstance(terminal_status_details, dict): - terminal_status = terminal_status_details.get("type") - response_terminal = payload_type == "response.done" or ( - payload_type == "response.listen" and session.active_response_id is not None - ) - can_promote_overlap = response_terminal and terminal_status not in {"cancelled", "failed"} - deferred_overlap_payload: dict[str, object] | None = None - continuous_input_crosses_terminal = ( - can_promote_overlap - and realtime_protocol is not None - and self._session_auto_responds(session) - and native.input_since_commit - and not native.deferred_response_create - ) - if continuous_input_crosses_terminal: - # Auto-response is a model-owned continuous stream. A response - # terminal advances the model turn without requiring the browser - # to close its Realtime input item. Keep any partial model unit so - # later PCM can complete it. Serving does not create or promote a - # new model generation here; Stage0 advances generation identity - # only after the model's terminal token. - session.reset_overlap_speech() - return True, None - realtime_input_still_open = ( - can_promote_overlap - and realtime_protocol is not None - and native.input_since_commit - and not native.deferred_response_create - ) - if realtime_input_still_open: - # response.done only closes the assistant response. It does not - # close the current Realtime input item, and it does not reserve - # or advance a model generation. - session.reset_overlap_speech() - return True, None - if can_promote_overlap and session.overlap_speech_ms > 0: - has_deferred_overlap = native.audio_buffer.has_pending() or native.committed_audio_payload is not None - should_promote_overlap = ( - session.state == DuplexSessionState.OPEN - and self._uses_native_input_append(session) - and has_deferred_overlap - and session.overlap_speech_ms > session.config.overlap_short_ack_ms - ) - if should_promote_overlap: - flushed_reserved_bytes = native.audio_buffer.pending_byte_count - deferred_overlap_payload = native.audio_buffer.flush( - chunk_period_ms=session.capabilities.chunk_period_ms or 1000 - ) - if native.committed_audio_payload is not None: - if deferred_overlap_payload is not None: - deferred_overlap_payload = self._merge_native_audio_payloads( - native.committed_audio_payload, - deferred_overlap_payload, - ) - else: - deferred_overlap_payload = native.committed_audio_payload - if self._session_auto_responds(session) and deferred_overlap_payload is not None: - deferred_overlap_payload = dict(deferred_overlap_payload) - deferred_overlap_payload["force_listen"] = False - if deferred_overlap_payload is not None: - native.retain_committed_audio( - deferred_overlap_payload, - operation_id=native.committed_audio_operation_id, - reserved_bytes=flushed_reserved_bytes, - ) - native.input_since_commit = deferred_overlap_payload is not None - if realtime_protocol is not None: - native.deferred_response_create = True - native.deferred_precreate_response = False - deferred_overlap_payload = None - else: - had_pending_overlap_audio = native.audio_buffer.has_pending() - native.audio_buffer.clear() - native.input_since_commit = False - native.speech_since_commit = False - if had_pending_overlap_audio and realtime_protocol is not None: - await realtime_protocol.discard_pending_input_audio(audio_end_ms=session.overlap_speech_ms) - if payload_type in {"audio.cancelled", "input.cancelled", "session.closed"}: - session.release_input_bytes(native.clear_committed_audio()) - - session.reset_overlap_speech() - if can_promote_overlap and native.deferred_response_create and native.committed_audio_payload is not None: - deferred_overlap_payload = native.committed_audio_payload - native.deferred_response_create = False - native.input_since_commit = False - native.speech_since_commit = False - return True, deferred_overlap_payload - - @staticmethod - def _advance_barge_in_epoch(session: DuplexSession) -> tuple[int, dict[str, int]]: - old_playback = session.playback.as_dict() - new_epoch = session.barge_in() - session.clear_playback_cursor() - return new_epoch, old_playback - - @staticmethod - def _commit_played_response_history( - session: DuplexSession, - response_id: str | None, - committed_ms: int, - ) -> None: - if not response_id or committed_ms < 0: - return - session.truncate_history_item( - f"item_{response_id}", - audio_end_ms=committed_ms, - playback=session.playback_for_response(response_id), - ) - - @staticmethod - def _should_commit_response_to_history(session: DuplexSession, response_id: str | None) -> bool: - if response_id is not None and response_id != session.active_response_id: - return True - mode = session.response_config.extra_body.get("realtime_response_conversation") - return not isinstance(mode, str) or mode.strip().lower() != "none" - - def _response_created_payload( - self, - session: DuplexSession, - response_id: str, - *, - epoch: int, - request_id: str | None = None, - ) -> dict[str, object]: - response_config = session.response_config - payload: dict[str, object] = { - "type": "response.created", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "modalities": list(response_config.modalities), - } - if request_id is not None: - payload["request_id"] = request_id - metadata = response_config.extra_body.get("realtime_response_metadata") - if not isinstance(metadata, dict): - metadata = response_config.extra_body.get("realtime_metadata") - if isinstance(metadata, dict): - payload["metadata"] = dict(metadata) - conversation = response_config.extra_body.get("realtime_response_conversation") - if isinstance(conversation, str): - payload["conversation"] = conversation - prompt = response_config.extra_body.get("realtime_response_prompt") - if isinstance(prompt, dict): - payload["prompt"] = dict(prompt) - return payload - - def _overlap_decision( - self, - session: DuplexSession, - event: dict[str, object], - payload: dict[str, object], - ) -> dict[str, object]: - """Classify input that arrives while assistant audio is active. - - This is a serving-side policy. The model still owns listen/speak - decisions for normal chunks; overlap policy only decides whether the - current assistant response should be interrupted before the new audio is - appended. - """ - duration_ms = self._input_audio_duration_ms(event, payload) - is_speech = self._input_looks_like_speech(event, payload, session=session) - if not session.capabilities.supports_barge_in and self._event_requests_barge_in(event): - return self._defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=is_speech) - explicit = event.get("overlap_action") or event.get("overlap") - if isinstance(explicit, str): - normalized = explicit.strip().lower() - if normalized in {"barge_in", "interrupt", "cancel"}: - return { - "action": "barge_in", - "reason": "client_overlap_action", - "duration_ms": duration_ms, - "buffer_audio": True, - } - if normalized in {"listen", "continue", "continue_output", "ack"}: - session.reset_overlap_speech() - return { - "action": "listen", - "reason": "client_overlap_action", - "duration_ms": duration_ms, - "buffer_audio": ( - normalized == "listen" and is_speech and duration_ms > session.config.overlap_short_ack_ms - ), - "defer_runtime_append": True, - } - if normalized in {"drop", "ignore", "silence"}: - session.reset_overlap_speech() - return { - "action": "drop", - "reason": "client_overlap_action", - "duration_ms": duration_ms, - "buffer_audio": False, - } - - if bool(event.get("force_barge_in", False)): - return { - "action": "barge_in", - "reason": "client_force_barge_in", - "duration_ms": duration_ms, - "buffer_audio": True, - } - if self._session_auto_responds(session): - # Full-duplex input remains model-owned while output is active. Feed - # complete model units into the existing Stage0 stream immediately; - # playback only controls history ACKs, not model admission. - if is_speech: - session.accumulate_overlap_speech(duration_ms) - vad_speech_started = self._vad_speech_started(event, payload) - if ( - self._uses_native_input_append(session) - and session.capabilities.supports_barge_in - and is_speech - and session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - and vad_speech_started is not False - ): - return { - "action": "barge_in", - "reason": ("server_vad_speech_started" if vad_speech_started is True else "barge_in_on_speech"), - "cancel_reason": "turn_detected" if vad_speech_started is True else "barge_in", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - } - return { - "action": "listen", - "reason": "auto_response_continuous", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": False, - "force_listen": event.get("force_listen") is True or payload.get("force_listen") is True, - "preserve_realtime_input": True, - } - if bool(event.get("force_listen", False)): - session.reset_overlap_speech() - return { - "action": "listen", - "reason": "client_force_listen", - "duration_ms": duration_ms, - "buffer_audio": is_speech, - "defer_runtime_append": True, - } - - policy = session.config.overlap_policy - if not is_speech: - if session.overlap_speech_ms <= 0: - session.reset_overlap_speech() - return { - "action": "drop", - "reason": "silence_or_noise", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": False, - } - - if self._is_short_ack_transcript_hint(event, payload): - session.reset_overlap_speech() - return { - "action": "listen", - "reason": "short_ack_transcript", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": False, - "defer_runtime_append": True, - } - - if policy == DuplexOverlapPolicy.LISTEN_ONLY.value: - session.accumulate_overlap_speech(duration_ms) - return { - "action": "listen", - "reason": "policy_listen_only", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - - if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value and not session.capabilities.supports_barge_in: - return self._defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=True) - - session.accumulate_overlap_speech(duration_ms) - if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: - vad_speech_started = self._vad_speech_started(event, payload) - if vad_speech_started is False: - return { - "action": "listen", - "reason": "server_vad_utterance_active", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": False, - "force_listen": True, - "preserve_realtime_input": True, - } - return { - "action": "barge_in", - "reason": ("server_vad_speech_started" if vad_speech_started is True else "policy_barge_in_on_speech"), - "cancel_reason": "turn_detected" if vad_speech_started is True else "barge_in", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - } - - if ( - duration_ms <= session.config.overlap_short_ack_ms - and session.overlap_speech_ms <= session.config.overlap_short_ack_ms - ): - return { - "action": "listen", - "reason": "short_ack", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - if session.overlap_speech_ms >= session.config.overlap_barge_in_ms: - if not session.capabilities.supports_barge_in: - return { - "action": "listen", - "reason": "barge_in_unsupported", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - return { - "action": "barge_in", - "reason": "long_overlap_speech", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - } - return { - "action": "listen", - "reason": "accumulating_overlap_speech", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - - @staticmethod - def _vad_speech_started( - event: Mapping[str, object], - payload: Mapping[str, object], - ) -> bool | None: - for source in (event, payload): - vad = source.get("vad") - if isinstance(vad, Mapping) and isinstance(vad.get("speech_started"), bool): - return bool(vad["speech_started"]) - return None - - @staticmethod - def _event_requests_barge_in(event: Mapping[str, object]) -> bool: - if event.get("force_barge_in") is True: - return True - explicit = event.get("overlap_action") or event.get("overlap") - return isinstance(explicit, str) and explicit.strip().lower() in { - "barge_in", - "interrupt", - "cancel", - } - - @staticmethod - def _defer_unsupported_barge_in( - session: DuplexSession, - *, - duration_ms: int, - is_speech: bool, - ) -> dict[str, object]: - if is_speech: - session.accumulate_overlap_speech(duration_ms) - return { - "action": "listen", - "reason": "barge_in_unsupported", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": is_speech, - "defer_runtime_append": True, - } - - @staticmethod - def _barge_in_unsupported_error(session: DuplexSession) -> dict[str, object]: - return { - "type": "error", - "session_id": session.session_id, - "code": "barge_in_unsupported", - "error": "Barge-in is not supported by this duplex model", - } - - @staticmethod - def _is_short_ack_transcript_hint(event: dict[str, object], payload: dict[str, object]) -> bool: - raw_text = event.get("transcript") or event.get("text") or payload.get("transcript") or payload.get("text") - if not isinstance(raw_text, str): - return False - normalized = raw_text.strip().lower() - if not normalized: - return False - compact = "".join(ch for ch in normalized if ch.isalnum() or "\u4e00" <= ch <= "\u9fff") - if compact in { - "嗯", - "嗯嗯", - "对", - "对的", - "好", - "好的", - "继续", - "继续说", - "可以", - "是的", - "yes", - "yeah", - "yep", - "ok", - "okay", - "continue", - "goon", - "right", - }: - return True - return normalized in {"go on", "keep going", "please continue"} - - @staticmethod - def _input_audio_duration_ms(event: dict[str, object], payload: dict[str, object]) -> int: - for key in ("duration_ms", "audio_duration_ms"): - value = event.get(key) - if isinstance(value, int | float): - return max(0, int(value)) - fmt = payload.get("format") - sample_rate_hz = payload.get("sample_rate_hz") - audio = payload.get("audio") - if fmt == "pcm_f32le" and isinstance(sample_rate_hz, int) and sample_rate_hz > 0 and isinstance(audio, str): - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return 0 - return int((len(raw) // 4) * 1000 / sample_rate_hz) - return 0 - - @staticmethod - def _merge_native_audio_payloads( - first: dict[str, object], - second: dict[str, object], - ) -> dict[str, object]: - if first.get("format") != "pcm_f32le" or second.get("format") != "pcm_f32le": - return second - first_rate = first.get("sample_rate_hz") - second_rate = second.get("sample_rate_hz") - if not isinstance(first_rate, int) or not isinstance(second_rate, int) or first_rate != second_rate: - return second - first_audio = first.get("audio") - second_audio = second.get("audio") - if not isinstance(first_audio, str) or not isinstance(second_audio, str): - return second - try: - first_raw = base64.b64decode(first_audio, validate=True) - second_raw = base64.b64decode(second_audio, validate=True) - except (binascii.Error, ValueError): - return second - merged = dict(second) - merged["audio"] = base64.b64encode(first_raw + second_raw).decode("ascii") - merged["sample_rate_hz"] = first_rate - merged_frames = [ - frame - for source in (first.get("video_frames"), second.get("video_frames")) - if isinstance(source, list) - for frame in source - if isinstance(frame, str) and frame - ] - if merged_frames: - merged["video_frames"] = merged_frames - else: - merged.pop("video_frames", None) - merged["force_listen"] = bool(first.get("force_listen", False)) or bool(second.get("force_listen", False)) - merged.pop("force_speak", None) - merged["is_speech"] = bool(first.get("is_speech", False)) or bool(second.get("is_speech", False)) - return merged - - @classmethod - def _should_force_listen_for_short_commit( - cls, - session: DuplexSession, - event: dict[str, object], - payload: dict[str, object], - ) -> bool: - """Keep very short committed Realtime chunks in listen mode. - - Realtime VAD can emit a commit for a short pause even when the user has - not actually yielded the turn. For MiniCPM-o native duplex, make that - policy explicit by steering the scheduler path to the model listen - token instead of letting a sub-second chunk start a response. - """ - if event.get("force_listen") is True or payload.get("force_listen") is True: - return True - if event.get("force_barge_in") is True: - return False - if event.get("response_create") is not True: - return False - duration_ms = cls._input_audio_duration_ms(event, payload) - return 0 < duration_ms <= session.config.overlap_short_ack_ms - - def _create_runtime_session_state(self) -> ServingRuntimeSessionState: - adapter = self._serving_runtime_adapter - if adapter is not None: - return adapter.create_session_state() - return TurnBasedServingSessionState() - - def _runtime_session_state(self, session: DuplexSession) -> ServingRuntimeSessionState: - adapter = self._serving_runtime_adapter - if adapter is not None: - return adapter.session_state(session.session_id) - state = self._serving_session_states.get(session.session_id) - if state is None: - state = TurnBasedServingSessionState() - self._serving_session_states[session.session_id] = state - return state - - def _require_serving_runtime_adapter(self) -> ServingRuntimeAdapter: - adapter = self._serving_runtime_adapter - if adapter is None: - raise RuntimeError("Model-native duplex runtime is not configured for this session") - return adapter - - def _should_force_listen_for_auto_response_overlap( - self, - session: DuplexSession, - event: dict[str, object], - payload: dict[str, object], - ) -> bool: - if not self._session_auto_responds(session): - return False - if event.get("force_barge_in") is True: - return False - return event.get("force_listen") is True or payload.get("force_listen") is True - - @staticmethod - def _assistant_playback_active(session: DuplexSession) -> bool: - return ( - session.config.playback_commit_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value - and session.playback.sent_ms > session.playback.committed_ms - ) - - @staticmethod - def _input_looks_like_speech( - event: dict[str, object], - payload: dict[str, object], - *, - session: DuplexSession, - ) -> bool: - # Prefer serving-derived fields in ``payload`` over optional client - # hints retained in ``event``. In server-VAD mode the pipeline, rather - # than the client or the RMS fallback below, owns speech classification. - for source in (payload, event): - for key in ("is_speech", "speech"): - value = source.get(key) - if isinstance(value, bool): - return value - vad = source.get("vad") - if isinstance(vad, dict): - value = vad.get("is_speech") - if isinstance(value, bool): - return value - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) >= 0.5 - probability = source.get("speech_probability") - if isinstance(probability, int | float): - return float(probability) >= 0.5 - - fmt = payload.get("format") - audio = payload.get("audio") - if fmt in {"pcm_f32le", "pcm16"} and isinstance(audio, str): - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return True - if fmt == "pcm_f32le": - if len(raw) < 4 or len(raw) % 4 != 0: - return True - samples = np.frombuffer(raw, dtype=np.float32) - else: - if len(raw) < 2 or len(raw) % 2 != 0: - return True - samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 - if samples.size == 0: - return False - rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) - return rms >= session.config.overlap_silence_rms - return True - - async def _emit_overlap_decision( - self, - send_json, - session: DuplexSession, - decision: dict[str, object], - ) -> None: - await send_json( - { - "type": "overlap.decision", - "session_id": session.session_id, - "epoch": session.epoch, - "policy": session.config.overlap_policy, - **decision, - } - ) - - async def _open_session( + async def _open( self, websocket: WebSocket, - send_json, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - attachment_send=None, - attachment_close=None, - ) -> _DuplexSessionHandshake | None: - raw = await self._receive_text( - websocket, - self._config_timeout_s, - realtime_protocol=realtime_protocol, - ) - if raw is None: - await send_json({"type": "error", "error": "Timeout waiting for session.create", "code": "config_timeout"}) - return None + envelope: RealtimeEnvelope, + session_payload: Mapping[str, object], + send_json: SendJson, + ) -> _Attachment | None: + # Any ``session_id`` / ``id`` in the payload is ignored: the server + # allocates the id and announces it in ``session.created``. try: - event = json.loads(raw) - except json.JSONDecodeError: - await send_json({"type": "error", "error": "Invalid JSON in session.create", "code": "invalid_json"}) + handle = await self._omni.open_session(session_payload) + except DuplexSessionError as exc: + await send_json(envelope.error_payload(exc.code, str(exc))) return None - if isinstance(event, dict) and event.get("type") == "session.resume": - if realtime_protocol is None: - await send_json( - { - "type": "error", - "error": "session.resume requires the Realtime protocol", - "code": "unsupported_session_resume", - } - ) - return None - return await self._resume_session_handshake( - event, - send_json=send_json, - realtime_protocol=realtime_protocol, - attachment_send=attachment_send, - attachment_close=attachment_close, - ) - if not isinstance(event, dict) or event.get("type") not in {"session.create", "open_session", "session.config"}: - await send_json( - { - "type": "error", - "error": f"Expected session.create, got: {event.get('type') if isinstance(event, dict) else None}", - "code": "bad_event", - } - ) + except (TypeError, ValueError) as exc: + await send_json(envelope.error_payload("invalid_session_config", str(exc))) return None - - try: - config = DuplexSessionConfig.from_event(event) - except ValueError as exc: - error: dict[str, object] = { - "type": "error", - "error": str(exc), - "code": "invalid_request_error", - } - event_id = event.get("_realtime_event_id") - if isinstance(event_id, str) and event_id: - error["event_id"] = event_id - await send_json(error) - return None - if config.idle_timeout_s == _DEFAULT_IDLE_TIMEOUT_S: - config.idle_timeout_s = self._idle_timeout_s - runtime_adapter = self._serving_runtime_adapter if self._uses_serving_runtime_adapter(config) else None - runtime_config: dict[str, object] = {} - if runtime_adapter is not None: - try: - runtime_config = await runtime_adapter.prepare_runtime_config( - config, - model_config=getattr(self._chat_service, "model_config", None), - ) - except ServingRuntimeConfigError as exc: - await send_json({"type": "error", "error": str(exc), "code": exc.code}) - return None - except ValueError as exc: - await send_json({"type": "error", "error": str(exc), "code": "unsupported_ref_audio_path"}) - return None - session_id = event.get("session_id") if isinstance(event.get("session_id"), str) else None - session = self._registry.create(config=config, session_id=session_id) - if runtime_adapter is not None: - session.replace_capabilities( - runtime_adapter.capabilities(max_sessions=self._duplex_session_config.max_sessions) - ) - session.replace_runtime_config(runtime_config) - session_payload = event.get("session") - self._resolve_server_vad_defaults( - session, config, session_payload if isinstance(session_payload, dict) else event - ) - event_id = event.get("_realtime_event_id") - return _DuplexSessionHandshake( - session=session, - event_id=event_id if isinstance(event_id, str) and event_id else None, + attachment_send, attachment_close = attachment_callbacks(websocket) + created = await self._attachment_registry.create( + handle.session_id, + send=attachment_send, + close=attachment_close, ) + resume_supported = bool(handle.capabilities.supports_session_resume) + attachment = _Attachment(handle=handle, generation=created.attachment_generation) + credentials: _ResumeCredentials | None = None + if resume_supported: + credentials = _ResumeCredentials( + attachment_generation=created.attachment_generation, + resume_token=created.resume_token.plaintext, + ) + self._start_pump(handle, credentials) + return attachment - async def _resume_session_handshake( + async def _resume( self, - event: dict[str, object], - *, - send_json, - realtime_protocol: NativeRealtimeSessionProtocol, - attachment_send, - attachment_close, - ) -> _DuplexSessionHandshake | None: - session_id = event.get("session_id") - incarnation = event.get("incarnation") - resume_token = event.get("resume_token") - last_received = event.get("last_received_server_event_seq", 0) - if ( - not isinstance(session_id, str) - or not session_id - or not isinstance(incarnation, int) - or incarnation < 0 - or not isinstance(resume_token, str) - or not resume_token - or not isinstance(last_received, int) - or last_received < 0 - ): + websocket: WebSocket, + envelope: RealtimeEnvelope, + event: Mapping[str, object], + send_json: SendJson, + ) -> _Attachment | None: + request = parse_resume_request(event) + if request is None: await send_json( - { - "type": "error", - "error": ( - "session.resume requires session_id, incarnation, resume_token, " - "and a non-negative event sequence" - ), - "code": "invalid_session_resume", - } + envelope.error_payload( + "invalid_session_resume", + "session.resume requires session_id, resume_token, and a non-negative event sequence", + ) ) return None - session = self._registry.get(session_id) - if session is None or session.state != DuplexSessionState.OPEN: + session_id = request.session_id + handle = self._omni.get_session(session_id) + if handle is None or handle.closed: await send_json( - { - "type": "error", - "error": f"Unknown or expired duplex session: {session_id}", - "code": "session_resume_expired", - } + envelope.error_payload("session_resume_expired", f"Unknown or expired duplex session: {session_id}") ) return None - if not session.capabilities.supports_session_resume: + if not handle.capabilities.supports_session_resume: await send_json( - { - "type": "error", - "error": f"Session does not support resume: {session_id}", - "code": "unsupported_session_resume", - } + envelope.error_payload("unsupported_session_resume", f"Session does not support resume: {session_id}") ) return None try: await self._attachment_registry.authenticate_resume( session_id, - incarnation=incarnation, - resume_token=resume_token, - last_received_server_event_seq=last_received, + resume_token=request.resume_token, + last_received_server_event_seq=request.last_received_server_event_seq, ) except InvalidResumeTokenError: - await send_json( - { - "type": "error", - "error": "Invalid duplex session resume token", - "code": "invalid_resume_token", - } - ) + await send_json(envelope.error_payload("invalid_resume_token", "Invalid duplex session resume token")) return None except DuplexJournalGapError: - await send_json( - { - "type": "session.resync_required", - "session_id": session_id, - "reason": "journal_gap", - } - ) + await send_json(SessionResyncRequired(session_id=session_id, reason="journal_gap").to_realtime()) return None except (KeyError, ValueError) as exc: - await send_json( - { - "type": "error", - "error": str(exc), - "code": "session_resume_conflict", - } - ) + await send_json(envelope.error_payload("session_resume_conflict", str(exc))) return None - - resume_runtime = getattr(self._chat_service.engine_client, "resume_duplex_session_async", None) - if not callable(resume_runtime): - await send_json( - { - "type": "error", - "error": "Duplex runtime does not expose session resume control", - "code": "runtime_resume_unsupported", - } - ) - return None - expected_generation = self._lease_generations.get(session_id, 0) try: - runtime_result = await resume_runtime( + await self._omni.resume_session( session_id, - fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - expected_lease_generation=expected_generation, - ) - lease_generation = self._runtime_lease_generation(runtime_result) - if lease_generation is None: - raise RuntimeError("runtime resume result omitted lease_generation") - except Exception as exc: - await send_json( - { - "type": "error", - "error": str(exc), - "code": "runtime_resume_failed", - } + expected_lease_generation=handle.lease_generation, ) + except DuplexSessionError as exc: + await send_json(envelope.error_payload("runtime_resume_failed", str(exc))) return None - if not callable(attachment_send) or not callable(attachment_close): - raise RuntimeError("session.resume requires transport attachment callbacks") + attachment_send, attachment_close = attachment_callbacks(websocket) - def activation_payload_factory(token, generation: int) -> dict[str, object]: - internal = { - "type": "session.resumed", - "session_id": session_id, - "incarnation": incarnation, - "attachment_generation": generation, - "resume_token": token.plaintext, - } - return realtime_protocol.encode_outbound_event(internal)[0] + def activation_payload_factory(token: ResumeToken, generation: int) -> dict[str, object]: + return SessionResumed( + session_id=session_id, + session=dict(handle.public_session), + attachment_generation=generation, + resume_token=token.plaintext, + ).to_realtime() try: resumed = await self._attachment_registry.resume( session_id, - incarnation=incarnation, - resume_token=resume_token, - last_received_server_event_seq=last_received, + resume_token=request.resume_token, + last_received_server_event_seq=request.last_received_server_event_seq, send=attachment_send, close=attachment_close, activation_payload_factory=activation_payload_factory, ) except Exception as exc: - # The engine-side CAS already advanced even if the transport - # vanished before it received the rotated token. Keep that - # generation so the registry's one-shot recovery token can retry. - self._lease_generations[session_id] = lease_generation - await send_json( - { - "type": "error", - "error": str(exc), - "code": "session_resume_conflict", - } - ) + # The engine lease was already resumed above, so a failed activation + # would leave the session attached to nothing -- unless another + # connection won the race and is attached right now. Detaching then + # would start the disconnect grace for the *winner*, which ordinary + # heartbeats do not clear. Roll back only what this attempt owns. + if not await self._attachment_registry.has_attachment(session_id): + with suppress(DuplexSessionError): + await self._omni.detach_session(session_id) + await send_json(envelope.error_payload("session_resume_conflict", str(exc))) return None - self._lease_generations[session_id] = lease_generation replaced = resumed.replaced_attachment if replaced is not None: - replaced_payload = realtime_protocol.encode_outbound_event( - { - "type": "session.replaced", - "session_id": session_id, - "attachment_generation": replaced.generation, - } - )[0] with suppress(Exception): - await replaced.send(replaced_payload) + await replaced.send( + SessionReplaced(session_id=session_id, attachment_generation=replaced.generation).to_realtime() + ) with suppress(Exception): await replaced.close("session_replaced") - return _DuplexSessionHandshake( - session=session, - resumed=True, - attachment_generation=resumed.attachment_generation, - ) + # A reconnect brings a fresh envelope carrying pcm16/16 kHz wire + # defaults. The negotiated input format is a wire default, not part of + # the public session object, so it has to be carried over explicitly: + # otherwise the first append that omits format/rate is decoded as pcm16. + remembered = self._input_defaults.get(session_id) + if remembered is not None: + envelope.defaults = remembered + self._start_pump(handle, None) + return _Attachment(handle=handle, generation=resumed.attachment_generation) - @classmethod - def _runtime_lease_generation(cls, result: object) -> int | None: - if isinstance(result, dict): - generation = result.get("lease_generation") - if isinstance(generation, int): - return generation - for key in ("stage_results", "result"): - generation = cls._runtime_lease_generation(result.get(key)) - if generation is not None: - return generation - return None - if isinstance(result, list | tuple): - for item in result: - generation = cls._runtime_lease_generation(item) - if generation is not None: - return generation - return None + # ------------------------------------------------------------------ # + # Outbound pump (session-scoped, survives reconnects) # + # ------------------------------------------------------------------ # - def _uses_serving_runtime_adapter(self, config: DuplexSessionConfig) -> bool: - adapter = self._serving_runtime_adapter - return adapter is not None and adapter.is_enabled(config) - - def _runtime_session_update_error( - self, - session: DuplexSession, - payload: dict[str, object], - ) -> dict[str, object] | None: - if not self._uses_native_input_append(session): - return None - try: - self._require_serving_runtime_adapter().validate_client_extra_body(payload.get("extra_body")) - except ServingRuntimeConfigError as exc: - return { - "type": "error", - "session_id": session.session_id, - "code": exc.code, - "error": str(exc), - } - return None - - def _runtime_config_for_session_update( - self, - session: DuplexSession, - candidate_config: DuplexSessionConfig, - ) -> dict[str, object]: - if not self._uses_native_input_append(session): - return dict(session.runtime_config) - return self._require_serving_runtime_adapter().runtime_config_for_update( - candidate_config, - dict(session.runtime_config), - ) - - def _resolve_server_vad_defaults( - self, - session: DuplexSession, - config: DuplexSessionConfig, - payload: dict[str, object], - ) -> None: - """Resolve omitted options without overriding explicit client policies.""" - if config.server_vad is None: + def _start_pump(self, handle: DuplexSessionHandle, credentials: _ResumeCredentials | None) -> None: + existing = self._pumps.get(handle.session_id) + if existing is not None and not existing.done(): return - if config.server_vad.interrupt_response is None: - config.server_vad = replace(config.server_vad, interrupt_response=self._uses_native_input_append(session)) - extra = payload.get("extra_body") - overlap_policy = extra.get("overlap_policy") if isinstance(extra, dict) else None - if not isinstance(overlap_policy, str): - overlap_policy = payload.get("overlap_policy") - config.overlap_policy = ( - overlap_policy - if isinstance(overlap_policy, str) - else "barge_in_on_speech" - if config.server_vad.interrupt_response - else "listen_only" + self._pumps[handle.session_id] = asyncio.create_task( + self._pump_events(handle, credentials), + name=f"duplex-session-pump-{handle.session_id}", ) - def _runtime_session_candidate_update_error( - self, - session: DuplexSession, - candidate_config: DuplexSessionConfig, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - ) -> dict[str, object] | None: - if candidate_config.server_vad is not None: - desired_policy = "barge_in_on_speech" if candidate_config.server_vad.interrupt_response else "listen_only" - if candidate_config.overlap_policy != desired_policy: - return { - "type": "error", - "code": "unsupported_turn_detection", - "error": f"overlap_policy conflicts with turn_detection; expected {desired_policy!r}", - "param": "overlap_policy", - } - if not self._uses_native_input_append(session): - if ( - candidate_config.server_vad is not None - and candidate_config.server_vad.min_speech_duration_ms is not None - ): - return { - "type": "error", - "code": "unsupported_turn_detection", - "error": "server_vad.min_speech_duration_ms is only supported by model-native duplex runtimes", - } - requests_interruption = ( - candidate_config.server_vad is not None and candidate_config.server_vad.interrupt_response - ) or candidate_config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - if requests_interruption: - return { - "type": "error", - "code": "server_vad_requires_native_duplex", - "error": ( - "Realtime server_vad interruption requires a model-native duplex runtime; " - "set interrupt_response=false for turn-based models" - ), - } - if candidate_config.server_vad is not None and realtime_protocol is not None: - session_payload = candidate_config.extra_body.get("realtime_session_payload") - audio_error = realtime_protocol.validate_server_vad_input_audio( - session_payload if isinstance(session_payload, dict) else {} - ) - if audio_error is not None: - return { - "type": "error", - "code": "unsupported_turn_detection", - "error": audio_error, - "param": "audio.input.format", - } - return None - if candidate_config.server_vad is not None: - if not candidate_config.server_vad.interrupt_response: - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": "Model-native duplex server_vad requires interrupt_response=true", - } - if candidate_config.server_vad.threshold <= SILERO_VAD_MIN_THRESHOLD: - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": ( - f"Model-native duplex server_vad.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD}" - ), - } - if not candidate_config.server_vad.create_response: - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": "Model-native duplex server_vad requires create_response=true", - } - if ( - candidate_config.instructions != session.config.instructions - and self._runtime_session_state(session).native_context_locked - ): - return { - "type": "error", - "session_id": session.session_id, - "code": "instructions_update_unsupported", - "error": "session.update cannot change instructions after the native duplex context is initialized", - } - if native_duplex_opt_in(session.config.extra_body) is not True: - return None - if not self._config_requests_audio_output(candidate_config): - return None - if "ref_audio_data" in session.runtime_config: - return None - return { - "type": "error", - "session_id": session.session_id, - "code": "ref_audio_required", - "error": "MiniCPM-o native duplex audio output requires ref_audio", - } + async def _pump_events(self, handle: DuplexSessionHandle, credentials: _ResumeCredentials | None) -> None: + session_id = handle.session_id + close_reason = "session_closed" + try: + async for event in handle.events(): + if isinstance(event, SessionCreated) and credentials is not None: + event = replace( + event, + attachment_generation=credentials.attachment_generation, + resume_token=credentials.resume_token, + ) + await self._send_event(session_id, event) + if isinstance(event, SessionClosed): + close_reason = event.reason or event.type + break + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Duplex event pump failed for %s: %s", session_id, exc) + finally: + self._pumps.pop(session_id, None) + self._resync_required_sessions.discard(session_id) + self._input_defaults.pop(session_id, None) + attachment = None + with suppress(Exception): + attachment = await self._attachment_registry.close(session_id) + if attachment is not None: + with suppress(Exception): + await attachment.close(close_reason) - @staticmethod - def _uses_native_input_append(session: DuplexSession) -> bool: - return ( - session.capabilities.implementation_level == "model_native_duplex" - and session.capabilities.supports_input_append - and native_duplex_opt_in(session.config.extra_body) is not False - ) + async def _drain_terminal_pump(self, attachment: _Attachment) -> None: + """Let the pump deliver ``session.closed`` before the endpoint returns. - @staticmethod - def _config_requests_audio_output(config: DuplexSessionConfig) -> bool: - return any(str(modality).lower() == "audio" for modality in config.modalities) + ``DuplexSessionHandle._deliver`` queues ``SessionClosed`` on the outbox + and marks the handle closed in the same synchronous step, so + ``handle.closed`` is already true while that event is still sitting in + the queue. ``_read_loop`` stops on exactly that flag: if the engine + closes the session between two reads, the loop returns, this endpoint + returns, and the ASGI server tears the socket down with the terminal + event unsent -- the client sees an abrupt close instead of + ``session.closed``. Waiting for the pump keeps the two in order; it is + the pump that sends the event and then closes with code 1000. - @staticmethod - def _native_stage0_request_id(session: DuplexSession, epoch: int) -> str: - return duplex_resource_request_id( - DuplexFence( - session.session_id, - epoch=epoch, - incarnation=session.incarnation, - ), - "stage0", - ) - - @staticmethod - def _session_auto_responds(session: DuplexSession) -> bool: - """Full-duplex / model-driven mode. - - When set, the server runs per-chunk speak-generation continuously (like - the official MiniCPM-o ``duplex_generate`` loop) instead of waiting for an - explicit ``response.create``: each ~chunk_period of appended audio is - emitted and fed to the stage0 stream so the model itself decides to speak - or listen. Signaled by the client via ``extra_body.auto_response`` (or - ``extra_body.full_duplex``). + Only for a closed session. A takeover or a resumable disconnect leaves + the pump running for the next attachment, and must not be waited on. """ - extra = getattr(session.config, "extra_body", None) - if not isinstance(extra, dict): - return False - return extra.get("auto_response") is True or extra.get("full_duplex") is True + if not attachment.handle.closed: + return + pump = self._pumps.get(attachment.handle.session_id) + if pump is None or pump.done(): + return + # shield: a drain timeout must not cancel a pump that is still the + # session's only writer. + with suppress(Exception): + await asyncio.wait_for(asyncio.shield(pump), _PUMP_DRAIN_TIMEOUT_S) - async def _receive_text( + async def _send_event(self, session_id: str, event: DuplexEvent) -> None: + payload = event.to_realtime() + journal = not isinstance(event, _UNJOURNALED_EVENTS) and session_id not in self._resync_required_sessions + try: + try: + await self._attachment_registry.send_event(session_id, payload, journal=journal) + except DuplexJournalOverflowError: + first_overflow = session_id not in self._resync_required_sessions + self._resync_required_sessions.add(session_id) + if first_overflow: + resync = SessionResyncRequired(session_id=session_id, reason="journal_overflow") + await self._attachment_registry.send_event(session_id, resync.to_realtime(), journal=False) + await self._attachment_registry.send_event(session_id, payload, journal=False) + except KeyError: + # Attachment already closed (takeover or teardown); the journal is gone. + pass + except asyncio.CancelledError: + raise + except Exception as exc: + # The socket died mid-send. The session is alive in the engine: + # treat it as a disconnect of the current attachment (the pump + # keeps journaling for a resume) instead of tearing the session down. + logger.info("Duplex transport send failed for %s: %s", session_id, exc) + await self._detach_current_attachment(session_id) + + async def _detach_current_attachment(self, session_id: str) -> None: + """Disconnect semantics for the socket currently attached to ``session_id``. + + The pump is session-scoped and outlives any one connection, so it can + only name the session: the registry detaches whichever socket is + attached, which is the one whose send just failed. + """ + handle = self._omni.get_session(session_id) + if handle is None or handle.closed: + return + if handle.capabilities.supports_session_resume: + if await self._attachment_registry.detach(session_id, attachment_generation=None): + with suppress(DuplexSessionError): + await self._omni.detach_session(session_id) + return + with suppress(DuplexSessionError): + await handle.close(reason="disconnect") + + # ------------------------------------------------------------------ # + # Inbound # + # ------------------------------------------------------------------ # + + async def _read_loop( self, websocket: WebSocket, - timeout_s: float, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - ) -> str | None: - try: - if realtime_protocol is not None: - return await asyncio.wait_for( - realtime_protocol.receive_internal_event_text(websocket), - timeout=max(0.1, timeout_s), - ) - return await asyncio.wait_for(websocket.receive_text(), timeout=max(0.1, timeout_s)) - except asyncio.TimeoutError: - return None - - @staticmethod - def _input_committed_payload( - session: DuplexSession, - committed: DuplexCommittedInput, - *, - item_id: object | None = None, - ) -> dict[str, object]: - payload: dict[str, object] = { - "type": "input.committed", - "session_id": session.session_id, - "turn_id": committed.turn_id, - "epoch": committed.epoch, - "history_len": len(session.history), - "message": committed.message, - } - if isinstance(item_id, str) and item_id: - payload["item_id"] = item_id - return payload - - @staticmethod - def _commit_native_audio_input( - session: DuplexSession, - *, - item_id: object | None = None, - transcript: object | None = None, - turn_id: int | None = None, - ) -> DuplexCommittedInput: - clean_transcript = transcript.strip() if isinstance(transcript, str) else None - committed = session.commit_native_audio_input( - transcript=clean_transcript or None, - turn_id=turn_id, - ) - if isinstance(item_id, str) and item_id: - session.register_history_item(item_id, committed.message) - return committed - - @staticmethod - def _native_audio_committed_payload( - session: DuplexSession, - *, - committed: DuplexCommittedInput | None = None, - item_id: object | None = None, - transcript: object | None = None, - ) -> dict[str, object]: - message = committed.message if committed is not None else None - if not isinstance(message, dict): - input_audio_part: dict[str, object] = { - "type": "audio_url", - "audio_url": {"url": "native-duplex:input-audio"}, - } - if isinstance(transcript, str) and transcript: - input_audio_part["transcript"] = transcript - message = { - "role": "user", - "content": [input_audio_part], - } - payload: dict[str, object] = { - "type": "input.committed", - "session_id": session.session_id, - "turn_id": committed.turn_id if committed is not None else session.turn_id, - "epoch": committed.epoch if committed is not None else session.epoch, - "history_len": len(session.history), - "native_audio": True, - "message": message, - } - if isinstance(transcript, str) and transcript: - payload["transcript"] = transcript - if isinstance(item_id, str) and item_id: - payload["item_id"] = item_id - return payload - - def _apply_session_update( - self, - session: DuplexSession, - payload: dict[str, object], - ) -> dict[str, object] | None: - model = payload.get("model") - audio_config = payload.get("audio") - audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None - audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - voice = payload.get("voice") - if not isinstance(voice, str) and isinstance(audio_output, dict): - voice = audio_output.get("voice") - if isinstance(model, str) and session.config.model is not None and model != session.config.model: - return { - "type": "error", - "session_id": session.session_id, - "code": "model_update_unsupported", - "error": "session.update cannot change model for an open realtime duplex session", - } - if isinstance(model, str) and session.config.model is None: - session.config.model = model - if isinstance(voice, str) and (session.playback.generated_ms > 0 or session.playback.sent_ms > 0): - return { - "type": "error", - "session_id": session.session_id, - "code": "voice_update_after_audio_unsupported", - "error": "session.update cannot change voice after audio output has started", - } - if isinstance(payload.get("ref_audio"), str): - return { - "type": "error", - "session_id": session.session_id, - "code": "ref_audio_update_unsupported", - "error": "session.update cannot change ref_audio after the native duplex runtime is open", - } - extra_body_payload = payload.get("extra_body") - if isinstance(extra_body_payload, dict): - # Fold the deprecated model-prefixed alias into the canonical - # native_duplex key before comparing or merging. - extra_body_payload = normalize_native_duplex_key(dict(extra_body_payload)) - payload = {**payload, "extra_body": extra_body_payload} - if isinstance(extra_body_payload, dict) and NATIVE_DUPLEX_KEY in extra_body_payload: - requested_native_duplex = extra_body_payload[NATIVE_DUPLEX_KEY] - if not isinstance(requested_native_duplex, bool): - return { - "type": "error", - "session_id": session.session_id, - "code": "invalid_request_error", - "error": "extra_body.native_duplex must be a boolean", - } - current_native_duplex = native_duplex_opt_in(session.config.extra_body) - if current_native_duplex is not None and extra_body_payload[NATIVE_DUPLEX_KEY] != current_native_duplex: - return { - "type": "error", - "session_id": session.session_id, - "code": "native_duplex_mode_update_unsupported", - "error": "session.update cannot change native_duplex after the session is created", - } - if isinstance(payload.get("instructions"), str): - session.config.instructions = str(payload["instructions"]) - elif "instructions" in payload and payload.get("instructions") is None: - session.config.instructions = None - if isinstance(voice, str): - session.config.voice = str(voice) - elif "voice" in payload and payload.get("voice") is None: - session.config.voice = None - response_format = payload.get("output_audio_format") or payload.get("response_format") - if response_format is None and isinstance(audio_config, dict): - if isinstance(audio_output, dict): - response_format = audio_output.get("format") - response_format, _ = NativeRealtimeSessionProtocol._parse_realtime_audio_format(response_format) - if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: - session.config.response_format = NativeRealtimeSessionProtocol._duplex_response_format(response_format) - if isinstance(payload.get("temperature"), int | float): - session.config.temperature = float(payload["temperature"]) - speed = payload.get("speed") - if not isinstance(speed, int | float) and isinstance(audio_output, dict): - speed = audio_output.get("speed") - if isinstance(speed, int | float): - session.config.speed = float(speed) - max_tokens = ( - payload.get("max_response_output_tokens") - if "max_response_output_tokens" in payload - else payload.get("max_output_tokens") - if "max_output_tokens" in payload - else payload.get("max_tokens") - ) - if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: - session.config.max_tokens = NativeRealtimeSessionProtocol.realtime_max_output_tokens(max_tokens) - if isinstance(payload.get("overlap_policy"), str): - session.config.overlap_policy = DuplexSessionConfig._normalize_overlap_policy( - str(payload["overlap_policy"]) - ) - if isinstance(payload.get("overlap_short_ack_ms"), int | float): - session.config.overlap_short_ack_ms = max(0, int(payload["overlap_short_ack_ms"])) - if isinstance(payload.get("overlap_barge_in_ms"), int | float): - session.config.overlap_barge_in_ms = max(0, int(payload["overlap_barge_in_ms"])) - if isinstance(payload.get("overlap_silence_rms"), int | float): - session.config.overlap_silence_rms = max(0.0, float(payload["overlap_silence_rms"])) - if isinstance(payload.get("playback_commit_policy"), str): - session.config.playback_commit_policy = DuplexSessionConfig._normalize_playback_commit_policy( - str(payload["playback_commit_policy"]) - ) - modalities = payload.get("modalities") or payload.get("output_modalities") - if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): - session.config.modalities = list(modalities) - if isinstance(payload.get("extra_body"), dict): - session.config.extra_body.update(payload["extra_body"]) - extra = payload["extra_body"] - if isinstance(extra.get("overlap_policy"), str): - session.config.overlap_policy = DuplexSessionConfig._normalize_overlap_policy( - str(extra["overlap_policy"]) - ) - if isinstance(extra.get("playback_commit_policy"), str): - session.config.playback_commit_policy = DuplexSessionConfig._normalize_playback_commit_policy( - str(extra["playback_commit_policy"]) - ) - if isinstance(payload.get("tools"), list): - session.config.extra_body["realtime_tools"] = payload["tools"] - elif "tools" in payload and payload.get("tools") is None: - session.config.extra_body.pop("realtime_tools", None) - if isinstance(payload.get("tool_choice"), str | dict): - session.config.extra_body["realtime_tool_choice"] = payload["tool_choice"] - elif "tool_choice" in payload and payload.get("tool_choice") is None: - session.config.extra_body.pop("realtime_tool_choice", None) - if isinstance(payload.get("metadata"), dict): - session.config.extra_body["realtime_metadata"] = dict(payload["metadata"]) - elif "metadata" in payload and payload.get("metadata") is None: - session.config.extra_body.pop("realtime_metadata", None) - if isinstance(payload.get("include"), list): - session.config.extra_body["realtime_include"] = list(payload["include"]) - elif "include" in payload and payload.get("include") is None: - session.config.extra_body.pop("realtime_include", None) - if isinstance(payload.get("prompt"), dict): - session.config.extra_body["realtime_prompt"] = dict(payload["prompt"]) - elif "prompt" in payload and payload.get("prompt") is None: - session.config.extra_body.pop("realtime_prompt", None) - input_audio_transcription = NativeRealtimeSessionProtocol._input_audio_transcription_config(payload) - if isinstance(input_audio_transcription, dict): - session.config.extra_body["realtime_input_audio_transcription"] = dict(input_audio_transcription) - elif "input_audio_transcription" in payload and payload.get("input_audio_transcription") is None: - session.config.extra_body.pop("realtime_input_audio_transcription", None) - if isinstance(payload.get("input_audio_noise_reduction"), dict): - session.config.extra_body["realtime_input_audio_noise_reduction"] = dict( - payload["input_audio_noise_reduction"] - ) - elif "input_audio_noise_reduction" in payload and payload.get("input_audio_noise_reduction") is None: - session.config.extra_body.pop("realtime_input_audio_noise_reduction", None) - if isinstance(audio_input, dict) and isinstance(audio_input.get("noise_reduction"), dict): - session.config.extra_body["realtime_input_audio_noise_reduction"] = dict(audio_input["noise_reduction"]) - elif isinstance(audio_input, dict) and audio_input.get("noise_reduction") is None: - session.config.extra_body.pop("realtime_input_audio_noise_reduction", None) - if isinstance(payload.get("audio"), dict): - session.config.extra_body["realtime_audio"] = dict(payload["audio"]) - elif "audio" in payload and payload.get("audio") is None: - session.config.extra_body.pop("realtime_audio", None) - if isinstance(payload.get("tracing"), str | dict): - session.config.extra_body["realtime_tracing"] = payload["tracing"] - elif "tracing" in payload and payload.get("tracing") is None: - session.config.extra_body.pop("realtime_tracing", None) - session.config.extra_body["realtime_session_payload"] = ( - NativeRealtimeSessionProtocol._json_safe_realtime_payload(payload) - ) - try: - turn_detection_configured, server_vad = parse_session_turn_detection(payload) - if turn_detection_configured: - vad_config = replace(session.config, server_vad=server_vad) - self._resolve_server_vad_defaults(session, vad_config, payload) - server_vad = vad_config.server_vad - if ( - not self._uses_native_input_append(session) - and session.turn_detection_config_locked - and (not session.config.turn_detection_configured or server_vad != session.config.server_vad) - ): - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": "turn_detection cannot be changed after the first audio append", - } - session.config.turn_detection_configured = True - session.config.server_vad = server_vad - session.config.overlap_policy = vad_config.overlap_policy - except ValueError as exc: - return { - "type": "error", - "session_id": session.session_id, - "code": "invalid_request_error", - "error": str(exc), - } - return None - - async def _configure_server_vad(self, session: DuplexSession) -> None: - pipeline = await self._prepare_server_vad_pipeline(session, session.config) - self._install_server_vad_pipeline(session, pipeline) - - def _install_server_vad_pipeline( - self, - session: DuplexSession, - pipeline: ServerVADPipeline | None, + envelope: RealtimeEnvelope, + attachment: _Attachment, + send_json: SendJson, ) -> None: - existing = self._server_vad_pipelines.pop(session.session_id, None) - if existing is not None: - existing.reset() - self._realtime_vad_metrics.session_finished() - if pipeline is None: - return - self._server_vad_pipelines[session.session_id] = pipeline - self._realtime_vad_metrics.session_started() - - async def _prepare_server_vad_pipeline( - self, - session: DuplexSession, - config: DuplexSessionConfig, - ) -> ServerVADPipeline | None: - if config.server_vad is None: - return None - backend = await asyncio.to_thread(self._server_vad_backend_provider.get) - pipeline_config = config.server_vad - if self._uses_native_input_append(session) and pipeline_config.min_speech_duration_ms is None: - pipeline_config = replace( - pipeline_config, - min_speech_duration_ms=SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS, - ) - return ServerVADPipeline(backend, pipeline_config) - - def _apply_response_create_options( - self, - session: DuplexSession, - payload: dict[str, object], - ) -> str | None: - """Reserve options that apply only to the next response lifecycle.""" - audio_config = payload.get("audio") - audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - if session.capabilities.implementation_level == "model_native_duplex": - nested_voice = audio_output.get("voice") if isinstance(audio_output, dict) else None - unsupported = ( - payload.get("instructions") is not None - or payload.get("voice") is not None - or nested_voice is not None - or payload.get("temperature") is not None - or any( - payload.get(field_name) is not None - for field_name in ("max_response_output_tokens", "max_output_tokens", "max_tokens") - ) - or payload.get("tools") is not None - or payload.get("tool_choice") is not None - ) - if unsupported: - return "unsupported_native_response_options" - - instructions = str(payload["instructions"]) if isinstance(payload.get("instructions"), str) else None - voice = payload.get("voice") - if not isinstance(voice, str) and isinstance(audio_output, dict): - voice = audio_output.get("voice") - voice = str(voice) if isinstance(voice, str) else None - response_format = payload.get("output_audio_format") or payload.get("response_format") - if response_format is None and isinstance(audio_config, dict): - if isinstance(audio_output, dict): - response_format = audio_output.get("format") - response_format, _ = NativeRealtimeSessionProtocol._parse_realtime_audio_format(response_format) - if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: - response_format = NativeRealtimeSessionProtocol._duplex_response_format(response_format) - else: - response_format = None - temperature = float(payload["temperature"]) if isinstance(payload.get("temperature"), int | float) else None - speed = payload.get("speed") - if not isinstance(speed, int | float) and isinstance(audio_output, dict): - speed = audio_output.get("speed") - speed = float(speed) if isinstance(speed, int | float) else None - max_tokens = ( - payload.get("max_response_output_tokens") - if "max_response_output_tokens" in payload - else payload.get("max_output_tokens") - if "max_output_tokens" in payload - else payload.get("max_tokens") - ) - if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: - max_tokens = NativeRealtimeSessionProtocol.realtime_max_output_tokens(max_tokens) - else: - max_tokens = None - modalities = payload.get("modalities") or payload.get("output_modalities") - if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): - modalities = tuple(modalities) - else: - modalities = None - response_extra: dict[str, object] = {} - conversation = payload.get("conversation") - if isinstance(conversation, str): - response_extra["realtime_response_conversation"] = conversation - metadata = payload.get("metadata") - if isinstance(metadata, dict): - response_extra["realtime_response_metadata"] = dict(metadata) - prompt = payload.get("prompt") - if isinstance(prompt, dict): - response_extra["realtime_response_prompt"] = dict(prompt) - if isinstance(payload.get("tools"), list): - response_extra["realtime_response_tools"] = payload["tools"] - if isinstance(payload.get("tool_choice"), str | dict): - response_extra["realtime_response_tool_choice"] = payload["tool_choice"] - extra_body = payload.get("extra_body") - if isinstance(extra_body, dict): - if session.capabilities.implementation_level == "model_native_duplex": - response_extra.update( - (key, value) - for key, value in extra_body.items() - if key not in self._require_serving_runtime_adapter().private_runtime_config_keys - ) - else: - response_extra.update(extra_body) - try: - session.reserve_response_options( - ResponseCreateOptions( - instructions=instructions, - voice=voice, - response_format=response_format, - temperature=temperature, - max_tokens=max_tokens, - speed=speed, - modalities=modalities, - extra_body=response_extra, - ) - ) - except RuntimeError: - return "response_already_active" - return None - - @staticmethod - def _realtime_item_to_history_message(item: object) -> dict[str, object] | None: - if not isinstance(item, dict): - return None - role = item.get("role") - if role not in {"system", "user", "assistant"}: - return None - content = item.get("content") - if isinstance(content, str): - text = content.strip() - return {"role": role, "content": text} if text else None - if not isinstance(content, list): - return None - text_chunks: list[str] = [] - audio_chunks: list[dict[str, object]] = [] - for part in content: - if not isinstance(part, dict): + handle = attachment.handle + idle_timeout_s = self._idle_timeout_s + config_timeout = handle.public_session.get("idle_timeout_s") if handle.public_session else None + if isinstance(config_timeout, int | float) and config_timeout > 0: + idle_timeout_s = float(config_timeout) + while not handle.closed: + raw = await receive_text_with_timeout(websocket, idle_timeout_s) + if not await self._attachment_registry.is_current_attachment(handle.session_id, attachment.generation): + # A newer connection took the session over; this socket is done. + return + if raw is None: + # Nothing from this socket for the idle window: same as a + # disconnect (engine grace for resumable sessions); the engine + # lease decides whether the session itself expires. + await self._on_disconnect(attachment) + return + if len(raw.encode("utf-8")) > MAX_EVENT_BYTES: + await send_json(envelope.error_payload("event_too_large", "Duplex event too large")) continue - part_type = part.get("type") - if part_type in {"input_text", "text", "output_text"} and isinstance(part.get("text"), str): - text_chunks.append(str(part["text"])) - elif part_type in {"input_audio", "audio"}: - audio = part.get("audio") or part.get("data") - fmt = part.get("format") if isinstance(part.get("format"), str) else "wav" - if isinstance(audio, str) and audio: - audio_chunks.append( - { - "type": "audio_url", - "audio_url": { - "url": f"data:audio/{fmt};base64,{audio}", - }, - } - ) - elif part_type in {"audio_transcript", "transcript"} and isinstance(part.get("text"), str): - text_chunks.append(str(part["text"])) - text = "".join(text_chunks).strip() - if audio_chunks: - content_items: list[dict[str, object]] = [] - if text: - content_items.append({"type": "text", "text": text}) - content_items.extend(audio_chunks) - return {"role": role, "content": content_items} - if text: - return {"role": role, "content": text} - return None - - async def _handle_playback_ack(self, session: DuplexSession, event: dict[str, object], send_json) -> None: - played_ms = event.get("played_ms", event.get("audio_ms", 0)) - committed_ms = event.get("committed_ms") - if not isinstance(played_ms, int | float): - await send_json({"type": "error", "error": "playback.ack requires played_ms", "code": "bad_event"}) - return - committed_cursor = int(committed_ms) if isinstance(committed_ms, int | float) else int(played_ms) - item_id = event.get("item_id") - response_id = event.get("response_id") - response_id = response_id if isinstance(response_id, str) and response_id else None - if not isinstance(item_id, str) or not item_id: - item_id = f"item_{response_id}" if response_id is not None else None - elif response_id is None and item_id.startswith("item_"): - response_id = item_id.removeprefix("item_") - if response_id is None and item_id is None and len(session.pending_history_item_ids) == 1: - item_id = next(iter(session.pending_history_item_ids)) - if item_id.startswith("item_"): - response_id = item_id.removeprefix("item_") - if response_id is None and item_id is None and session.active_response_id is not None: - response_id = session.active_response_id - item_id = f"item_{response_id}" - if response_id is not None: - expected_item_id = f"item_{response_id}" - if item_id is None: - item_id = expected_item_id - elif item_id != expected_item_id: - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_item_mismatch", - "error": "playback.ack item_id must match item_.", - } - ) - return - if not session.has_assistant_response_item(response_id, item_id): - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_item_not_found", - "error": f"No assistant response item is registered for {response_id}.", - } - ) - return - if session.playback_ack_is_too_late(response_id, item_id): - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_ack_too_late", - "error": "playback.ack arrived after a later user input was committed.", - } - ) - return - # Reserve the response's current history position before any later - # input commit can append a user item. A 0 ms ACK is intentional: - # the response may be active but have no audio delta yet. - session.reserve_history_item(item_id) - elif item_id is not None: - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_item_not_found", - "error": "playback.ack requires a response-owned assistant item.", - } - ) - return - hard_truncate = event.get("truncate") is True - if hard_truncate: - playback = session.acknowledge_playback( - int(played_ms), - committed_cursor, - response_id=response_id, - ) - playback = session.truncate_playback_commit( - committed_cursor, - response_id=response_id, - ) - else: - playback = session.acknowledge_playback( - int(played_ms), - committed_cursor, - response_id=response_id, - ) - committed_history = False - if isinstance(item_id, str) and item_id: - expected_item_id = f"item_{response_id}" if response_id is not None else None - if ( - expected_item_id == item_id - and item_id not in session.history_item_ids - and item_id not in session.pending_history_item_ids - ): - # ACK_ONLY responses are normally registered as pending when - # the response ends. Recover the response-local item if that - # registration was lost across the response/ack boundary. - session.register_history_item(item_id, None) - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - elif session.pending_history_item_ids: - # A plain playback ack has no OpenAI item id. Commit the only - # uncommitted assistant candidate if the session has an unambiguous - # pending response; otherwise wait for conversation.item.truncate. - pending_ids = list(session.pending_history_item_ids) - if len(pending_ids) == 1: - item_id = pending_ids[0] - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - elif session.active_response_id is not None: - item_id = f"item_{session.active_response_id}" - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - elif session.last_assistant_full_message is not None: - if item_id is None and session.history_item_ids: - assistant_item_ids = [ - known_item_id - for known_item_id, message in session.history_item_ids.items() - if message.get("role") == "assistant" - ] - if len(assistant_item_ids) == 1: - item_id = assistant_item_ids[0] - if isinstance(item_id, str) and item_id: - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - await send_json( - { - "type": "playback.acknowledged", - "session_id": session.session_id, - "epoch": session.epoch, - "item_id": item_id, - "played_ms": int(played_ms), - "committed_ms": committed_cursor, - "truncate": event.get("truncate") is True, - "playback": playback.as_dict(), - "history_committed": committed_history, - } - ) - if committed_history and committed_cursor >= max(playback.sent_ms, playback.generated_ms): - session.release_response_playback(response_id) - session.release_response_history_snapshot(response_id) - - async def _cancel_active_response( - self, - session: DuplexSession, - active_task: asyncio.Task[None] | None, - send_json, - *, - reason: str, - notify: bool = True, - ) -> bool: - has_running_task = active_task is not None and not active_task.done() - if not has_running_task and session.active_request_id is None and session.active_response_id is None: - return False - - old_epoch = session.epoch - old_request_id = session.active_request_id - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - committed_message = session.end_response( - commit_text=self._should_commit_response_to_history(session, old_response_id), - playback_commit_policy=DuplexPlaybackCommitPolicy.ACK_ONLY.value, - ) - if old_response_id is not None: - item_id = f"item_{old_response_id}" - if committed_message is not None: - session.register_history_item(item_id, committed_message) - elif committed_ms > 0 and not session.playback_ack_is_too_late(old_response_id, item_id): - session.truncate_history_item(item_id, audio_end_ms=committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - if old_request_id is not None: - # Epoch-scoped request ids prevent state reuse, but explicitly - # release projector/parser cursors so cancelled epochs do not - # accumulate until the whole session closes. - if self._serving_runtime_adapter is not None: - self._serving_runtime_adapter.data_plane.close_stream(old_request_id) - await self._abort_request_background( - session, - old_request_id, - send_json, - notify=notify, - ) - if has_running_task and active_task is not None: - active_task.cancel() try: - await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25) - except asyncio.TimeoutError: - pass - if notify: - await send_json( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - return True + payload = json.loads(raw) + except json.JSONDecodeError: + await send_json(envelope.error_payload("invalid_json", "Invalid JSON event")) + continue + if not isinstance(payload, dict) or not isinstance(payload.get("type"), str): + await send_json(envelope.error_payload("bad_event", "Duplex event must be a JSON object with a type")) + continue + if envelope.is_envelope_event(payload): + await self._handle_envelope_event(attachment, envelope, payload, send_json) + continue + await self._submit_wire_event(attachment, envelope, payload, send_json) - async def _abort_request_background( + async def _handle_envelope_event( self, - session: DuplexSession, - request_id: str, - send_json, - *, - notify: bool, + attachment: _Attachment, + envelope: RealtimeEnvelope, + payload: dict[str, object], + send_json: SendJson, + ) -> None: + handle = attachment.handle + event_type = payload.get("type") + event_id = payload.get("event_id") + if event_type == "session.event_ack": + acknowledged = payload.get("server_event_seq") + if not isinstance(acknowledged, int) or acknowledged < 0: + message = "session.event_ack requires a non-negative server_event_seq" + await send_json(envelope.error_payload("invalid_event_ack", message, event_id=event_id)) + return + try: + await self._attachment_registry.acknowledge(handle.session_id, acknowledged) + except ValueError as exc: + await send_json(envelope.error_payload("invalid_event_ack", str(exc), event_id=event_id)) + return + message = "session.resume must be the first message of a new connection" + await send_json(envelope.error_payload("unsupported_session_resume", message, event_id=event_id)) + + async def _submit_wire_event( + self, + attachment: _Attachment, + envelope: RealtimeEnvelope, + payload: dict[str, object], + send_json: SendJson, ) -> None: try: - abort_internal = getattr(self._chat_service.engine_client, "_abort_internal_requests", None) - if callable(abort_internal): - result = abort_internal([request_id]) - else: - result = self._chat_service.engine_client.abort([request_id]) - if inspect.isawaitable(result): - await result - except Exception as exc: - logger.exception("Failed to abort duplex request %s: %s", request_id, exc) - if notify and session.state != DuplexSessionState.CLOSED: - await self._send_runtime_error(send_json, "runtime_abort_failed", exc, session=session) + command = envelope.translate(payload) + except DuplexCommandError as exc: + await send_json(envelope.command_error_payload(exc)) + return + # ``translate`` folds a session.update's audio settings into the + # envelope; remember them so a later reconnect starts where this left off. + self._input_defaults[attachment.handle.session_id] = envelope.defaults + await self._submit_command(attachment, envelope, command, send_json) - async def _cancel_pending_input(self, session: DuplexSession, send_json, *, reason: str) -> None: - cancelled = session.cancel_pending_input() - self._advance_barge_in_epoch(session) - await send_json( - { - "type": "input.cancelled", - "session_id": session.session_id, - "reason": reason, - "epoch": session.epoch, - "cancelled": cancelled, - } - ) + async def _submit_command( + self, + attachment: _Attachment, + envelope: RealtimeEnvelope, + command: DuplexCommand, + send_json: SendJson, + ) -> None: + try: + await attachment.handle.submit(command) + except DuplexSessionError as exc: + await send_json(envelope.error_payload(exc.code, str(exc), event_id=command.event_id)) + + # ------------------------------------------------------------------ # + # Disconnect # + # ------------------------------------------------------------------ # + + async def _on_disconnect(self, attachment: _Attachment) -> None: + handle = attachment.handle + if handle.closed: + return + if handle.capabilities.supports_session_resume: + if await self._attachment_registry.detach(handle.session_id, attachment_generation=attachment.generation): + # Engine-owned disconnect grace: expiry arrives as session.expired. + with suppress(DuplexSessionError): + await self._omni.detach_session(handle.session_id) + # Otherwise a newer connection already took the session over; the + # replaced socket's disconnect must not touch it. + return + with suppress(DuplexSessionError): + await handle.close(reason="disconnect") diff --git a/vllm_omni/entrypoints/duplex/session_attachment.py b/vllm_omni/entrypoints/duplex/session_attachment.py index 242169a6f..f2dc09f1c 100644 --- a/vllm_omni/entrypoints/duplex/session_attachment.py +++ b/vllm_omni/entrypoints/duplex/session_attachment.py @@ -180,7 +180,6 @@ class DuplexEventJournal: @dataclass(frozen=True) class DuplexSessionAttachmentCreated: session_id: str - incarnation: int attachment_generation: int resume_token: ResumeToken = field(repr=False) @@ -188,7 +187,6 @@ class DuplexSessionAttachmentCreated: @dataclass(frozen=True) class DuplexSessionResumeResult: session_id: str - incarnation: int attachment_generation: int resume_token: ResumeToken = field(repr=False) replay_entries: tuple[JournalEntry, ...] = () @@ -198,14 +196,12 @@ class DuplexSessionResumeResult: @dataclass class _DuplexSessionAttachmentState: session_id: str - incarnation: int credential: DuplexResumeCredential journal: DuplexEventJournal attachment: DuplexTransportAttachment | None attachment_generation: int outbound_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) recovery_token_digest: bytes | None = field(default=None, repr=False) - grace_task: asyncio.Task[None] | None = field(default=None, repr=False) class DuplexSessionAttachmentRegistry: @@ -214,18 +210,14 @@ class DuplexSessionAttachmentRegistry: *, replay_ttl_s: float, replay_max_bytes_per_session: int, - disconnect_grace_s: float = 30.0, clock: Callable[[], float] | None = None, ) -> None: if replay_ttl_s <= 0: raise ValueError("replay_ttl_s must be positive") if replay_max_bytes_per_session <= 0: raise ValueError("replay_max_bytes_per_session must be positive") - if disconnect_grace_s <= 0: - raise ValueError("disconnect_grace_s must be positive") self._replay_ttl_s = replay_ttl_s self._replay_max_bytes_per_session = replay_max_bytes_per_session - self._disconnect_grace_s = disconnect_grace_s self._clock = clock or time.monotonic self._sessions: dict[str, _DuplexSessionAttachmentState] = {} self._lock = asyncio.Lock() @@ -237,7 +229,6 @@ class DuplexSessionAttachmentRegistry: self, session_id: str, *, - incarnation: int, send: Callable[[dict[str, object]], Awaitable[None]], close: Callable[[str], Awaitable[None]], ) -> DuplexSessionAttachmentCreated: @@ -248,7 +239,6 @@ class DuplexSessionAttachmentRegistry: generation = 1 self._sessions[session_id] = _DuplexSessionAttachmentState( session_id=session_id, - incarnation=incarnation, credential=DuplexResumeCredential.from_token(token), journal=DuplexEventJournal( max_bytes=self._replay_max_bytes_per_session, @@ -264,7 +254,6 @@ class DuplexSessionAttachmentRegistry: ) return DuplexSessionAttachmentCreated( session_id=session_id, - incarnation=incarnation, attachment_generation=generation, resume_token=token, ) @@ -309,33 +298,39 @@ class DuplexSessionAttachmentRegistry: async with self._lock: return self._require(session_id).journal.acknowledge(sequence) - async def detach( - self, - session_id: str, - *, - attachment_generation: int, - on_grace_expired: Callable[[], Awaitable[None]] | None = None, - ) -> bool: + async def detach(self, session_id: str, *, attachment_generation: int | None = None) -> bool: + """Drop the current transport; the engine lease owns the disconnect grace. + + ``attachment_generation`` names the connection asking to detach, so a + socket that already lost a takeover cannot detach the winner. ``None`` + means "whichever connection is attached right now" and is for callers + that only know the session (the outbound pump, whose send just failed). + + Returns whether this call is the one that detached: an already-detached + session answers ``False`` so a second disconnect signal for the same + socket cannot restart the engine's disconnect grace window. + """ async with self._lock: state = self._sessions.get(session_id) - if state is None or state.attachment_generation != attachment_generation: + if state is None or state.attachment is None: + return False + if attachment_generation is not None and state.attachment_generation != attachment_generation: return False state.attachment = None - if state.grace_task is not None: - state.grace_task.cancel() - state.grace_task = ( - asyncio.create_task( - self._run_disconnect_grace( - state, - attachment_generation=attachment_generation, - callback=on_grace_expired, - ) - ) - if on_grace_expired is not None - else None - ) return True + async def has_attachment(self, session_id: str) -> bool: + """Whether some connection is attached right now, whoever it is. + + A resume that failed to activate has no generation of its own, so it + cannot ask ``is_current_attachment``. What it needs to know before + rolling the engine lease back into its disconnect grace is only whether + it would be rolling back somebody else's live attachment. + """ + async with self._lock: + state = self._sessions.get(session_id) + return state is not None and state.attachment is not None + async def is_current_attachment(self, session_id: str, attachment_generation: int) -> bool: async with self._lock: state = self._sessions.get(session_id) @@ -349,25 +344,19 @@ class DuplexSessionAttachmentRegistry: self, session_id: str, *, - incarnation: int, resume_token: str, last_received_server_event_seq: int, ) -> None: """Validate transport credentials before any engine resume control.""" async with self._lock: state = self._require(session_id) - self._validate_resume_identity( - state, - incarnation=incarnation, - resume_token=resume_token, - ) + self._validate_resume_identity(state, resume_token=resume_token) state.journal.replay_after(last_received_server_event_seq) async def resume( self, session_id: str, *, - incarnation: int, resume_token: str, last_received_server_event_seq: int, send: Callable[[dict[str, object]], Awaitable[None]], @@ -380,18 +369,11 @@ class DuplexSessionAttachmentRegistry: async with self._lock: if self._sessions.get(session_id) is not state: raise KeyError(f"unknown duplex attachment session: {session_id}") - used_recovery = self._validate_resume_identity( - state, - incarnation=incarnation, - resume_token=resume_token, - ) + used_recovery = self._validate_resume_identity(state, resume_token=resume_token) replay_entries = state.journal.replay_after(last_received_server_event_seq) accepted_token_digest = ( state.recovery_token_digest if used_recovery else bytes(state.credential.token_digest) ) - if state.grace_task is not None: - state.grace_task.cancel() - state.grace_task = None state.recovery_token_digest = None rotated_token = state.credential.rotate() replaced = state.attachment @@ -407,53 +389,53 @@ class DuplexSessionAttachmentRegistry: await send(dict(activation_payload_factory(rotated_token, attachment_generation))) for entry in replay_entries: await send(dict(entry.payload)) - except Exception: - async with self._lock: - if ( - self._sessions.get(session_id) is state - and state.attachment_generation == attachment_generation - ): - state.attachment = None - state.recovery_token_digest = accepted_token_digest + except (Exception, asyncio.CancelledError): + # CancelledError derives from BaseException, so an + # ``except Exception`` here would let a cancellation during + # activation or replay skip the rollback: the dead + # attachment would stay current while the token it rotated + # past is gone, leaving the session attached to an unusable + # transport with no way back. Both failures take the same + # path deliberately, so they cannot drift apart. + await self._rollback_resume( + session_id, + state=state, + attachment_generation=attachment_generation, + accepted_token_digest=accepted_token_digest, + ) raise return DuplexSessionResumeResult( session_id=session_id, - incarnation=incarnation, attachment_generation=attachment_generation, resume_token=rotated_token, replay_entries=replay_entries, replaced_attachment=replaced, ) + async def _rollback_resume( + self, + session_id: str, + *, + state: _DuplexSessionAttachmentState, + attachment_generation: int, + accepted_token_digest: bytes, + ) -> None: + """Undo a resume whose activation never reached the client. + + Generation-checked: a later resume that already took over must keep its + attachment. Restoring ``recovery_token_digest`` is what lets the caller + retry with the token it presented, which ``resume`` had rotated past. + """ + async with self._lock: + if self._sessions.get(session_id) is state and state.attachment_generation == attachment_generation: + state.attachment = None + state.recovery_token_digest = accepted_token_digest + async def close(self, session_id: str) -> DuplexTransportAttachment | None: async with self._lock: state = self._sessions.pop(session_id, None) - if state is not None and state.grace_task is not None: - state.grace_task.cancel() return state.attachment if state is not None else None - async def _run_disconnect_grace( - self, - state: _DuplexSessionAttachmentState, - *, - attachment_generation: int, - callback: Callable[[], Awaitable[None]] | None, - ) -> None: - try: - await asyncio.sleep(self._disconnect_grace_s) - except asyncio.CancelledError: - return - async with self._lock: - if ( - self._sessions.get(state.session_id) is not state - or state.attachment is not None - or state.attachment_generation != attachment_generation - ): - return - state.grace_task = None - if callback is not None: - await callback() - def _require(self, session_id: str) -> _DuplexSessionAttachmentState: state = self._sessions.get(session_id) if state is None: @@ -461,14 +443,7 @@ class DuplexSessionAttachmentRegistry: return state @staticmethod - def _validate_resume_identity( - state: _DuplexSessionAttachmentState, - *, - incarnation: int, - resume_token: str, - ) -> bool: - if incarnation != state.incarnation: - raise ValueError(f"duplex attachment incarnation mismatch: expected {state.incarnation}, got {incarnation}") + def _validate_resume_identity(state: _DuplexSessionAttachmentState, *, resume_token: str) -> bool: if state.credential.verify(resume_token): return False recovery_digest = state.recovery_token_digest diff --git a/vllm_omni/entrypoints/duplex/session_runner.py b/vllm_omni/entrypoints/duplex/session_runner.py deleted file mode 100644 index 90251e998..000000000 --- a/vllm_omni/entrypoints/duplex/session_runner.py +++ /dev/null @@ -1,2520 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import binascii -import json -import time -import uuid -from collections.abc import Callable -from contextlib import suppress -from copy import deepcopy - -import numpy as np -import pybase64 as base64 -from fastapi import WebSocket, WebSocketDisconnect -from vllm.logger import init_logger - -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.duplex.audio import ( - convert_input_audio_with_rate, - validate_input_sample_rate_hz, -) -from vllm_omni.entrypoints.duplex.commit_policy import ( - CommitAction, - CommitSnapshot, - decide_commit_action, -) -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexOverlapPolicy, - DuplexPlaybackCommitPolicy, - DuplexSession, - DuplexSessionState, - DuplexTurnEventType, -) -from vllm_omni.entrypoints.duplex.realtime_session import ( - NativeRealtimeSessionProtocol, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - PcmAppendReservation, - ServingRuntimeConfigError, - ServingRuntimeSessionState, - payload_turn_id, -) -from vllm_omni.entrypoints.duplex.server_vad import ServerVADPipeline -from vllm_omni.entrypoints.duplex.session_attachment import ( - DuplexJournalOverflowError, -) -from vllm_omni.entrypoints.duplex.websocket import ( - DuplexWebSocketActor, - is_input_event, - normalize_duplex_input_event, -) - -logger = init_logger(__name__) - -_MAX_EVENT_BYTES = 15 * 1024 * 1024 - - -class DuplexSessionRunnerMixin: - """Run one ordered WebSocket mailbox for a duplex session.""" - - async def handle_session( - self, - websocket: WebSocket, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - ) -> None: - await websocket.accept() - session: DuplexSession | None = None - actor = DuplexWebSocketActor( - websocket, - current_epoch=lambda: session.epoch if session is not None else None, - session_closed=lambda: session is not None and session.state == DuplexSessionState.CLOSED, - outbound_protocol=realtime_protocol, - ) - runtime_opened = False - runtime_closed = False - attachment_ready = False - attachment_generation: int | None = None - resume_credential_delivered = False - transport_detached = False - pending_turn_reservations = 0 - - async def attachment_send(payload: dict[str, object]) -> None: - nonlocal resume_credential_delivered - try: - await websocket.send_json(payload) - except RuntimeError as exc: - message = str(exc) - if "after sending 'websocket.close'" in message or "response already completed" in message: - raise WebSocketDisconnect(code=1006) from exc - raise - if payload.get("type") == "session.created" and isinstance(payload.get("resume_token"), str): - resume_credential_delivered = True - - async def attachment_close(reason: str) -> None: - close = getattr(websocket, "close", None) - if callable(close): - await close(code=1000, reason=reason) - - async def send_attachment_payload( - payload: dict[str, object], *, journal: bool, on_accepted: Callable[[], None] | None = None - ) -> None: - assert session is not None - session_id = session.session_id - should_journal = journal and session_id not in self._resync_required_sessions - try: - await self._attachment_registry.send_event( - session_id, - payload, - journal=should_journal, - on_accepted=on_accepted, - ) - except DuplexJournalOverflowError: - first_overflow = session_id not in self._resync_required_sessions - self._resync_required_sessions.add(session_id) - if first_overflow: - resync_payload: dict[str, object] = { - "type": "session.resync_required", - "session_id": session_id, - "reason": "journal_overflow", - } - if realtime_protocol is not None: - resync_payload = realtime_protocol.encode_outbound_event(resync_payload)[0] - await self._attachment_registry.send_event( - session_id, - resync_payload, - journal=False, - ) - await self._attachment_registry.send_event( - session_id, - payload, - journal=False, - on_accepted=on_accepted, - ) - - if realtime_protocol is not None: - - async def send_realtime_raw(payload: dict[str, object]) -> None: - if attachment_ready and session is not None: - await send_attachment_payload( - payload, - journal=payload.get("type") not in {"session.created", "session.resumed"}, - ) - return - raw_payload = dict(payload) - raw_payload["_realtime_raw"] = True - await actor.send_json(raw_payload) - - realtime_protocol.bind_sender(send_realtime_raw) - native: ServingRuntimeSessionState = self._create_runtime_session_state() - - def begin_close(reason: str) -> None: - actor.closing = True - actor.close_reason = reason - if session is not None: - session.mark_closing() - - event_emit_lock = asyncio.Lock() - - async def send_outbound(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - if not attachment_ready or session is None: - await actor.send_json(dict(payload)) - if on_accepted is not None: - on_accepted() - return - if realtime_protocol is None: - await send_attachment_payload( - payload, - journal=payload.get("type") not in {"session.created", "session.resumed"}, - on_accepted=on_accepted, - ) - return - for projected in realtime_protocol.encode_outbound_event(payload): - await send_attachment_payload( - projected, - journal=projected.get("type") not in {"session.created", "session.resumed"}, - # Item/part announcements may precede the audio. Only the - # actual audio event establishes acceptance of this chunk. - on_accepted=on_accepted if projected.get("type") == "response.audio.delta" else None, - ) - - async def emit_event(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - deferred_precreate_response = False - async with event_emit_lock: - accepted, deferred_overlap_payload = await self._apply_outbound_session_event( - payload, - session=session, - actor=actor, - native=native, - realtime_protocol=realtime_protocol, - ) - if not accepted: - return - await send_outbound(payload, on_accepted=on_accepted) - if deferred_overlap_payload is not None: - deferred_precreate_response = native.deferred_precreate_response - native.deferred_precreate_response = False - if deferred_overlap_payload is not None and session is not None and not actor.closing: - await start_native_append( - deferred_overlap_payload, - final=True, - precreate_response=deferred_precreate_response, - operation_id=native.committed_audio_operation_id, - retained_committed_payload=( - deferred_overlap_payload if native.committed_audio_payload is deferred_overlap_payload else None - ), - ) - - writer_task = asyncio.create_task(actor.writer_loop(), name="duplex-session-writer") - reader_task: asyncio.Task[None] | None = None - - async def read_event_loop() -> None: - nonlocal pending_turn_reservations - assert session is not None - try: - while not actor.closing: - raw = await self._receive_text( - websocket, - session.config.idle_timeout_s, - realtime_protocol=realtime_protocol, - ) - if raw is None: - await actor.enqueue_event({"type": "__timeout__"}) - return - if len(raw.encode("utf-8")) > _MAX_EVENT_BYTES: - await emit_event( - {"type": "error", "error": "Duplex event too large", "code": "event_too_large"} - ) - continue - try: - event = json.loads(raw) - except json.JSONDecodeError: - await emit_event({"type": "error", "error": "Invalid JSON event", "code": "invalid_json"}) - continue - if not isinstance(event, dict): - await emit_event( - { - "type": "error", - "error": "Duplex event must be a JSON object", - "code": "bad_event", - } - ) - continue - event = normalize_duplex_input_event(event) - event_type = event.get("type") - if not isinstance(event_type, str): - await emit_event( - {"type": "error", "error": "Duplex event missing string type", "code": "bad_event"} - ) - continue - if event_type in {"input.commit", "input_audio_buffer.commit"}: - if not session.reserve_pending_turn( - limit=self._duplex_session_config.max_pending_turns_per_session - ): - await emit_event( - { - "type": "error", - "error": "Duplex session has too many pending input turns", - "code": "input_backpressure", - } - ) - continue - event["_duplex_pending_turn_reserved"] = True - pending_turn_reservations += 1 - if is_input_event(event_type) and native_response_in_progress(): - event["_duplex_overlap_candidate"] = True - await actor.enqueue_event(event) - except WebSocketDisconnect: - await actor.enqueue_event({"type": "__disconnect__"}) - except Exception as exc: - await emit_event({"type": "error", "error": str(exc), "code": "realtime_input_failed"}) - await actor.enqueue_event({"type": "__disconnect__"}) - - async def next_actor_event() -> dict[str, object]: - nonlocal pending_turn_reservations, transport_detached - event = await actor.next_event() - if event.pop("_duplex_pending_turn_reserved", False) and session is not None: - session.release_pending_turn() - pending_turn_reservations = max(0, pending_turn_reservations - 1) - if ( - attachment_ready - and attachment_generation is not None - and session is not None - and not await self._attachment_registry.is_current_attachment( - session.session_id, - attachment_generation, - ) - ): - transport_detached = True - return {"type": "__replaced_attachment__"} - return event - - def native_response_in_progress() -> bool: - if session is None: - return False - if session.active_response_id is not None: - return True - if ( - session.config.playback_commit_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value - and session.playback.sent_ms > session.playback.committed_ms - ): - return True - if actor.active_response_task is not None and not actor.active_response_task.done(): - return True - if actor.has_response_bound_append_tasks(): - return True - return False - - def start_chat_response() -> None: - assert session is not None - - async def run_and_notify() -> None: - task = asyncio.current_task() - try: - await self._run_response(session, emit_event) - finally: - await actor.enqueue_event( - { - "type": "__chat_response_finished__", - "_response_task": task, - } - ) - - actor.active_response_task = asyncio.create_task(run_and_notify()) - - def activate_waiting_server_vad_turn() -> None: - assert session is not None - pending_turn = session.activate_server_vad_turn() - if pending_turn is not None and pending_turn.create_response: - start_chat_response() - - def correlate_realtime_error( - payload: dict[str, object], - event_id: object, - ) -> dict[str, object]: - if isinstance(event_id, str) and event_id: - payload["event_id"] = event_id - return payload - - async def reject_and_clear_server_vad_input( - pipeline: ServerVADPipeline, - *, - code: str, - error: str, - metric_reason: str, - event_id: object, - ) -> None: - assert session is not None - session.cancel_pending_input() - pipeline.reset() - self._realtime_vad_metrics.error(metric_reason) - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": code, - "error": error, - }, - event_id, - ) - ) - await emit_event( - { - "type": "input_audio_buffer.cleared", - "session_id": session.session_id, - } - ) - - async def commit_server_vad_turn( - *, - item_id: str, - create_response: bool, - retain_input_bytes: int, - event_id: object = None, - ) -> None: - assert session is not None - # Retain ownership until the actor consumes the terminal event so - # queued audio cannot overtake pending-turn activation. - active_response = actor.active_response_task is not None - if active_response and session.pending_server_vad_turn is not None: - released = session.discard_uncommitted_server_vad_utterance() - session.release_input_bytes(released) - self._realtime_vad_metrics.error("pending_turn_backpressure") - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "input_backpressure", - "error": "Only one committed Server VAD turn may wait for the active response", - }, - event_id, - ) - ) - return - if not session.stage_server_vad_audio_for_commit(): - return - committed = session.commit_user_input( - append_history=not active_response, - retain_input_bytes=retain_input_bytes, - ) - if committed is None: - return - if active_response: - session.stage_server_vad_turn( - committed, - create_response=create_response, - item_id=item_id, - ) - else: - session.register_history_item(item_id, committed.message) - await emit_event( - self._input_committed_payload( - session, - committed, - item_id=item_id, - ) - ) - if create_response and not active_response: - start_chat_response() - - def clear_completed_pending_silence() -> None: - task = native.pending_silence_task - if task is not None and task.done(): - native.pending_silence_task = None - native.pending_silence_owner_id = None - - def mark_pending_silence_superseded() -> None: - task = native.pending_silence_task - if task is None: - return - if task.done(): - native.pending_silence_task = None - # Do not cancel the task or clear native_append_tail here. A - # silence append may already have reached the Engine, and local - # cancellation cannot retract that RPC. The sequencer preserves - # wire order; before_append will skip silence that has not started. - native.pending_silence_owner_id = None - - def real_native_input_waiting() -> bool: - clear_completed_pending_silence() - return ( - native.audio_buffer.has_pending() - or native.audio_buffer.has_reserved() - or actor.has_queued_input_events() - ) - - async def start_native_append( - payload: object, - *, - final: bool, - precreate_response: bool = False, - pcm_reservation: PcmAppendReservation | None = None, - operation_id: str | None = None, - retained_committed_payload: dict[str, object] | None = None, - silence_continuation: bool = False, - before_append=None, - ) -> asyncio.Task[bool] | None: - if session is None: - return - if not silence_continuation: - mark_pending_silence_superseded() - append_epoch = session.epoch - append_turn_id = payload_turn_id(payload) - if append_turn_id is None: - append_turn_id = session.turn_id - request_id = self._native_stage0_request_id(session, append_epoch) - if final or precreate_response: - session.bind_request(request_id) - if precreate_response: - session.bind_response_turn(append_turn_id) - if precreate_response and session.active_response_id is None: - response_id = session.begin_response(turn_id=append_turn_id) - await emit_event( - self._response_created_payload( - session, - response_id, - epoch=append_epoch, - ) - ) - precreated_response_id = session.active_response_id if precreate_response else None - - def _discard_retained_committed_audio() -> None: - if ( - retained_committed_payload is not None - and native.committed_audio_payload is retained_committed_payload - ): - session.release_input_bytes(native.clear_committed_audio()) - - async def _run() -> bool: - nonlocal runtime_closed - try: - session.mark_model_turn_request_started(append_turn_id, time.monotonic()) - append_ok, emitted_response = await self._append_runtime_input( - session, - payload, - operation_id=(pcm_reservation.operation_id if pcm_reservation is not None else operation_id), - final=final, - send_json=emit_event, - mode="append_audio_chunk", - expected_epoch=append_epoch, - ) - if append_ok: - native.native_context_locked = True - if pcm_reservation is not None: - pcm_reservation.commit() - session.release_input_bytes(pcm_reservation.byte_count) - if ( - retained_committed_payload is not None - and native.committed_audio_payload is retained_committed_payload - ): - session.release_input_bytes(native.clear_committed_audio()) - else: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - if ( - not append_ok - and precreated_response_id is not None - and session.active_response_id == precreated_response_id - ): - session.end_response(commit_text=False) - await emit_event( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": precreated_response_id, - "epoch": session.epoch, - "committed": False, - "status": "failed", - "status_details": { - "type": "failed", - "reason": "runtime_append_failed", - }, - "playback": session.playback.as_dict(), - } - ) - if not append_ok and session.state == DuplexSessionState.CLOSED: - runtime_closed = True - return False - if not emitted_response and session.epoch == append_epoch: - if session.active_request_id == self._native_stage0_request_id(session, append_epoch): - session.clear_request(request_id) - if final: - await emit_event( - self._turn_controller.signal(session, DuplexTurnEventType.USER_STARTED.value) - ) - return append_ok - except asyncio.CancelledError: - if pcm_reservation is not None: - pcm_reservation.rollback() - raise - except Exception as exc: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - logger.exception("Native duplex append task failed: %s", exc) - await self._send_runtime_error(emit_event, "runtime_append_task_failed", exc, session=session) - if session.state != DuplexSessionState.CLOSED: - begin_close("runtime_append_task_failed") - if await self._close_runtime_session( - session, - reason="runtime_append_task_failed", - send_json=emit_event, - ): - runtime_closed = True - session.close() - await emit_event( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": "runtime_append_task_failed", - } - ) - return False - - async def _run_in_wire_order(predecessor: asyncio.Task[bool] | None) -> bool: - if predecessor is not None: - try: - predecessor_ok = await predecessor - except asyncio.CancelledError: - current = asyncio.current_task() - if current is not None and current.cancelling(): - raise - predecessor_ok = False - except Exception: - predecessor_ok = False - if not predecessor_ok: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - return False - if actor.closing or runtime_closed or session.state != DuplexSessionState.OPEN: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - return False - if before_append is not None and not before_append(): - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - return True - if pcm_reservation is not None and not pcm_reservation.active: - _discard_retained_committed_audio() - return False - return await _run() - - def _release_cancelled_retained_audio(done: asyncio.Task[bool]) -> None: - if done.cancelled(): - _discard_retained_committed_audio() - return - try: - append_ok = done.result() - except Exception: - append_ok = False - if not append_ok: - _discard_retained_committed_audio() - - predecessor = actor.native_append_tail - if predecessor is not None and predecessor.done(): - try: - predecessor_ok = predecessor.result() - except (asyncio.CancelledError, Exception): - predecessor_ok = False - if not predecessor_ok: - # Appends already queued behind a failed predecessor keep - # their captured dependency and stop. A later wire command - # is an explicit retry and starts a new chain. - predecessor = None - task = asyncio.create_task(_run_in_wire_order(predecessor)) - task.add_done_callback(_release_cancelled_retained_audio) - actor.native_append_tail = task - actor.track_append_task( - task, - epoch=append_epoch, - mode="append_audio_chunk", - final=final, - response_bound=final or precreate_response, - ) - if silence_continuation: - native.pending_silence_task = task - - def _clear_done_pending_silence(done: asyncio.Task[bool]) -> None: - if native.pending_silence_task is done: - native.pending_silence_task = None - native.pending_silence_owner_id = None - - task.add_done_callback(_clear_done_pending_silence) - # Let this wire-order effect start before the next mailbox event can - # cancel it. Later native appends still serialize on predecessor. - await asyncio.sleep(0) - return task - - async def schedule_native_silence_continuation( - payload: object, - *, - request_id: str, - owner_id: str, - response_id: str | None, - response_owned: bool, - expected_epoch: int | None, - expected_incarnation: int, - expected_model_turn_id: int | None, - send_json, - ) -> bool: - del send_json - if session is None: - return False - clear_completed_pending_silence() - pending_silence = native.pending_silence_task - if pending_silence is not None and not pending_silence.done(): - if pending_silence is asyncio.current_task(): - return False - try: - if not await pending_silence: - return False - except asyncio.CancelledError: - current = asyncio.current_task() - if current is not None and current.cancelling(): - raise - return False - except Exception: - return False - clear_completed_pending_silence() - pending_silence = native.pending_silence_task - if pending_silence is not None and not pending_silence.done(): - return False - append_tail = actor.native_append_tail - if (append_tail is None or append_tail.done()) and real_native_input_waiting(): - return False - continuation_delay_s = max( - 0.0, - float(session.capabilities.chunk_period_ms or 1000) / 1000.0, - ) - if continuation_delay_s > 0: - await asyncio.sleep(continuation_delay_s) - if ( - actor.native_append_tail is not append_tail - or ((append_tail is None or append_tail.done()) and real_native_input_waiting()) - or self._native_silence_continuation_is_stale( - session, - request_id=request_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_incarnation=expected_incarnation, - expected_model_turn_id=expected_model_turn_id, - ) - ): - return False - - def _still_valid() -> bool: - return not real_native_input_waiting() and not self._native_silence_continuation_is_stale( - session, - request_id=request_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_incarnation=expected_incarnation, - expected_model_turn_id=expected_model_turn_id, - ) - - native.pending_silence_owner_id = owner_id - task = await start_native_append( - payload, - final=False, - silence_continuation=True, - before_append=_still_valid, - ) - if task is None: - native.pending_silence_owner_id = None - return False - return True - - async def wait_for_native_append_tail() -> bool: - predecessor = actor.native_append_tail - if predecessor is None: - return True - try: - return await predecessor - except asyncio.CancelledError: - current = asyncio.current_task() - if current is not None and current.cancelling(): - raise - return False - except Exception: - # The append task already emitted its runtime error. Preserve - # wire order before the update path decides whether to continue. - return False - - async def start_runtime_append( - payload: object, - *, - final: bool, - mode: str = "append_tokens", - ) -> None: - """Schedule non-native runtime appends without blocking WS input. - - Model-native appends use ``start_native_append`` because they can - create data-plane response streams. Generic runtime appends still - need the same control-plane isolation so cancel/close can preempt a - slow engine append. - """ - if session is None: - return - append_epoch = session.epoch - - async def _run() -> None: - nonlocal runtime_closed - try: - append_ok, emitted_response = await self._append_runtime_input( - session, - payload, - final=final, - send_json=emit_event, - mode=mode, - expected_epoch=append_epoch, - ) - if not append_ok: - if session.state == DuplexSessionState.CLOSED: - runtime_closed = True - return - if not emitted_response and session.epoch == append_epoch: - await emit_event(self._turn_controller.signal(session, DuplexTurnEventType.USER_STARTED.value)) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.exception("Duplex runtime append task failed: %s", exc) - await self._send_runtime_error(emit_event, "runtime_append_task_failed", exc, session=session) - - task = asyncio.create_task(_run()) - actor.track_append_task( - task, - epoch=append_epoch, - mode=mode, - final=final, - response_bound=final, - ) - - try: - handshake = await self._open_session( - websocket, - emit_event, - realtime_protocol=realtime_protocol, - attachment_send=attachment_send, - attachment_close=attachment_close, - ) - if handshake is None: - return - session = handshake.session - if initial_update_error := self._runtime_session_candidate_update_error( - session, session.config, realtime_protocol=realtime_protocol - ): - await emit_event( - correlate_realtime_error( - initial_update_error, - handshake.event_id, - ) - ) - return - if handshake.resumed: - native = self._serving_session_states[session.session_id] - actor.tasks = self._session_tasks[session.session_id] - persisted_protocol = self._realtime_protocols.get(session.session_id) - if persisted_protocol is None: - raise RuntimeError(f"Missing Realtime protocol state for resumed session {session.session_id}") - realtime_protocol = persisted_protocol - realtime_protocol.bind_sender(send_realtime_raw) - attachment_generation = handshake.attachment_generation - attachment_ready = True - resume_credential_delivered = True - runtime_opened = True - self._ensure_lifecycle_listener() - reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") - else: - self._serving_session_states[session.session_id] = native - self._session_tasks[session.session_id] = actor.tasks - if realtime_protocol is not None: - self._realtime_protocols[session.session_id] = realtime_protocol - native.silence_continuation_scheduler = schedule_native_silence_continuation - if realtime_protocol is not None: - realtime_protocol.bind_native_input_append(self._uses_native_input_append(session)) - session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value - if not handshake.resumed: - open_result = await self._open_runtime_session(session, emit_event) - if open_result is False: - return - runtime_opened = True - try: - await self._configure_server_vad(session) - except Exception as exc: - logger.exception("Server VAD initialization failed: %s", exc) - self._realtime_vad_metrics.error("initialization") - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "server_vad_initialization_failed", - "error": str(exc), - }, - handshake.event_id, - ) - ) - return - created_attachment = await self._attachment_registry.create( - session.session_id, - incarnation=session.incarnation, - send=attachment_send, - close=attachment_close, - ) - attachment_generation = created_attachment.attachment_generation - attachment_ready = True - self._lease_generations[session.session_id] = 0 - self._ensure_lifecycle_listener() - created_payload: dict[str, object] = { - "type": "session.created", - "session": session.as_public_dict(), - } - if session.capabilities.supports_session_resume: - created_payload.update( - { - "incarnation": session.incarnation, - "attachment_generation": attachment_generation, - "resume_token": created_attachment.resume_token.plaintext, - } - ) - if isinstance(open_result, dict): - created_payload["runtime_control"] = self._redact_runtime_control_result(open_result) - if realtime_protocol is not None: - realtime_protocol.commit_realtime_turn_detection_update(session.config.server_vad) - await emit_event(created_payload) - resume_credential_delivered = session.capabilities.supports_session_resume - reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") - - while True: - event = await next_actor_event() - event_type = event.get("type") - - if event_type == "__replaced_attachment__": - return - - if event_type == "__chat_response_finished__": - completed_task = event.get("_response_task") - if completed_task is not actor.active_response_task: - continue - actor.active_response_task = None - activate_waiting_server_vad_turn() - continue - - if event_type == "__timeout__": - begin_close("timeout") - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - await actor.cancel_append_tasks() - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="timeout", - notify=True, - ) - actor.active_response_task = None - runtime_closed = await self._close_runtime_session( - session, - reason="timeout", - send_json=emit_event, - ) - if not runtime_closed: - return - session.close() - await emit_event( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": "timeout", - } - ) - return - - if event_type == "__disconnect__": - transport_detached = True - return - - if not isinstance(event_type, str): - continue - - if event_type == "session.heartbeat": - touch_session = getattr(self._chat_service.engine_client, "touch_duplex_session_async", None) - if not callable(touch_session): - await emit_event( - { - "type": "error", - "error": "Duplex runtime does not expose heartbeat control", - "code": "runtime_touch_unsupported", - } - ) - continue - try: - await touch_session( - session.session_id, - fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - activity=DuplexLeaseActivity.HEARTBEAT, - ) - except Exception as exc: - await emit_event( - { - "type": "error", - "error": str(exc), - "code": "runtime_touch_failed", - } - ) - continue - await emit_event( - { - "type": "session.heartbeat_ack", - "session_id": session.session_id, - } - ) - continue - - if event_type == "session.event_ack": - acknowledged = event.get("server_event_seq") - if not isinstance(acknowledged, int) or acknowledged < 0: - await emit_event( - { - "type": "error", - "error": "session.event_ack requires a non-negative server_event_seq", - "code": "invalid_event_ack", - } - ) - continue - try: - await self._attachment_registry.acknowledge(session.session_id, acknowledged) - except ValueError as exc: - await emit_event( - { - "type": "error", - "error": str(exc), - "code": "invalid_event_ack", - } - ) - continue - - if event_type == "turn.signal" and event.get("event") in { - "input.cancel", - "response.cancel", - "barge_in", - }: - payload = event.get("payload") - normalized_event = dict(payload) if isinstance(payload, dict) else {} - normalized_event.update(event) - normalized_event["type"] = event["event"] - event = normalized_event - event_type = str(event["type"]) - - if event_type == "session.close": - with suppress(asyncio.TimeoutError): - await asyncio.wait_for(actor.output_queue.join(), timeout=1.0) - begin_close("session_close") - if session.state == DuplexSessionState.CLOSED: - runtime_closed = True - return - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - await actor.cancel_append_tasks() - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="session_close", - ) - runtime_closed = await self._close_runtime_session( - session, - reason="session_close", - send_json=emit_event, - ) - actor.active_response_task = None - if not runtime_closed: - return - await emit_event({"type": "session.closed", "session_id": session.session_id}) - session.close() - return - - if event_type == "input_audio_buffer.clear": - native.audio_buffer.clear() - session.release_all_input_bytes() - server_vad_pipeline = self._server_vad_pipelines.get(session.session_id) - if server_vad_pipeline is not None: - server_vad_pipeline.reset() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - cancelled = session.cancel_pending_input() - await emit_event( - { - "type": "input_audio_buffer.cleared", - "session_id": session.session_id, - "epoch": session.epoch, - "drained_input_events": 0, - "cancelled": cancelled, - } - ) - continue - - if event_type == "barge_in" and not session.capabilities.supports_barge_in: - await emit_event(self._barge_in_unsupported_error(session)) - continue - - if event_type in {"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}: - cancel_reason = ( - "output_audio_buffer_clear" - if event_type == "output_audio_buffer.clear" - else "client_cancelled" - if event_type == "response.cancel" - else "barge_in" - ) - cancelled_fence = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - if event_type == "response.cancel": - requested_response_id = event.get("response_id") - has_active_response_work = native_response_in_progress() - if isinstance(requested_response_id, str) and ( - session.active_response_id is None or requested_response_id != session.active_response_id - ): - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": "response_not_active", - "error": f"Response is not active: {requested_response_id}", - } - ) - continue - if event_type == "response.cancel": - if not has_active_response_work: - if realtime_protocol is not None and isinstance(requested_response_id, str): - continue - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": "response_not_active", - "error": "response.cancel requires an active response", - } - ) - continue - had_native_unbuffered_append = ( - self._uses_native_input_append(session) - and native.input_since_commit - and not native.audio_buffer.has_pending() - ) - playback_was_active = self._assistant_playback_active(session) - if event_type in {"input.cancel", "barge_in"}: - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - had_native_append = await actor.cancel_append_tasks( - response_bound_only=event_type in {"response.cancel", "output_audio_buffer.clear"}, - ) - if event_type == "response.cancel": - session.release_input_bytes(native.clear_committed_audio()) - had_native_stream = native.data_plane_task is not None - cancelled = await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason=cancel_reason, - ) - had_native_stream = await self._cancel_native_data_plane_stream(session) or had_native_stream - if not cancelled and (had_native_stream or had_native_append or had_native_unbuffered_append): - old_epoch = session.epoch - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and playback_was_active: - old_epoch = session.epoch - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, session.last_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": session.last_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and event_type == "response.cancel": - old_epoch = session.epoch - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and event_type == "output_audio_buffer.clear": - old_playback = session.playback.as_dict() - committed_ms = session.playback.committed_ms - session.clear_playback_cursor() - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": session.active_response_id, - "reason": cancel_reason, - "cancelled_epoch": session.epoch, - "epoch": session.epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - actor.active_response_task = None - activate_waiting_server_vad_turn() - continue - if not cancelled: - await self._cancel_pending_input(session, emit_event, reason="barge_in") - if not await self._signal_runtime_session( - session, - "barge_in", - emit_event, - fence=cancelled_fence, - next_fence=( - DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - if session.epoch > cancelled_fence.epoch - else None - ), - ): - continue - actor.active_response_task = None - activate_waiting_server_vad_turn() - continue - - if event_type == "turn.signal": - turn_event = event.get("event") - if isinstance(turn_event, str): - if turn_event == "barge_in" and not session.capabilities.supports_barge_in: - await emit_event(self._barge_in_unsupported_error(session)) - continue - if turn_event == "session.update": - realtime_event_id = event.get("realtime_event_id") - - async def emit_update_event(update_event: dict[str, object]) -> None: - if isinstance(realtime_event_id, str) and realtime_event_id: - update_event = {**update_event, "realtime_event_id": realtime_event_id} - await emit_event(update_event) - - def reject_update() -> None: - if realtime_protocol is not None: - realtime_protocol.reject_realtime_turn_detection_update() - - payload = event.get("payload") - if not isinstance(payload, dict): - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": "bad_event", - "error": "session.update requires a session payload", - } - ) - reject_update() - continue - if not await wait_for_native_append_tail(): - await emit_update_event( - { - "type": "error", - "code": "session_update_aborted", - "error": "session.update was not applied because the preceding append failed", - } - ) - reject_update() - continue - if ( - realtime_protocol is not None - and not self._uses_native_input_append(session) - and session.config.server_vad is not None - and session.turn_detection_config_locked - and realtime_protocol.pending_realtime_input_audio_config_changed() - ): - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_audio_format", - "error": ( - "session.update cannot change Server VAD input format or sample rate " - "after the first audio append" - ), - "param": "audio.input.format", - } - ) - reject_update() - continue - runtime_update_error = self._runtime_session_update_error(session, payload) - if runtime_update_error is not None: - await emit_update_event(runtime_update_error) - reject_update() - continue - previous_config = session.config - candidate_config = deepcopy(previous_config) - session.replace_config(candidate_config) - try: - update_error = self._apply_session_update(session, payload) - finally: - session.replace_config(previous_config) - server_vad_changed = candidate_config.server_vad != previous_config.server_vad - if update_error is not None: - await emit_update_event(update_error) - reject_update() - continue - runtime_update_error = self._runtime_session_candidate_update_error( - session, - candidate_config, - realtime_protocol=realtime_protocol, - ) - if runtime_update_error is not None: - await emit_update_event(runtime_update_error) - reject_update() - continue - try: - candidate_runtime_config = self._runtime_config_for_session_update( - session, - candidate_config, - ) - except ServingRuntimeConfigError as exc: - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": exc.code, - "error": str(exc), - } - ) - reject_update() - continue - candidate_server_vad_pipeline = None - if server_vad_changed: - try: - candidate_server_vad_pipeline = await self._prepare_server_vad_pipeline( - session, - candidate_config, - ) - except Exception as exc: - logger.exception("Server VAD session update failed: %s", exc) - self._realtime_vad_metrics.error("initialization") - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": "server_vad_initialization_failed", - "error": str(exc), - } - ) - reject_update() - continue - runtime_updated = await self._signal_runtime_session( - session, - turn_event, - emit_update_event, - session_config=candidate_config.as_dict(), - runtime_config=candidate_runtime_config, - ) - if not runtime_updated: - reject_update() - continue - session.replace_config(candidate_config) - session.replace_runtime_config(candidate_runtime_config) - if server_vad_changed: - self._install_server_vad_pipeline( - session, - candidate_server_vad_pipeline, - ) - if realtime_protocol is not None: - realtime_protocol.bind_native_input_append(self._uses_native_input_append(session)) - realtime_protocol.commit_realtime_turn_detection_update(session.config.server_vad) - await emit_update_event( - { - "type": "session.updated", - "session": session.as_public_dict(), - } - ) - continue - if turn_event == "conversation.item.create": - payload = event.get("payload") - item = payload.get("item") if isinstance(payload, dict) else None - item_type = item.get("type") if isinstance(item, dict) else None - prepare_function_output = getattr( - self._serving_runtime_adapter, - "runtime_config_for_function_output", - None, - ) - if item_type == "function_call_output" and callable(prepare_function_output): - if not await wait_for_native_append_tail(): - continue - try: - candidate_runtime_config = prepare_function_output( - item, - dict(session.runtime_config), - ) - except ServingRuntimeConfigError as exc: - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": exc.code, - "error": str(exc), - } - ) - continue - if not await self._signal_runtime_session( - session, - "session.update", - emit_event, - runtime_config=candidate_runtime_config, - ): - continue - session.replace_runtime_config(candidate_runtime_config) - await emit_event( - { - "type": "conversation.item.created", - "session_id": session.session_id, - "item": item, - "created": True, - } - ) - await self._maybe_continue_native_response( - emit_event, - session=session, - expected_epoch=session.epoch, - expected_model_turn_id=session.turn_id, - ) - continue - history_item = payload.get("history_item") if isinstance(payload, dict) else None - if not isinstance(history_item, dict): - history_item = item - message = self._realtime_item_to_history_message(history_item) - item_id = item.get("id") if isinstance(item, dict) else None - if message is not None: - session.append_history_message(message) - session.register_history_item(item_id if isinstance(item_id, str) else None, message) - await emit_event( - { - "type": "conversation.item.created", - "session_id": session.session_id, - "item": item, - "created": message is not None, - } - ) - continue - if turn_event == "conversation.item.delete": - payload = event.get("payload") - item_id = payload.get("item_id") if isinstance(payload, dict) else None - deleted = session.delete_history_item(item_id) if isinstance(item_id, str) else False - await emit_event( - { - "type": "conversation.item.deleted", - "session_id": session.session_id, - "item_id": item_id, - "deleted": deleted, - } - ) - continue - if turn_event == "conversation.item.truncate": - payload = event.get("payload") - item_id = payload.get("item_id") if isinstance(payload, dict) else None - audio_end_ms = payload.get("audio_end_ms") if isinstance(payload, dict) else None - truncated = ( - session.truncate_history_item( - item_id, - audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, - hard=True, - ) - if isinstance(item_id, str) - else False - ) - await emit_event( - { - "type": "conversation.item.truncated", - "session_id": session.session_id, - "item_id": item_id, - "content_index": ( - payload.get("content_index", 0) if isinstance(payload, dict) else 0 - ), - "audio_end_ms": audio_end_ms, - "truncated": truncated, - } - ) - continue - await emit_event(self._turn_controller.signal(session, turn_event, event)) - else: - await emit_event({"type": "error", "error": "turn.signal requires event", "code": "bad_event"}) - continue - - if event_type == "playback.ack": - await self._handle_playback_ack(session, event, emit_event) - continue - - if event_type == "input.text.append": - session.mark_user_input_activity() - text = event.get("text") - if not isinstance(text, str): - await emit_event( - { - "type": "error", - "error": "input.text.append requires text", - "code": "bad_event", - } - ) - continue - if self._uses_native_input_append(session): - await emit_event( - { - "type": "error", - "error": "The selected native duplex runtime accepts audio append only", - "code": "native_text_append_unsupported", - } - ) - continue - else: - session.append_text(text) - if session.capabilities.supports_input_append: - await start_runtime_append(text, final=False, mode="append_tokens") - continue - - if event_type == "input_audio_buffer.append": - session.mark_user_input_activity() - audio = event.get("audio") or event.get("data") - if not isinstance(audio, str): - await emit_event( - { - "type": "error", - "error": "input.audio.append requires audio", - "code": "bad_event", - } - ) - continue - if not session.capabilities.supports_barge_in and self._event_requests_barge_in(event): - await emit_event(self._barge_in_unsupported_error(session)) - event = dict(event) - event.pop("force_barge_in", None) - for key in ("overlap_action", "overlap"): - value = event.get(key) - if isinstance(value, str) and value.strip().lower() in { - "barge_in", - "interrupt", - "cancel", - }: - event.pop(key, None) - fmt = event.get("format") if isinstance(event.get("format"), str) else "pcm16" - default_sample_rate_hz = 16000 - sr_raw = event.get("sample_rate_hz") or event.get("sample_rate") - sample_rate_hz = sr_raw if isinstance(sr_raw, int | float) else default_sample_rate_hz - native_input = self._uses_native_input_append(session) - server_vad_config = session.config.server_vad - turn_based_server_vad = server_vad_config is not None and not native_input - if not turn_based_server_vad: - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz, - ) - except ValueError as exc: - await emit_event({"type": "error", "error": str(exc), "code": "bad_event"}) - continue - if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le"}: - await emit_event( - { - "type": "error", - "error": "input_audio_buffer.append pcm16 audio could not be decoded", - "code": "bad_audio", - } - ) - continue - if server_vad_config is None: - session.lock_turn_detection_config() - force_listen = bool(event.get("force_listen", False)) - payload = { - "type": "audio", - "audio": audio, - "format": fmt, - "sample_rate_hz": sample_rate_hz, - "force_listen": force_listen, - } - video_frames = event.get("video_frames") - if isinstance(video_frames, list): - frames = [frame for frame in video_frames if isinstance(frame, str) and frame] - if frames: - payload["video_frames"] = frames - if server_vad_config is not None: - realtime_event_id = event.get("_realtime_event_id") - pipeline = self._server_vad_pipelines.get(session.session_id) - if pipeline is None: - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "internal_error", - "error": "server_vad is configured but its detector pipeline is unavailable", - }, - realtime_event_id, - ) - ) - continue - speech_was_active = pipeline.speech_active - reserved_bytes = 0 - previous_scratch_bytes = pipeline.scratch_bytes - try: - if native_input: - if fmt != "pcm_f32le" or sample_rate_hz != pipeline.sample_rate_hz: - raise ValueError("native server_vad input must be mono PCM float32 at 16000 Hz") - raw_audio = base64.b64decode(audio, validate=True) - if len(raw_audio) % np.dtype(" 0 - else: - if fmt != "pcm16": - raise ValueError("server_vad input must be mono PCM16") - source_sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) - if source_sample_rate_hz not in {16_000, 24_000}: - raise ValueError("server_vad PCM16 input sample rate must be 16000 or 24000 Hz") - if ( - pipeline.source_sample_rate_hz is not None - and source_sample_rate_hz != pipeline.source_sample_rate_hz - ): - raise ValueError( - "server_vad input sample rate cannot change within a continuous audio stream" - ) - try: - raw_audio = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError( - "server_vad input audio must be valid base64-encoded PCM16" - ) from exc - if len(raw_audio) % np.dtype(" 0 - except (binascii.Error, ValueError) as exc: - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "bad_audio", - "error": str(exc), - "param": "audio", - }, - realtime_event_id, - ) - ) - continue - if not native_input: - normalized_sample_upper_bound = ( - source_samples - if source_sample_rate_hz == pipeline.sample_rate_hz - else ((source_samples + 2) // 3) * 2 - ) - reserved_bytes = normalized_sample_upper_bound * np.dtype(np.float32).itemsize - if not session.reserve_input_bytes( - reserved_bytes, - limit=self._duplex_session_config.max_pending_input_bytes_per_session, - ): - await reject_and_clear_server_vad_input( - pipeline, - code="input_backpressure", - error="Server VAD input exceeds the per-session input limit", - metric_reason="input_backpressure", - event_id=realtime_event_id, - ) - continue - try: - vad_batch = ( - await pipeline.push(normalized_audio) - if native_input - else await pipeline.push_pcm16( - raw_audio, - source_sample_rate_hz=source_sample_rate_hz, - ) - ) - except Exception as exc: - logger.exception("Server VAD inference failed: %s", exc) - await reject_and_clear_server_vad_input( - pipeline, - code="server_vad_inference_failed", - error=str(exc), - metric_reason="inference", - event_id=realtime_event_id, - ) - continue - if not native_input: - retained_delta = ( - sum(frame.samples.nbytes for frame in vad_batch.frames) - + pipeline.scratch_bytes - - previous_scratch_bytes - ) - session.release_input_bytes(max(0, reserved_bytes - retained_delta)) - self._realtime_vad_metrics.observe_inference(vad_batch.inference_ms) - prefix_samples = round(server_vad_config.prefix_padding_ms * pipeline.sample_rate_hz / 1000) - speech_started = False - speech_stopped = False - for frame_index, vad_frame in enumerate(vad_batch.frames): - decision = vad_frame.decision - if not native_input: - released_bytes = session.append_server_vad_frame( - vad_frame.samples, - speech_started=decision.speech_started, - speech_stopped=decision.speech_stopped, - prefix_samples=prefix_samples, - ) - session.release_input_bytes(released_bytes) - if decision.speech_started: - speech_started = True - item_id = f"item_{uuid.uuid4().hex}" - session.begin_server_vad_speech(item_id) - await emit_event( - { - "type": "input_audio_buffer.speech_started", - "session_id": session.session_id, - "item_id": item_id, - "audio_start_ms": decision.audio_start_ms or 0, - } - ) - if decision.speech_stopped: - speech_stopped = True - item_id = session.finish_server_vad_speech() or f"item_{uuid.uuid4().hex}" - await emit_event( - { - "type": "input_audio_buffer.speech_stopped", - "session_id": session.session_id, - "item_id": item_id, - "audio_end_ms": decision.audio_end_ms or 0, - } - ) - if decision.endpoint_delay_ms is not None: - self._realtime_vad_metrics.observe_endpoint_delay(decision.endpoint_delay_ms) - if not native_input: - remaining_frame_bytes = sum( - frame.samples.nbytes for frame in vad_batch.frames[frame_index + 1 :] - ) - await commit_server_vad_turn( - item_id=item_id, - create_response=server_vad_config.create_response, - retain_input_bytes=remaining_frame_bytes + pipeline.scratch_bytes, - event_id=realtime_event_id, - ) - if accepted_audio: - session.lock_turn_detection_config() - if not native_input: - continue - speech_active = pipeline.speech_active - contained_speech = speech_was_active or speech_started or speech_stopped or speech_active - payload["is_speech"] = contained_speech - payload["vad"] = { - "backend": "silero", - "is_speech": contained_speech, - "speech_active": speech_active, - "speech_started": speech_started, - "speech_stopped": speech_stopped, - } - if speech_active: - payload["force_listen"] = True - # Speech/silence tag for the Stage0 turn-ended latch. - payload["is_speech"] = self._input_looks_like_speech(event, payload, session=session) - defer_native_append = False - buffer_overlap_audio = True - if native_input: - mark_pending_silence_superseded() - overlap_active = native_response_in_progress() and ( - not self._session_auto_responds(session) - or ( - session.capabilities.supports_barge_in - and ( - session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - or self._event_requests_barge_in(event) - ) - ) - ) - if overlap_active: - decision = self._overlap_decision(session, event, payload) - await self._emit_overlap_decision(emit_event, session, decision) - action = decision.get("action") - if action == "drop": - if realtime_protocol is not None: - await realtime_protocol.discard_pending_input_audio( - audio_end_ms=self._input_audio_duration_ms(event, payload) - ) - continue - if action == "listen": - buffer_overlap_audio = bool(decision.get("buffer_audio", True)) - defer_native_append = bool(decision.get("defer_runtime_append", True)) - if ( - not buffer_overlap_audio - and realtime_protocol is not None - and decision.get("preserve_realtime_input") is not True - ): - await realtime_protocol.discard_pending_input_audio( - audio_end_ms=self._input_audio_duration_ms(event, payload) - ) - if decision.get("force_listen", True) is True: - payload["force_listen"] = True - else: - event["force_barge_in"] = True - cancelled_fence = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - playback_was_active = self._assistant_playback_active(session) - buffer_overlap_audio = True - defer_native_append = False - native.audio_buffer.clear_force_listen() - session.reset_overlap_speech() - native.input_since_commit = False - native.speech_since_commit = False - await actor.cancel_append_tasks() - had_native_stream = native.data_plane_task is not None - cancel_reason = str(decision.get("cancel_reason") or "barge_in") - cancelled = await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason=cancel_reason, - ) - had_native_stream = ( - await self._cancel_native_data_plane_stream(session) or had_native_stream - ) - if not cancelled and had_native_stream: - old_epoch = session.epoch - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and playback_was_active: - old_epoch = session.epoch - committed_ms = session.playback.committed_ms - self._commit_played_response_history( - session, session.last_response_id, committed_ms - ) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": session.last_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if session.epoch > cancelled_fence.epoch: - if not await self._signal_runtime_session( - session, - "barge_in", - emit_event, - fence=cancelled_fence, - next_fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - ): - continue - actor.active_response_task = None - elif not self._session_auto_responds(session) and not self._input_looks_like_speech( - event, payload, session=session - ): - # Turn-mode only: skip silent chunks so they don't open a - # response. In auto-respond (full-duplex) mode the model owns - # the speak/listen decision and MUST receive silence units -- - # the official model typically starts speaking during the - # silence after a question. - # No response_id: a silence-skip listen can arrive - # while a response is still streaming and must not - # read as that response's terminal decision. - await emit_event( - { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": "silence_or_noise", - } - ) - continue - if self._should_force_listen_for_auto_response_overlap(session, event, payload): - # Auto-response keeps a long-lived native Stage0 stream. - # While assistant audio is still active, silence from the - # browser should advance the model in listen mode rather - # than letting the same stream start another assistant - # segment. Speech/barge-in chunks remain model-owned. - payload["force_listen"] = True - if not buffer_overlap_audio: - continue - session.mark_user_input_activity() - native.input_since_commit = True - native.speech_since_commit = native.speech_since_commit or self._input_looks_like_speech( - event, payload, session=session - ) - try: - raw_audio_bytes = self._native_audio_payload_size_bytes(payload) - if not session.reserve_input_bytes( - raw_audio_bytes, - limit=self._duplex_session_config.max_pending_input_bytes_per_session, - ): - await emit_event( - { - "type": "error", - "error": "Duplex session pending input exceeds server limit", - "code": "input_backpressure", - } - ) - continue - allow_emit = not defer_native_append and ( - realtime_protocol is None - or event_type != "input_audio_buffer.append" - # Full-duplex: emit each ~chunk_period of audio so the model - # runs per-chunk generation (speak/listen) without an explicit - # response.create, matching the official duplex_generate loop. - or self._session_auto_responds(session) - ) - pcm_reservation = native.audio_buffer.prepare_append( - payload, - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - allow_emit=allow_emit, - ) - except ValueError as exc: - session.release_input_bytes(raw_audio_bytes) - await emit_event({"type": "error", "error": str(exc), "code": "bad_event"}) - continue - if pcm_reservation is None: - continue - if pcm_reservation.byte_count == 0: - session.release_input_bytes(raw_audio_bytes) - payload = pcm_reservation.payload - else: - session.append_audio(audio, fmt=fmt, sample_rate_hz=sample_rate_hz) - if native_input: - await start_native_append( - payload, - final=False, - pcm_reservation=pcm_reservation, - ) - continue - if session.capabilities.supports_input_append: - await start_runtime_append(payload, final=False, mode="append_audio_chunk") - continue - - if event_type in {"input.commit", "input_audio_buffer.commit", "response.create"}: - if ( - event_type == "response.create" - and session.config.server_vad is not None - and actor.active_response_task is not None - ): - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": "response_already_active", - "error": "response.create must wait for the active response to complete", - } - ) - continue - item_id = event.get("item_id") - realtime_validated_audio_commit = ( - event_type == "input_audio_buffer.commit" and isinstance(item_id, str) and bool(item_id) - ) - if ( - self._uses_native_input_append(session) - and event_type in {"input.commit", "input_audio_buffer.commit"} - and not await wait_for_native_append_tail() - ): - continue - if event_type == "input_audio_buffer.commit" and event.get("is_speech") is False: - native.input_since_commit = False - native.speech_since_commit = False - native.audio_buffer.clear() - session.release_all_input_bytes() - native.clear_committed_audio() - await emit_event( - { - "type": "input.committed", - "session_id": session.session_id, - "item_id": item_id, - "turn_id": session.turn_id, - "epoch": session.epoch, - "empty": True, - "is_speech": False, - "no_response": True, - } - ) - # No response_id: this silence listen answers an empty - # commit, it is not a response terminal. - await emit_event( - { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": "silence_or_noise", - } - ) - continue - should_create_response = ( - event_type == "response.create" - or bool(event.get("response_create", event_type == "input.commit")) - or (event_type == "input_audio_buffer.commit" and self._session_auto_responds(session)) - ) - precreate_response_requested = event_type == "response.create" or bool( - event.get("response_create", event_type == "input.commit") - ) - if event_type == "response.create": - response_payload = event.get("response") - if isinstance(response_payload, dict): - response_options_error = self._apply_response_create_options(session, response_payload) - if response_options_error is not None: - error_message = ( - "The selected native duplex runtime does not support generation " - "overrides for instructions, voice, temperature, max tokens, tools, or " - "tool_choice." - if response_options_error == "unsupported_native_response_options" - else "response.create cannot reserve options while another response is active." - ) - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": response_options_error, - "error": error_message, - } - ) - continue - if self._uses_native_input_append(session) and event_type == "input_audio_buffer.commit": - has_pending_native_audio = ( - native.input_since_commit - or native.audio_buffer.has_pending() - or native.committed_audio_payload is not None - or realtime_validated_audio_commit - ) - if ( - not has_pending_native_audio - and not actor.native_append_tasks - and native.data_plane_task is None - ): - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "input_audio_buffer_empty", - "error": "input_audio_buffer.commit requires a non-empty input audio buffer.", - } - ) - continue - commit_action = decide_commit_action( - CommitSnapshot( - auto_responds=self._session_auto_responds(session), - speech_since_commit=native.speech_since_commit, - active_response_id=session.active_response_id, - overlap_speech_ms=session.overlap_speech_ms, - native_response_in_progress=native_response_in_progress(), - playback_active=self._assistant_playback_active(session), - ) - ) - if commit_action is CommitAction.DEFER_ACTIVE_RESPONSE: - if session.overlap_speech_ms <= session.config.overlap_short_ack_ms: - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - if realtime_protocol is not None: - await realtime_protocol.discard_pending_input_audio( - audio_end_ms=session.overlap_speech_ms - ) - await emit_event( - { - "type": "input.committed", - "session_id": session.session_id, - "item_id": item_id, - "turn_id": session.turn_id, - "epoch": session.epoch, - "empty": True, - "is_speech": False, - "overlap_ack": True, - "no_response": True, - } - ) - session.reset_overlap_speech() - session.discard_response_options() - continue - - commit_reservation = native.audio_buffer.prepare_commit( - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - ) - deferred_payload = commit_reservation.payload - if deferred_payload is None: - commit_reservation.commit() - else: - if native.committed_audio_payload is not None: - deferred_payload = self._merge_native_audio_payloads( - native.committed_audio_payload, - deferred_payload, - ) - native.retain_committed_audio( - deferred_payload, - operation_id=commit_reservation.operation_id, - reserved_bytes=commit_reservation.byte_count, - ) - commit_reservation.commit() - native.deferred_response_create = should_create_response - native.deferred_precreate_response = precreate_response_requested - native.input_since_commit = False - native.speech_since_commit = False - committed = self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - committed_payload = self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - committed_payload["overlap_deferred"] = True - committed_payload["response_create_deferred"] = should_create_response - await emit_event(committed_payload) - continue - if commit_action is CommitAction.START_AUTO_RESPONSE: - commit_reservation = native.audio_buffer.prepare_commit( - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - ) - committed_input = commit_reservation.payload - final_payload = committed_input - if native.committed_audio_payload is not None: - if final_payload is not None: - final_payload = self._merge_native_audio_payloads( - native.committed_audio_payload, - final_payload, - ) - else: - final_payload = native.committed_audio_payload - commit_reservation.commit() - if final_payload is not None: - native.retain_committed_audio( - final_payload, - operation_id=commit_reservation.operation_id, - reserved_bytes=commit_reservation.byte_count, - ) - native.deferred_response_create = False - native.input_since_commit = False - native.speech_since_commit = False - data_plane_turn_id = session.turn_id - committed = self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - turn_id=data_plane_turn_id, - ) - await emit_event( - self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - ) - if final_payload is not None: - await start_native_append( - { - **final_payload, - "duplex_turn_id": data_plane_turn_id, - }, - final=True, - precreate_response=False, - operation_id=commit_reservation.operation_id, - retained_committed_payload=final_payload, - ) - continue - if self._uses_native_input_append(session) and event_type == "response.create": - if ( - native_response_in_progress() - or actor.native_append_tasks - or native.data_plane_task is not None - ): - if session.active_response_id is None and ( - session.active_request_id is not None - or actor.native_append_tasks - or native.data_plane_task is not None - ): - continue - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "response_already_active", - "error": "response.create cannot start while another response is active.", - } - ) - session.discard_response_options() - continue - if native.committed_audio_payload is not None: - committed_payload = native.committed_audio_payload - operation_id = native.committed_audio_operation_id - if operation_id is None: - operation_id = uuid.uuid4().hex - native.committed_audio_operation_id = operation_id - await start_native_append( - committed_payload, - final=True, - precreate_response=True, - operation_id=operation_id, - retained_committed_payload=committed_payload, - ) - continue - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "response_create_without_input", - "error": "Native duplex response.create requires committed audio input.", - } - ) - session.discard_response_options() - continue - if self._uses_native_input_append(session) and not native_response_in_progress(): - commit_reservation = ( - native.audio_buffer.prepare_commit( - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - ) - if event_type in {"input.commit", "input_audio_buffer.commit"} - else None - ) - flushed_buffer_reserved_bytes = ( - native.audio_buffer.pending_byte_count if commit_reservation is None else 0 - ) - flushed = ( - commit_reservation.payload - if commit_reservation is not None - else native.audio_buffer.flush(chunk_period_ms=session.capabilities.chunk_period_ms or 1000) - ) - if native.committed_audio_payload is not None: - if flushed is not None: - flushed = self._merge_native_audio_payloads( - native.committed_audio_payload, - flushed, - ) - else: - flushed = native.committed_audio_payload - if commit_reservation is not None: - commit_reservation.commit() - if flushed is not None: - if self._should_force_listen_for_short_commit(session, event, flushed): - flushed = dict(flushed) - flushed["force_listen"] = True - native.input_since_commit = False - committed = self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - await emit_event( - self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - ) - if should_create_response: - native.retain_committed_audio( - flushed, - operation_id=( - commit_reservation.operation_id - if commit_reservation is not None - else uuid.uuid4().hex - ), - reserved_bytes=( - commit_reservation.byte_count - if commit_reservation is not None - else flushed_buffer_reserved_bytes - ), - ) - await start_native_append( - flushed, - final=True, - precreate_response=True, - operation_id=native.committed_audio_operation_id, - retained_committed_payload=flushed, - ) - else: - native.retain_committed_audio( - flushed, - operation_id=( - commit_reservation.operation_id - if commit_reservation is not None - else uuid.uuid4().hex - ), - reserved_bytes=( - commit_reservation.byte_count - if commit_reservation is not None - else flushed_buffer_reserved_bytes - ), - ) - native.deferred_response_create = False - continue - native_had_uncommitted_audio = self._uses_native_input_append(session) and ( - native.input_since_commit - or native.audio_buffer.has_pending() - or native.committed_audio_payload is not None - or realtime_validated_audio_commit - ) - committed = session.commit_user_input() - if self._uses_native_input_append(session) and event_type in { - "input_audio_buffer.commit", - "input.commit", - }: - native.input_since_commit = False - native.speech_since_commit = False - if committed is None and event_type != "response.create": - if self._uses_native_input_append(session): - committed = ( - self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - if native_had_uncommitted_audio - else None - ) - await emit_event( - self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - ) - else: - await emit_event( - { - "type": "input.committed", - "session_id": session.session_id, - "item_id": item_id, - "empty": True, - } - ) - continue - if committed is not None: - item_id = event.get("item_id") - if isinstance(item_id, str): - session.register_history_item(item_id, committed.message) - await emit_event( - self._input_committed_payload( - session, - committed, - item_id=item_id, - ) - ) - if not should_create_response: - continue - if actor.active_response_task is not None and not actor.active_response_task.done(): - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="new_response", - ) - if session.config.server_vad is not None and not self._uses_native_input_append(session): - start_chat_response() - else: - actor.active_response_task = asyncio.create_task(self._run_response(session, emit_event)) - continue - - await emit_event( - { - "type": "error", - "error": f"Unknown duplex event: {event_type}", - "code": "unknown_event", - } - ) - - except WebSocketDisconnect: - logger.info("Duplex session disconnected") - transport_detached = True - except Exception as exc: - logger.exception("Duplex session failed: %s", exc) - with suppress(Exception): - await emit_event({"type": "error", "error": str(exc), "code": "internal_error"}) - finally: - if realtime_protocol is not None: - realtime_protocol.reject_realtime_turn_detection_update() - if session is not None: - if reader_task is not None and not reader_task.done(): - reader_task.cancel() - with suppress(asyncio.CancelledError): - await reader_task - while pending_turn_reservations > 0: - session.release_pending_turn() - pending_turn_reservations -= 1 - resumable_detach = ( - transport_detached - and runtime_opened - and not runtime_closed - and session.state == DuplexSessionState.OPEN - and session.capabilities.supports_session_resume - and attachment_generation is not None - and resume_credential_delivered - ) - if resumable_detach: - - async def cancel_orphan_response_after_grace() -> None: - tasks = self._session_tasks.get(session.session_id) - current_session = self._registry.get(session.session_id) - if tasks is None or current_session is not session or session.state != DuplexSessionState.OPEN: - return - await tasks.cancel_append_tasks(response_bound_only=True) - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - tasks.active_response_task, - emit_event, - reason="disconnect_grace_expired", - notify=False, - ) - tasks.active_response_task = None - - detached = await self._attachment_registry.detach( - session.session_id, - attachment_generation=attachment_generation, - on_grace_expired=cancel_orphan_response_after_grace, - ) - touch_session = getattr(self._chat_service.engine_client, "touch_duplex_session_async", None) - if detached and callable(touch_session): - with suppress(Exception): - await touch_session( - session.session_id, - fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - activity=DuplexLeaseActivity.DETACH, - ) - else: - begin_close(actor.close_reason or "disconnect") - await actor.cancel_append_tasks() - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="disconnect", - notify=False, - ) - if runtime_opened and not runtime_closed and session.state != DuplexSessionState.CLOSED: - await self._close_runtime_session(session, reason="disconnect") - self._cleanup_duplex_session_state(session) - self._registry.close(session.session_id) - self._session_tasks.pop(session.session_id, None) - self._realtime_protocols.pop(session.session_id, None) - self._lease_generations.pop(session.session_id, None) - self._resync_required_sessions.discard(session.session_id) - with suppress(Exception): - await self._attachment_registry.close(session.session_id) - self._stop_lifecycle_listener_if_idle() - await actor.close_writer() - with suppress(Exception): - await asyncio.wait_for(actor.output_queue.join(), timeout=2.0) - if not writer_task.done(): - writer_task.cancel() - with suppress(asyncio.CancelledError): - await writer_task diff --git a/vllm_omni/entrypoints/duplex/warmup.py b/vllm_omni/entrypoints/duplex/warmup.py index 5051a1d31..2c138096d 100644 --- a/vllm_omni/entrypoints/duplex/warmup.py +++ b/vllm_omni/entrypoints/duplex/warmup.py @@ -43,10 +43,11 @@ async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None: model_name = served else: model_name = args.model - adapter = getattr(getattr(app.state, "openai_serving_duplex", None), "_serving_runtime_adapter", None) - frame_samples = int(getattr(adapter, "silence_continuation_samples", 16000)) + # One silence unit as the engine-side plugin defines it (DuplexOmniEngine.plugin). + plugin = getattr(getattr(app.state.engine_client, "engine", None), "plugin", None) + frame_samples = int(getattr(plugin, "silence_continuation_samples", 16000)) silence = base64.b64encode(bytes(frame_samples * 4)).decode("ascii") - from vllm_omni.experimental.fullduplex.client import build_realtime_url + from vllm_omni.clients.duplex import build_realtime_url url = ( build_realtime_url(f"ws://127.0.0.1:{args.port}/v1/realtime", model_name, autostart=False) diff --git a/vllm_omni/entrypoints/duplex/websocket.py b/vllm_omni/entrypoints/duplex/websocket.py index d51a7ed57..8ec9cc780 100644 --- a/vllm_omni/entrypoints/duplex/websocket.py +++ b/vllm_omni/entrypoints/duplex/websocket.py @@ -1,266 +1,60 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""WebSocket transport helpers for the duplex Realtime handler. + +The ordered session mailbox and the per-session task handles that used to be +defined here now live engine-side (``vllm_omni.engine.duplex.session.runner``). +""" + from __future__ import annotations import asyncio -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import Any +from collections.abc import Awaitable, Callable +from contextlib import suppress -from fastapi import WebSocketDisconnect +from fastapi import WebSocket, WebSocketDisconnect -INPUT_EVENTS = frozenset( - { - "input.text.append", - "input_audio_buffer.append", - "input.commit", - "input_audio_buffer.commit", - "response.create", - } -) +#: Largest client frame accepted before it is parsed (base64 audio + video frames). +MAX_EVENT_BYTES = 15 * 1024 * 1024 -_INPUT_EVENT_ALIASES = { - "signal_turn": "turn.signal", - "close_session": "session.close", - "audio.playback_ack": "playback.ack", - "input_text.append": "input.text.append", - "push_text": "input.text.append", - "input.audio.append": "input_audio_buffer.append", - "push_chunk": "input_audio_buffer.append", -} -_WAV_AUDIO_ALIASES = frozenset({"input.audio.append", "push_chunk"}) +SendJson = Callable[[dict[str, object]], Awaitable[None]] +CloseFn = Callable[[str], Awaitable[None]] -def normalize_duplex_input_event(event: dict[str, object]) -> dict[str, object]: - """Normalize compatibility aliases before an event enters the mailbox.""" - event_type = event.get("type") - if not isinstance(event_type, str): - return event - canonical_type = _INPUT_EVENT_ALIASES.get(event_type) - if canonical_type is None: - return event - normalized = dict(event) - normalized["type"] = canonical_type - if event_type in _WAV_AUDIO_ALIASES: - normalized.setdefault("format", "wav") - return normalized +async def receive_text_with_timeout(websocket: WebSocket, timeout_s: float | None) -> str | None: + """Receive one text frame; ``None`` on idle timeout. Disconnects propagate.""" + if timeout_s is None or timeout_s <= 0: + return await websocket.receive_text() + try: + return await asyncio.wait_for(websocket.receive_text(), timeout=timeout_s) + except TimeoutError: + return None -MODEL_OUTPUT_EVENTS = frozenset( - { - "response.created", - "response.listen", - "response.speak", - "response.output_item.added", - "response.content_part.added", - "response.output_audio.delta", - "response.output_audio.done", - "response.output_text.delta", - "response.output_text.done", - "response.text.delta", - "response.text.done", - "response.message", - "response.output_item.done", - "response.content_part.done", - "response.done", - "runtime.control", - } -) -DOMAIN_TERMINAL_EVENTS = frozenset( - { - "response.done", - "response.listen", - "audio.cancelled", - "input.cancelled", - "session.closed", - } -) +def attachment_callbacks(websocket: WebSocket) -> tuple[SendJson, CloseFn]: + """``send`` / ``close`` callables for a ``DuplexSessionAttachmentRegistry`` attachment. + A send on a socket that already closed surfaces as ``WebSocketDisconnect`` + so the caller's disconnect path runs instead of a generic failure. + """ -def is_input_event(event_type: object) -> bool: - return isinstance(event_type, str) and event_type in INPUT_EVENTS - - -@dataclass(frozen=True, slots=True) -class DuplexAppendTaskMeta: - epoch: int - mode: str - final: bool - response_bound: bool - - -@dataclass -class DuplexSessionTasks: - """Connection-independent task handles for one resumable session.""" - - native_append_tasks: dict[asyncio.Task[bool], DuplexAppendTaskMeta] = field(default_factory=dict) - native_append_tail: asyncio.Task[bool] | None = None - active_response_task: asyncio.Task[None] | None = None - - def track_append_task( - self, - task: asyncio.Task[bool], - *, - epoch: int, - mode: str, - final: bool, - response_bound: bool, - ) -> None: - self.native_append_tasks[task] = DuplexAppendTaskMeta(epoch, mode, final, response_bound) - task.add_done_callback(self.native_append_tasks.pop) - - def has_response_bound_append_tasks(self) -> bool: - return any(meta.response_bound for meta in self.native_append_tasks.values()) - - async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: - tasks = [ - task for task, meta in self.native_append_tasks.items() if not response_bound_only or meta.response_bound - ] - if not tasks: - return False - cancelled_tail = self.native_append_tail if self.native_append_tail in tasks else None - for task in tasks: - task.cancel() + async def attachment_send(payload: dict[str, object]) -> None: try: - await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) - except TimeoutError: - pass - if cancelled_tail is not None and self.native_append_tail is cancelled_tail: - self.native_append_tail = None - return True + await websocket.send_json(payload) + except RuntimeError as exc: + message = str(exc) + if "after sending 'websocket.close'" in message or "response already completed" in message: + raise WebSocketDisconnect(code=1006) from exc + raise + + async def attachment_close(reason: str) -> None: + close = getattr(websocket, "close", None) + if callable(close): + with suppress(Exception): + await close(code=1000, reason=reason) + + return attachment_send, attachment_close -@dataclass -class DuplexWebSocketActor: - """Own ordered WebSocket I/O queues, but no domain identity.""" - - websocket: Any - current_epoch: Callable[[], int | None] | None = None - session_closed: Callable[[], bool] | None = None - output_queue: asyncio.Queue[dict[str, object] | None] = field(default_factory=asyncio.Queue) - mailbox: asyncio.Queue[dict[str, object]] = field(default_factory=asyncio.Queue) - outbound_protocol: Any | None = None - tasks: DuplexSessionTasks = field(default_factory=DuplexSessionTasks) - closing: bool = False - close_reason: str | None = None - stale_output_dropped: int = 0 - _queued_input_events: int = 0 - - @property - def native_append_tasks(self) -> dict[asyncio.Task[bool], DuplexAppendTaskMeta]: - return self.tasks.native_append_tasks - - @property - def active_response_task(self) -> asyncio.Task[None] | None: - return self.tasks.active_response_task - - @active_response_task.setter - def active_response_task(self, task: asyncio.Task[None] | None) -> None: - self.tasks.active_response_task = task - - @property - def native_append_tail(self) -> asyncio.Task[bool] | None: - return self.tasks.native_append_tail - - @native_append_tail.setter - def native_append_tail(self, task: asyncio.Task[bool] | None) -> None: - self.tasks.native_append_tail = task - - async def enqueue_event(self, event: dict[str, object]) -> None: - if is_input_event(event.get("type")): - self._queued_input_events += 1 - await self.mailbox.put(event) - - async def next_event(self) -> dict[str, object]: - event = await self.mailbox.get() - if is_input_event(event.get("type")): - self._queued_input_events = max(0, self._queued_input_events - 1) - self.mailbox.task_done() - return event - - def has_queued_input_events(self) -> bool: - return self._queued_input_events > 0 - - async def send_json(self, payload: dict[str, object]) -> None: - await self.output_queue.put(payload) - - async def close_writer(self) -> None: - await self.output_queue.put(None) - - async def writer_loop(self) -> None: - while True: - payload = await self.output_queue.get() - try: - if payload is None: - return - raw_realtime = payload.pop("_realtime_raw", False) is True - if not raw_realtime and self._is_stale_model_output(payload): - self.stale_output_dropped += 1 - continue - try: - if raw_realtime: - await self.websocket.send_json(payload) - elif self.outbound_protocol is not None: - for projected in self.outbound_protocol.encode_outbound_event(payload): - await self.websocket.send_json(projected) - else: - await self.websocket.send_json(payload) - except (WebSocketDisconnect, RuntimeError): - return - finally: - self.output_queue.task_done() - - def _is_stale_model_output(self, payload: dict[str, object]) -> bool: - event_type = payload.get("type") - if event_type in DOMAIN_TERMINAL_EVENTS: - return False - if event_type not in MODEL_OUTPUT_EVENTS: - return False - if self.closing and event_type not in {"response.listen", "runtime.control"}: - return True - expected_epoch = self.current_epoch() if self.current_epoch is not None else None - if ( - self.session_closed is not None - and self.session_closed() - and event_type - not in { - "response.listen", - "runtime.control", - } - ): - return True - epoch = payload.get("epoch") - return isinstance(epoch, int) and isinstance(expected_epoch, int) and epoch != expected_epoch - - def track_append_task( - self, - task: asyncio.Task[bool], - *, - epoch: int, - mode: str, - final: bool, - response_bound: bool, - ) -> None: - self.tasks.track_append_task( - task, - epoch=epoch, - mode=mode, - final=final, - response_bound=response_bound, - ) - - def has_response_bound_append_tasks(self) -> bool: - return self.tasks.has_response_bound_append_tasks() - - async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: - return await self.tasks.cancel_append_tasks(timeout_s, response_bound_only=response_bound_only) - - -__all__ = [ - "DOMAIN_TERMINAL_EVENTS", - "DuplexAppendTaskMeta", - "DuplexSessionTasks", - "DuplexWebSocketActor", - "is_input_event", -] +__all__ = ["MAX_EVENT_BYTES", "CloseFn", "SendJson", "attachment_callbacks", "receive_text_with_timeout"] diff --git a/vllm_omni/entrypoints/duplex_omni.py b/vllm_omni/entrypoints/duplex_omni.py new file mode 100644 index 000000000..968b69f74 --- /dev/null +++ b/vllm_omni/entrypoints/duplex_omni.py @@ -0,0 +1,476 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""DuplexOmni: the Python API for full-duplex models. + +Sessions run inside the engine (``DuplexSessionRunner`` on the orchestrator +loop of ``DuplexOrchestrator``). This class opens / resumes / closes them and +pipes typed ``DuplexCommand`` objects in and typed ``DuplexEvent`` objects out +through a ``DuplexSessionHandle``. It holds no session state beyond the +handle registry, so the websocket handler and ``InlineDuplexClient`` are both +thin consumers of the same surface. + +Example:: + + omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True) + async with await omni.open_session({"modalities": ["audio", "text"]}) as session: + async def consume(): + async for event in session.events(): + if isinstance(event, AudioDelta): + play(event.audio) + task = asyncio.create_task(consume()) + await session.submit(AppendAudio(audio=pcm_bytes, format="pcm16", sample_rate_hz=16000)) + await session.submit(Commit()) +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import suppress +from typing import Any +from uuid import uuid4 + +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands as duplex_commands +from vllm_omni.engine.duplex.commands import DuplexCommand +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig, ResponseCreateOptions +from vllm_omni.engine.duplex.events import DuplexEvent, SessionClosed +from vllm_omni.engine.duplex.messages import ( + DuplexControlResultMessage, + DuplexSessionError, + DuplexSessionEventMessage, +) +from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine +from vllm_omni.entrypoints.async_omni import AsyncOmni + +logger = init_logger(__name__) + +_DEFAULT_CONTROL_TIMEOUT_S = 10.0 + + +class DuplexSessionHandle: + """Client-side view of one engine-resident duplex session. + + Connection-independent: a websocket may detach and a later one resume the + same handle. ``events()`` is single-consumer at any given time but may be + re-entered after the previous iterator was closed (resume). + """ + + def __init__(self, omni: DuplexOmni, session_id: str) -> None: + self._omni = omni + #: Allocated by ``DuplexOmni.open_session``; unique for the engine's lifetime. + self.session_id = session_id + self.capabilities: DuplexCapabilities = DuplexCapabilities() + self.public_session: dict[str, object] = {} + self.lease_generation: int = 0 + self._outbox: asyncio.Queue[DuplexEvent | None] = asyncio.Queue() + self._closed = False + self._close_reason: str | None = None + self._closed_event = asyncio.Event() + self._consumer_active = False + + # ---- state ---- + + @property + def closed(self) -> bool: + return self._closed + + @property + def close_reason(self) -> str | None: + return self._close_reason + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"DuplexSessionHandle(session_id={self.session_id!r}, closed={self._closed})" + + # ---- input ---- + + async def submit(self, command: DuplexCommand) -> None: + """Enqueue one command in caller order; rejections arrive as ``ErrorEvent`` on ``events()``.""" + if self._closed: + raise DuplexSessionError( + f"duplex session {self.session_id} is closed", code="session_closed", session_id=self.session_id + ) + await self._omni.engine.submit_command_async(self.session_id, command) + + async def append_audio( + self, + audio: bytes, + *, + format: str = "pcm16", + sample_rate_hz: int | None = None, + is_speech: bool | None = None, + video_frames: Sequence[str] | None = None, + duration_ms: int | None = None, + audio_end_ms: int | None = None, + hints: Mapping[str, object] | None = None, + ) -> None: + await self.submit( + duplex_commands.AppendAudio( + audio=audio, + format=format, + sample_rate_hz=sample_rate_hz, + is_speech=is_speech, + video_frames=tuple(video_frames or ()), + duration_ms=duration_ms, + audio_end_ms=audio_end_ms, + hints=dict(hints or {}), + ) + ) + + async def append_text(self, text: str) -> None: + await self.submit(duplex_commands.AppendText(text=text)) + + async def commit( + self, + *, + final: bool = True, + create_response: bool | None = None, + is_speech: bool | None = None, + ) -> None: + await self.submit(duplex_commands.Commit(final=final, create_response=create_response, is_speech=is_speech)) + + async def create_response(self, options: ResponseCreateOptions | Mapping[str, object] | None = None) -> None: + if isinstance(options, ResponseCreateOptions): + payload: dict[str, object] = { + key: value + for key, value in ( + ("instructions", options.instructions), + ("voice", options.voice), + ("output_audio_format", options.response_format), + ("temperature", options.temperature), + ("max_output_tokens", options.max_tokens), + ("speed", options.speed), + ("modalities", list(options.modalities) if options.modalities is not None else None), + ) + if value is not None + } + if options.extra_body: + payload["extra_body"] = dict(options.extra_body) + else: + payload = dict(options or {}) + await self.submit(duplex_commands.CreateResponse(options=payload)) + + async def clear_input(self) -> None: + await self.submit(duplex_commands.ClearInput()) + + async def cancel_input(self) -> None: + await self.submit(duplex_commands.CancelInput()) + + async def cancel_response(self, response_id: str | None = None) -> None: + await self.submit(duplex_commands.CancelResponse(response_id=response_id)) + + async def barge_in(self) -> None: + await self.submit(duplex_commands.BargeIn()) + + async def clear_output_audio(self, response_id: str | None = None) -> None: + await self.submit(duplex_commands.ClearOutputAudio(response_id=response_id)) + + async def signal_turn(self, event: str, payload: Mapping[str, object] | None = None) -> None: + await self.submit(duplex_commands.SignalTurn(event=event, signal_payload=dict(payload or {}))) + + async def update(self, session_patch: Mapping[str, object]) -> None: + await self.submit(duplex_commands.UpdateSession(patch=dict(session_patch))) + + async def ack_playback( + self, + played_ms: int, + *, + response_id: str | None = None, + item_id: str | None = None, + committed_ms: int | None = None, + ) -> None: + await self.submit( + duplex_commands.AckPlayback( + played_ms=int(played_ms), + committed_ms=int(committed_ms) if committed_ms is not None else None, + response_id=response_id, + item_id=item_id, + ) + ) + + async def heartbeat(self) -> None: + await self.submit(duplex_commands.Heartbeat()) + + async def create_item(self, item: Mapping[str, object], *, previous_item_id: str | None = None) -> None: + await self.submit(duplex_commands.CreateItem(item=dict(item), previous_item_id=previous_item_id)) + + async def delete_item(self, item_id: str) -> None: + await self.submit(duplex_commands.DeleteItem(item_id=item_id)) + + async def truncate_item(self, item_id: str, *, audio_end_ms: int, content_index: int = 0) -> None: + await self.submit( + duplex_commands.TruncateItem(item_id=item_id, audio_end_ms=int(audio_end_ms), content_index=content_index) + ) + + async def close(self, *, reason: str = "client_close", timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None: + if self._closed: + return + await self._omni.close_session(self.session_id, reason=reason, timeout=timeout) + + # ---- output ---- + + async def events(self) -> AsyncIterator[DuplexEvent]: + """Ordered public events; ends after ``session.closed`` / ``session.expired``.""" + if self._consumer_active: + raise RuntimeError(f"duplex session {self.session_id} already has an active events() consumer") + self._consumer_active = True + try: + while True: + if self._closed and self._outbox.empty(): + return + event = await self._outbox.get() + if event is None: + if self._closed and self._outbox.empty(): + return + continue + yield event + if event.is_terminal: + return + finally: + self._consumer_active = False + + async def wait_closed(self) -> str: + await self._closed_event.wait() + return self._close_reason or "closed" + + async def __aenter__(self) -> DuplexSessionHandle: + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + try: + await self.close() + except DuplexSessionError: + pass + + # ---- internals used by DuplexOmni ---- + + def _adopt(self, result: DuplexControlResultMessage) -> None: + if result.capabilities is not None: + self.capabilities = result.capabilities + if result.public_session: + self.public_session = dict(result.public_session) + if result.lease_generation is not None: + self.lease_generation = int(result.lease_generation) + + def _deliver(self, event: DuplexEvent) -> None: + self._outbox.put_nowait(event) + if isinstance(event, SessionClosed): # SessionExpired is a SessionClosed + self._mark_closed(event.reason or event.type) + + def _mark_closed(self, reason: str) -> None: + if self._closed: + return + self._closed = True + self._close_reason = reason + self._closed_event.set() + self._outbox.put_nowait(None) + + +class DuplexOmni(AsyncOmni): + """Async Python API for full-duplex models (see module docstring). + + Construct it like ``AsyncOmni``. The pipeline must declare + ``duplex_plugin`` and the deploy config ``session_mode: duplex``; the + engine raises at startup otherwise. There is no ``generate()``: duplex + serving is session-only. + """ + + engine: DuplexOmniEngine + + def _create_engine(self, **engine_kwargs: Any) -> DuplexOmniEngine: + from vllm_omni.entrypoints.duplex.audio_encoding import encode_audio + + return DuplexOmniEngine(duplex_audio_encoder=encode_audio, **engine_kwargs) + + def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None: + super().__init__(*args, model=model, **kwargs) + self._handles: dict[str, DuplexSessionHandle] = {} + + # ---- deployment facts ---- + + @property + def duplex_session_config(self) -> DuplexSessionRuntimeConfig: + return self.engine.duplex_session_config + + @property + def duplex_capabilities(self) -> DuplexCapabilities: + return self.engine.duplex_capabilities + + @property + def sessions(self) -> Mapping[str, DuplexSessionHandle]: + return dict(self._handles) + + def active_session_count(self) -> int: + return sum(1 for handle in self._handles.values() if not handle.closed) + + # ---- session lifecycle ---- + + @staticmethod + def _resolve_session_config( + config: DuplexSessionConfig | Mapping[str, object] | None, + *, + model: str, + ) -> DuplexSessionConfig: + """Build the normalized session config the engine will own (the one parse/normalize site). + + A ``session_id`` / ``id`` key inside a Realtime session object is + ignored: Realtime clients echo the session object back, and the server + allocates every session id. + """ + if config is None: + resolved = DuplexSessionConfig(model=model) + elif isinstance(config, DuplexSessionConfig): + resolved = config + elif isinstance(config, Mapping): + resolved = DuplexSessionConfig.from_realtime(config, model=model) + else: + raise TypeError(f"unsupported duplex session config: {type(config).__name__}") + if resolved.model is None: + resolved.model = model + return resolved.normalized() + + async def open_session( + self, + config: DuplexSessionConfig | Mapping[str, object] | None = None, + *, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexSessionHandle: + """Open an engine-resident session and return its handle. + + The session id is always allocated here (``duplex-``; never + reused, so the id alone identifies a session) and the handle is + registered before the open RPC, so the first event (``session.created``, + which announces the id) can never arrive before it exists. + """ + session_id = f"duplex-{uuid4().hex}" + session_config = self._resolve_session_config(config, model=self.model) + handle = DuplexSessionHandle(self, session_id) + self._handles[session_id] = handle + self._final_output_handler() + try: + result = await self.engine.open_session_async(session_id, session_config, timeout=timeout) + except BaseException: + if self._handles.get(session_id) is handle: + self._handles.pop(session_id, None) + # Dropping the handle only forgets the id here. The engine may still + # be admitting this session -- a control timeout unregisters the RPC + # waiter but cannot cancel the manager's open -- so a late admission + # would hold an admission slot and a Stage0 reservation until idle + # expiry. Send the compensating close; the manager applies it + # immediately if the session already landed, and records it for the + # open to honour on arrival if not. + with suppress(Exception): + await self.engine.close_session_async(session_id, reason="open_abandoned", timeout=timeout) + raise + handle._adopt(result) + return handle + + def get_session(self, session_id: str) -> DuplexSessionHandle | None: + return self._handles.get(session_id) + + async def resume_session( + self, + session_id: str, + *, + expected_lease_generation: int, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexSessionHandle: + """Engine lease resume (CAS on the lease generation); returns the existing handle.""" + handle = self._require_handle(session_id) + result = await self.engine.resume_session_async( + session_id, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ) + # The result carries the engine's current public session (state, + # epoch, turn ...), which a resumed client must see, not the open-time snapshot. + handle._adopt(result) + return handle + + async def touch_session( + self, + session_id: str, + *, + activity: str = "heartbeat", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> None: + self._require_handle(session_id) + await self.engine.touch_session_async(session_id, activity=activity, timeout=timeout) + + async def detach_session(self, session_id: str, *, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None: + """Start the engine-owned disconnect grace; expiry arrives as ``session.expired``.""" + await self.touch_session(session_id, activity="detach", timeout=timeout) + + async def close_session( + self, + session_id: str, + *, + reason: str = "client_close", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> None: + handle = self._handles.get(session_id) + if handle is None or handle.closed: + return + try: + await self.engine.close_session_async(session_id, reason=reason, timeout=timeout) + except Exception: + # The engine refused or lost the close: the handle is unusable + # either way, and it must not shadow the id in ``sessions``. + handle._mark_closed(reason) + self._handles.pop(session_id, None) + raise + # The manager emits session.closed after the stage cleanup, but the + # event travels on the output queue: wait for it so ``events()`` ends + # with the typed SessionClosed rather than an abrupt stop. + try: + await asyncio.wait_for(handle.wait_closed(), timeout=timeout) + except asyncio.TimeoutError: + logger.warning("[DuplexOmni] session.closed for %s did not arrive in time", session_id) + handle._mark_closed(reason) + self._handles.pop(session_id, None) + + async def close_all_sessions(self, *, reason: str = "shutdown") -> None: + for session_id in list(self._handles): + try: + await self.close_session(session_id, reason=reason) + except Exception: + logger.exception("[DuplexOmni] failed to close session %s", session_id) + + def _require_handle(self, session_id: str) -> DuplexSessionHandle: + handle = self._handles.get(session_id) + if handle is None or handle.closed: + raise DuplexSessionError( + f"unknown or closed duplex session: {session_id}", code="unknown_session", session_id=session_id + ) + return handle + + # ---- engine output routing ---- + + def _route_engine_message(self, msg: object) -> bool: + if not isinstance(msg, DuplexSessionEventMessage): + return False + handle = self._handles.get(msg.session_id) + if handle is None: + logger.debug("[DuplexOmni] dropping event for unknown session %s", msg.session_id) + return True + handle._deliver(msg.event) + if handle.closed: + self._handles.pop(msg.session_id, None) + return True + + def _on_engine_dead(self, error: str) -> None: + for handle in list(self._handles.values()): + handle._mark_closed(f"engine_dead: {error}") + self._handles.clear() + + def shutdown(self, timeout: float | None = None) -> None: + for handle in list(self._handles.values()): + handle._mark_closed("shutdown") + self._handles.clear() + super().shutdown(timeout) + + +__all__ = ["DuplexOmni", "DuplexSessionHandle"] diff --git a/vllm_omni/entrypoints/duplex_request_client.py b/vllm_omni/entrypoints/duplex_request_client.py deleted file mode 100644 index d9e70f672..000000000 --- a/vllm_omni/entrypoints/duplex_request_client.py +++ /dev/null @@ -1,383 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import time -from collections.abc import Callable, MutableMapping -from typing import Protocol - -from vllm_omni.engine.duplex.contracts import ( - duplex_data_plane_request_info, - duplex_resource_request_id, -) -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.messages import ErrorMessage, OutputMessage -from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import get_duplex_output_decision - - -class DuplexEnginePort(Protocol): - """Engine control methods required by the experimental request client. - - The protocol keeps request preregistration and output collection testable - without coupling this component to ``AsyncOmniEngine``. - """ - - async def open_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def append_duplex_input_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def signal_duplex_turn_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def close_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def touch_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def resume_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - -class DuplexRequestOutputPort: - """Narrow output-processing surface owned by the AsyncOmni entrypoint.""" - - def __init__( - self, - *, - request_states: MutableMapping[str, ClientRequestState], - num_stages: int, - log_stats: bool, - start_output_handler: Callable[[], object], - process_single_result: Callable[..., OmniRequestOutput | None], - ) -> None: - self.request_states = request_states - self.num_stages = num_stages - self.log_stats = log_stats - self._start_output_handler = start_output_handler - self._process_single_result = process_single_result - - def start_output_handler(self) -> None: - self._start_output_handler() - - def process_single_result(self, *args, **kwargs) -> OmniRequestOutput | None: - return self._process_single_result(*args, **kwargs) - - -class DuplexRequestClient: - """Owns duplex request identity, preregistration and output collection.""" - - def __init__(self, engine: DuplexEnginePort, output_port: DuplexRequestOutputPort) -> None: - self.engine = engine - self.output_port = output_port - - @staticmethod - def request_info(result: dict[str, object]) -> tuple[str | None, int | None]: - return duplex_data_plane_request_info(result) - - async def open( - self, - session_id: str, - *, - session_mode: str, - capabilities: dict[str, object] | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - # Lifecycle expiry messages share the engine output channel. Start its - # sole dispatcher before open so an idle session can be reaped even if - # it never submits an append request. - self.output_port.start_output_handler() - kwargs: dict[str, object] = { - "session_mode": session_mode, - "capabilities": capabilities, - "session_config": session_config, - "fence": fence, - "timeout": timeout, - } - if runtime_config is not None: - kwargs["runtime_config"] = runtime_config - return await self.engine.open_duplex_session_async(session_id, **kwargs) - - async def append( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None, - final: bool, - expected_epoch: int | None, - fence: DuplexFence, - timeout: float | None, - collect_outputs: bool, - ) -> dict[str, object]: - kwargs: dict[str, object] = { - "mode": mode, - "payload": payload, - "final": final, - "expected_epoch": expected_epoch, - "timeout": timeout, - "fence": fence, - } - if operation_id is not None: - kwargs["operation_id"] = operation_id - - expected_request_id = duplex_resource_request_id(fence, "stage0") - request_state = self.output_port.request_states.get(expected_request_id) - created_request_state = request_state is None - if request_state is None: - request_state = ClientRequestState(expected_request_id) - self.output_port.request_states[expected_request_id] = request_state - if request_state.metrics is None: - wall_start_ts = time.time() - request_state.metrics = OrchestratorMetrics( - self.output_port.num_stages, - self.output_port.log_stats, - wall_start_ts, - max(0, self.output_port.num_stages - 1), - ) - request_state.request_arrival_ts = wall_start_ts - self.output_port.start_output_handler() - try: - result = await self.engine.append_duplex_input_async(session_id, **kwargs) - except BaseException: - if created_request_state and request_state.queue.empty(): - self.output_port.request_states.pop(expected_request_id, None) - raise - - request_id, response_stage_id = duplex_data_plane_request_info(result) - if request_id is None: - if created_request_state and request_state.queue.empty(): - self.output_port.request_states.pop(expected_request_id, None) - return result - if request_id != expected_request_id: - if created_request_state and request_state.queue.empty(): - self.output_port.request_states.pop(expected_request_id, None) - raise RuntimeError( - f"duplex data-plane request id mismatch: expected {expected_request_id!r}, got {request_id!r}" - ) - if not collect_outputs: - return result - outputs = await self.collect_outputs( - request_id, - request_state, - response_stage_id=response_stage_id, - timeout=timeout, - ) - if outputs: - result = dict(result) - result["data_plane_outputs"] = outputs - return result - - async def collect_registered_outputs( - self, - request_id: str, - *, - response_stage_id: int | None, - timeout: float | None, - ) -> list[OmniRequestOutput]: - self.output_port.start_output_handler() - request_state = self.output_port.request_states.get(request_id) - if request_state is None: - return [] - return await self.collect_outputs( - request_id, - request_state, - response_stage_id=response_stage_id, - timeout=timeout, - ) - - async def signal( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - timeout: float | None, - ) -> dict[str, object]: - kwargs: dict[str, object] = {"event": event, "fence": fence, "timeout": timeout} - if next_fence is not None: - kwargs["next_fence"] = next_fence - if session_config is not None: - kwargs["session_config"] = session_config - if runtime_config is not None: - kwargs["runtime_config"] = runtime_config - result = await self.engine.signal_duplex_turn_async(session_id, **kwargs) - if event in {"barge_in", "input.cancel", "response.cancel"}: - self.output_port.request_states.pop(duplex_resource_request_id(fence, "stage0"), None) - return result - - async def close( - self, - session_id: str, - *, - reason: str, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - result = await self.engine.close_duplex_session_async( - session_id, - reason=reason, - fence=fence, - timeout=timeout, - ) - self.output_port.request_states.pop(duplex_resource_request_id(fence, "stage0"), None) - return result - - async def touch( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None, - ) -> dict[str, object]: - return await self.engine.touch_duplex_session_async( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ) - - async def resume( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None, - ) -> dict[str, object]: - return await self.engine.resume_duplex_session_async( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ) - - async def collect_outputs( - self, - request_id: str, - request_state: ClientRequestState, - *, - response_stage_id: int | None, - timeout: float | None, - ) -> list[OmniRequestOutput]: - deadline = None if timeout is None else time.monotonic() + timeout - wall_start_ts = request_state.request_arrival_ts or time.time() - final_stage_id = response_stage_id if response_stage_id is not None else max(0, self.output_port.num_stages - 1) - num_stages = max(self.output_port.num_stages, final_stage_id + 1) - metrics = getattr(request_state, "metrics", None) or OrchestratorMetrics( - num_stages, - self.output_port.log_stats, - wall_start_ts, - final_stage_id, - ) - stage_event_cursor = metrics.stage_event_cursor(request_id) - request_start_ts = {request_id: wall_start_ts} - outputs: list[OmniRequestOutput] = [] - while True: - remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) - if remaining == 0.0: - break - try: - message = await asyncio.wait_for(request_state.queue.get(), timeout=remaining) - except asyncio.TimeoutError: - break - if isinstance(message, ErrorMessage): - raise RuntimeError(message.error) - if not isinstance(message, OutputMessage): - continue - engine_outputs = message.engine_outputs - is_response_stage = response_stage_id is None or message.stage_id >= response_stage_id - is_direct_response = self.is_direct_response(engine_outputs) - output_to_collect = None - if isinstance(engine_outputs, OmniRequestOutput): - output_to_collect = engine_outputs - if message.metrics is not None: - existing_metrics = dict(output_to_collect.metrics) - metrics.process_stage_metrics( - result={"metrics": message.metrics}, - stage_type="", - stage_id=message.stage_id, - req_id=request_id, - engine_outputs=engine_outputs, - finished=message.finished, - final_output_type=engine_outputs.final_output_type, - output_to_yield=output_to_collect, - event_cursor=stage_event_cursor, - ) - generated_metrics = dict(output_to_collect.metrics) - existing_stage_metrics = existing_metrics.get("stage_metrics") - generated_stage_metrics = generated_metrics.get("stage_metrics") - if isinstance(existing_stage_metrics, dict) and isinstance(generated_stage_metrics, dict): - generated_metrics["stage_metrics"] = { - **existing_stage_metrics, - **generated_stage_metrics, - } - output_to_collect.metrics = { - **existing_metrics, - **generated_metrics, - } - elif is_response_stage: - output_to_collect = self.output_port.process_single_result( - message, - message.stage_id, - metrics, - request_start_ts, - wall_start_ts, - final_stage_id, - stage_event_cursor, - ) - if output_to_collect is not None and (is_response_stage or is_direct_response): - if message.finished: - output_to_collect.finished = True - outputs.append(output_to_collect) - if message.finished or outputs: - break - return outputs - - @classmethod - def is_direct_response(cls, output: object) -> bool: - multimodal_output = cls.multimodal_output(output) - if not multimodal_output: - return False - return multimodal_output.get("duplex_direct_response") is True or multimodal_output.get( - "duplex_native_decision" - ) in {"listen", "speak"} - - @classmethod - def multimodal_output(cls, output: object) -> dict[str, object]: - if isinstance(output, OmniRequestOutput): - decision = get_duplex_output_decision(output) - if decision is not None: - return dict(decision.metadata) - private_output = getattr(output, "_multimodal_output", None) - multimodal_output = output.multimodal_output - if isinstance(multimodal_output, dict) and multimodal_output: - return multimodal_output - inner_output = getattr(output, "request_output", None) - if inner_output is not None and inner_output is not output: - inner_multimodal_output = cls.multimodal_output(inner_output) - if inner_multimodal_output: - return inner_multimodal_output - if isinstance(private_output, dict): - return private_output - return multimodal_output if isinstance(multimodal_output, dict) else {} - multimodal_output = getattr(output, "multimodal_output", None) - if isinstance(multimodal_output, dict): - return multimodal_output - outputs = getattr(output, "outputs", None) - completion = outputs[0] if isinstance(outputs, list) and outputs else None - multimodal_output = getattr(completion, "multimodal_output", None) if completion is not None else None - return multimodal_output if isinstance(multimodal_output, dict) else {} - - -__all__ = ["DuplexEnginePort", "DuplexRequestClient", "DuplexRequestOutputPort"] diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index cf4da3b7a..ca9fc0196 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -7,12 +7,13 @@ import copy import time import uuid from collections.abc import Callable, Generator, Iterable, Sequence -from typing import TYPE_CHECKING, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, overload from tqdm.auto import tqdm from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import OutputMessage from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni_base import OmniBase @@ -28,6 +29,9 @@ logger = init_logger(__name__) class Omni(OmniBase): """Synchronous entrypoint for offline generation.""" + def _create_engine(self, **engine_kwargs: Any) -> AsyncOmniEngine: + return AsyncOmniEngine(**engine_kwargs) + def _maybe_force_final_only_for_llm_stages( self, sampling_params_list: Sequence[OmniSamplingParams], diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index ead0e6f1d..79217ca8b 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -19,13 +19,13 @@ from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.stage_config import merge_sampling_constraints -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, OutputMessage, StageMetricsMessage, ) +from vllm_omni.engine.omni_engine_base import OmniEngineBase from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e @@ -66,7 +66,7 @@ class OmniEngineDeadError(EngineDeadError): self.error_stage_id = error_stage_id -def _weak_shutdown_engine(engine: AsyncOmniEngine) -> None: +def _weak_shutdown_engine(engine: OmniEngineBase) -> None: """Best-effort engine cleanup for GC finalization.""" try: engine.shutdown() @@ -206,13 +206,13 @@ class OmniBase(PDDisaggregationMixin): self.tts_batch_max_items: int = kwargs.pop("tts_batch_max_items", 32) logger.info("[%s] Initializing with model %s", self.__class__.__name__, model) - # Construct transfer_metrics first so we can hand it to AsyncOmniEngine + # Construct transfer_metrics first so we can hand it to the engine # (which forwards it to the Orchestrator background thread for - # TX-side emit; see Orchestrator._forward_to_next_stage). + # TX-side emit; see OrchestratorBase._forward_to_next_stage). self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats) self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) st = time.time() - self.engine = AsyncOmniEngine( + self.engine = self._create_engine( model=model, init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, @@ -224,8 +224,8 @@ class OmniBase(PDDisaggregationMixin): self._shutdown_called = False self._weak_finalizer = weakref.finalize(self, _weak_shutdown_engine, self.engine) et = time.time() - logger.info("[%s] AsyncOmniEngine initialized in %.2f seconds", self.__class__.__name__, et - st) - # Authoritative: ``AsyncOmniEngine`` resolves (pipeline + deploy YAML + + logger.info("[%s] %s initialized in %.2f seconds", self.__class__.__name__, type(self.engine).__name__, et - st) + # Authoritative: the engine resolves (pipeline + deploy YAML + # CLI overrides) through ``StageConfigFactory`` and stores the final # value on ``engine.async_chunk``; mirror it here so ``--no-async-chunk`` # (explicit ``False``) is not fallen-back-through by ``or``. @@ -253,6 +253,10 @@ class OmniBase(PDDisaggregationMixin): # PD disaggregation state (detects if a prefill/decode stage pair is configured) self._init_pd_state() + def _create_engine(self, **engine_kwargs: Any) -> OmniEngineBase: + """Construct this entrypoint's engine (``AsyncOmniEngine`` for turn-based use).""" + raise NotImplementedError + @property def num_stages(self) -> int: return self.engine.num_stages diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 36b7544ee..376a9ccd2 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -87,9 +87,9 @@ from vllm_omni.config.endpoint_policy import ( shutdown_unsupported_routes, ) from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni -from vllm_omni.entrypoints.duplex.capability import should_enable_duplex_endpoint from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime +from vllm_omni.entrypoints.duplex_omni import DuplexOmni from vllm_omni.entrypoints.openai import app_state as openai_app_state from vllm_omni.entrypoints.openai.app_state import ( ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, @@ -440,6 +440,23 @@ async def build_async_omni( yield async_omni +def _is_duplex_model(model: str, kwargs: dict[str, Any]) -> bool: + """Whether ``model``'s pipeline declares a ``duplex_plugin`` (served by ``DuplexOmni``). + + Resolution errors propagate: a duplex model whose pipeline or deploy + config cannot be resolved must fail startup rather than silently start a + turn-based server. + """ + from vllm_omni.config.config_factory import StageConfigFactory + + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=bool(kwargs.get("trust_remote_code")), + deploy_config_path=kwargs.get("deploy_config"), + ) + return bool(pipeline_config is not None and getattr(pipeline_config, "duplex_plugin", None)) + + @asynccontextmanager async def build_async_omni_from_stage_config( args: TrackingNamespace, @@ -501,7 +518,12 @@ async def build_async_omni_from_stage_config( kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) - async_omni = AsyncOmni(model=model, **kwargs) + if _is_duplex_model(model, kwargs): + # A duplex model is always served in duplex mode: sessions over + # /v1/realtime?duplex=1, no turn-based HTTP routes. + async_omni = DuplexOmni(model=model, **kwargs) + else: + async_omni = AsyncOmni(model=model, **kwargs) # # Don't keep the dummy data in memory # await async_llm.reset_mm_cache() @@ -512,6 +534,134 @@ async def build_async_omni_from_stage_config( async_omni.shutdown() +async def _init_duplex_app_state( + engine_client: DuplexOmni, + state: State, + args: Namespace, + base_model_paths: list[BaseModelPath], + vllm_config: Any, + request_logger: RequestLogger | None, +) -> None: + """Minimal app state for a duplex server: only the surfaces a duplex session backs.""" + state.vllm_config = vllm_config + state.diffusion_engine = None + state.openai_serving_models = OpenAIServingModels( + engine_client=engine_client, # type: ignore[arg-type] + base_model_paths=base_model_paths, + lora_modules=None, + ) + state.serving_tokenization = None + state.serving_tokens = None + # Replaced by the chat init below when the model serves chat. + state.online_renderer = None + for attribute in ( + "openai_serving_chat_batch", + "openai_serving_completion", + "openai_serving_responses", + "openai_serving_embedding", + "openai_serving_pooling", + "openai_serving_classification", + "openai_serving_scores", + "openai_serving_transcription", + "openai_serving_translation", + "openai_serving_speech", + "openai_serving_audio_generate", + "openai_serving_video", + "openai_streaming_speech", + "openai_streaming_video", + "openai_streaming_video_output", + "openai_serving_realtime", + "openai_serving_realtime_robot", + "anthropic_serving_messages", + ): + setattr(state, attribute, None) + state.openai_serving_duplex = OmniDuplexSessionHandler(duplex_omni=engine_client) + # One engine, both surfaces. ``DuplexOmni`` extends ``AsyncOmni``, so the + # ordinary chat service runs on it unchanged: a chat request is a turn-based + # request on the same pipeline, not a session, and costs no admission slot. + state.openai_serving_chat = await _init_duplex_chat(engine_client, state, args, request_logger) + state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) + state.server_load_metrics = 0 + if state.openai_serving_chat is not None: + logger.info( + "Duplex mode: serving %s over /v1/realtime?duplex=1 and /v1/chat/completions", + engine_client.model, + ) + else: + logger.info( + "Duplex mode: serving %s over /v1/realtime?duplex=1 only " + "(/v1/chat/completions unavailable: the model does not declare supports_chat_completions)", + engine_client.model, + ) + + +async def _init_duplex_chat( + engine_client: DuplexOmni, + state: State, + args: Namespace, + request_logger: RequestLogger | None, +) -> OmniOpenAIServingChat | None: + """The ordinary chat service, on a duplex engine, when the model allows it. + + Gated on ``DuplexCapabilities.supports_chat_completions`` so the decision + stays the model's: a duplex model that should not answer chat requests says + so in its plugin, and the route reports "not available" rather than + answering badly. ``endpoint_restrictions`` remains the per-deployment + opt-out on top of this. + """ + if not engine_client.duplex_capabilities.supports_chat_completions: + return None + supported_tasks: set[str] = {"generate"} + if hasattr(engine_client, "get_supported_tasks"): + supported_tasks = set(await engine_client.get_supported_tasks()) + if "generate" not in supported_tasks: + return None + + resolved_chat_template = load_chat_template(args.chat_template) + if resolved_chat_template is None: + try: + tokenizer = await engine_client.get_tokenizer() + except Exception as exc: + logger.debug("Could not inspect tokenizer chat_template before duplex chat init: %s", exc) + tokenizer = None + if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: + resolved_chat_template = _load_model_chat_template_json(args.model) + + state.online_renderer = OnlineRenderer( + model_config=engine_client.model_config, + renderer=engine_client.renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + ) + return OmniOpenAIServingChat( + engine_client=engine_client, + models=state.openai_serving_models, + response_role=args.response_role, + online_renderer=state.online_renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=args.default_chat_template_kwargs, + trust_request_chat_template=args.trust_request_chat_template, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, + ) + + async def omni_init_app_state( engine_client: EngineClient, state: State, @@ -562,6 +712,13 @@ async def omni_init_app_state( state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None model_name = served_model_names[0] if served_model_names else args.model + # Duplex mode: a duplex model is served through DuplexOmni only. Sessions + # run over /v1/realtime?duplex=1; every turn-based HTTP route reports + # "not available". + if isinstance(engine_client, DuplexOmni): + await _init_duplex_app_state(engine_client, state, args, base_model_paths, vllm_config, request_logger) + return + # Pure Diffusion mode: use simplified initialization logic if is_pure_diffusion: state.vllm_config = None @@ -942,17 +1099,6 @@ async def omni_init_app_state( else None ) state.openai_serving_duplex = None - if state.openai_serving_chat is not None and should_enable_duplex_endpoint( - state.stage_configs, - config_path=getattr(engine_client, "config_path", None) or getattr(args, "deploy_config", None), - ): - state.openai_serving_duplex = OmniDuplexSessionHandler( - chat_service=state.openai_serving_chat, - served_model_name=model_name, - log_stats=state.log_stats, - duplex_session_config=getattr(engine_client, "duplex_session_config", None), - serving_runtime_adapter_path=getattr(engine_client, "duplex_serving_adapter_path", None), - ) state.openai_serving_realtime = OpenAIServingRealtime( engine_client=engine_client, models=state.openai_serving_models, @@ -1477,14 +1623,7 @@ async def streaming_video_output(websocket: WebSocket): @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" - # Hold real clients until the startup duplex warmup finishes (the warmup - # connection marks itself with vllm_omni_warmup=1 and passes through). - warmup_done = getattr(websocket.app.state, "duplex_warmup_done", None) - if warmup_done is not None and not warmup_done.is_set() and websocket.query_params.get("vllm_omni_warmup") != "1": - try: - await asyncio.wait_for(warmup_done.wait(), timeout=120) - except (TimeoutError, asyncio.TimeoutError): - logger.warning("Duplex warmup still running after 120 s; admitting the client anyway.") + await _wait_for_duplex_warmup(websocket) duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None) duplex_query = websocket.query_params.get("duplex") use_duplex_realtime = duplex_handler is not None and ( @@ -1504,6 +1643,32 @@ async def realtime_websocket(websocket: WebSocket): await connection.handle_connection() +async def _wait_for_duplex_warmup(websocket: WebSocket) -> None: + """Hold real clients until the startup duplex warmup finishes. + + The warmup connection marks itself with ``vllm_omni_warmup=1`` and passes through. + """ + warmup_done = getattr(websocket.app.state, "duplex_warmup_done", None) + if warmup_done is not None and not warmup_done.is_set() and websocket.query_params.get("vllm_omni_warmup") != "1": + try: + await asyncio.wait_for(warmup_done.wait(), timeout=120) + except (TimeoutError, asyncio.TimeoutError): + logger.warning("Duplex warmup still running after 120 s; admitting the client anyway.") + + +@router.websocket("/v1/duplex") +async def duplex_websocket(websocket: WebSocket): + """Alias of ``/v1/realtime?duplex=1``: the same Realtime duplex session protocol.""" + await _wait_for_duplex_warmup(websocket) + handler = getattr(websocket.app.state, "openai_serving_duplex", None) + if handler is None: + await websocket.accept() + await websocket.send_json({"type": "error", "error": "Duplex API is not available", "code": "unsupported"}) + await websocket.close() + return + await handler.handle_realtime_session(websocket) + + @router.websocket("/v1/realtime/robot/openpi") async def realtime_robot_openpi(websocket: WebSocket): """WebSocket endpoint for robot policy inference via OpenPI messages.""" @@ -1534,18 +1699,6 @@ async def realtime_robot_openpi(websocket: WebSocket): await connection.handle_connection() -@router.websocket("/v1/duplex") -async def duplex_websocket(websocket: WebSocket): - """WebSocket endpoint for vLLM-Omni duplex session control.""" - handler = getattr(websocket.app.state, "openai_serving_duplex", None) - if handler is None: - await websocket.accept() - await websocket.send_json({"type": "error", "error": "Duplex API is not available", "code": "unsupported"}) - await websocket.close() - return - await handler.handle_session(websocket) - - # RL Rollout serving (RFC #3747, P0) diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index a13318ad4..b82842c42 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -21,7 +21,6 @@ LONGEVITY: See ``openai/README.md`` (utils vs helpers, no overlap). """ -import math from io import BytesIO import numpy as np @@ -36,6 +35,9 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( CreateAudio, ) +# Re-exported: serving_speech and the audio tests import it from here. +from vllm_omni.utils.audio_resample import StreamingAudioResampler # noqa: F401 + try: import soundfile except ImportError: @@ -44,155 +46,6 @@ except ImportError: logger = init_logger(__name__) -class StreamingAudioResampler: - """Stateful polyphase resampler for streaming mono float audio. - - Retains filter state so output is invariant to input chunk boundaries. - """ - - _half_filter_width = 10 - _kaiser_beta = 5.0 - _max_polyphase_factor = 2_048 - - def __init__(self, source_rate: int, target_rate: int): - if source_rate <= 0 or target_rate <= 0: - raise ValueError("Audio sample rates must be positive") - self.source_rate = source_rate - self.target_rate = target_rate - rate_gcd = math.gcd(source_rate, target_rate) - self._up = target_rate // rate_gcd - self._down = source_rate // rate_gcd - if max(self._up, self._down) > self._max_polyphase_factor: - raise ValueError("source and target sample rates have an unsupported resampling ratio") - self._half_len, self._phase_kernels = self._design_polyphase_filter() - self._history_samples = self._phase_kernels.shape[1] - 1 - self.reset() - - def _design_polyphase_filter(self) -> tuple[int, np.ndarray]: - if self._up == self._down: - return 0, np.ones((1, 1), dtype=np.float32) - - max_rate = max(self._up, self._down) - half_len = self._half_filter_width * max_rate - offsets = np.arange(-half_len, half_len + 1, dtype=np.float64) - cutoff = 1.0 / max_rate - taps = cutoff * np.sinc(cutoff * offsets) - taps *= np.kaiser(taps.size, self._kaiser_beta) - taps /= np.sum(taps) - taps *= self._up - taps = np.ascontiguousarray(taps, dtype=np.float32) - - phases = [taps[phase :: self._up] for phase in range(self._up)] - phase_width = max(phase.size for phase in phases) - kernels = np.zeros((self._up, phase_width), dtype=np.float32) - for phase_index, phase in enumerate(phases): - kernels[phase_index, -phase.size :] = phase[::-1] - return half_len, kernels - - @property - def scratch_bytes(self) -> int: - pending_samples = self._ceil_div(self._input_samples * self._up, self._down) - self._output_samples - return max(0, pending_samples) * np.dtype(np.float32).itemsize - - @staticmethod - def _ceil_div(numerator: int, denominator: int) -> int: - return -(-numerator // denominator) - - def _render_outputs( - self, - combined: np.ndarray, - *, - first_output: int, - output_count: int, - input_start: int, - ) -> np.ndarray: - if output_count <= 0: - return np.empty(0, dtype=np.float32) - - output_indexes = np.arange(first_output, first_output + output_count, dtype=np.int64) - filter_indexes = output_indexes * self._down + self._half_len - source_indexes = filter_indexes // self._up - phases = filter_indexes % self._up - windows = np.lib.stride_tricks.sliding_window_view( - combined, - self._phase_kernels.shape[1], - ) - selected_windows = windows[source_indexes - input_start] - return np.sum( - selected_windows * self._phase_kernels[phases], - axis=1, - dtype=np.float32, - ) - - def _push(self, chunk: np.ndarray) -> np.ndarray: - if chunk.size == 0: - return np.empty(0, dtype=np.float32) - if self._flushed: - raise RuntimeError("cannot process audio after the resampler has been finalized; call reset first") - - combined = np.concatenate((self._history, chunk)) - next_input_samples = self._input_samples + chunk.size - stable_output_samples = max( - 0, - self._ceil_div(next_input_samples * self._up - self._half_len, self._down), - ) - output = self._render_outputs( - combined, - first_output=self._output_samples, - output_count=stable_output_samples - self._output_samples, - input_start=self._input_samples, - ) - if self._history_samples: - self._history = combined[-self._history_samples :].copy() - self._input_samples = next_input_samples - self._output_samples = stable_output_samples - return output - - def _flush(self) -> np.ndarray: - if self._flushed: - return np.empty(0, dtype=np.float32) - - total_output_samples = self._ceil_div(self._input_samples * self._up, self._down) - output_count = total_output_samples - self._output_samples - if output_count: - last_output = total_output_samples - 1 - last_source_index = (last_output * self._down + self._half_len) // self._up - right_padding = max(0, last_source_index - self._input_samples + 1) - combined = np.concatenate((self._history, np.zeros(right_padding, dtype=np.float32))) - output = self._render_outputs( - combined, - first_output=self._output_samples, - output_count=output_count, - input_start=self._input_samples, - ) - else: - output = np.empty(0, dtype=np.float32) - self._output_samples = total_output_samples - self._flushed = True - return output - - def process(self, audio: np.ndarray, *, final: bool = False) -> np.ndarray: - chunk = np.asarray(audio, dtype=np.float32) - if chunk.ndim != 1: - raise ValueError(f"Streaming audio resampling only supports mono audio, got shape {chunk.shape}") - chunk = np.ascontiguousarray(chunk, dtype=np.float32) - output = self._push(chunk) - if not final: - return output - tail = self._flush() - if not output.size: - return tail - if not tail.size: - return output - return np.concatenate((output, tail)) - - def reset(self) -> None: - self._history = np.zeros(self._history_samples, dtype=np.float32) - self._input_samples = 0 - self._output_samples = 0 - self._flushed = False - - class AudioMixin: """Mixin class to add audio-related utilities.""" diff --git a/vllm_omni/experimental/fullduplex/README.md b/vllm_omni/experimental/fullduplex/README.md index b6c918fd6..b710924f1 100644 --- a/vllm_omni/experimental/fullduplex/README.md +++ b/vllm_omni/experimental/fullduplex/README.md @@ -16,16 +16,20 @@ video_stacking.py camera-frame tiling for omni duplex video input To run JoyVL, see [`recipes/JD/JoyAI-VL-Interaction.md`](../../../recipes/JD/JoyAI-VL-Interaction.md). -The MiniCPM-o 4.5 and PersonaPlex native full-duplex runtimes graduated out -of this package. They now live in the stable tree: +The native full-duplex runtimes graduated out of this package. They now live +in the stable tree, where sessions are engine-resident (RFC +[vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)): ```text -vllm_omni/engine/duplex/ engine control plane, sessions, leases -vllm_omni/entrypoints/duplex/ WebSocket serving and Realtime projection -vllm_omni/entrypoints/duplex_request_client.py request/output lifecycle -vllm_omni/model_executor/models/minicpmo_4_5/duplex/ MiniCPM adapter -vllm_omni/model_executor/models/personaplex/duplex/ PersonaPlex adapter -vllm_omni/model_executor/models/nemotron_voicechat/duplex/ Nemotron VoiceChat adapter +vllm_omni/engine/duplex/ sessions, manager, runner, leases, typed contract +vllm_omni/engine/duplex_omni_engine.py DuplexOmniEngine (session message surface) +vllm_omni/engine/duplex_orchestrator.py DuplexOrchestrator (hosts the session manager) +vllm_omni/entrypoints/duplex_omni.py DuplexOmni + DuplexSessionHandle (Python API) +vllm_omni/entrypoints/duplex/ WebSocket transport (attachments, resume, replay) +vllm_omni/clients/duplex.py DuplexClient / DuplexClientBase +vllm_omni/model_executor/models/minicpmo_4_5/duplex/ MiniCPM-o 4.5 DuplexModelPlugin +vllm_omni/model_executor/models/personaplex/duplex/ PersonaPlex (pre-framework, not ported yet) +vllm_omni/model_executor/models/nemotron_voicechat/duplex/ Nemotron VoiceChat (pre-framework, not ported yet) vllm_omni/model_executor/duplex_sampling.py AR-runner sampling hook helper vllm_omni/outputs/duplex.py typed output decision envelope ``` @@ -51,6 +55,9 @@ only model policy. 3. Promote a helper from a model package up into `core/` only once a second model actually needs it. -For production serving, prefer the stable plugin seams instead -(`duplex_serving_adapter` / `duplex_runtime_extension` dotted strings in the -model's `pipeline.py`), as MiniCPM-o 4.5 and PersonaPlex do. +For production serving, prefer the stable plugin seam instead: implement one +`vllm_omni.engine.duplex.plugin.DuplexModelPlugin` and name it in the model's +`pipeline.py` as `duplex_plugin`, as MiniCPM-o 4.5 does. PersonaPlex and +Nemotron VoiceChat still carry their pre-framework duplex code and are ported +to the plugin contract in follow-up PRs. The contract is documented in +[`docs/design/fullduplex.md`](../../../docs/design/fullduplex.md). diff --git a/vllm_omni/model_executor/duplex_sampling.py b/vllm_omni/model_executor/duplex_sampling.py index 7dbc185d1..dd48431f6 100644 --- a/vllm_omni/model_executor/duplex_sampling.py +++ b/vllm_omni/model_executor/duplex_sampling.py @@ -19,7 +19,6 @@ class DuplexSamplingRow: row_idx: int request_id: str session_id: str | None - incarnation: int seq: int | None payload: dict[str, object] | None max_tokens: int | None @@ -73,10 +72,6 @@ class DuplexSamplingHelper: session_id = duplex.get("session_id") if not isinstance(session_id, str) or not session_id: session_id = None - try: - incarnation = int(duplex.get("incarnation", 0)) - except (TypeError, ValueError): - incarnation = 0 try: seq = int(duplex.get("seq")) except (TypeError, ValueError): @@ -95,7 +90,6 @@ class DuplexSamplingHelper: row_idx=row_idx, request_id=req_id, session_id=session_id, - incarnation=incarnation, seq=seq, payload=payload, max_tokens=max_tokens if max_tokens > 0 else None, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py index b681ad51e..e38cc12ef 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from .adapter import MiniCPMO45ClientRuntimeConfigError, MiniCPMO45NativeDuplexServingAdapter from .input import MiniCPMO45PcmAppendBuffer +from .plugin import MiniCPMO45ClientRuntimeConfigError, MiniCPMO45DuplexPlugin from .policy import MiniCPMO45DuplexPolicy from .stage0 import MiniCPMO45Stage0DuplexRuntime __all__ = [ - "MiniCPMO45DuplexPolicy", - "MiniCPMO45NativeDuplexServingAdapter", "MiniCPMO45ClientRuntimeConfigError", + "MiniCPMO45DuplexPlugin", + "MiniCPMO45DuplexPolicy", "MiniCPMO45PcmAppendBuffer", "MiniCPMO45Stage0DuplexRuntime", ] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py deleted file mode 100644 index 758beafc9..000000000 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py +++ /dev/null @@ -1,335 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import base64 -from copy import deepcopy -from typing import Any - -import numpy as np -from vllm.multimodal.media import MediaConnector - -from vllm_omni.entrypoints.duplex.protocol import DuplexSessionConfig, native_duplex_opt_in -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - ServingRuntimeConfigError, - reject_changed_runtime_value, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy - - -class MiniCPMO45ClientRuntimeConfigError(ServingRuntimeConfigError): - pass - - -class MiniCPMO45NativeDuplexServingAdapter: - """Serving-side MiniCPM-o 4.5 native duplex session preparation. - - The generic duplex WebSocket handler should not let client-supplied local - paths reach workers. This adapter follows the existing media connector - boundary: serving resolves client media URIs, then workers receive only - normalized PCM payloads. Server-owned model assets remain local paths here. - """ - - PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( - { - "duplex_stage_sampling_params", - "duplex_stage_max_tokens", - "duplex_stage0_max_tokens", - "duplex_scheduler_token_id", - "duplex_first_append_context_tokens", - "ref_audio_data", - "ref_audio_format", - "ref_audio_sample_rate_hz", - "initial_user_text", - } - ) - - @classmethod - def is_enabled(cls, config: DuplexSessionConfig) -> bool: - return native_duplex_opt_in(config.extra_body) is True - - @classmethod - def validate_client_config(cls, config: DuplexSessionConfig) -> None: - cls.validate_client_extra_body(config.extra_body) - - @classmethod - def validate_client_extra_body(cls, extra_body: object) -> None: - if not isinstance(extra_body, dict): - return - private_keys = sorted(cls.PRIVATE_RUNTIME_CONFIG_KEYS.intersection(extra_body)) - if private_keys: - raise MiniCPMO45ClientRuntimeConfigError( - "native duplex runtime configuration is server-owned: " + ", ".join(private_keys) - ) - - @classmethod - def runtime_config_for_update( - cls, - config: DuplexSessionConfig, - current: object, - ) -> dict[str, object]: - runtime_config = deepcopy(dict(current)) if isinstance(current, dict) else {} - reject_changed_runtime_value( - config.instructions, - runtime_config.get("instructions"), - message="instructions cannot be changed after the session is created", - code="instructions_update_unsupported", - error_cls=MiniCPMO45ClientRuntimeConfigError, - ) - stage_max_tokens = runtime_config.get("duplex_stage_max_tokens") - stage_max_tokens = deepcopy(stage_max_tokens) if isinstance(stage_max_tokens, dict) else {} - stage_max_tokens["0"] = ( - config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 - ) - stage_max_tokens.setdefault("1", 8192) - runtime_config["duplex_stage_max_tokens"] = stage_max_tokens - - stage_sampling = runtime_config.get("duplex_stage_sampling_params") - stage_sampling = deepcopy(stage_sampling) if isinstance(stage_sampling, dict) else {} - stage0 = stage_sampling.get("0") - stage0 = deepcopy(stage0) if isinstance(stage0, dict) else {} - stage0["temperature"] = config.temperature if config.temperature is not None else 0.7 - stage_sampling["0"] = stage0 - runtime_config["duplex_stage_sampling_params"] = stage_sampling - return runtime_config - - @classmethod - async def prepare_runtime_config(cls, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]: - extra_body = dict(config.extra_body) - if any(key in extra_body for key in ("ref_audio_path", "tts_ref_audio_path")): - raise ValueError("ref_audio_path is not accepted by native duplex; use ref_audio URI instead") - cls.validate_client_config(config) - runtime_config: dict[str, object] = {"instructions": config.instructions} - initial_user_text = extra_body.pop("duplex_initial_user_text", None) - if isinstance(initial_user_text, str) and initial_user_text: - runtime_config["initial_user_text"] = initial_user_text - cls._apply_default_scheduler_policy(runtime_config, config=config, model_config=model_config) - - ref_audio = config.ref_audio - extra_ref_audio = extra_body.get("ref_audio") - if ref_audio is None and isinstance(extra_ref_audio, str): - extra_body.pop("ref_audio") - ref_audio = extra_ref_audio - extra_tts_ref_audio = extra_body.get("tts_ref_audio") - if ref_audio is None and isinstance(extra_tts_ref_audio, str): - extra_body.pop("tts_ref_audio") - ref_audio = extra_tts_ref_audio - - if ref_audio is None: - if any(str(modality).lower() == "audio" for modality in config.modalities): - raise MiniCPMO45ClientRuntimeConfigError( - "MiniCPM-o native duplex audio output requires ref_audio", - code="ref_audio_required", - ) - cls._apply_first_append_context_tokens( - runtime_config, - model_config=model_config, - instructions=config.instructions, - initial_user_text=initial_user_text, - ref_sample_count=None, - ) - config.extra_body = extra_body - return runtime_config - else: - wav_np, sr = await cls.resolve_ref_audio(ref_audio, model_config=model_config) - - wav_np = cls.normalize_ref_audio(wav_np, int(sr), target_sr=16000) - # Trim to a whole number of pooled audio embeddings (100 ms frames) so - # the first-append scheduler reserve can count them exactly. - usable = (len(wav_np) // MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN) * ( - MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN - ) - wav_np = wav_np[:usable] - ref_audio_bytes = np.ascontiguousarray(wav_np, dtype=np.float32).tobytes() - runtime_config["ref_audio_data"] = base64.b64encode(ref_audio_bytes).decode("ascii") - runtime_config["ref_audio_format"] = "pcm_f32le" - runtime_config["ref_audio_sample_rate_hz"] = 16000 - cls._apply_first_append_context_tokens( - runtime_config, - model_config=model_config, - instructions=config.instructions, - initial_user_text=initial_user_text, - ref_sample_count=len(wav_np), - ) - config.extra_body = extra_body - config.ref_audio = None - return runtime_config - - @classmethod - def _apply_default_scheduler_policy( - cls, - runtime_config: dict[str, object], - *, - config: DuplexSessionConfig, - model_config: Any, - ) -> None: - stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 - runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192} - stage0_params: dict[str, object] = { - "temperature": config.temperature if config.temperature is not None else 0.7, - "top_p": 0.8, - "top_k": 20, - "repetition_penalty": 1.05, - } - stop_token_ids = cls._native_stage0_stop_token_ids(model_config) - if stop_token_ids: - stage0_params["stop_token_ids"] = stop_token_ids - # Stage 1 keeps the deploy YAML's codec knobs, minus upstream's - # min_new_token=50: that floor is for whole-utterance chat TTS, while a - # duplex chunk is 26 codec samples, so check_stop would never release - # the Talker and Thinker would stall on the next model turn. - runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params, "1": {"min_tokens": 0}} - scheduler_token_id = cls._native_scheduler_token_id(model_config) - if scheduler_token_id is not None: - runtime_config["duplex_scheduler_token_id"] = scheduler_token_id - - @classmethod - def _apply_first_append_context_tokens( - cls, - runtime_config: dict[str, object], - *, - model_config: Any, - instructions: object, - initial_user_text: object, - ref_sample_count: int | None, - ) -> None: - """Precompute the exact session-context token count for the engine. - - The first data-plane append carries the system template and optional - reference-audio embeddings ahead of the first unit. The engine - reserves scheduler slots from this count; an inexact count turns into - pad embeddings inside the model KV (surplus) or truncated context - (deficit), so it is computed with the same template and pooling math - the worker uses. - """ - if "duplex_first_append_context_tokens" in runtime_config: - return - tokenizer = cls._load_native_tokenizer(model_config) - if tokenizer is None: - return - prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts( - instructions, - ref_sample_count is not None, - initial_user_text, - ) - try: - prefix_ids = tokenizer.encode(prefix, add_special_tokens=False) - suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) - except Exception: - return - ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0) - runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) - - @staticmethod - def _native_stage0_stop_token_ids(model_config: Any) -> list[int]: - tokenizer = MiniCPMO45NativeDuplexServingAdapter._load_native_tokenizer(model_config) - if tokenizer is None: - return [] - out: list[int] = [] - stop_token_fields = ( - "chunk_eos_token_id", - "chunk_tts_eos_token_id", - "listen_token_id", - "turn_eos_token_id", - ) - for field in stop_token_fields: - token = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS[field] - token_id = MiniCPMO45NativeDuplexServingAdapter._convert_token_to_id(tokenizer, token) - if token_id is not None and token_id not in out: - out.append(token_id) - return out - - @staticmethod - def _native_scheduler_token_id(model_config: Any) -> int | None: - tokenizer = MiniCPMO45NativeDuplexServingAdapter._load_native_tokenizer(model_config) - if tokenizer is None: - return None - scheduler_tokens = ( - MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS["unit_token_id"], - MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS["audio_placeholder_token_id"], - ) - for token in scheduler_tokens: - token_id = MiniCPMO45NativeDuplexServingAdapter._convert_token_to_id(tokenizer, token) - if token_id is not None: - return token_id - eos_id = getattr(tokenizer, "eos_token_id", None) - if eos_id is None: - return None - try: - return int(eos_id) - except (TypeError, ValueError): - return None - - @staticmethod - def _load_native_tokenizer(model_config: Any) -> Any | None: - model_path = getattr(model_config, "model", None) - if not isinstance(model_path, str) or not model_path: - return None - try: - from transformers import AutoTokenizer - - return AutoTokenizer.from_pretrained( - model_path, - trust_remote_code=True, - local_files_only=True, - ) - except Exception: - return None - - @staticmethod - def _convert_token_to_id(tokenizer: Any, token: str) -> int | None: - convert = getattr(tokenizer, "convert_tokens_to_ids", None) - value = None - if callable(convert): - value = convert(token) - if isinstance(value, list): - value = value[0] if len(value) == 1 else None - if value is None: - token_id = -1 - else: - try: - token_id = int(value) - except (TypeError, ValueError): - token_id = -1 - unk_token_id = getattr(tokenizer, "unk_token_id", None) - if token_id >= 0 and token_id != unk_token_id: - return token_id - encode = getattr(tokenizer, "encode", None) - if callable(encode): - try: - ids = list(encode(token, add_special_tokens=False)) - except TypeError: - ids = list(encode(token)) - if len(ids) == 1: - try: - token_id = int(ids[0]) - except (TypeError, ValueError): - token_id = -1 - if token_id >= 0 and token_id != unk_token_id: - return token_id - return None - - @staticmethod - async def resolve_ref_audio(ref_audio: str, *, model_config: Any) -> tuple[np.ndarray, int]: - connector = MediaConnector( - allowed_local_media_path=getattr(model_config, "allowed_local_media_path", None), - allowed_media_domains=getattr(model_config, "allowed_media_domains", None), - ) - wav_np, sr = await connector.fetch_audio_async(ref_audio) - return np.asarray(wav_np, dtype=np.float32), int(sr) - - @staticmethod - def normalize_ref_audio(wav_np: np.ndarray, sample_rate: int, *, target_sr: int) -> np.ndarray: - wav_np = np.asarray(wav_np, dtype=np.float32) - if wav_np.ndim > 1: - wav_np = wav_np.mean(axis=-1) - wav_np = wav_np.reshape(-1) - if sample_rate <= 0 or sample_rate == target_sr or wav_np.size == 0: - return wav_np.astype(np.float32, copy=False) - import torch - import torchaudio - - audio = torch.from_numpy(wav_np).to(dtype=torch.float32).unsqueeze(0) - resampled = torchaudio.functional.resample(audio, int(sample_rate), int(target_sr)) - return resampled.squeeze(0).cpu().numpy().astype(np.float32, copy=False) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py index 668405892..d850f4c3c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py @@ -3,7 +3,12 @@ from __future__ import annotations -from vllm_omni.entrypoints.duplex.protocol import DuplexCapabilities +from vllm_omni.engine.duplex.config import DuplexCapabilities + +#: Silence units a text-seeded turn is given to generate on. The model speaks +#: per audio unit, so a seeded turn still needs a clock; three units is what +#: the benchmark harness uses for the same purpose. +MINICPMO45_TEXT_TURN_PRIMING_UNITS = 3 def minicpmo45_native_capabilities(*, max_sessions: int = 1) -> DuplexCapabilities: @@ -31,11 +36,14 @@ def minicpmo45_native_capabilities(*, max_sessions: int = 1) -> DuplexCapabiliti supports_session_resume=True, session_admission_mode="engine_managed", supports_audio_truncate=True, + # The session template ends at the assistant turn, so seeded text puts + # the model in position to answer; it still generates per audio unit, + # hence the priming units. + supports_chat_completions=True, # via initial_user_text + text_turn_priming_units=MINICPMO45_TEXT_TURN_PRIMING_UNITS, requires_model_runner_kv=True, requires_native_stage_role=True, - implementation_level="model_native_duplex", adapter_patterns=["scheduler_data_plane"], - input_modes=["append_audio_chunk"], signal_sources=["model_native", "client_event", "server_policy"], stage_handoff_transport="scheduler_data_plane", chunk_period_ms=1000, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py index 9a4bcf9b0..5b6b5c909 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py @@ -3,14 +3,12 @@ from __future__ import annotations -from typing import Any - -def patch_minicpmo_remote_config(config: Any) -> None: +def patch_minicpmo_remote_config(config: object) -> None: """Fill config fields expected by the vLLM MiniCPM implementation.""" seen: set[int] = set() - def patch_cfg(cfg: Any) -> None: + def patch_cfg(cfg: object) -> None: if cfg is None or id(cfg) in seen: return seen.add(id(cfg)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py index d4063ff3f..7bf51abfd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py @@ -12,6 +12,7 @@ from vllm.logger import init_logger from vllm_omni.engine.duplex.contracts import ( duplex_resource_request_belongs_to_session, ) +from vllm_omni.engine.duplex.plugin import DuplexDataPlane from vllm_omni.outputs.duplex import get_duplex_output_decision logger = init_logger(__name__) @@ -44,6 +45,9 @@ class _TurnState: @dataclass(slots=True) class _RequestState: audio_offset: int = 0 + #: A snapshot below ``audio_offset`` that has not been accepted as a restart + #: yet. See :meth:`MiniCPMO45DataPlaneSession.slice_cumulative_audio`. + audio_restart_candidate: int | None = None uses_segment_text_metadata: bool = False pending_audio_without_text: list[dict[str, object]] = field(default_factory=list) terminal: bool = False @@ -53,7 +57,7 @@ class _RequestState: return self.turns.setdefault(turn_id, _TurnState()) -class MiniCPMO45DataPlaneSession: +class MiniCPMO45DataPlaneSession(DuplexDataPlane): """MiniCPM output projector and request/turn cursor owner. The scheduler request owns cumulative Stage1 audio, while model turns own @@ -487,13 +491,33 @@ class MiniCPMO45DataPlaneSession: prev_samples = state.audio_offset if prev_samples <= 0: state.audio_offset = num_samples + state.audio_restart_candidate = None return audio_data if num_samples == prev_samples: return None if num_samples < prev_samples: + # Below the cursor means one of two things, and they are not + # distinguishable from this snapshot alone: the producer restarted + # its buffer, or this is a unit that carries no cumulative audio at + # all (a continuation placeholder arrives as a single sample). + # Rewinding the cursor for the second case is expensive and silent: + # the next real snapshot is then sliced from ~0 and the whole + # session's audio goes out again as one delta. So defer the + # decision — hold the candidate and leave the cursor where it is. + candidate = state.audio_restart_candidate + if candidate is None or num_samples <= candidate: + state.audio_restart_candidate = num_samples + return audio_data + # A second below-cursor snapshot that grew from the first: the + # producer really did restart, and this is its delta. + state.audio_offset = num_samples + state.audio_restart_candidate = num_samples + prev_samples = candidate + else: + # Above the cursor: whatever the dip was, this stream never + # restarted, so it is sliced from the cursor as usual. + state.audio_restart_candidate = None state.audio_offset = num_samples - return audio_data - state.audio_offset = num_samples try: import torch diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py index 1b30ec4c2..4feb8cb97 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py @@ -6,9 +6,11 @@ from __future__ import annotations import base64 import binascii from dataclasses import dataclass -from typing import Any import numpy as np +from numpy.typing import NDArray + +from vllm_omni.engine.duplex.plugin import PcmAppendBuffer, PcmAppendReservation @dataclass(frozen=True, slots=True) @@ -18,7 +20,7 @@ class _PcmSpan: is_speech: bool -class MiniCPMO45PcmAppendReservation: +class MiniCPMO45PcmAppendReservation(PcmAppendReservation): __slots__ = ( "_active", "_force_listen", @@ -74,7 +76,7 @@ class MiniCPMO45PcmAppendReservation: self._owner._rollback_reservation(self) -def validate_native_ref_audio_config(session_config: dict[str, Any]) -> None: +def validate_native_ref_audio_config(session_config: dict[str, object]) -> None: extra_body = session_config.get("extra_body") if not isinstance(extra_body, dict): extra_body = {} @@ -84,7 +86,7 @@ def validate_native_ref_audio_config(session_config: dict[str, Any]) -> None: raise ValueError("native duplex ref_audio_path is not accepted; resolve ref_audio in serving first") -def decode_native_ref_audio_from_config(session_config: dict[str, Any]) -> np.ndarray | None: +def decode_native_ref_audio_from_config(session_config: dict[str, object]) -> NDArray[np.float32] | None: validate_native_ref_audio_config(session_config) extra_body = session_config.get("extra_body") if not isinstance(extra_body, dict): @@ -106,7 +108,7 @@ def decode_native_ref_audio_from_config(session_config: dict[str, Any]) -> np.nd return np.frombuffer(raw, dtype=" None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py new file mode 100644 index 000000000..673640461 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py @@ -0,0 +1,773 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""MiniCPM-o 4.5 full-duplex model plugin: engine policy and session policy in one class. + +Engine policy (scheduler token budget, resumable Stage0 prompt, listen +decision) and session policy (client config validation, reference-audio +resolution, per-stage sampling knobs) both run engine-side, so one class owns +them. Client-supplied media URIs are resolved here through the media +connector; workers only ever receive normalized PCM payloads. +""" + +from __future__ import annotations + +import asyncio +import base64 +from base64 import b64decode +from binascii import Error as BinasciiError +from collections.abc import Mapping +from copy import deepcopy +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +from numpy.typing import NDArray +from vllm.multimodal.media import MediaConnector +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexOutputAction, + DuplexOutputDecision, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexModelPlugin, + DuplexRuntimeConfigError, + EncodeAudio, + reject_changed_runtime_value, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( + minicpmo45_native_capabilities, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( + MiniCPMO45DataPlaneContext, + MiniCPMO45DataPlaneSession, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( + MiniCPMO45ServingSessionState, +) + +if TYPE_CHECKING: + import torch + from transformers import PreTrainedTokenizerBase + from vllm.config import ModelConfig + +_DUPLEX_CHUNK_SAMPLES = 16000 +_DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600 +# + 64 resampler embeddings + per frame (max_slice_nums=1), +# matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME. +_DUPLEX_VISION_TOKENS_PER_FRAME = 66 +# Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD +# sliced (1 source + 2 patches on 960x540) and the composite is not. +_DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 + +PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( + { + "duplex_stage_sampling_params", + "duplex_stage_max_tokens", + "duplex_stage0_max_tokens", + "duplex_scheduler_token_id", + "duplex_first_append_context_tokens", + "ref_audio_data", + "ref_audio_format", + "ref_audio_sample_rate_hz", + "initial_user_text", + } +) + + +class MiniCPMO45ClientRuntimeConfigError(DuplexRuntimeConfigError): + pass + + +# ---- engine policy helpers: scheduler token budget ---- + + +def _duplex_frame_count(payload: object) -> int: + if not isinstance(payload, dict): + return 0 + frames = payload.get("video_frames") + if not isinstance(frames, list): + return 0 + return sum(1 for frame in frames if isinstance(frame, str) and frame) + + +def _duplex_vision_tokens(payload: object) -> int: + """Scheduler slots for this append's camera track. + + Audio is never stacked: a unit still carries one second of soundtrack. + ``stack_frames`` only adds a second *image*. Official HD on that pair is + ``[2, 1]``, so the base frame reserves three 66-token blocks and every + extra frame reserves one. + """ + count = _duplex_frame_count(payload) + if count <= 0: + return 0 + if count >= 2: + return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME + return count * _DUPLEX_VISION_TOKENS_PER_FRAME + + +def _duplex_pcm_sample_count(payload: object) -> int | None: + if not isinstance(payload, dict): + return None + audio = payload.get("audio") or payload.get("data") + if payload.get("format") != "pcm_f32le" or not isinstance(audio, str): + return None + try: + raw = b64decode(audio, validate=True) + except (BinasciiError, ValueError): + return None + return len(raw) // 4 + + +def duplex_payload_is_exact_chunks(payload: object) -> bool: + sample_count = _duplex_pcm_sample_count(payload) + return sample_count is not None and sample_count != 0 and sample_count % _DUPLEX_CHUNK_SAMPLES == 0 + + +def duplex_first_append_unit_count(payload: object) -> int | None: + sample_count = _duplex_pcm_sample_count(payload) + if not sample_count or sample_count % _DUPLEX_CHUNK_SAMPLES != 0: + return None + return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1) + + +def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: + vision_tokens = _duplex_vision_tokens(payload) + sample_count = _duplex_pcm_sample_count(payload) + if sample_count is None: + return max(1, int(default)) + vision_tokens + sample_count = max(1, sample_count) + if sample_count % _DUPLEX_CHUNK_SAMPLES == 0: + units = sample_count // _DUPLEX_CHUNK_SAMPLES + return units * (2 + _DUPLEX_CHUNK_SAMPLES // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN) + vision_tokens + return max(16, min(768, sample_count // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8)) + vision_tokens + + +def duplex_first_append_context_reserve(runtime_config: object) -> int: + if not isinstance(runtime_config, dict): + return 48 + exact = runtime_config.get("duplex_first_append_context_tokens") + if isinstance(exact, int) and exact >= 0: + return exact + reserve = 48 + ref = runtime_config.get("ref_audio_data") + if isinstance(ref, str) and ref: + try: + raw = b64decode(ref, validate=True) + except (BinasciiError, ValueError): + raw = b"" + if raw: + reserve += max(0, (len(raw) // 4) // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8) + return reserve + + +def _duplex_force_listen_count(extra_body: object) -> int: + raw = extra_body.get("force_listen_count") if isinstance(extra_body, dict) else None + try: + return 0 if raw is None else max(0, int(raw)) + except (TypeError, ValueError): + return 0 + + +def build_duplex_data_plane_prompt( + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, object], + runtime_config: dict[str, object], + seq: int, + turn_seq: int, + payload: object, + final: bool, +) -> dict[str, object]: + token_budget = duplex_scheduler_token_budget(payload) + if seq <= 1: + context_reserve = duplex_first_append_context_reserve(runtime_config) + token_budget += context_reserve + first_units = duplex_first_append_unit_count(payload) + if first_units is not None: + token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) + if seq > 1 and duplex_payload_is_exact_chunks(payload): + token_budget += 1 + if final and duplex_payload_is_exact_chunks(payload): + token_budget += 12 + extra_body = session_config.get("extra_body") + raw_token_id = runtime_config.get("duplex_scheduler_token_id") + try: + token_id = 0 if raw_token_id is None else max(0, int(raw_token_id)) + except (TypeError, ValueError): + token_id = 0 + force_listen_count = _duplex_force_listen_count(extra_body) + if ( + force_listen_count > 0 + and turn_seq <= force_listen_count + and isinstance(payload, dict) + and payload.get("force_listen") is not True + ): + payload = {**payload, "force_listen": True} + return { + "prompt_token_ids": [token_id] * token_budget, + "model_intermediate_buffer": { + "request_id": request_id, + "global_request_id": [fence.session_id], + "duplex": { + "fence": fence, + "session_id": fence.session_id, + "epoch": fence.epoch, + "seq": seq, + "turn_id": fence.turn_id, + "turn_seq": turn_seq, + "mode": "append_audio_chunk", + "payload": payload, + "final": final, + "data_plane": True, + "session_config": dict(session_config), + "runtime_config": dict(runtime_config), + "scheduler_token_budget": token_budget, + "scheduler_token_id": token_id, + }, + }, + } + + +# ---- engine policy helpers: listen decision ---- + + +def _coerce_int(value: object) -> int | None: + detach = getattr(value, "detach", None) + if callable(detach): + try: + flat: torch.Tensor = detach().cpu().reshape(-1) + if flat.numel() == 0: + return None + value = flat[0].item() + except Exception: + return None + try: + return int(cast(Any, value)) # Any: duck-typed scalar (int/float/str/tensor item) + except (TypeError, ValueError): + return None + + +def _coerce_int_list(value: object) -> list[int]: + if value is None: + return [] + if hasattr(value, "detach"): + try: + value = value.detach().cpu().reshape(-1).tolist() + except Exception: + return [] + if not isinstance(value, (list, tuple)): + return [] + return [token_id for item in value if (token_id := _coerce_int(item)) is not None] + + +def _first_completion(output: object) -> object | None: + outputs = getattr(output, "outputs", None) + return outputs[0] if isinstance(outputs, list) and outputs else None + + +def _multimodal_output(output: object, completion: object | None) -> dict[str, object]: + metadata = getattr(output, "multimodal_output", None) + if isinstance(metadata, dict): + return metadata + metadata = getattr(completion, "multimodal_output", None) if completion is not None else None + return metadata if isinstance(metadata, dict) else {} + + +def _special_token_ids(metadata: dict[str, object]) -> dict[str, int]: + sources: list[object] = [metadata.get("special_token_ids"), metadata.get("meta")] + sources.append( + { + key.removeprefix("meta."): value + for key, value in metadata.items() + if isinstance(key, str) and key.startswith("meta.") + } + ) + token_ids: dict[str, int] = {} + for source in sources: + if not isinstance(source, dict): + continue + for key, value in source.items(): + token_id = _coerce_int(value) + if isinstance(key, str) and token_id is not None and token_id >= 0: + token_ids[key] = token_id + return token_ids + + +def _completion_token_ids(completion: object | None) -> list[int]: + if completion is None: + return [] + for attribute in ("token_ids", "cumulative_token_ids"): + token_ids = _coerce_int_list(getattr(completion, attribute, None)) + if token_ids: + return token_ids + return [] + + +def _stage_config_value(runtime_config: dict[str, object], key: str, stage_id: int) -> object | None: + raw = runtime_config.get(key) + if isinstance(raw, dict): + value = raw.get(stage_id) + return raw.get(str(stage_id)) if value is None else value + if isinstance(raw, (list, tuple)) and stage_id < len(raw): + return raw[stage_id] + return None + + +# ---- session policy helpers: tokenizer / reference audio ---- + + +def _load_tokenizer(model_config: ModelConfig | None) -> PreTrainedTokenizerBase | None: + model_path = getattr(model_config, "model", None) + if not isinstance(model_path, str) or not model_path: + return None + try: + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained( + model_path, + trust_remote_code=True, + local_files_only=True, + ) + except Exception: + return None + + +def _convert_token_to_id(tokenizer: PreTrainedTokenizerBase, token: str) -> int | None: + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + value = None + if callable(convert): + value = convert(token) + if isinstance(value, list): + value = value[0] if len(value) == 1 else None + if value is None: + token_id = -1 + else: + try: + token_id = int(value) + except (TypeError, ValueError): + token_id = -1 + unk_token_id = getattr(tokenizer, "unk_token_id", None) + if token_id >= 0 and token_id != unk_token_id: + return token_id + encode = getattr(tokenizer, "encode", None) + if callable(encode): + try: + ids = list(encode(token, add_special_tokens=False)) + except TypeError: + ids = list(encode(token)) + if len(ids) == 1: + try: + token_id = int(ids[0]) + except (TypeError, ValueError): + token_id = -1 + if token_id >= 0 and token_id != unk_token_id: + return token_id + return None + + +def _stage0_stop_token_ids(tokenizer: PreTrainedTokenizerBase | None) -> list[int]: + if tokenizer is None: + return [] + out: list[int] = [] + stop_token_fields = ( + "chunk_eos_token_id", + "chunk_tts_eos_token_id", + "listen_token_id", + "turn_eos_token_id", + ) + for field in stop_token_fields: + token = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS[field] + token_id = _convert_token_to_id(tokenizer, token) + if token_id is not None and token_id not in out: + out.append(token_id) + return out + + +def _scheduler_token_id(tokenizer: PreTrainedTokenizerBase | None) -> int | None: + if tokenizer is None: + return None + scheduler_tokens = ( + MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS["unit_token_id"], + MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS["audio_placeholder_token_id"], + ) + for token in scheduler_tokens: + token_id = _convert_token_to_id(tokenizer, token) + if token_id is not None: + return token_id + eos_id = getattr(tokenizer, "eos_token_id", None) + if eos_id is None: + return None + try: + return int(eos_id) + except (TypeError, ValueError): + return None + + +async def resolve_ref_audio(ref_audio: str, *, model_config: ModelConfig | None) -> tuple[NDArray[np.float32], int]: + connector = MediaConnector( + allowed_local_media_path=getattr(model_config, "allowed_local_media_path", None), + allowed_media_domains=getattr(model_config, "allowed_media_domains", None), + ) + wav_np, sr = await connector.fetch_audio_async(ref_audio) + return np.asarray(wav_np, dtype=np.float32), int(sr) + + +def normalize_ref_audio(wav_np: NDArray[np.float32], sample_rate: int, *, target_sr: int) -> NDArray[np.float32]: + wav_np = np.asarray(wav_np, dtype=np.float32) + if wav_np.ndim > 1: + wav_np = wav_np.mean(axis=-1) + wav_np = wav_np.reshape(-1) + if sample_rate <= 0 or sample_rate == target_sr or wav_np.size == 0: + return wav_np.astype(np.float32, copy=False) + import torch + import torchaudio + + audio = torch.from_numpy(wav_np).to(dtype=torch.float32).unsqueeze(0) + resampled = torchaudio.functional.resample(audio, int(sample_rate), int(target_sr)) + return resampled.squeeze(0).cpu().numpy().astype(np.float32, copy=False) + + +def _apply_first_append_context_tokens( + runtime_config: dict[str, object], + *, + tokenizer: PreTrainedTokenizerBase | None, + instructions: object, + initial_user_text: object, + ref_sample_count: int | None, +) -> None: + """Precompute the exact session-context token count for the engine. + + The first data-plane append carries the system template and optional + reference-audio embeddings ahead of the first unit. The engine reserves + scheduler slots from this count; an inexact count turns into pad + embeddings inside the model KV (surplus) or truncated context (deficit), + so it is computed with the same template and pooling math the worker uses. + """ + if "duplex_first_append_context_tokens" in runtime_config or tokenizer is None: + return + prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts( + instructions, + ref_sample_count is not None, + initial_user_text, + ) + try: + prefix_ids = tokenizer.encode(prefix, add_special_tokens=False) + suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) + except Exception: + return + ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0) + runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) + + +def _apply_default_scheduler_policy( + runtime_config: dict[str, object], + *, + config: DuplexSessionConfig, + tokenizer: PreTrainedTokenizerBase | None, +) -> None: + stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 + runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192} + stage0_params: dict[str, object] = { + "temperature": config.temperature if config.temperature is not None else 0.7, + "top_p": 0.8, + "top_k": 20, + "repetition_penalty": 1.05, + } + stop_token_ids = _stage0_stop_token_ids(tokenizer) + if stop_token_ids: + stage0_params["stop_token_ids"] = stop_token_ids + # Stage 1 keeps the deploy YAML's codec knobs, minus upstream's + # min_new_token=50: that floor is for whole-utterance chat TTS, while a + # duplex chunk is 26 codec samples, so check_stop would never release + # the Talker and Thinker would stall on the next model turn. + runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params, "1": {"min_tokens": 0}} + scheduler_token_id = _scheduler_token_id(tokenizer) + if scheduler_token_id is not None: + runtime_config["duplex_scheduler_token_id"] = scheduler_token_id + + +class MiniCPMO45DuplexPlugin(DuplexModelPlugin): + """MiniCPM-owned sampling policy, append planning, session state and output projection.""" + + plugin_id = "minicpmo45" + private_runtime_config_keys = PRIVATE_RUNTIME_CONFIG_KEYS + silence_continuation_samples = 16000 + + def __init__(self, encode_audio: EncodeAudio) -> None: + super().__init__(encode_audio) + self.data_plane = MiniCPMO45DataPlaneSession(encode_audio) + # The tokenizer is only needed for special-token ids and the session + # context token count; it is the same for every session, so load it + # once per model path, off the orchestrator loop. + self._tokenizers: dict[str, PreTrainedTokenizerBase | None] = {} + self._tokenizer_lock = asyncio.Lock() + + async def _tokenizer_for(self, model_config: ModelConfig | None) -> PreTrainedTokenizerBase | None: + model_path = getattr(model_config, "model", None) + key = model_path if isinstance(model_path, str) else "" + if key in self._tokenizers: + return self._tokenizers[key] + async with self._tokenizer_lock: + if key not in self._tokenizers: + self._tokenizers[key] = await asyncio.to_thread(_load_tokenizer, model_config) + return self._tokenizers[key] + + # ---- engine policy (the resumable Stage0 request) ---- + + def configure_sampling_params( + self, + *, + runtime_config: dict[str, object], + defaults: tuple[object, ...], + ) -> tuple[object, ...]: + configured: list[object] = [] + for stage_id, default in enumerate(defaults): + max_tokens = _coerce_int(_stage_config_value(runtime_config, "duplex_stage_max_tokens", stage_id)) + raw_overrides = _stage_config_value(runtime_config, "duplex_stage_sampling_params", stage_id) + overrides = dict(raw_overrides) if isinstance(raw_overrides, dict) else {} + if not isinstance(default, SamplingParams) or (not overrides and (max_tokens is None or max_tokens <= 0)): + configured.append(default) + continue + params = default.clone() + if max_tokens is not None and max_tokens > 0: + params.max_tokens = max_tokens + for name, value in overrides.items(): + if not hasattr(params, name): + continue + setattr(params, name, value) + if name == "stop_token_ids": + all_stop_token_ids = getattr(params, "_all_stop_token_ids", None) + if isinstance(all_stop_token_ids, set): + all_stop_token_ids.update(int(token_id) for token_id in value) + configured.append(params) + return tuple(configured) + + def plan_append( + self, + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, object], + runtime_config: dict[str, object], + seq: int, + turn_seq: int, + payload: object, + final: bool, + sampling_params: object, + ) -> DuplexAppendPlan: + del sampling_params + return DuplexAppendPlan( + prompt=build_duplex_data_plane_prompt( + request_id=request_id, + fence=fence, + session_config=session_config, + runtime_config=runtime_config, + seq=seq, + turn_seq=turn_seq, + payload=payload, + final=final, + ) + ) + + def decide_output( + self, + *, + stage_id: int, + final_stage_id: int, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: dict[str, object], + output: object, + ) -> DuplexOutputDecision | None: + if stage_id >= final_stage_id or not segment_finished: + return None + + completion = _first_completion(output) + output_metadata = _multimodal_output(output, completion) + special_token_ids = _special_token_ids(segment_output_metadata) + special_token_ids.update(_special_token_ids(output_metadata)) + listen_id = special_token_ids.get("listen_token_id") + if listen_id is None: + return None + + stop_reason = getattr(completion, "stop_reason", None) if completion is not None else None + token_ids = _completion_token_ids(completion) or list(segment_token_ids) + if _coerce_int(stop_reason) != listen_id and (not token_ids or token_ids[-1] != listen_id): + return None + + metadata = dict(output_metadata) + for key, value in special_token_ids.items(): + metadata.setdefault(f"meta.{key}", value) + metadata.update( + { + "duplex_direct_response": True, + "duplex_native_decision": "listen", + "model_listen": True, + "listen_source": "model_listen", + } + ) + return DuplexOutputDecision( + action=DuplexOutputAction.DIRECT_RESPONSE, + metadata=metadata, + ) + + # ---- session policy ---- + + def create_session_state(self) -> MiniCPMO45ServingSessionState: + return MiniCPMO45ServingSessionState() + + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: + return minicpmo45_native_capabilities(max_sessions=max_sessions) + + def validate_client_extra_body(self, extra_body: object) -> None: + if not isinstance(extra_body, dict): + return + private_keys = sorted(PRIVATE_RUNTIME_CONFIG_KEYS.intersection(extra_body)) + if private_keys: + raise MiniCPMO45ClientRuntimeConfigError( + "duplex runtime configuration is server-owned: " + ", ".join(private_keys) + ) + + async def prepare_runtime_config( + self, config: DuplexSessionConfig, *, model_config: ModelConfig | None + ) -> dict[str, object]: + extra_body = dict(config.extra_body) + if any(key in extra_body for key in ("ref_audio_path", "tts_ref_audio_path")): + raise MiniCPMO45ClientRuntimeConfigError( + "ref_audio_path is not accepted by duplex sessions; use ref_audio URI instead", + code="unsupported_ref_audio_path", + ) + runtime_config: dict[str, object] = {"instructions": config.instructions} + # ``duplex_initial_user_text`` is the older extra_body spelling and + # still works; the session field is the framework-level one. + initial_user_text = extra_body.pop("duplex_initial_user_text", None) + if not (isinstance(initial_user_text, str) and initial_user_text): + initial_user_text = config.initial_user_text + if isinstance(initial_user_text, str) and initial_user_text: + runtime_config["initial_user_text"] = initial_user_text + tokenizer = await self._tokenizer_for(model_config) + _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer) + + ref_audio = config.ref_audio + extra_ref_audio = extra_body.get("ref_audio") + if ref_audio is None and isinstance(extra_ref_audio, str): + extra_body.pop("ref_audio") + ref_audio = extra_ref_audio + extra_tts_ref_audio = extra_body.get("tts_ref_audio") + if ref_audio is None and isinstance(extra_tts_ref_audio, str): + extra_body.pop("tts_ref_audio") + ref_audio = extra_tts_ref_audio + + if ref_audio is None: + if any(str(modality).lower() == "audio" for modality in config.modalities): + raise MiniCPMO45ClientRuntimeConfigError( + "MiniCPM-o duplex audio output requires ref_audio", + code="ref_audio_required", + ) + _apply_first_append_context_tokens( + runtime_config, + tokenizer=tokenizer, + instructions=config.instructions, + initial_user_text=initial_user_text, + ref_sample_count=None, + ) + config.extra_body = extra_body + return runtime_config + + wav_np, sr = await resolve_ref_audio(ref_audio, model_config=model_config) + # torchaudio resampling is CPU work: keep it off the orchestrator loop. + wav_np = await asyncio.to_thread(normalize_ref_audio, wav_np, int(sr), target_sr=16000) + # Trim to a whole number of pooled audio embeddings (100 ms frames) so + # the first-append scheduler reserve can count them exactly. + usable = (len(wav_np) // MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN) * ( + MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN + ) + wav_np = wav_np[:usable] + ref_audio_bytes = np.ascontiguousarray(wav_np, dtype=np.float32).tobytes() + runtime_config["ref_audio_data"] = base64.b64encode(ref_audio_bytes).decode("ascii") + runtime_config["ref_audio_format"] = "pcm_f32le" + runtime_config["ref_audio_sample_rate_hz"] = 16000 + _apply_first_append_context_tokens( + runtime_config, + tokenizer=tokenizer, + instructions=config.instructions, + initial_user_text=initial_user_text, + ref_sample_count=len(wav_np), + ) + config.extra_body = extra_body + config.ref_audio = None + return runtime_config + + def runtime_config_for_update( + self, + config: DuplexSessionConfig, + current: Mapping[str, object], + ) -> dict[str, object]: + runtime_config = deepcopy(dict(current)) + reject_changed_runtime_value( + config.instructions, + runtime_config.get("instructions"), + message="instructions cannot be changed after the session is created", + code="instructions_update_unsupported", + error_cls=MiniCPMO45ClientRuntimeConfigError, + ) + stage_max_tokens = runtime_config.get("duplex_stage_max_tokens") + stage_max_tokens = deepcopy(stage_max_tokens) if isinstance(stage_max_tokens, dict) else {} + stage_max_tokens["0"] = ( + config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 + ) + stage_max_tokens.setdefault("1", 8192) + runtime_config["duplex_stage_max_tokens"] = stage_max_tokens + + stage_sampling = runtime_config.get("duplex_stage_sampling_params") + stage_sampling = deepcopy(stage_sampling) if isinstance(stage_sampling, dict) else {} + stage0 = stage_sampling.get("0") + stage0 = deepcopy(stage0) if isinstance(stage0, dict) else {} + stage0["temperature"] = config.temperature if config.temperature is not None else 0.7 + stage_sampling["0"] = stage0 + runtime_config["duplex_stage_sampling_params"] = stage_sampling + return runtime_config + + def data_plane_context( + self, + *, + epoch: int, + turn_id: int, + active_response_turn_id: int | None, + active_response_id: str | None, + auto_responds: bool, + response_format: str, + speed: float | None, + modalities: tuple[str, ...], + ) -> MiniCPMO45DataPlaneContext: + return MiniCPMO45DataPlaneContext( + epoch=epoch, + turn_id=turn_id, + active_response_turn_id=active_response_turn_id, + active_response_id=active_response_id, + auto_responds=auto_responds, + response_format=response_format, + speed=speed, + modalities=modalities, + ) + + +__all__ = [ + "PRIVATE_RUNTIME_CONFIG_KEYS", + "MiniCPMO45ClientRuntimeConfigError", + "MiniCPMO45DuplexPlugin", + "build_duplex_data_plane_prompt", + "duplex_first_append_context_reserve", + "duplex_first_append_unit_count", + "duplex_payload_is_exact_chunks", + "duplex_scheduler_token_budget", + "normalize_ref_audio", + "resolve_ref_audio", +] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py index 9a2643521..0332c81c0 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py @@ -3,7 +3,10 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase class MiniCPMO45DuplexPolicy: @@ -89,7 +92,7 @@ class MiniCPMO45DuplexPolicy: } @classmethod - def token_ids_from_tokenizer(cls, tokenizer: Any) -> dict[str, int]: + def token_ids_from_tokenizer(cls, tokenizer: PreTrainedTokenizerBase) -> dict[str, int]: convert = getattr(tokenizer, "convert_tokens_to_ids", None) def token_id(token: str) -> int: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py deleted file mode 100644 index d5e98d723..000000000 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py +++ /dev/null @@ -1,439 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import math -from base64 import b64decode -from binascii import Error as BinasciiError -from typing import Any, cast - -from vllm.sampling_params import SamplingParams - -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexOutputAction, - DuplexOutputDecision, -) - -_DUPLEX_CHUNK_SAMPLES = 16000 -_DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600 -# + 64 resampler embeddings + per frame (max_slice_nums=1), -# matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME. -_DUPLEX_VISION_TOKENS_PER_FRAME = 66 -# Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD -# sliced (1 source + up to 2 patches, e.g. on 960x540) and the composite is -# not. The patch count depends on the frame size (see -# ``_duplex_hd_slice_count``); the constant is the fallback when the frame -# header cannot be read. -_DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 -_DUPLEX_HD_MAX_SLICE_NUMS = 2 -_DUPLEX_SCALE_RESOLUTION = 448 - - -def _duplex_frames(payload: object) -> list[str]: - if not isinstance(payload, dict): - return [] - frames = payload.get("video_frames") - if not isinstance(frames, list): - return [] - return [frame for frame in frames if isinstance(frame, str) and frame] - - -def _duplex_frame_count(payload: object) -> int: - return len(_duplex_frames(payload)) - - -def _duplex_frame_size(frame_b64: str) -> tuple[int, int] | None: - """Pixel size of a base64 JPEG/PNG frame from its header, or ``None``.""" - from io import BytesIO - - try: - from PIL import Image - - with Image.open(BytesIO(b64decode(frame_b64, validate=True))) as image: - width, height = image.size - except Exception: # noqa: BLE001 - Stage0 rejects bad frames with a reason - return None - if width <= 0 or height <= 0: - return None - return int(width), int(height) - - -def _duplex_hd_slice_count(image_size: tuple[int, int], max_slice_nums: int) -> int: - """Number of HD patches ``MiniCPMVImageProcessor.get_sliced_grid`` adds. - - Port of the official grid selection (``scale_resolution=448``): an image - whose area fits in one 448x448 tile is not sliced at all, otherwise the - grid closest to the image aspect ratio among ``multiple-1 .. multiple+1`` - splits is used. Stage0 runs the same processor, so the scheduler budget - and the worker-built embeddings agree. - """ - width, height = image_size - ratio = width * height / (_DUPLEX_SCALE_RESOLUTION * _DUPLEX_SCALE_RESOLUTION) - multiple = min(math.ceil(ratio), max_slice_nums) - if multiple <= 1: - return 0 - log_ratio = math.log(width / height) - best_grid = (1, 1) - min_error = float("inf") - for split_grids_nums in (multiple - 1, multiple, multiple + 1): - if split_grids_nums == 1 or split_grids_nums > max_slice_nums: - continue - for m in range(1, split_grids_nums + 1): - if split_grids_nums % m: - continue - grid = (m, split_grids_nums // m) - error = abs(log_ratio - math.log(grid[0] / grid[1])) - if error < min_error: - best_grid = grid - min_error = error - return best_grid[0] * best_grid[1] - - -def _duplex_vision_tokens(payload: object) -> int: - """Scheduler slots for this append's camera track. - - Audio is never stacked: a unit still carries one second of soundtrack. - ``stack_frames`` only adds a second *image*. Official HD on that pair is - ``max_slice_nums=[2, 1]``: the base frame keeps its source block plus the - HD patches the processor cuts for its size (two for a 960x540 camera - frame, none for anything that fits in one 448x448 tile) and every extra - frame is one block. A single frame is never sliced. The count must match - the embeddings Stage0 builds exactly: surplus slots become pad - embeddings inside the KV. - """ - frames = _duplex_frames(payload) - count = len(frames) - if count <= 0: - return 0 - if count == 1: - return _DUPLEX_VISION_TOKENS_PER_FRAME - base_size = _duplex_frame_size(frames[0]) - if base_size is None: - base_blocks = _DUPLEX_HD_SLICES_PER_BASE_FRAME - else: - base_blocks = 1 + _duplex_hd_slice_count(base_size, _DUPLEX_HD_MAX_SLICE_NUMS) - return (base_blocks + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME - - -def _duplex_pcm_sample_count(payload: object) -> int | None: - if not isinstance(payload, dict): - return None - audio = payload.get("audio") or payload.get("data") - if payload.get("format") != "pcm_f32le" or not isinstance(audio, str): - return None - try: - raw = b64decode(audio, validate=True) - except (BinasciiError, ValueError): - return None - return len(raw) // 4 - - -def duplex_payload_is_exact_chunks(payload: object) -> bool: - sample_count = _duplex_pcm_sample_count(payload) - return sample_count is not None and sample_count != 0 and sample_count % _DUPLEX_CHUNK_SAMPLES == 0 - - -def duplex_first_append_unit_count(payload: object) -> int | None: - sample_count = _duplex_pcm_sample_count(payload) - if not sample_count or sample_count % _DUPLEX_CHUNK_SAMPLES != 0: - return None - return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1) - - -def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: - vision_tokens = _duplex_vision_tokens(payload) - sample_count = _duplex_pcm_sample_count(payload) - if sample_count is None: - return max(1, int(default)) + vision_tokens - sample_count = max(1, sample_count) - if sample_count % _DUPLEX_CHUNK_SAMPLES == 0: - units = sample_count // _DUPLEX_CHUNK_SAMPLES - return units * (2 + _DUPLEX_CHUNK_SAMPLES // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN) + vision_tokens - return max(16, min(768, sample_count // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8)) + vision_tokens - - -def duplex_first_append_context_reserve(runtime_config: object) -> int: - if not isinstance(runtime_config, dict): - return 48 - exact = runtime_config.get("duplex_first_append_context_tokens") - if isinstance(exact, int) and exact >= 0: - return exact - reserve = 48 - ref = runtime_config.get("ref_audio_data") - if isinstance(ref, str) and ref: - try: - raw = b64decode(ref, validate=True) - except (BinasciiError, ValueError): - raw = b"" - if raw: - reserve += max(0, (len(raw) // 4) // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8) - return reserve - - -def _duplex_force_listen_count(extra_body: object) -> int: - raw = extra_body.get("force_listen_count") if isinstance(extra_body, dict) else None - try: - return 0 if raw is None else max(0, int(raw)) - except (TypeError, ValueError): - return 0 - - -def build_duplex_data_plane_prompt( - *, - request_id: str, - fence: DuplexFence, - session_config: dict[str, Any], - runtime_config: dict[str, Any], - seq: int, - turn_seq: int, - mode: DuplexInputMode, - payload: object, - final: bool, -) -> dict[str, Any]: - token_budget = duplex_scheduler_token_budget(payload) - if seq <= 1: - context_reserve = duplex_first_append_context_reserve(runtime_config) - token_budget += context_reserve - first_units = duplex_first_append_unit_count(payload) - if first_units is not None: - token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) - if seq > 1 and duplex_payload_is_exact_chunks(payload): - token_budget += 1 - if final and duplex_payload_is_exact_chunks(payload): - token_budget += 12 - extra_body = session_config.get("extra_body") - raw_token_id = runtime_config.get("duplex_scheduler_token_id") - try: - token_id = 0 if raw_token_id is None else max(0, int(raw_token_id)) - except (TypeError, ValueError): - token_id = 0 - force_listen_count = _duplex_force_listen_count(extra_body) - if ( - force_listen_count > 0 - and turn_seq <= force_listen_count - and isinstance(payload, dict) - and payload.get("force_listen") is not True - ): - payload = {**payload, "force_listen": True} - return { - "prompt_token_ids": [token_id] * token_budget, - "model_intermediate_buffer": { - "request_id": request_id, - "global_request_id": [fence.session_id], - "duplex": { - "fence": fence, - "session_id": fence.session_id, - "incarnation": fence.incarnation, - "epoch": fence.epoch, - "seq": seq, - "turn_id": fence.turn_id, - "response_seq": fence.response_seq, - "turn_seq": turn_seq, - "mode": mode.value, - "payload": payload, - "final": final, - "data_plane": True, - "session_config": dict(session_config), - "runtime_config": dict(runtime_config), - "scheduler_token_budget": token_budget, - "scheduler_token_id": token_id, - }, - }, - } - - -def _coerce_int(value: object) -> int | None: - detach = getattr(value, "detach", None) - if callable(detach): - try: - flat: Any = detach().cpu().reshape(-1) - if flat.numel() == 0: - return None - value = flat[0].item() - except Exception: - return None - try: - return int(cast(Any, value)) - except (TypeError, ValueError): - return None - - -def _coerce_int_list(value: object) -> list[int]: - if value is None: - return [] - if hasattr(value, "detach"): - try: - value = value.detach().cpu().reshape(-1).tolist() - except Exception: - return [] - if not isinstance(value, (list, tuple)): - return [] - return [token_id for item in value if (token_id := _coerce_int(item)) is not None] - - -def _first_completion(output: object) -> object | None: - outputs = getattr(output, "outputs", None) - return outputs[0] if isinstance(outputs, list) and outputs else None - - -def _multimodal_output(output: object, completion: object | None) -> dict[str, Any]: - metadata = getattr(output, "multimodal_output", None) - if isinstance(metadata, dict): - return metadata - metadata = getattr(completion, "multimodal_output", None) if completion is not None else None - return metadata if isinstance(metadata, dict) else {} - - -def _special_token_ids(metadata: dict[str, Any]) -> dict[str, int]: - sources: list[object] = [metadata.get("special_token_ids"), metadata.get("meta")] - sources.append( - { - key.removeprefix("meta."): value - for key, value in metadata.items() - if isinstance(key, str) and key.startswith("meta.") - } - ) - token_ids: dict[str, int] = {} - for source in sources: - if not isinstance(source, dict): - continue - for key, value in source.items(): - token_id = _coerce_int(value) - if isinstance(key, str) and token_id is not None and token_id >= 0: - token_ids[key] = token_id - return token_ids - - -def _completion_token_ids(completion: object | None) -> list[int]: - if completion is None: - return [] - for attribute in ("token_ids", "cumulative_token_ids"): - token_ids = _coerce_int_list(getattr(completion, attribute, None)) - if token_ids: - return token_ids - return [] - - -def _stage_config_value(runtime_config: dict[str, Any], key: str, stage_id: int) -> object | None: - raw = runtime_config.get(key) - if isinstance(raw, dict): - value = raw.get(stage_id) - return raw.get(str(stage_id)) if value is None else value - if isinstance(raw, (list, tuple)) and stage_id < len(raw): - return raw[stage_id] - return None - - -class MiniCPMO45DuplexRuntimeExtension: - def configure_sampling_params( - self, - *, - runtime_config: dict[str, Any], - defaults: tuple[object, ...], - ) -> tuple[object, ...]: - configured: list[object] = [] - for stage_id, default in enumerate(defaults): - max_tokens = _coerce_int(_stage_config_value(runtime_config, "duplex_stage_max_tokens", stage_id)) - raw_overrides = _stage_config_value(runtime_config, "duplex_stage_sampling_params", stage_id) - overrides = dict(raw_overrides) if isinstance(raw_overrides, dict) else {} - if not isinstance(default, SamplingParams) or (not overrides and (max_tokens is None or max_tokens <= 0)): - configured.append(default) - continue - params = default.clone() - if max_tokens is not None and max_tokens > 0: - params.max_tokens = max_tokens - for name, value in overrides.items(): - if not hasattr(params, name): - continue - setattr(params, name, value) - if name == "stop_token_ids": - all_stop_token_ids = getattr(params, "_all_stop_token_ids", None) - if isinstance(all_stop_token_ids, set): - all_stop_token_ids.update(int(token_id) for token_id in value) - configured.append(params) - return tuple(configured) - - def plan_append( - self, - *, - request_id: str, - fence: DuplexFence, - session_config: dict[str, Any], - runtime_config: dict[str, Any], - seq: int, - turn_seq: int, - mode: DuplexInputMode, - payload: object, - final: bool, - sampling_params: object, - ) -> DuplexAppendPlan: - del sampling_params - return DuplexAppendPlan( - prompt=build_duplex_data_plane_prompt( - request_id=request_id, - fence=fence, - session_config=session_config, - runtime_config=runtime_config, - seq=seq, - turn_seq=turn_seq, - mode=mode, - payload=payload, - final=final, - ) - ) - - def decide_output( - self, - *, - stage_id: int, - final_stage_id: int, - segment_finished: bool, - segment_token_ids: tuple[int, ...], - segment_output_metadata: dict[str, Any], - output: object, - ) -> DuplexOutputDecision | None: - if stage_id >= final_stage_id or not segment_finished: - return None - - completion = _first_completion(output) - output_metadata = _multimodal_output(output, completion) - special_token_ids = _special_token_ids(segment_output_metadata) - special_token_ids.update(_special_token_ids(output_metadata)) - listen_id = special_token_ids.get("listen_token_id") - if listen_id is None: - return None - - stop_reason = getattr(completion, "stop_reason", None) if completion is not None else None - token_ids = _completion_token_ids(completion) or list(segment_token_ids) - if _coerce_int(stop_reason) != listen_id and (not token_ids or token_ids[-1] != listen_id): - return None - - metadata = dict(output_metadata) - for key, value in special_token_ids.items(): - metadata.setdefault(f"meta.{key}", value) - metadata.update( - { - "duplex_direct_response": True, - "duplex_native_decision": "listen", - "model_listen": True, - "listen_source": "model_listen", - } - ) - return DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata=metadata, - ) - - -__all__ = [ - "MiniCPMO45DuplexRuntimeExtension", - "build_duplex_data_plane_prompt", - "duplex_first_append_context_reserve", - "duplex_first_append_unit_count", - "duplex_payload_is_exact_chunks", - "duplex_scheduler_token_budget", -] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py deleted file mode 100644 index ea51662a9..000000000 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py +++ /dev/null @@ -1,102 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -from collections.abc import Callable, Mapping -from typing import Any - -from vllm_omni.entrypoints.duplex.protocol import DuplexCapabilities -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.adapter import ( - MiniCPMO45NativeDuplexServingAdapter, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( - minicpmo45_native_capabilities, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( - MiniCPMO45DataPlaneContext, - MiniCPMO45DataPlaneSession, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( - MiniCPMO45ServingSessionState, -) - -EncodeAudio = Callable[[object, int, str, float | None], str | None] - - -class MiniCPMO45ServingRuntimeAdapter: - """MiniCPM-owned serving state, input packing, and output projection.""" - - adapter_id = "minicpmo45" - clean_response_done_prefix = "" - interrupted_tts_prefix = "" - private_runtime_config_keys = MiniCPMO45NativeDuplexServingAdapter.PRIVATE_RUNTIME_CONFIG_KEYS - - def __init__(self, encode_audio: EncodeAudio) -> None: - self.session_states: dict[str, MiniCPMO45ServingSessionState] = {} - self.data_plane = MiniCPMO45DataPlaneSession(encode_audio) - - def create_session_state(self) -> MiniCPMO45ServingSessionState: - return MiniCPMO45ServingSessionState() - - def session_state(self, session_id: str) -> MiniCPMO45ServingSessionState: - state = self.session_states.get(session_id) - if state is None: - state = self.create_session_state() - self.session_states[session_id] = state - return state - - def remove_session_state(self, session_id: str) -> None: - self.session_states.pop(session_id, None) - - @staticmethod - def is_enabled(config: object) -> bool: - return MiniCPMO45NativeDuplexServingAdapter.is_enabled(config) # type: ignore[arg-type] - - @staticmethod - def capabilities(*, max_sessions: int) -> DuplexCapabilities: - return minicpmo45_native_capabilities(max_sessions=max_sessions) - - @staticmethod - def validate_client_extra_body(extra_body: object) -> None: - MiniCPMO45NativeDuplexServingAdapter.validate_client_extra_body(extra_body) - - @staticmethod - async def prepare_runtime_config(config: object, *, model_config: Any) -> dict[str, object]: - return await MiniCPMO45NativeDuplexServingAdapter.prepare_runtime_config( - config, # type: ignore[arg-type] - model_config=model_config, - ) - - @staticmethod - def runtime_config_for_update( - config: object, - current: Mapping[str, object], - ) -> dict[str, object]: - return MiniCPMO45NativeDuplexServingAdapter.runtime_config_for_update( - config, # type: ignore[arg-type] - dict(current), - ) - - @staticmethod - def data_plane_context( - *, - epoch: int, - turn_id: int, - active_response_turn_id: int | None, - active_response_id: str | None, - auto_responds: bool, - response_format: str, - speed: float | None, - modalities: tuple[str, ...], - ) -> MiniCPMO45DataPlaneContext: - return MiniCPMO45DataPlaneContext( - epoch=epoch, - turn_id=turn_id, - active_response_turn_id=active_response_turn_id, - active_response_id=active_response_id, - auto_responds=auto_responds, - response_format=response_format, - speed=speed, - modalities=modalities, - ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py index 40c42cb81..49d0caab6 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py @@ -4,34 +4,31 @@ from __future__ import annotations import asyncio -from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from vllm_omni.engine.duplex.plugin import DuplexModelSessionState from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import ( MiniCPMO45PcmAppendBuffer, ) @dataclass(slots=True) -class MiniCPMO45ServingSessionState: - """Mutable serving state owned by one MiniCPM duplex session.""" +class MiniCPMO45ServingSessionState(DuplexModelSessionState): + """Mutable model-owned state of one MiniCPM duplex session (owned by the session runner).""" audio_buffer: MiniCPMO45PcmAppendBuffer = field(default_factory=MiniCPMO45PcmAppendBuffer) input_since_commit: bool = False speech_since_commit: bool = False - native_context_locked: bool = False + context_locked: bool = False committed_audio_payload: dict[str, object] | None = None committed_audio_operation_id: str | None = None committed_audio_reserved_bytes: int = 0 deferred_response_create: bool = False deferred_precreate_response: bool = False - data_plane_task: asyncio.Task[None] | None = None - data_plane_restart_requested: bool = False continuation_owner_id: str | None = None continuation_units: int = 0 pending_silence_task: asyncio.Task[bool] | None = None pending_silence_owner_id: str | None = None - silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None = None def retain_committed_audio( self, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py index 6d9800025..bc13b58f2 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py @@ -6,15 +6,22 @@ from __future__ import annotations import base64 import copy import time +from collections.abc import Callable from contextlib import suppress from dataclasses import dataclass, field from threading import Lock -from typing import Any +from typing import TYPE_CHECKING, Any, cast import numpy as np +from numpy.typing import NDArray from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy +if TYPE_CHECKING: + import torch + from PIL import Image + from torch import nn + _MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() @@ -23,17 +30,19 @@ _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() @dataclass class _MiniCPMO45Stage0SessionState: session_id: str + #: Any: MiniCPM-o remote-code processor (transformers, no importable type). streaming_processor: Any | None = None - audio_buffer: np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32)) + audio_buffer: NDArray[np.float32] = field(default_factory=lambda: np.array([], dtype=np.float32)) audio_chunk_idx: int = 0 - context_embeds: list[Any] = field(default_factory=list) + context_embeds: list[torch.Tensor] = field(default_factory=list) context_token_ids: list[int] = field(default_factory=list) current_turn_ended: bool = True prepared_append_identity: tuple[int | None, int] | None = None - prepared_inputs_embeds: Any | None = None + prepared_inputs_embeds: torch.Tensor | None = None prepared_input_token_ids: list[int] = field(default_factory=list) - prepared_result: dict[str, Any] = field(default_factory=dict) - audio_past_key_values: Any | None = None + prepared_result: dict[str, object] = field(default_factory=dict) + #: Opaque transformers cache, handed back to the model unchanged. + audio_past_key_values: object | None = None pending_terminator_token: int | None = None last_terminator_token: int | None = None pending_speech_context: bool = False @@ -59,7 +68,7 @@ class MiniCPMO45Stage0DuplexRuntime: image_start_token_id: int = -1 image_end_token_id: int = -1 - def __init__(self, stage_model: Any, *, model_path: str | None = None, device: str = "cuda") -> None: + def __init__(self, stage_model: nn.Module, *, model_path: str | None = None, device: str = "cuda") -> None: self.stage_model = stage_model self.model_path = model_path self.device = device @@ -85,7 +94,7 @@ class MiniCPMO45Stage0DuplexRuntime: def _configure_streaming_processor( self, state: _MiniCPMO45Stage0SessionState | None = None, - ) -> Any | None: + ) -> Any | None: # Any: MiniCPM-o remote-code processor processor = self.processor if processor is None: return None @@ -129,9 +138,9 @@ class MiniCPMO45Stage0DuplexRuntime: def _prepare_session_context( self, state: _MiniCPMO45Stage0SessionState, - session_config: dict[str, Any], + session_config: dict[str, object], *, - runtime_config: dict[str, Any] | None = None, + runtime_config: dict[str, object] | None = None, ) -> None: if not self._stage_runtime_ready(): return @@ -163,14 +172,14 @@ class MiniCPMO45Stage0DuplexRuntime: def _stage_prefill_embeddings_only( self, state: _MiniCPMO45Stage0SessionState, - audio_waveform: Any, + audio_waveform: NDArray[np.float32] | None, *, - video_frames: list[Any] | None = None, + video_frames: list[Image.Image] | None = None, epoch: int | None = None, seq: int | None = None, is_speech: bool = False, final: bool = False, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build scheduler-owned Stage0 input embeddings for one audio append. Unlike the legacy worker-control path, this method never calls an eager @@ -217,7 +226,7 @@ class MiniCPMO45Stage0DuplexRuntime: # Omni duplex: encode this append's camera frames so the first unit can # carry them, mirroring official streaming_prefill (feed , then # every frame_list entry as its own block, then the audio). - frame_blocks: list[Any] = [] + frame_blocks: list[list[torch.Tensor]] = [] if video_frames: self._require_vision_token_ids() encoded_frames = self._stage_vision_embeddings(video_frames) @@ -227,7 +236,7 @@ class MiniCPMO45Stage0DuplexRuntime: return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") frame_blocks = encoded_frames - embed_parts: list[Any] = [] + embed_parts: list[torch.Tensor] = [] token_ids: list[int] = [] if state.audio_chunk_idx == 0 and state.context_embeds: embed_parts.extend(state.context_embeds) @@ -336,7 +345,7 @@ class MiniCPMO45Stage0DuplexRuntime: return result @staticmethod - def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, Any]: + def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, object]: return { "success": success, "prefill_success": success, @@ -347,14 +356,14 @@ class MiniCPMO45Stage0DuplexRuntime: } @staticmethod - def _as_2d_tensor(value: Any) -> Any: + def _as_2d_tensor(value: torch.Tensor) -> torch.Tensor: if value.ndim == 1: return value.unsqueeze(0) if value.ndim == 3 and value.shape[0] == 1: return value.squeeze(0) return value - def _embed_token(self, token_id: int) -> Any: + def _embed_token(self, token_id: int) -> torch.Tensor: import torch token = torch.tensor([int(token_id)], dtype=torch.long, device=self._model_device()) @@ -362,7 +371,7 @@ class MiniCPMO45Stage0DuplexRuntime: embeds = embedder(token) return self._as_2d_tensor(embeds) - def _token_embedding_dtype(self) -> Any: + def _token_embedding_dtype(self) -> torch.dtype | None: """dtype of the decoder token embeddings, the unit's reference dtype. Official ``get_vllm_embedding`` casts vision hidden states to the token @@ -376,7 +385,7 @@ class MiniCPMO45Stage0DuplexRuntime: return dtype return getattr(next(self.thinker.parameters()), "dtype", None) - def _token_embedder(self) -> Any: + def _token_embedder(self) -> Callable[[torch.Tensor], torch.Tensor]: nested_embed = getattr(getattr(getattr(self.thinker, "llm", None), "model", None), "embed_tokens", None) if callable(nested_embed): return nested_embed @@ -391,7 +400,7 @@ class MiniCPMO45Stage0DuplexRuntime: return embedder raise AttributeError("MiniCPM-o stage0 model does not expose token embeddings") - def _model_device(self) -> Any: + def _model_device(self) -> torch.device | str: try: return next(self.thinker.parameters()).device except Exception: @@ -402,14 +411,14 @@ class MiniCPMO45Stage0DuplexRuntime: pass return self.device - def _streaming_chunk_size(self, processor: Any | None = None) -> int: + def _streaming_chunk_size(self, processor: Any | None = None) -> int: # Any: remote-code processor processor = processor or self.processor get_chunk = getattr(processor, "get_streaming_chunk_size", None) if callable(get_chunk): return int(get_chunk()) return 16000 - def _sample_rate(self, processor: Any | None = None) -> int: + def _sample_rate(self, processor: Any | None = None) -> int: # Any: remote-code processor processor = processor or self.processor return int( self._stage_param( @@ -418,7 +427,7 @@ class MiniCPMO45Stage0DuplexRuntime: ) ) - def _first_chunk_samples(self, default_chunk_size: int, processor: Any | None = None) -> int: + def _first_chunk_samples(self, default_chunk_size: int, processor: Any | None = None) -> int: # Any: processor processor = processor or self.processor if getattr(processor, "_streaming_mel_processor", None) is None: return default_chunk_size @@ -427,7 +436,7 @@ class MiniCPMO45Stage0DuplexRuntime: def _pad_first_audio_chunk_if_needed( self, state: _MiniCPMO45Stage0SessionState, - processor: Any | None = None, + processor: Any | None = None, # Any: remote-code processor ) -> None: if state.audio_chunk_idx != 0 or len(state.audio_buffer) == 0: return @@ -440,7 +449,7 @@ class MiniCPMO45Stage0DuplexRuntime: padding: np.ndarray = np.zeros(first_chunk_samples - len(state.audio_buffer), dtype=np.float32) state.audio_buffer = np.concatenate([padding, state.audio_buffer]) - def _stage_param(self, name: str, default: Any) -> Any: + def _stage_param(self, name: str, default: object) -> Any: # Any: attribute read off the remote-code model for target in (self.stage_model, self.thinker, getattr(self.thinker, "llm", None)): value = getattr(target, name, None) if value is not None: @@ -455,7 +464,7 @@ class MiniCPMO45Stage0DuplexRuntime: chunk_idx: int, default_chunk_size: int, *, - processor: Any | None = None, + processor: Any | None = None, # Any: remote-code processor ) -> int: processor = processor or self.processor if chunk_idx != 0: @@ -472,11 +481,11 @@ class MiniCPMO45Stage0DuplexRuntime: def _process_streaming_audio( self, - audio_chunk: Any, + audio_chunk: NDArray[np.float32], chunk_idx: int, *, - processor: Any | None = None, - ) -> Any: + processor: Any | None = None, # Any: remote-code processor + ) -> Any: # Any: transformers BatchFeature (or the dict fallback below) processor = processor or self.processor process = getattr(processor, "process_audio_streaming", None) if callable(process): @@ -488,10 +497,10 @@ class MiniCPMO45Stage0DuplexRuntime: def _stage_audio_embeddings( self, - batch_feature: Any, + batch_feature: Any, # Any: transformers BatchFeature *, state: _MiniCPMO45Stage0SessionState | None = None, - ) -> Any | None: + ) -> torch.Tensor | None: if hasattr(batch_feature, "to"): batch_feature = batch_feature.to(self.device) self._ensure_dynamic_cache_compat() @@ -534,17 +543,17 @@ class MiniCPMO45Stage0DuplexRuntime: self.thinker.audio_past_key_values = previous_audio_past_key_values @staticmethod - def _decode_ref_audio_from_session_config(session_config: dict[str, Any]) -> Any | None: + def _decode_ref_audio_from_session_config(session_config: dict[str, object]) -> NDArray[np.float32] | None: from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import decode_native_ref_audio_from_config return decode_native_ref_audio_from_config({"extra_body": session_config}) def _stage_ref_audio_embeddings( self, - ref_audio: Any, + ref_audio: NDArray[np.float32], *, state: _MiniCPMO45Stage0SessionState | None = None, - ) -> Any | None: + ) -> torch.Tensor | None: process_audio = getattr(self.processor, "process_audio", None) if callable(process_audio): batch_feature = process_audio([ref_audio]) @@ -606,16 +615,16 @@ class MiniCPMO45Stage0DuplexRuntime: DynamicCache.get_usable_length = get_usable_length # type: ignore[attr-defined] @staticmethod - def _cat_nested_tensors(value: Any) -> Any | None: + def _cat_nested_tensors(value: object) -> torch.Tensor | None: import torch - tensors = [] + tensors: list[torch.Tensor] = [] - def collect(item: Any) -> None: + def collect(item: object) -> None: if item is None: return if hasattr(item, "detach"): - tensors.append(item) + tensors.append(cast("torch.Tensor", item)) return if isinstance(item, dict): for child in item.values(): @@ -655,7 +664,7 @@ class MiniCPMO45Stage0DuplexRuntime: } @staticmethod - def _load_processor_from_path(model_path: str | None) -> Any | None: + def _load_processor_from_path(model_path: str | None) -> Any | None: # Any: remote-code AutoProcessor if not model_path: return None try: @@ -752,7 +761,7 @@ class MiniCPMO45Stage0DuplexRuntime: return self.stage_padding_token_id() @staticmethod - def _decode_audio_payload(payload: dict[str, Any]) -> Any: + def _decode_audio_payload(payload: dict[str, object]) -> NDArray[np.float32]: audio = payload.get("audio") or payload.get("data") if not isinstance(audio, str): raise ValueError("audio append payload requires base64 audio") @@ -762,7 +771,7 @@ class MiniCPMO45Stage0DuplexRuntime: return np.frombuffer(base64.b64decode(audio), dtype=np.float32) @staticmethod - def _decode_video_frames_payload(payload: dict[str, Any]) -> list[Any]: + def _decode_video_frames_payload(payload: dict[str, object]) -> list[Image.Image]: """Decode omni-duplex camera frames (base64 JPEG/PNG) to PIL images.""" frames = payload.get("video_frames") if not isinstance(frames, list) or not frames: @@ -771,7 +780,7 @@ class MiniCPMO45Stage0DuplexRuntime: from PIL import Image - decoded: list[Any] = [] + decoded: list[Image.Image] = [] for frame_b64 in frames: if not isinstance(frame_b64, str) or not frame_b64: continue @@ -826,7 +835,7 @@ class MiniCPMO45Stage0DuplexRuntime: return [1] return [2] + [1] * (frame_count - 1) - def _encode_processed_vision(self, processed: Any) -> list[Any] | None: + def _encode_processed_vision(self, processed: Any) -> list[torch.Tensor] | None: # Any: BatchFeature targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)) for target in targets: if target is None: @@ -842,8 +851,8 @@ class MiniCPMO45Stage0DuplexRuntime: device, dtype = vpm_param.device, vpm_param.dtype pixel_nested = processed["pixel_values"] tgt_nested = processed["tgt_sizes"] - flat_pixels: list[Any] = [] - flat_tgt: list[Any] = [] + flat_pixels: list[torch.Tensor] = [] + flat_tgt: list[torch.Tensor] = [] for image_slices, image_tgt in zip(pixel_nested, tgt_nested): for slice_pixels in image_slices: flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) @@ -858,7 +867,7 @@ class MiniCPMO45Stage0DuplexRuntime: return None expected = MiniCPMO45DuplexPolicy.VISION_EMBEDS_PER_FRAME embed_dtype = self._token_embedding_dtype() - out: list[Any] = [] + out: list[torch.Tensor] = [] for block in hidden: block_2d = self._as_2d_tensor(block) if int(block_2d.shape[0]) != expected: @@ -869,7 +878,7 @@ class MiniCPMO45Stage0DuplexRuntime: return out return None - def _stage_vision_embeddings(self, frames: list[Any]) -> list[list[Any]] | None: + def _stage_vision_embeddings(self, frames: list[Image.Image]) -> list[list[torch.Tensor]] | None: """Encode camera frames for omni duplex via the loaded vision tower. Official ``streaming_prefill`` encodes each ``frame_list`` entry as its @@ -883,7 +892,7 @@ class MiniCPMO45Stage0DuplexRuntime: process_image = getattr(self.processor, "process_image", None) if not callable(process_image): return None - out: list[list[Any]] = [] + out: list[list[torch.Tensor]] = [] for frame, max_slices in zip(frames, self._official_max_slice_nums(len(frames))): try: processed = process_image([frame], max_slice_nums=max_slices) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 60252e78c..b1c8db89c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -157,15 +157,13 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp del sampling_metadata self._minicpmo45_active_duplex_rows = [row.row_idx for row in rows] self._minicpmo45_duplex_row_sessions = { - row.row_idx: (row.session_id, row.incarnation) for row in rows if row.session_id is not None + row.row_idx: row.session_id for row in rows if row.session_id is not None } request_sessions = getattr(self, "_minicpmo45_duplex_request_sessions", None) if not isinstance(request_sessions, dict): request_sessions = {} self._minicpmo45_duplex_request_sessions = request_sessions - request_sessions.update( - {row.request_id: (row.session_id, row.incarnation) for row in rows if row.session_id is not None} - ) + request_sessions.update({row.request_id: row.session_id for row in rows if row.session_id is not None}) self._minicpmo45_duplex_row_payloads = {row.row_idx: row.payload for row in rows if row.payload is not None} self._minicpmo45_duplex_row_max_tokens = { row.row_idx: row.max_tokens for row in rows if row.max_tokens is not None @@ -199,7 +197,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp force_listen = payload.get("force_listen") is True is_speech = payload.get("is_speech") segment_key = (row.request_id, row.seq if row.seq is not None else -1) - session_key = (row.session_id, row.incarnation) if row.session_id is not None else None + session_key = row.session_id if turn_eos_id >= 0 and session_key is not None: state = helper_sessions.get(session_key) if isinstance(helper_sessions, dict) else None pending_speech_context = ( @@ -326,16 +324,12 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp helper = self._duplex_data_plane_helper() session_id = str(duplex.get("session_id") or "") - try: - incarnation = int(duplex.get("incarnation", 0)) - except (TypeError, ValueError): - incarnation = 0 payload = duplex.get("payload") if not session_id or not isinstance(payload, dict): embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": "bad_duplex_payload"}} - session_key = (session_id, incarnation) + session_key = session_id state = helper.sessions.get(session_key) if state is None: from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index e48328d0b..e75393845 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -24,13 +24,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( model_type="minicpmo_4_5", default_deploy_config_name="minicpmo_4_5.yaml", model_arch="MiniCPMO45OmniForConditionalGeneration", - duplex_runtime_extension=( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime.MiniCPMO45DuplexRuntimeExtension" - ), - duplex_serving_adapter=( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter" - ), - duplex_control_enabled=True, + duplex_plugin="vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin.MiniCPMO45DuplexPlugin", # MiniCPM-o 4.5's HF config.json reports `model_type="minicpmo"` and # `architectures=["MiniCPMO"]` — both shared verbatim with older MiniCPM-o # 1.0 / 2.6 checkpoints. The only field distinguishing the generations is diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index 80995f92e..1c374be85 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -690,9 +690,6 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | session_id = duplex_state.get("session_id") if isinstance(session_id, str) and session_id: metadata["session_id"] = session_id - incarnation = duplex_state.get("incarnation") - if isinstance(incarnation, int): - metadata["incarnation"] = incarnation epoch = duplex_state.get("epoch") if isinstance(epoch, int): metadata["epoch"] = epoch diff --git a/vllm_omni/utils/audio_resample.py b/vllm_omni/utils/audio_resample.py new file mode 100644 index 000000000..fc6230f6e --- /dev/null +++ b/vllm_omni/utils/audio_resample.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Streaming resampling, shared by the engine and the entrypoints. + +Lives here rather than beside the OpenAI audio mixin because the duplex +engine needs it too, and engine code must not import the entrypoints layer. +Pure numpy on purpose: importing this must not pull in torch or the API +server. +""" + +from __future__ import annotations + +import math + +import numpy as np + +__all__ = ["StreamingAudioResampler"] + + +class StreamingAudioResampler: + """Stateful polyphase resampler for streaming mono float audio. + + Retains filter state so output is invariant to input chunk boundaries. + """ + + _half_filter_width = 10 + _kaiser_beta = 5.0 + max_polyphase_factor = 2_048 + + def __init__(self, source_rate: int, target_rate: int): + if source_rate <= 0 or target_rate <= 0: + raise ValueError("Audio sample rates must be positive") + self.source_rate = source_rate + self.target_rate = target_rate + rate_gcd = math.gcd(source_rate, target_rate) + self._up = target_rate // rate_gcd + self._down = source_rate // rate_gcd + if max(self._up, self._down) > self.max_polyphase_factor: + raise ValueError("source and target sample rates have an unsupported resampling ratio") + self._half_len, self._phase_kernels = self._design_polyphase_filter() + self._history_samples = self._phase_kernels.shape[1] - 1 + self.reset() + + def _design_polyphase_filter(self) -> tuple[int, np.ndarray]: + if self._up == self._down: + return 0, np.ones((1, 1), dtype=np.float32) + + max_rate = max(self._up, self._down) + half_len = self._half_filter_width * max_rate + offsets = np.arange(-half_len, half_len + 1, dtype=np.float64) + cutoff = 1.0 / max_rate + taps = cutoff * np.sinc(cutoff * offsets) + taps *= np.kaiser(taps.size, self._kaiser_beta) + taps /= np.sum(taps) + taps *= self._up + taps = np.ascontiguousarray(taps, dtype=np.float32) + + phases = [taps[phase :: self._up] for phase in range(self._up)] + phase_width = max(phase.size for phase in phases) + kernels = np.zeros((self._up, phase_width), dtype=np.float32) + for phase_index, phase in enumerate(phases): + kernels[phase_index, -phase.size :] = phase[::-1] + return half_len, kernels + + @property + def scratch_bytes(self) -> int: + pending_samples = self._ceil_div(self._input_samples * self._up, self._down) - self._output_samples + return max(0, pending_samples) * np.dtype(np.float32).itemsize + + @staticmethod + def _ceil_div(numerator: int, denominator: int) -> int: + return -(-numerator // denominator) + + def _render_outputs( + self, + combined: np.ndarray, + *, + first_output: int, + output_count: int, + input_start: int, + ) -> np.ndarray: + if output_count <= 0: + return np.empty(0, dtype=np.float32) + + output_indexes = np.arange(first_output, first_output + output_count, dtype=np.int64) + filter_indexes = output_indexes * self._down + self._half_len + source_indexes = filter_indexes // self._up + phases = filter_indexes % self._up + windows = np.lib.stride_tricks.sliding_window_view( + combined, + self._phase_kernels.shape[1], + ) + selected_windows = windows[source_indexes - input_start] + return np.sum( + selected_windows * self._phase_kernels[phases], + axis=1, + dtype=np.float32, + ) + + def _push(self, chunk: np.ndarray) -> np.ndarray: + if chunk.size == 0: + return np.empty(0, dtype=np.float32) + if self._flushed: + raise RuntimeError("cannot process audio after the resampler has been finalized; call reset first") + + combined = np.concatenate((self._history, chunk)) + next_input_samples = self._input_samples + chunk.size + stable_output_samples = max( + 0, + self._ceil_div(next_input_samples * self._up - self._half_len, self._down), + ) + output = self._render_outputs( + combined, + first_output=self._output_samples, + output_count=stable_output_samples - self._output_samples, + input_start=self._input_samples, + ) + if self._history_samples: + self._history = combined[-self._history_samples :].copy() + self._input_samples = next_input_samples + self._output_samples = stable_output_samples + return output + + def _flush(self) -> np.ndarray: + if self._flushed: + return np.empty(0, dtype=np.float32) + + total_output_samples = self._ceil_div(self._input_samples * self._up, self._down) + output_count = total_output_samples - self._output_samples + if output_count: + last_output = total_output_samples - 1 + last_source_index = (last_output * self._down + self._half_len) // self._up + right_padding = max(0, last_source_index - self._input_samples + 1) + combined = np.concatenate((self._history, np.zeros(right_padding, dtype=np.float32))) + output = self._render_outputs( + combined, + first_output=self._output_samples, + output_count=output_count, + input_start=self._input_samples, + ) + else: + output = np.empty(0, dtype=np.float32) + self._output_samples = total_output_samples + self._flushed = True + return output + + def process(self, audio: np.ndarray, *, final: bool = False) -> np.ndarray: + chunk = np.asarray(audio, dtype=np.float32) + if chunk.ndim != 1: + raise ValueError(f"Streaming audio resampling only supports mono audio, got shape {chunk.shape}") + chunk = np.ascontiguousarray(chunk, dtype=np.float32) + output = self._push(chunk) + if not final: + return output + tail = self._flush() + if not output.size: + return tail + if not tail.size: + return output + return np.concatenate((output, tail)) + + def reset(self) -> None: + self._history = np.zeros(self._history_samples, dtype=np.float32) + self._input_samples = 0 + self._output_samples = 0 + self._flushed = False From 507cb1d834db66f2dbb776602e82c62d58e92992 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:47:31 +0800 Subject: [PATCH 106/209] [Attention] [P1] Keep H3 VSA policy in the model and share sparse primitives (#7535) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- .../feature/attention_backend_selection.md | 38 ++- .../attention_backends/fastvideo_vsa.md | 78 ++++-- pyproject.toml | 5 + .../attention/test_block_sparse_ops.py | 126 +++++++++ .../diffusion/attention/test_fastvideo_vsa.py | 164 +---------- .../attention/test_impl_overrides.py | 136 +++++++++ .../cache/test_teacache_extractors.py | 10 +- .../minimax_h3/test_minimax_h3_modulation.py | 2 +- .../models/minimax_h3/test_vsa_layout.py | 221 +++++++++++++++ tests/helpers/runtime.py | 62 ++++- .../tests/test_runtime_server_ports.py | 109 ++++++++ .../attention/backends/fastvideo_vsa.py | 258 +----------------- vllm_omni/diffusion/attention/layer.py | 34 ++- .../diffusion/attention/ops/block_sparse.py | 115 ++++++++ .../diffusion/cache/teacache/extractors.py | 2 +- .../indexed_modulation.py} | 4 +- .../models/minimax_h3/attention/__init__.py | 2 + .../minimax_h3/attention/fastvideo_h3.py | 209 ++++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 15 +- 19 files changed, 1117 insertions(+), 473 deletions(-) create mode 100644 tests/diffusion/attention/test_block_sparse_ops.py create mode 100644 tests/diffusion/attention/test_impl_overrides.py create mode 100644 tests/diffusion/models/minimax_h3/test_vsa_layout.py create mode 100644 tests/helpers/tests/test_runtime_server_ports.py create mode 100644 vllm_omni/diffusion/attention/ops/block_sparse.py rename vllm_omni/diffusion/{attention/ops/minimax_h3_modulation.py => layers/indexed_modulation.py} (98%) create mode 100644 vllm_omni/diffusion/models/minimax_h3/attention/__init__.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md index 41cbaf3ff..3ae1115b1 100644 --- a/docs/design/feature/attention_backend_selection.md +++ b/docs/design/feature/attention_backend_selection.md @@ -76,6 +76,34 @@ preserve common `self` or `cross` policies. A model also decides whether its path is compatible with automatic TRTLLM selection, because only the model knows whether masking and packing satisfy the kernel contract. +### Model-owned implementation specializations + +Use `Attention(..., impl_overrides={backend_name: implementation_class})` +when a model needs different computation within an existing backend contract. +The model defines the implementation in its own directory and passes it when +constructing the layer: + +```python +# MyModelVSAImpl is a model-owned subclass of FastVideoVSAImpl. +self.attention = Attention( + num_heads=num_heads, + head_size=head_size, + softmax_scale=head_size**-0.5, + causal=False, + role="self", + impl_overrides={"FASTVIDEO_VSA": MyModelVSAImpl}, +) +``` + +The override is applied after role/platform selection, keyed by the selected +backend's `get_name()`. It must subclass the **selected implementation**; +incompatible platform implementations raise `TypeError`. Unselected overrides +have no effect. Backend capabilities, constructor options, and shared +parallel dispatch remain unchanged and must be honored by the specialization. + +Use a new backend for a new provider/capability contract. `custom_attention` +instead owns communication and requires `skip_sequence_parallel=True`. + ## Adding or changing a backend An implementation change is complete only when it: @@ -91,8 +119,9 @@ An implementation change is complete only when it: ## FastVideo VSA model metadata -FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in -the original latent video grid. Wan integrations therefore attach the +FastWan's text-to-video and image-to-video modes both use `Wan22Pipeline`. +It operates on a flattened DiT sequence but partitions tokens in the original +latent video grid. Wan integrations therefore attach the post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The backend validates that `T * H * W` equals the sequence length, derives the runtime block count, and selects the configured top-k key/value blocks for @@ -104,3 +133,8 @@ the checkpoint contains those weights. Checkpoints without the projection do not allocate or execute it. When top-k selects every block, native checkpoints route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path to retain checkpoint semantics. + +MiniMax-H3 supplies `MiniMaxH3VSAImpl` through `impl_overrides`. +`models/minimax_h3/attention/fastvideo_h3.py` owns its prefix layout, video tiling, and +learned compression gate. Pooling, block-map construction, and tile64 provider +calls are shared through `attention/ops/block_sparse.py`. diff --git a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md index c03b9ffde..df6e34b0c 100644 --- a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md +++ b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md @@ -1,24 +1,33 @@ # FastVideo VSA -FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's -self-attention by partitioning the post-patch latent grid into spatiotemporal +FastVideo Variable Sparse Attention (VSA) partitions the post-patch latent grid into spatiotemporal blocks. For every query block, VSA scores the key/value blocks and computes attention only against the selected top-k blocks. -The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope. +## Supported models -MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route -through the same backend. See the -[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry, -supported topologies, and commands. +| Model / checkpoint | Required adapter | Tasks | Sequence parallelism | +| --- | --- | --- | --- | +| `FastVideo/FastWan2.2-TI2V-5B-Diffusers` | None | T2V, I2V through `Wan22Pipeline` | Disabled | +| `MiniMaxAI/MiniMax-H3` | [FastH3 VSA (recipe)](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter) | T2VA | Disabled or pure Ulysses | +| Wan I2V-14B, S2V, VACE | — | Unsupported | — | -VSA is a CUDA-only, explicitly selected backend. It requires the -`fastvideo-kernel` package and currently supports non-causal self-attention -with equal query and key/value sequence lengths. Unsupported shapes, masks, -dtypes, sequence-parallel execution, or kernel failures fall back to -`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active -sequence-parallel context is one of those fallbacks; the H3 route supports pure -Ulysses and rejects ring or all-gather sequence parallelism at startup. +## Installation + +Install vLLM-Omni from this checkout with the `vsa` extra: + +```bash +uv pip install -e '.[vsa]' +``` + +The extra installs the tested kernel dependency automatically. Prebuilt kernels require +Linux, Python 3.12, and glibc 2.34 or newer (x86-64 or aarch64). The full +FastVideo framework and provider environment variables are not required. + +`FASTVIDEO_VSA` selects the attention algorithm. On SM120, the H3 integration +currently executes FastVideo's 64-token Triton block-sparse kernel; it does not +dispatch to FlashInfer. The `vsa` extra installs `fastvideo-kernel==0.3.4` for +this execution path. ## Enable the backend @@ -27,24 +36,18 @@ flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained for every query block: ```bash -vllm-omni serve \ +vllm serve --omni \ --diffusion-attention-backend FASTVIDEO_VSA \ --fastvideo-vsa-topk 64 ``` -The backwards-compatible environment variable selects the backend with the -default `topk=64`: +

+Alternative configuration formats + +Equivalent structured configuration: ```bash -export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA -vllm-omni serve -``` - -To tune top-k, pass the CLI backend and top-k flags together as shown above, -or use the equivalent structured configuration: - -```bash -vllm-omni serve \ +vllm serve --omni \ --diffusion-attention-config \ '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}' ``` @@ -73,8 +76,15 @@ stages: fastvideo_vsa_topk: 64 ``` +
+ ## Choose top-k +H3 uses 64-token video blocks and keeps +all prefix blocks; see the [FastH3 VSA recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter). + +### Wan top-k behavior + At runtime the backend logs the sequence shape and derived block count: ```text @@ -117,8 +127,16 @@ selecting VSA does not turn a native checkpoint into a distilled model. ## Verify routing and fallback -Check the startup and first-forward logs instead of assuming that selecting -the backend guarantees sparse execution: +VSA is a CUDA-only, explicitly selected backend. It requires the +`fastvideo-kernel` package and currently supports non-causal self-attention +with equal query and key/value sequence lengths. Unsupported shapes, masks, +dtypes, sequence-parallel execution, or kernel failures fall back to +`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active +sequence-parallel context is one of those fallbacks; the H3 route supports pure +Ulysses and rejects ring or all-gather sequence parallelism at startup. +H3 accelerator faults propagate instead of attempting dense recovery. + +On the Wan route, check the startup and first-forward logs: - `route=VSA` means top-k block selection is active. - `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all @@ -126,6 +144,10 @@ the backend guarantees sparse execution: - `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA executed; the warning includes the reason. +For H3, check `FastH3 adapter active` at startup and +`FASTVIDEO_VSA H3 routing` during DiT execution, as described in the +[model recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter). + The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks, standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask, and no active sequence-parallel context. The MiniMax-H3 route uses 64-token diff --git a/pyproject.toml b/pyproject.toml index fc1d5573b..434c4305c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,6 +123,11 @@ fa4 = [ "flash-attn-4[cu13]==4.0.0b18", ] +# FastVideo VSA kernels for the supported Linux/Python 3.12 wheel environment. +vsa = [ + "fastvideo-kernel==0.3.4", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", diff --git a/tests/diffusion/attention/test_block_sparse_ops.py b/tests/diffusion/attention/test_block_sparse_ops.py new file mode 100644 index 000000000..845f09b46 --- /dev/null +++ b/tests/diffusion/attention/test_block_sparse_ops.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import sys +import types + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from vllm_omni.diffusion.attention.ops.block_sparse import ( + build_prefix_dense_block_map, + fastvideo_block_sparse_attn_bshd, + mean_pool_tiles, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.mark.parametrize("prefix_blocks,topk", [(0, 1), (2, 0), (2, 1), (2, 3), (2, 5)]) +def test_block_map_keeps_prefix_dense_and_selects_non_prefix_keys(prefix_blocks, topk): + sparse_blocks = 3 + blocks = prefix_blocks + sparse_blocks + scores = torch.arange(2 * blocks * blocks, dtype=torch.float32).reshape(1, 2, blocks, blocks) + block_map = build_prefix_dense_block_map(scores, prefix_blocks, sparse_blocks, topk) + expected = torch.zeros_like(scores, dtype=torch.bool) + # Scores increase with key index, so top-k always selects the last keys. + for query_block in range(blocks): + for key_block in range(blocks): + if ( + query_block < prefix_blocks + or key_block < prefix_blocks + or key_block >= blocks - min(topk, sparse_blocks) + ): + expected[..., query_block, key_block] = True + assert torch.equal(block_map, expected) + + +@pytest.mark.parametrize("block_size", [4, 32, 64]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_pooling_uses_valid_tile_sizes_and_fp32_accumulation(block_size, dtype): + sizes = torch.tensor([block_size, block_size - 1, 0], dtype=torch.int32) + tiles = torch.arange(2 * 3 * block_size * 2 * 3).reshape(2, 3, block_size, 2, 3).to(dtype) + tiles[:, 1, -1] = 0 + tiles[:, 2] = 0 + expected = torch.stack( + [tiles[:, 0].float().mean(dim=1), tiles[:, 1, :-1].float().mean(dim=1), torch.zeros(2, 2, 3)], dim=2 + ) + result = mean_pool_tiles(tiles.flatten(1, 2), sizes, block_size) + assert result.dtype == torch.float32 + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +def test_provider_wrapper_restores_transport_rows_and_bshd_layout(monkeypatch): + monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False) + calls = [] + + def block_sparse_attn(q, k, v, block_map, sizes): + calls.append((q.shape, block_map.shape, sizes.tolist())) + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert sizes.dtype == torch.int32 + return q + k + v, None + + module = types.ModuleType("fastvideo_kernel.block_sparse_attn") + setattr(module, "block_sparse_attn", block_sparse_attn) + monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel")) + monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module) + query = torch.randn(1, 4 * 64, 2, 8) + block_map = torch.ones(1, 2, 4, 4, dtype=torch.bool) + sizes = torch.tensor([64, 17, 64, 0]) + output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3) + expected = torch.zeros_like(query) + expected[:, : 3 * 64] = query[:, : 3 * 64] * 3 + torch.testing.assert_close(output, expected) + assert calls == [(torch.Size([1, 2, 3 * 64, 8]), torch.Size([1, 2, 3, 3]), [64, 17, 64])] + + +def test_custom_op_fake_preserves_metadata_without_loading_provider(monkeypatch): + monkeypatch.setitem(sys.modules, "fastvideo_kernel", None) + with FakeTensorMode(): + query = torch.empty(1, 4 * 64, 2, 8, dtype=torch.bfloat16) + block_map = torch.empty(1, 2, 4, 4, dtype=torch.bool) + sizes = torch.empty(4, dtype=torch.int32) + output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3) + assert output.shape == query.shape + assert output.dtype == query.dtype + assert output.device == query.device + + +@pytest.mark.parametrize("error_type", [ImportError, RuntimeError, torch.AcceleratorError]) +def test_native_provider_failure_only_recovers_without_accelerator_fault(monkeypatch, mocker, error_type): + monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") + error = error_type("native provider failure") + native = mocker.Mock() + native.is_supported.return_value = True + native.block_sparse_attn_sm100a.side_effect = error + provider = types.ModuleType("fastvideo_kernel") + setattr(provider, "block_sparse_attn_sm100a", native) + index = types.ModuleType("fastvideo_kernel.triton_kernels.index") + setattr( + index, + "map_to_index", + mocker.Mock(return_value=(torch.zeros(1, dtype=torch.int32), torch.ones(1, dtype=torch.int32))), + ) + triton = types.ModuleType("fastvideo_kernel.block_sparse_attn") + fallback = mocker.Mock(side_effect=lambda q, *args: (q.clone(), None)) + setattr(triton, "block_sparse_attn", fallback) + monkeypatch.setitem(sys.modules, "fastvideo_kernel", provider) + monkeypatch.setitem( + sys.modules, "fastvideo_kernel.triton_kernels", types.ModuleType("fastvideo_kernel.triton_kernels") + ) + monkeypatch.setitem(sys.modules, "fastvideo_kernel.triton_kernels.index", index) + monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", triton) + query = torch.arange(128 * 2 * 8, dtype=torch.float32).reshape(1, 128, 2, 8) + block_map = torch.ones(1, 2, 2, 2, dtype=torch.bool) + sizes = torch.tensor([64, 64], dtype=torch.int32) + + if error_type is torch.AcceleratorError: + with pytest.raises(torch.AcceleratorError, match="native provider failure") as raised: + fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2) + assert raised.value is error + fallback.assert_not_called() + else: + output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2) + torch.testing.assert_close(output, query, rtol=0, atol=0) + fallback.assert_called_once() diff --git a/tests/diffusion/attention/test_fastvideo_vsa.py b/tests/diffusion/attention/test_fastvideo_vsa.py index 1bd7192ad..e5c854d8d 100644 --- a/tests/diffusion/attention/test_fastvideo_vsa.py +++ b/tests/diffusion/attention/test_fastvideo_vsa.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -import math import sys import types from typing import Any @@ -12,14 +11,10 @@ import torch from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, PackedPaddingMetadata, - VideoTokenLayout, - VideoTokenSpan, ) from vllm_omni.diffusion.attention.backends.fastvideo_vsa import ( FastVideoVSABackend, FastVideoVSAImpl, - _build_h3_block_map, - _get_h3_tile_metadata, ) from vllm_omni.diffusion.attention.backends.registry import ( DiffusionAttentionBackendEnum, @@ -218,144 +213,8 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks(): assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2" -def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None): - extra: dict[str, Any] = { - "vsa_h3_prefix_segments": prefix_segments, - } - if gate is not None: - extra["gate_compress"] = gate - return AttentionMetadata( - packed_padding=packed_padding, - video_layout=VideoTokenLayout( - used_len=sum(prefix_segments) + math.prod(video_shape), - video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),), - ), - extra=extra, - ) - - -def _fake_block_sparse_kernel(monkeypatch, calls): - """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA.""" - - def block_sparse_attn(q, k, v, block_map, variable_block_sizes): - calls["block_map"] = block_map.clone() - calls["variable_block_sizes"] = variable_block_sizes.tolist() - calls["q_shape"] = tuple(q.shape) - blocks = block_map.shape[-1] - assert blocks * 64 == q.shape[2] - # Expand the per-block map to a per-token mask so the stand-in honours - # the sparsity contract the real kernel implements, and drop the rows - # each partial block pads with, which the real kernel skips through - # variable_block_sizes. - token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2) - within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None] - token_mask = token_mask & within.reshape(1, 1, 1, -1) - out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask) - return out.to(q.dtype), None - - module = types.ModuleType("fastvideo_kernel.block_sparse_attn") - setattr(module, "block_sparse_attn", block_sparse_attn) - monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel")) - monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module) - - -def _h3_impl(**backend_kwargs): - kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False} - kwargs.update(backend_kwargs) - return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs) - - -def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch): - calls: dict[str, Any] = {} - _fake_block_sparse_kernel(monkeypatch, calls) - impl = _h3_impl(topk=64) - prefix_segments, video_shape = (5, 70), (2, 4, 4) - seq_len = sum(prefix_segments) + 2 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) - - output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - - assert output.shape == query.shape - # topk covers every video block, so the block map is dense and the tiled - # kernel must reproduce plain attention over the same rows. - reference = torch.nn.functional.scaled_dot_product_attention( - *(x.float().transpose(1, 2) for x in (query, query, query)) - ).transpose(1, 2) - torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2) - # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile. - assert calls["variable_block_sizes"] == [5, 64, 6, 32] - - -def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch): - calls: dict[str, Any] = {} - _fake_block_sparse_kernel(monkeypatch, calls) - impl = _h3_impl(topk=1) - prefix_segments, video_shape = (5,), (8, 4, 4) - seq_len = 5 + 8 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) - - impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - - # 1 prefix block + 2 video blocks. The odd count is padded to an even one - # for the paired-CTA kernel contract, and the Triton route slices that - # transport-only partner off again before launching. - block_map = calls["block_map"] - assert block_map.shape[-1] == 3 - assert calls["variable_block_sizes"] == [5, 64, 64] - assert block_map[:, :, :1, :].all(), "prefix queries stay dense" - assert block_map[..., :1].all(), "prefix keys are exempt from selection" - assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query" - - # The same odd geometry must still untile back onto the original rows. - dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - reference = torch.nn.functional.scaled_dot_product_attention( - *(x.float().transpose(1, 2) for x in (query, query, query)) - ).transpose(1, 2) - torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2) - - -def test_h3_forward_applies_the_learned_compression_gate(monkeypatch): - _fake_block_sparse_kernel(monkeypatch, {}) - impl = _h3_impl(topk=64) - prefix_segments, video_shape = (5,), (2, 4, 4) - seq_len = 5 + 2 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) - - without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - zero_gate = torch.zeros_like(query) - with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate)) - live_gate = torch.full_like(query, 0.5) - with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate)) - - torch.testing.assert_close(with_zero_gate.float(), without_gate.float()) - assert not torch.allclose(with_live_gate.float(), without_gate.float()) - - -def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch): - _fake_block_sparse_kernel(monkeypatch, {}) - impl = _h3_impl(topk=64) - prefix_segments, video_shape = (5,), (2, 4, 4) - valid = 5 + 2 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16) - padding = PackedPaddingMetadata( - q_length=valid, - kv_length=valid, - cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), - cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), - ) - - output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding)) - - assert output.shape == query.shape - assert torch.count_nonzero(output[:, valid:]) == 0 - - def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch): - # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips + # The backend advertises supports_packed_mask_free, so a model may skip # building the padding mask. Every fallback out of the VSA path must honour # that contract itself; SDPA reads attn_mask and nothing else. impl = FastVideoVSAImpl( @@ -374,6 +233,8 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch): cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), ) + monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda) + baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={})) perturbed_input = query.clone() perturbed_input[:, valid:] += 100.0 @@ -384,22 +245,3 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch): assert baseline.shape == query.shape torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid]) assert torch.count_nonzero(baseline[:, valid:]) == 0 - - -def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d(): - partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( - (5, 70, 9), (5, 6, 6), torch.device("cpu") - ) - assert sizes[:4].tolist() == [5, 64, 6, 9] - assert prefix_blocks == 4 - assert video_blocks == 8 - assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6 - assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum()) - - -def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt(): - scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5) - block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1) - assert block_map[:, :, :2].all() - assert block_map[..., :2].all() - assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all() diff --git a/tests/diffusion/attention/test_impl_overrides.py b/tests/diffusion/attention/test_impl_overrides.py new file mode 100644 index 000000000..f7755de5c --- /dev/null +++ b/tests/diffusion/attention/test_impl_overrides.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.diffusion.attention import layer +from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl +from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend +from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention +from vllm_omni.diffusion.data import AttentionSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _SpecializedImpl(FastVideoVSAImpl): + pass + + +@pytest.fixture +def local_attention(monkeypatch): + monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None) + monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + + +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("override", [False, True]) +def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override): + spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None + select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec)) + attention = layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + role="video.self", + role_category="self", + qkv_layout="BSND", + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None, + ) + + assert attention.attn_backend is FastVideoVSABackend + assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl) + assert attention.attention.topk == (7 if explicit else 64) + assert attention.attention.qkv_layout == "BSND" + assert attention.backend_pref == "FASTVIDEO_VSA" + assert attention.backend_explicit is explicit + assert attention.attn_spec is spec + assert isinstance(attention.parallel_strategy, NoParallelAttention) + select.assert_called_once_with( + role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False + ) + + +def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention): + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None)) + attention = layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl}, + ) + assert attention.attn_backend is SDPABackend + assert type(attention.attention) is SDPABackend.get_impl_cls() + + +def test_specialization_does_not_bypass_selection_errors(mocker, local_attention): + mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable")) + with pytest.raises(ImportError, match="optional kernel unavailable"): + layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl}, + ) + + +def test_specialization_must_extend_selected_implementation(mocker, local_attention): + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None)) + with pytest.raises(TypeError, match="must subclass the selected implementation"): + layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()}, + ) + + +def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention): + from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import MiniMaxH3VSAImpl + + class PlatformBackend(FastVideoVSABackend): + @classmethod + def supports_packed_mask_free(cls): + return False + + @staticmethod + def get_supported_head_sizes(): + return [128] + + # A registry/platform may refine capabilities without changing the + # implementation. The model must not replace those capability decisions. + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None)) + attention = layer.Attention( + num_heads=2, + head_size=128, + softmax_scale=128**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl}, + ) + assert attention.attn_backend is PlatformBackend + assert not attention.attn_backend.supports_packed_mask_free() + assert attention.attn_backend.get_supported_head_sizes() == [128] + assert type(attention.attention) is MiniMaxH3VSAImpl + + +def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention): + class PlatformImpl(FastVideoVSAImpl): + pass + + class PlatformBackend(FastVideoVSABackend): + @staticmethod + def get_impl_cls(): + return PlatformImpl + + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None)) + with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"): + layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl}, + ) diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index d821541d0..55b92b820 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -68,12 +68,12 @@ class BaseExtractorTest(ABC): pass @abstractmethod - def get_module(self): + def get_module(self, module, /): """Return model module instance.""" pass @abstractmethod - def get_sample_inputs(self): + def get_sample_inputs(self, sample_inputs, /): """Return sample inputs for model.""" pass @@ -737,10 +737,10 @@ class TestMiniMaxH3Extractor(BaseExtractorTest): monkeypatch, ): """Strict SP must not mix local TeaCache state with gathered block rows.""" - from vllm_omni.diffusion.attention.ops import minimax_h3_modulation from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.layers import indexed_modulation seq_len = sample_inputs["x"].shape[1] local_len = seq_len // 2 @@ -844,13 +844,13 @@ class TestMiniMaxH3Extractor(BaseExtractorTest): ) monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group) - original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_ + original_indexed_scale_shift = indexed_modulation.indexed_scale_shift_ def capture_indexed_scale_shift(hidden, shift, scale, indices): captured_modulation_indices.append(indices.clone()) return original_indexed_scale_shift(hidden, shift, scale, indices) - monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift) + monkeypatch.setattr(indexed_modulation, "indexed_scale_shift_", capture_indexed_scale_shift) local_img_pos = torch.tensor([0, 1]) local_audio_pos = torch.tensor([2, 3]) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index e10fbe56a..7fde01354 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -6,7 +6,7 @@ import torch from vllm.triton_utils import HAS_TRITON from tests.helpers.mark import hardware_marks -from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( +from vllm_omni.diffusion.layers.indexed_modulation import ( _MAX_1D_GRID_SIZE, _iter_row_chunks, _launch_row_chunks, diff --git a/tests/diffusion/models/minimax_h3/test_vsa_layout.py b/tests/diffusion/models/minimax_h3/test_vsa_layout.py new file mode 100644 index 000000000..c556543d0 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_vsa_layout.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import math +import sys +import types +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, + PackedPaddingMetadata, + VideoTokenLayout, + VideoTokenSpan, +) +from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import ( + MiniMaxH3VSAImpl, + _get_h3_tile_metadata, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None): + extra: dict[str, Any] = { + "vsa_h3_prefix_segments": prefix_segments, + } + if gate is not None: + extra["gate_compress"] = gate + return AttentionMetadata( + packed_padding=packed_padding, + video_layout=VideoTokenLayout( + used_len=sum(prefix_segments) + math.prod(video_shape), + video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),), + ), + extra=extra, + ) + + +def _fake_block_sparse_kernel(monkeypatch, calls): + """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA.""" + + def block_sparse_attn(q, k, v, block_map, variable_block_sizes): + calls["block_map"] = block_map.clone() + calls["variable_block_sizes"] = variable_block_sizes.tolist() + calls["q_shape"] = tuple(q.shape) + blocks = block_map.shape[-1] + assert blocks * 64 == q.shape[2] + # Expand the per-block map to a per-token mask so the stand-in honours + # the sparsity contract the real kernel implements, and drop the rows + # each partial block pads with, which the real kernel skips through + # variable_block_sizes. + token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2) + within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None] + token_mask = token_mask & within.reshape(1, 1, 1, -1) + out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask) + return out.to(q.dtype), None + + module = types.ModuleType("fastvideo_kernel.block_sparse_attn") + setattr(module, "block_sparse_attn", block_sparse_attn) + monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel")) + monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module) + + +def _h3_impl(**backend_kwargs): + kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False} + kwargs.update(backend_kwargs) + return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs) + + +def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch): + calls: dict[str, Any] = {} + _fake_block_sparse_kernel(monkeypatch, calls) + impl = _h3_impl(topk=64) + prefix_segments, video_shape = (5, 70), (2, 4, 4) + seq_len = sum(prefix_segments) + 2 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) + + output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + + assert output.shape == query.shape + # topk covers every video block, so the block map is dense and the tiled + # kernel must reproduce plain attention over the same rows. + reference = torch.nn.functional.scaled_dot_product_attention( + *(x.float().transpose(1, 2) for x in (query, query, query)) + ).transpose(1, 2) + torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2) + # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile. + assert calls["variable_block_sizes"] == [5, 64, 6, 32] + + +def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch): + calls: dict[str, Any] = {} + _fake_block_sparse_kernel(monkeypatch, calls) + impl = _h3_impl(topk=1) + prefix_segments, video_shape = (5,), (8, 4, 4) + seq_len = 5 + 8 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) + + impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + + # 1 prefix block + 2 video blocks. The odd count is padded to an even one + # for the paired-CTA kernel contract, and the Triton route slices that + # transport-only partner off again before launching. + block_map = calls["block_map"] + assert block_map.shape[-1] == 3 + assert calls["variable_block_sizes"] == [5, 64, 64] + assert block_map[:, :, :1, :].all(), "prefix queries stay dense" + assert block_map[..., :1].all(), "prefix keys are exempt from selection" + assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query" + + # The same odd geometry must still untile back onto the original rows. + dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + reference = torch.nn.functional.scaled_dot_product_attention( + *(x.float().transpose(1, 2) for x in (query, query, query)) + ).transpose(1, 2) + torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2) + + +def test_h3_forward_applies_the_learned_compression_gate(monkeypatch): + _fake_block_sparse_kernel(monkeypatch, {}) + impl = _h3_impl(topk=64) + prefix_segments, video_shape = (5,), (2, 4, 4) + seq_len = 5 + 2 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) + + without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + zero_gate = torch.zeros_like(query) + with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate)) + live_gate = torch.full_like(query, 0.5) + with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate)) + + torch.testing.assert_close(with_zero_gate.float(), without_gate.float()) + assert not torch.allclose(with_live_gate.float(), without_gate.float()) + + +def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch): + _fake_block_sparse_kernel(monkeypatch, {}) + impl = _h3_impl(topk=64) + prefix_segments, video_shape = (5,), (2, 4, 4) + valid = 5 + 2 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16) + padding = PackedPaddingMetadata( + q_length=valid, + kv_length=valid, + cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), + ) + + output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding)) + + assert output.shape == query.shape + assert torch.count_nonzero(output[:, valid:]) == 0 + + +def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d(): + partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( + (5, 70, 9), (5, 6, 6), torch.device("cpu") + ) + assert sizes[:4].tolist() == [5, 64, 6, 9] + assert prefix_blocks == 4 + assert video_blocks == 8 + assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6 + assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum()) + + +@pytest.mark.parametrize("fallback_on_error", [False, True]) +def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error): + impl = _h3_impl(fallback_on_error=fallback_on_error) + query = torch.randn(1, 37, 2, 8) + metadata = _h3_metadata((5,), (2, 4, 4)) + + def fail(*args, **kwargs): + raise RuntimeError("kernel failure") + + monkeypatch.setattr(impl, "_forward_h3", fail) + monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda) + if fallback_on_error: + output = impl.forward_cuda(query, query, query, metadata) + expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + else: + with pytest.raises(RuntimeError, match="kernel failure"): + impl.forward_cuda(query, query, query, metadata) + + +def test_h3_accelerator_failure_is_not_swallowed(monkeypatch): + impl = _h3_impl(fallback_on_error=True) + query = torch.randn(1, 37, 2, 8) + + def fail(*args, **kwargs): + raise torch.AcceleratorError("device failure") + + monkeypatch.setattr(impl, "_forward_h3", fail) + with pytest.raises(torch.AcceleratorError, match="device failure"): + impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4))) + + +def test_h3_without_layout_delegates_to_shared_backend(monkeypatch): + from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl + + impl = _h3_impl() + query = torch.randn(1, 37, 2, 8) + metadata = AttentionMetadata() + expected = torch.empty_like(query) + calls = [] + + def forward(self, q, k, v, attn_metadata): + calls.append((self, q, k, v, attn_metadata)) + return expected + + monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward) + assert impl.forward_cuda(query, query, query, metadata) is expected + assert len(calls) == 1 + assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata))) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 66e0cefab..0632db55d 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -23,6 +23,7 @@ from typing import Any, NamedTuple import numpy as np import psutil import yaml +from filelock import FileLock, Timeout from vllm import TextPrompt from vllm.logger import init_logger @@ -170,9 +171,54 @@ class OmniServer: self.use_omni = use_omni self.proc: subprocess.Popen | None = None self.host = "127.0.0.1" + self._auto_port = port is None + self._port_lock: FileLock | None = None self.port = get_open_port() if port is None else port + def _reserve_port(self) -> None: + # bind(0)/close does not reserve a port across concurrent pytest workers. + # Keep an advisory lock until teardown, including the long import phase + # before the subprocess binds its HTTP socket. Never unlink lock files: + # another worker may already have the same inode open. + for attempt in range(128): + if attempt: + self.port = get_open_port(self.host) + lock_path = Path(tempfile.gettempdir()) / f"vllm-omni-test-port-{os.getuid()}-{self.port}.lock" + lock = FileLock(lock_path) + try: + lock.acquire(timeout=0) + except Timeout as error: + if not self._auto_port: + raise RuntimeError(f"HTTP test port {self.port} is already reserved") from error + continue + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind((self.host, self.port)) + except OSError as error: + lock.release() + if not self._auto_port or error.errno != errno.EADDRINUSE: + raise + continue + self._port_lock = lock + return + raise RuntimeError("Could not reserve an HTTP test port after 128 attempts") + + def _owns_listening_port(self) -> bool: + assert self.proc is not None + try: + parent = psutil.Process(self.proc.pid) + processes = [parent, *parent.children(recursive=True)] + return any( + conn.status == psutil.CONN_LISTEN and conn.laddr.port == self.port + for process in processes + for conn in process.net_connections(kind="tcp") + ) + except psutil.NoSuchProcess: + # A child can exit between enumeration and inspecting its sockets. + return False + def _start_server(self) -> None: + self._reserve_port() env = os.environ.copy() if self.env_dict is not None: env.update(self.env_dict) @@ -208,7 +254,10 @@ class OmniServer: raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.") with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.settimeout(1) - if sock.connect_ex((self.host, self.port)) == 0: + if sock.connect_ex((self.host, self.port)) == 0 and self._owns_listening_port(): + ret = self.proc.poll() + if ret is not None: + raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.") startup_s = time.perf_counter() - startup_t0 if self.log_stats: print( @@ -368,9 +417,14 @@ class OmniServer: return self def __exit__(self, exc_type, exc_val, exc_tb): - if self.proc: - self._kill_process_tree(self.proc.pid) - cleanup_test_environment() + try: + if self.proc: + self._kill_process_tree(self.proc.pid) + finally: + if self._port_lock is not None: + self._port_lock.release() + self._port_lock = None + cleanup_test_environment() class OmniServerStageCli(OmniServer): diff --git a/tests/helpers/tests/test_runtime_server_ports.py b/tests/helpers/tests/test_runtime_server_ports.py new file mode 100644 index 000000000..1d8c0b2f1 --- /dev/null +++ b/tests/helpers/tests/test_runtime_server_ports.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""A listening port alone must not make a different test's server ready.""" + +import socket +import subprocess +import sys + +import pytest +from filelock import FileLock + +from tests.helpers import runtime + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_server_does_not_accept_another_process_listener(monkeypatch): + monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None) + server = runtime.OmniServer("fake-model", []) + popen = subprocess.Popen + + # Occupy the chosen port after allocation, while the requested subprocess + # is still importing. The old TCP-only probe returns ready for this socket. + with socket.socket() as foreign_listener: + + def launch(*args, **kwargs): + foreign_listener.bind((server.host, server.port)) + foreign_listener.listen() + return popen([sys.executable, "-c", "import time; time.sleep(0.2); raise SystemExit(7)"]) + + monkeypatch.setattr(runtime.subprocess, "Popen", launch) + with pytest.raises(RuntimeError, match="exited with code 7"): + with server: + pass + # Failed __enter__ must also release the lease for a later server. + retry = runtime.OmniServer("retry-model", [], port=server.port) + try: + retry._reserve_port() + finally: + retry.__exit__(None, None, None) + + +def test_auto_ports_are_reserved_before_subprocesses_bind(monkeypatch, tmp_path): + monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None) + monkeypatch.setattr(runtime.tempfile, "gettempdir", lambda: str(tmp_path)) + # Both constructors see the same free port before either server has bound. + first = runtime.OmniServer("first-model", []) + candidate = first.port + monkeypatch.setattr(runtime, "get_open_port", lambda *args: candidate) + second = runtime.OmniServer("second-model", []) + first._reserve_port() + with socket.socket() as occupied, socket.socket() as probe: + occupied.bind((first.host, candidate)) + probe.bind((first.host, 0)) + alternative = probe.getsockname()[1] + monkeypatch.setattr(runtime, "get_open_port", lambda *args: alternative) + try: + second._reserve_port() + assert first.port == candidate + assert second.port == alternative + # Check OS-level lock exclusion, not only the Python object's state. + lock_path = next(tmp_path.glob(f"*-{candidate}.lock")) + contender = FileLock(lock_path) + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys\nfrom filelock import FileLock, Timeout\n" + "try:\n FileLock(sys.argv[1]).acquire(timeout=0)\n" + "except Timeout:\n sys.exit(23)\n", + str(lock_path), + ], + timeout=10, + check=False, + ) + assert result.returncode == 23 + explicit = runtime.OmniServer("explicit-model", [], port=candidate) + with pytest.raises(RuntimeError, match="already reserved"): + with explicit: + pass + finally: + first.__exit__(None, None, None) + second.__exit__(None, None, None) + with contender.acquire(timeout=0): + pass + + +@pytest.mark.parametrize("in_child", [False, True]) +def test_server_accepts_its_own_listener(monkeypatch, in_child): + monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None) + server = runtime.OmniServer("fake-model", []) + popen = subprocess.Popen + + def launch(*args, **kwargs): + code = ( + "import socket,time; s=socket.socket(); " + f"s.bind(({server.host!r}, {server.port})); s.listen(); time.sleep(30)" + ) + if in_child: + code = f"import subprocess,sys; child=subprocess.Popen([sys.executable,'-c',{code!r}]); child.wait()" + return popen([sys.executable, "-c", code]) + + monkeypatch.setattr(runtime.subprocess, "Popen", launch) + with server: + assert server.proc is not None + assert server.proc.poll() is None + with socket.create_connection((server.host, server.port), timeout=1): + pass diff --git a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py index 465912d39..cbab82c9a 100644 --- a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py +++ b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py @@ -6,7 +6,6 @@ from __future__ import annotations import functools import importlib.util import math -import os from collections.abc import Mapping from typing import Any @@ -88,77 +87,6 @@ if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"): _fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd -# H3 needs an explicit per-query block map: prefix queries are dense, while -# video queries select prefix + top-k video tiles. The generic -# video_sparse_attn() entry point cannot express that contract. -if not hasattr(torch.ops.vllm_omni, "fastvideo_h3_vsa_bhsd"): - - @torch.library.custom_op("vllm_omni::fastvideo_h3_vsa_bhsd", mutates_args=()) - def _fastvideo_h3_vsa_bhsd_op( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, - logical_blocks: int, - ) -> torch.Tensor: - q = query.transpose(1, 2).contiguous() - k = key.transpose(1, 2).contiguous() - v = value.transpose(1, 2).contiguous() - - # Official FastVideo's measured FastH3 route opts into this native - # Blackwell forward. Wheels without the extension (including SM103 - # builds today) retain the corrected explicit-mask Triton route. - if os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1": - try: - from fastvideo_kernel import block_sparse_attn_sm100a - from fastvideo_kernel.triton_kernels.index import map_to_index - - if block_sparse_attn_sm100a.is_supported(q, variable_block_sizes): - q2k_idx, q2k_num = map_to_index(block_map) - out, _ = block_sparse_attn_sm100a.block_sparse_attn_sm100a( - q, - k, - v, - q2k_idx.to(torch.int32).contiguous(), - q2k_num.to(torch.int32).contiguous(), - variable_block_sizes.to(torch.int32).contiguous(), - need_lse=False, - ) - return out.transpose(1, 2).contiguous() - except (ImportError, RuntimeError) as exc: - # Opting in explicitly and then silently getting a different - # numeric path is worse than the slower route it lands on. - logger.warning_once( - "FASTVIDEO_VSA_SM100A=1 requested but the native Blackwell forward is " - "unavailable (%s); using the Triton block-sparse route instead.", - exc, - ) - - from fastvideo_kernel.block_sparse_attn import block_sparse_attn - - logical_len = logical_blocks * 64 - out, _ = block_sparse_attn( - q[:, :, :logical_len].contiguous(), - k[:, :, :logical_len].contiguous(), - v[:, :, :logical_len].contiguous(), - block_map[..., :logical_blocks, :logical_blocks].contiguous(), - variable_block_sizes[:logical_blocks].to(torch.int32).contiguous(), - ) - out = out.transpose(1, 2).contiguous() - if out.shape[1] != query.shape[1]: - out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, query.shape[1] - out.shape[1])) - return out - - @_fastvideo_h3_vsa_bhsd_op.register_fake - def _(query, key, value, block_map, variable_block_sizes, logical_blocks): - del key, value, block_map, variable_block_sizes, logical_blocks - return torch.empty_like(query) - - -_fastvideo_h3_vsa_bhsd_op = torch.ops.vllm_omni.fastvideo_h3_vsa_bhsd - - @functools.lru_cache(maxsize=32) def _get_tile_partition_indices( dit_seq_shape: tuple[int, int, int], @@ -226,81 +154,6 @@ def _get_tile_metadata( return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index -@functools.lru_cache(maxsize=32) -def _get_h3_tile_metadata( - prefix_segments: tuple[int, ...], - video_shape: tuple[int, int, int], - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: - """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles.""" - block_size = (4, 4, 4) - block_elements = 64 - prefix_len = sum(prefix_segments) - prefix_sizes: list[int] = [] - for segment in prefix_segments: - full, remainder = divmod(segment, block_elements) - prefix_sizes.extend([block_elements] * full) - if remainder: - prefix_sizes.append(remainder) - - video_indices = _get_tile_partition_indices(video_shape, block_size, device) + prefix_len - video_sizes = _construct_variable_block_sizes(video_shape, block_size, device) - partition = torch.cat([torch.arange(prefix_len, device=device, dtype=torch.long), video_indices]) - sizes = torch.cat([torch.tensor(prefix_sizes, device=device, dtype=torch.int32), video_sizes.to(torch.int32)]) - non_pad = _get_non_pad_index(sizes, block_elements) - untile = non_pad[torch.argsort(partition)] - total = prefix_len + math.prod(video_shape) - if int(sizes.sum()) != total or untile.numel() != total: - raise ValueError( - f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, " - f"sizes_sum={int(sizes.sum())}, total={total}" - ) - return partition, sizes, non_pad, untile, len(prefix_sizes), int(video_sizes.numel()) - - -def _get_h3_layout( - attn_metadata: AttentionMetadata | None, -) -> tuple[tuple[int, ...], tuple[int, int, int], int] | None: - if attn_metadata is None or attn_metadata.video_layout is None: - return None - prefix = attn_metadata.extra.get("vsa_h3_prefix_segments") - if not isinstance(prefix, (tuple, list)): - return None - target = next( - (span for span in reversed(attn_metadata.video_layout.video_spans) if span.role == "target"), - None, - ) - if target is None: - return None - return tuple(int(x) for x in prefix if int(x) > 0), target.latent_grid, target.start - - -def _pool_h3_tiles(x: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor: - batch, seq_len, heads, dim = x.shape - blocks = seq_len // 64 - pooled = x.view(batch, blocks, 64, heads, dim).sum(dim=2, dtype=torch.float32) - pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1) - return pooled.permute(0, 2, 1, 3) - - -def _build_h3_block_map( - scores: torch.Tensor, - num_prefix_blocks: int, - num_video_blocks: int, - topk: int, -) -> torch.Tensor: - """Prefix K/V are exempt and prefix queries stay dense, as in FastVideo.""" - keep_video = min(topk, num_video_blocks) - if keep_video == num_video_blocks: - return torch.ones_like(scores, dtype=torch.bool) - block_map = torch.zeros_like(scores, dtype=torch.bool) - indices = scores[..., num_prefix_blocks:].topk(keep_video, dim=-1).indices + num_prefix_blocks - block_map.scatter_(-1, indices, True) - block_map[..., :num_prefix_blocks] = True - block_map[:, :, :num_prefix_blocks, :] = True - return block_map - - def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None: if attn_metadata is None: return None @@ -343,7 +196,9 @@ class FastVideoVSABackend(AttentionBackend): if importlib.util.find_spec("fastvideo_kernel") is None: raise ImportError( "FASTVIDEO_VSA requires the optional fastvideo-kernel package " - "(the installation must provide the 'fastvideo_kernel' Python module)." + "included in vllm-omni[vsa]. Install with `uv pip install 'vllm-omni[vsa]'` " + "(from source: `uv pip install -e '.[vsa]'`). " + "Prebuilt kernels require Linux, Python 3.12 and glibc >= 2.34." ) @staticmethod @@ -504,93 +359,6 @@ class FastVideoVSAImpl(AttentionImpl): return "sequence parallel context is active" return None - def _forward_h3( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - layout = _get_h3_layout(attn_metadata) - if layout is None: - raise ValueError("incomplete VSA-H3 layout metadata") - prefix_segments, video_shape, target_start = layout - if sum(prefix_segments) != target_start: - raise ValueError(f"VSA-H3 prefix segments sum to {sum(prefix_segments)}, target starts at {target_start}") - expected = target_start + math.prod(video_shape) - if query.shape[1] != expected: - raise ValueError(f"VSA-H3 layout has {expected} rows but attention received {query.shape[1]}") - if query.dtype not in (torch.float16, torch.bfloat16): - raise ValueError(f"VSA-H3 requires fp16/bf16 tensors, got {query.dtype}") - gate = _get_gate_compress(attn_metadata) - if gate is not None: - # H3 pads the packed document to 64 rows, while VSA operates on - # the valid prefix. Validate/slice before launching attention so a - # metadata error can never trigger an unsafe dense fallback after - # an asynchronous custom kernel. - if gate.shape[0] != query.shape[0] or gate.shape[2:] != query.shape[2:] or gate.shape[1] < query.shape[1]: - raise ValueError(f"gate_compress shape {gate.shape} cannot cover query shape {query.shape}") - gate = gate[:, : query.shape[1]] - - partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( - prefix_segments, video_shape, query.device - ) - logical_blocks = int(sizes.numel()) - # The native sm100a kernel assigns pairs of query blocks to CTAs. Its - # contract requires an even block count; the synthetic partner is - # transport-only and is removed before returning. - pair_pad = logical_blocks % 2 - kernel_blocks = logical_blocks + pair_pad - target_shape = (query.shape[0], kernel_blocks * 64, query.shape[2], query.shape[3]) - q_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype) - k_tiled = torch.zeros_like(q_tiled) - v_tiled = torch.zeros_like(q_tiled) - q_tiled[:, non_pad] = query[:, partition] - k_tiled[:, non_pad] = key[:, partition] - v_tiled[:, non_pad] = value[:, partition] - - q_pool = _pool_h3_tiles(q_tiled[:, : logical_blocks * 64], sizes) - k_pool = _pool_h3_tiles(k_tiled[:, : logical_blocks * 64], sizes) - scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * self.softmax_scale - block_map = _build_h3_block_map(scores, prefix_blocks, video_blocks, self.topk) - kernel_sizes = sizes - if pair_pad: - block_map = torch.nn.functional.pad(block_map, (0, 1, 0, 1), value=False) - kernel_sizes = torch.nn.functional.pad(sizes, (0, 1), value=0) - - logger.info_once( - "FASTVIDEO_VSA H3 routing: seq_len=%d, prefix_segments=%s, video_shape=%s, " - "prefix_blocks=%d, video_blocks=%d, topk=%d, kernel_blocks=%d", - query.shape[1], - prefix_segments, - video_shape, - prefix_blocks, - video_blocks, - min(self.topk, video_blocks), - kernel_blocks, - ) - output = _fastvideo_h3_vsa_bhsd_op( - q_tiled.contiguous(), - k_tiled.contiguous(), - v_tiled.contiguous(), - block_map.contiguous(), - kernel_sizes.contiguous(), - logical_blocks, - )[:, : logical_blocks * 64] - - if gate is not None: - gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * 64]) - gate_tiled[:, non_pad] = gate[:, partition] - v_pool = _pool_h3_tiles(v_tiled[:, : logical_blocks * 64], sizes) - compressed = torch.matmul(torch.softmax(scores, dim=-1), v_pool) - compressed = compressed.permute(0, 2, 1, 3).to(output.dtype) - output = ( - output.view(output.shape[0], logical_blocks, 64, output.shape[2], output.shape[3]) - + compressed.unsqueeze(2) - * gate_tiled.view(gate_tiled.shape[0], logical_blocks, 64, gate_tiled.shape[2], gate_tiled.shape[3]) - ).view_as(output) - return output[:, untile].contiguous() - def forward_cuda( self, query: torch.Tensor, @@ -611,26 +379,6 @@ class FastVideoVSAImpl(AttentionImpl): key = key[:, :valid_seq_len] value = value[:, :valid_seq_len] - if attn_metadata is not None and _get_h3_layout(attn_metadata) is not None: - try: - output = self._forward_h3(query, key, value, attn_metadata) - if valid_seq_len == original_seq_len: - return output - restored = torch.zeros_like(original_query) - restored[:, :valid_seq_len] = output - return restored - except Exception as exc: - # A CUDA fault poisons the process context; attempting SDPA - # afterwards obscures the original kernel failure and cannot - # recover the request. - if isinstance(exc, torch.AcceleratorError): - raise - if not self.fallback_on_error: - raise - return self._fallback( - original_query, original_key, original_value, attn_metadata, f"VSA-H3 kernel failed: {exc}" - ) - reason = self._fallback_reason(query, key, value, attn_metadata) if reason is not None: return self._fallback(original_query, original_key, original_value, attn_metadata, reason) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4c588af0c..2e0b2604d 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -7,7 +7,9 @@ # https://github.com/feifeibear/long-context-attention/blob/main/yunchang/attention/layer.py +from collections.abc import Mapping from dataclasses import replace +from typing import cast import torch import torch.nn as nn @@ -16,7 +18,7 @@ from vllm.logger import init_logger from vllm.model_executor.models.utils import extract_layer_index from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention @@ -83,6 +85,10 @@ class Attention(nn.Module): custom_attention: nn.Module | None = None, # Preserve dense FP32 inference for models opting into CUDA auto fallback. allow_fp32_fallback: bool = False, + # Model-owned implementation subclasses, keyed by backend name. + # Keep the platform-selected backend and its capabilities unchanged. + # Each override must preserve the selected implementation's contract. + impl_overrides: Mapping[str, type[AttentionImpl]] | None = None, ): super().__init__() @@ -181,8 +187,18 @@ class Attention(nn.Module): self.backend_pref = attn_backend_cls.get_name() logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref) - self.attn_backend = attn_backend_cls + self.attn_backend: type[AttentionBackend] | None = attn_backend_cls self.attn_impl_cls = self.attn_backend.get_impl_cls() + if impl_overrides is not None: + override = impl_overrides.get(attn_backend_cls.get_name()) + if override is not None: + if not issubclass(override, self.attn_impl_cls): + raise TypeError( + f"Attention implementation override {override.__qualname__} must subclass " + f"the selected implementation {self.attn_impl_cls.__qualname__} " + f"for backend {attn_backend_cls.__qualname__}" + ) + self.attn_impl_cls = override self.attention = self.attn_impl_cls( num_heads=num_heads, head_size=head_size, @@ -197,7 +213,7 @@ class Attention(nn.Module): ) # Compatibility kernels run inside shared dispatch, between the # parallel strategy's input preparation and output restoration. - self.sdpa_fallback = SDPABackend.get_impl_cls()( + self.sdpa_fallback: AttentionImpl | None = SDPABackend.get_impl_cls()( num_heads=num_heads, head_size=head_size, softmax_scale=softmax_scale, @@ -299,6 +315,7 @@ class Attention(nn.Module): def _init_kv_cache_quantization(self, config) -> None: if config is None or self._has_custom_attention: return + assert self.attn_backend is not None dtype = getattr(config, "diffusion_kv_cache_dtype", None) if dtype == "auto": dtype = None @@ -400,6 +417,7 @@ class Attention(nn.Module): ) use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False): + assert self.attn_backend is not None raise NotImplementedError( f"Diffusion paged KV requires an Omni backend with paged support; " f"selected {self.attn_backend.get_name()}" @@ -496,7 +514,7 @@ class Attention(nn.Module): def _run_local_attention(self, query, key, value, attn_metadata): if self._has_custom_attention: - return self.attention(query, key, value, attn_metadata) + return cast(nn.Module, self.attention)(query, key, value, attn_metadata) self._assert_metadata_compatible(attn_metadata) @@ -506,7 +524,7 @@ class Attention(nn.Module): and query.dtype == torch.float32 and query.ndim == 4 and not self.backend_explicit - and self.attn_backend.get_name() == "FLASH_ATTN" + and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN" and ( attn_metadata is None or ( @@ -519,7 +537,7 @@ class Attention(nn.Module): ) ): logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.") - return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata) in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile # The startup KV-capacity profile needs tensor shapes, not a paged @@ -531,14 +549,14 @@ class Attention(nn.Module): self._scheduler_paged_kv and self.paged_kv_cache_role is not None and in_kv_memory_profile - and self.attn_backend.get_name() == "FLASH_ATTN" + and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN" and not current_omni_platform.supports_diffusion_dense_flash_attention() ): logger.warning_once( "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. " "Formal paged requests still use the platform-native paged attention backend." ) - return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata) return self.attention.forward(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/ops/block_sparse.py b/vllm_omni/diffusion/attention/ops/block_sparse.py new file mode 100644 index 000000000..aeec14699 --- /dev/null +++ b/vllm_omni/diffusion/attention/ops/block_sparse.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared pooling, routing and provider calls for explicit block-sparse attention.""" + +import os + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +if not hasattr(torch.ops.vllm_omni, "fastvideo_block_sparse_attn_bshd"): + + @torch.library.custom_op("vllm_omni::fastvideo_block_sparse_attn_bshd", mutates_args=()) + def fastvideo_block_sparse_attn_bshd( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + logical_blocks: int, + ) -> torch.Tensor: + """Run explicit block-sparse attention on BSHD tensors in 64-row blocks. + + Inputs and the block map may include transport-only blocks after + ``logical_blocks``. The Triton provider sees only logical blocks and + the wrapper restores the original row count. Native provider selection + and its supported-input checks remain owned by FastVideo. + """ + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + + # Prefer the explicitly selected native provider when it supports + # these tensors; retain the existing Triton provider otherwise. + if os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1": + try: + from fastvideo_kernel import block_sparse_attn_sm100a + from fastvideo_kernel.triton_kernels.index import map_to_index + + if block_sparse_attn_sm100a.is_supported(q, variable_block_sizes): + q2k_idx, q2k_num = map_to_index(block_map) + out, _ = block_sparse_attn_sm100a.block_sparse_attn_sm100a( + q, + k, + v, + q2k_idx.to(torch.int32).contiguous(), + q2k_num.to(torch.int32).contiguous(), + variable_block_sizes.to(torch.int32).contiguous(), + need_lse=False, + ) + return out.transpose(1, 2).contiguous() + except torch.AcceleratorError: + # A device fault can poison the context; do not launch another kernel. + raise + except (ImportError, RuntimeError) as exc: + logger.warning_once( + "FASTVIDEO_VSA_SM100A=1 requested but the native Blackwell forward is " + "unavailable (%s); using the Triton block-sparse route instead.", + exc, + ) + + from fastvideo_kernel.block_sparse_attn import block_sparse_attn + + logical_len = logical_blocks * 64 + out, _ = block_sparse_attn( + q[:, :, :logical_len].contiguous(), + k[:, :, :logical_len].contiguous(), + v[:, :, :logical_len].contiguous(), + block_map[..., :logical_blocks, :logical_blocks].contiguous(), + variable_block_sizes[:logical_blocks].to(torch.int32).contiguous(), + ) + out = out.transpose(1, 2).contiguous() + if out.shape[1] != query.shape[1]: + out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, query.shape[1] - out.shape[1])) + return out + + @fastvideo_block_sparse_attn_bshd.register_fake + def _(query, key, value, block_map, variable_block_sizes, logical_blocks): + del key, value, block_map, variable_block_sizes, logical_blocks + return torch.empty_like(query) + + +fastvideo_block_sparse_attn_bshd = torch.ops.vllm_omni.fastvideo_block_sparse_attn_bshd + + +def mean_pool_tiles(x: torch.Tensor, sizes: torch.Tensor, block_size: int) -> torch.Tensor: + """Mean-pool zero-padded BSHD tiles with FP32 accumulation. + + ``sizes`` gives the number of valid rows in each tile. + """ + batch, seq_len, heads, dim = x.shape + blocks = seq_len // block_size + pooled = x.view(batch, blocks, block_size, heads, dim).sum(dim=2, dtype=torch.float32) + pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1) + return pooled.permute(0, 2, 1, 3) + + +def build_prefix_dense_block_map( + scores: torch.Tensor, + num_prefix_blocks: int, + num_sparse_blocks: int, + topk: int, +) -> torch.Tensor: + """Keep prefix queries/keys dense and select top-k non-prefix keys per row.""" + keep_sparse = min(topk, num_sparse_blocks) + if keep_sparse == num_sparse_blocks: + return torch.ones_like(scores, dtype=torch.bool) + block_map = torch.zeros_like(scores, dtype=torch.bool) + indices = scores[..., num_prefix_blocks:].topk(keep_sparse, dim=-1).indices + num_prefix_blocks + block_map.scatter_(-1, indices, True) + block_map[..., :num_prefix_blocks] = True + block_map[:, :, :num_prefix_blocks, :] = True + return block_map diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 3aca638ec..a7539b34a 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1262,7 +1262,7 @@ def extract_minimax_h3_context( preprocessing via ``_embed``, cacheable ``blocks`` loop, and ``final_layer`` postprocessing with row selection and update masks. """ - from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_ + from vllm_omni.diffusion.layers.indexed_modulation import indexed_scale_shift_ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( _FORWARD_SUPPORTED_KWARGS, MINIMAX_H3_ADALN_MODALITY_NUM, diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/layers/indexed_modulation.py similarity index 98% rename from vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py rename to vllm_omni/diffusion/layers/indexed_modulation.py index 20327b7e9..385ce1a1b 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/layers/indexed_modulation.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Fused MiniMax H3 modulation with FP32 accumulation.""" +"""Fused indexed modulation with FP32 accumulation.""" from __future__ import annotations @@ -244,7 +244,7 @@ def rms_norm_indexed_scale_shift( indices: torch.Tensor, eps: float, ) -> torch.Tensor: - """Fuse H3 RMSNorm with its indexed AdaLN affine transform.""" + """Fuse RMSNorm with an indexed affine transform.""" if x.is_cpu: input_dtype = x.dtype normalized = x.float() diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py b/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py b/vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py new file mode 100644 index 000000000..35b87a58b --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax-H3 VSA prefix routing, tile layout, and learned compression gate.""" + +from __future__ import annotations + +import functools +import math + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.fastvideo_vsa import ( + FastVideoVSAImpl, + _construct_variable_block_sizes, + _get_gate_compress, + _get_non_pad_index, + _get_tile_partition_indices, +) +from vllm_omni.diffusion.attention.ops.block_sparse import ( + build_prefix_dense_block_map, + fastvideo_block_sparse_attn_bshd, + mean_pool_tiles, +) + +logger = init_logger(__name__) + + +@functools.lru_cache(maxsize=32) +def _get_h3_tile_metadata( + prefix_segments: tuple[int, ...], + video_shape: tuple[int, int, int], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles.""" + block_size = (4, 4, 4) + block_elements = 64 + prefix_len = sum(prefix_segments) + prefix_sizes: list[int] = [] + for segment in prefix_segments: + full, remainder = divmod(segment, block_elements) + prefix_sizes.extend([block_elements] * full) + if remainder: + prefix_sizes.append(remainder) + + video_indices = _get_tile_partition_indices(video_shape, block_size, device) + prefix_len + video_sizes = _construct_variable_block_sizes(video_shape, block_size, device) + partition = torch.cat([torch.arange(prefix_len, device=device, dtype=torch.long), video_indices]) + sizes = torch.cat([torch.tensor(prefix_sizes, device=device, dtype=torch.int32), video_sizes.to(torch.int32)]) + non_pad = _get_non_pad_index(sizes, block_elements) + untile = non_pad[torch.argsort(partition)] + total = prefix_len + math.prod(video_shape) + if int(sizes.sum()) != total or untile.numel() != total: + raise ValueError( + f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, " + f"sizes_sum={int(sizes.sum())}, total={total}" + ) + return partition, sizes, non_pad, untile, len(prefix_sizes), int(video_sizes.numel()) + + +def _get_h3_layout( + attn_metadata: AttentionMetadata | None, +) -> tuple[tuple[int, ...], tuple[int, int, int], int] | None: + if attn_metadata is None or attn_metadata.video_layout is None: + return None + prefix = attn_metadata.extra.get("vsa_h3_prefix_segments") + if not isinstance(prefix, (tuple, list)): + return None + target = next( + (span for span in reversed(attn_metadata.video_layout.video_spans) if span.role == "target"), + None, + ) + if target is None: + return None + return tuple(int(x) for x in prefix if int(x) > 0), target.latent_grid, target.start + + +class MiniMaxH3VSAImpl(FastVideoVSAImpl): + """Apply H3 tile64 routing after shared parallel attention dispatch.""" + + def _forward_h3( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + layout = _get_h3_layout(attn_metadata) + if layout is None: + raise ValueError("incomplete VSA-H3 layout metadata") + prefix_segments, video_shape, target_start = layout + if sum(prefix_segments) != target_start: + raise ValueError(f"VSA-H3 prefix segments sum to {sum(prefix_segments)}, target starts at {target_start}") + expected = target_start + math.prod(video_shape) + if query.shape[1] != expected: + raise ValueError(f"VSA-H3 layout has {expected} rows but attention received {query.shape[1]}") + if query.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"VSA-H3 requires fp16/bf16 tensors, got {query.dtype}") + gate = _get_gate_compress(attn_metadata) + if gate is not None: + # H3 pads the packed document to 64 rows, while VSA operates on + # the valid prefix. Validate/slice before launching attention so a + # metadata error can never trigger an unsafe dense fallback after + # an asynchronous custom kernel. + if gate.shape[0] != query.shape[0] or gate.shape[2:] != query.shape[2:] or gate.shape[1] < query.shape[1]: + raise ValueError(f"gate_compress shape {gate.shape} cannot cover query shape {query.shape}") + gate = gate[:, : query.shape[1]] + + partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( + prefix_segments, video_shape, query.device + ) + logical_blocks = int(sizes.numel()) + # The native sm100a kernel assigns pairs of query blocks to CTAs. Its + # contract requires an even block count; the synthetic partner is + # transport-only and is removed before returning. + pair_pad = logical_blocks % 2 + kernel_blocks = logical_blocks + pair_pad + target_shape = (query.shape[0], kernel_blocks * 64, query.shape[2], query.shape[3]) + q_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype) + k_tiled = torch.zeros_like(q_tiled) + v_tiled = torch.zeros_like(q_tiled) + q_tiled[:, non_pad] = query[:, partition] + k_tiled[:, non_pad] = key[:, partition] + v_tiled[:, non_pad] = value[:, partition] + + q_pool = mean_pool_tiles(q_tiled[:, : logical_blocks * 64], sizes, block_size=64) + k_pool = mean_pool_tiles(k_tiled[:, : logical_blocks * 64], sizes, block_size=64) + scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * self.softmax_scale + block_map = build_prefix_dense_block_map(scores, prefix_blocks, video_blocks, self.topk) + kernel_sizes = sizes + if pair_pad: + block_map = torch.nn.functional.pad(block_map, (0, 1, 0, 1), value=False) + kernel_sizes = torch.nn.functional.pad(sizes, (0, 1), value=0) + + logger.info_once( + "FASTVIDEO_VSA H3 routing: seq_len=%d, prefix_segments=%s, video_shape=%s, " + "prefix_blocks=%d, video_blocks=%d, topk=%d, kernel_blocks=%d", + query.shape[1], + prefix_segments, + video_shape, + prefix_blocks, + video_blocks, + min(self.topk, video_blocks), + kernel_blocks, + ) + output = fastvideo_block_sparse_attn_bshd( + q_tiled.contiguous(), + k_tiled.contiguous(), + v_tiled.contiguous(), + block_map.contiguous(), + kernel_sizes.contiguous(), + logical_blocks, + )[:, : logical_blocks * 64] + + if gate is not None: + gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * 64]) + gate_tiled[:, non_pad] = gate[:, partition] + v_pool = mean_pool_tiles(v_tiled[:, : logical_blocks * 64], sizes, block_size=64) + compressed = torch.matmul(torch.softmax(scores, dim=-1), v_pool) + compressed = compressed.permute(0, 2, 1, 3).to(output.dtype) + output = ( + output.view(output.shape[0], logical_blocks, 64, output.shape[2], output.shape[3]) + + compressed.unsqueeze(2) + * gate_tiled.view(gate_tiled.shape[0], logical_blocks, 64, gate_tiled.shape[2], gate_tiled.shape[3]) + ).view_as(output) + return output[:, untile].contiguous() + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + if _get_h3_layout(attn_metadata) is None: + return super().forward_cuda(query, key, value, attn_metadata) + + original_query, original_key, original_value = query, key, value + original_seq_len = query.shape[1] + valid_seq_len = original_seq_len + if attn_metadata is not None and attn_metadata.packed_padding is not None: + valid_seq_len = attn_metadata.packed_padding.q_length + if attn_metadata.packed_padding.kv_length != valid_seq_len: + return self._fallback( + original_query, original_key, original_value, attn_metadata, "packed Q/KV lengths must match" + ) + query = query[:, :valid_seq_len] + key = key[:, :valid_seq_len] + value = value[:, :valid_seq_len] + + try: + output = self._forward_h3(query, key, value, attn_metadata) + if valid_seq_len == original_seq_len: + return output + restored = torch.zeros_like(original_query) + restored[:, :valid_seq_len] = output + return restored + except Exception as exc: + # A CUDA fault poisons the process context; attempting SDPA + # afterwards obscures the original kernel failure and cannot + # recover the request. + if isinstance(exc, torch.AcceleratorError): + raise + if not self.fallback_on_error: + raise + return self._fallback( + original_query, original_key, original_value, attn_metadata, f"VSA-H3 kernel failed: {exc}" + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index f0839429b..e26493931 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -33,12 +33,6 @@ from vllm_omni.diffusion.attention.backends.abstract import ( VideoTokenLayout, ) from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( - indexed_gate, - indexed_gate_rms_norm_scale_shift, - indexed_scale_shift_, - rms_norm_indexed_scale_shift, -) from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -50,6 +44,12 @@ from vllm_omni.diffusion.forward_context import ( ) from vllm_omni.diffusion.layers.activation import SiluAndMul from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope +from vllm_omni.diffusion.layers.indexed_modulation import ( + indexed_gate, + indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, + rms_norm_indexed_scale_shift, +) from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract @@ -442,6 +442,8 @@ class MiniMaxH3Attention(nn.Module): self._gate_hidden_size = arch.hidden_size self._gate_quant_config = quant_config self._gate_prefix = f"{prefix}.to_gate_compress" + from .attention.fastvideo_h3 import MiniMaxH3VSAImpl + self.attention = Attention( num_heads=self.num_heads, num_kv_heads=self.num_kv_heads, @@ -454,6 +456,7 @@ class MiniMaxH3Attention(nn.Module): role_category=role_category, skip_sequence_parallel=skip_sequence_parallel, prefix=prefix, + impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl}, ) def enable_vsa_gate(self) -> None: From 5f7f62e13cfb88bc1add00db5d3ac2e9ef1328f5 Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 16 Sep 2026 16:24:31 +0800 Subject: [PATCH 107/209] [Bugfix][MiniCPM-o] Stabilize code2wav CFM CUDA graph caching with reference audio normalization and mel-frame bucketing (#7416) Signed-off-by: y_null Co-authored-by: amy-why-3459 --- tests/config/test_config_factory.py | 6 +- .../minicpmo_4_5/test_cfm_graph_bucketing.py | 162 ++++++++++++++ .../minicpmo_4_5/test_code2wav_batching.py | 102 ++++++++- .../minicpmo_4_5/test_cuda_graph_wrapper.py | 44 +++- .../test_reference_audio_normalization.py | 204 ++++++++++++++++++ vllm_omni/deploy/minicpmo_4_5.yaml | 2 +- vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 2 +- vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 2 +- vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 2 +- .../models/minicpmo_4_5/batched_token2wav.py | 147 +++++++++++++ .../models/minicpmo_4_5/cuda_graph_wrapper.py | 33 ++- .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 143 +++++++++++- .../npu/models/minicpmo_4_5_code2wav.py | 14 ++ 13 files changed, 838 insertions(+), 25 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index d1dc0fd6d..3699eeda7 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1328,7 +1328,11 @@ class TestDeployConfigLoading: stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.session_mode == "duplex" - assert deploy.active_stream_window == max_sessions + # ``0`` disables the limiter: the configured session/stage capacity is + # what bounds the pipeline. Any positive window serializes concurrent + # audio first-packet generation (measured: TTFP -23% at 0 vs 4 on + # L20X, 128 requests / concurrency 8). + assert deploy.active_stream_window == 0 assert deploy.duplex_session.max_sessions == max_sessions assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py new file mode 100644 index 000000000..c88fc6d5b --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the CFM mel-frame bucketing decision. + +Steady-state chunk lengths vary per request; without bucketing every length +becomes its own CUDA-graph capture shape (428 captures / 13 flushes in the +#6628 regression). ``_cfm_pad_frames`` aligns the frame axis onto a grid so +steady-state calls share one capture shape; ``_decode_cfm`` trims the output +back. These tests pin the decision math on CPU, no CUDA required. +""" + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( + _cfm_pad_frames, + _zero_padded_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_pads_partial_chunk_up_to_bucket(): + # 50 frames on a 16-frame grid -> 14 padding frames. + assert _cfm_pad_frames(mel_frames=50, offset=0, noise_capacity=1000, bucket_frames=16, disabled=False) == 14 + + +def test_aligned_chunk_needs_no_padding(): + assert _cfm_pad_frames(mel_frames=64, offset=0, noise_capacity=1000, bucket_frames=16, disabled=False) == 0 + + +def test_ragged_valid_lengths_path_disables_bucketing(): + """The ragged per-request cache path must never be padded.""" + for ragged in (True,): + for mel in (50, 64, 96): + assert ( + _cfm_pad_frames( + mel_frames=mel, + offset=0, + noise_capacity=1000, + bucket_frames=16, + disabled=ragged, + ) + == 0 + ) + + +def test_disabled_bucketing_returns_zero(): + """bucket_frames <= 1 (or wrapper absent) means plain eager behavior.""" + for bucket in (0, 1): + assert ( + _cfm_pad_frames( + mel_frames=50, + offset=0, + noise_capacity=1000, + bucket_frames=bucket, + disabled=False, + ) + == 0 + ) + + +def test_padding_never_overflows_noise_buffer(): + """Padding past the decoder's noise buffer must fall back to zero. + + The padded call slices x from decoder.rand_noise[offset:end]; exceeding + its width would crash. Bucketing is best-effort, never required. + """ + # 96 frames + 16 pad = 112 > 100 capacity -> give up, pad 0. + assert _cfm_pad_frames(mel_frames=96, offset=0, noise_capacity=100, bucket_frames=16, disabled=False) == 0 + # Same mel length fits when un-padded: 96 <= 100. + assert _cfm_pad_frames(mel_frames=96, offset=0, noise_capacity=100, bucket_frames=1, disabled=False) == 0 + # Cache offset eats into the capacity: 64 + 0 pad would fit, 64 + 48 pad + # would not, so the pad must be dropped. + assert _cfm_pad_frames(mel_frames=64, offset=50, noise_capacity=110, bucket_frames=16, disabled=False) == 0 + # The original chunk fits while the padded one does not: 50 + 50 = 100 + # <= 110, but 100 + 14 = 114 > 110, so the pad must be dropped. + assert _cfm_pad_frames(mel_frames=50, offset=50, noise_capacity=110, bucket_frames=16, disabled=False) == 0 + # Same numbers with room for the padding: 114 <= 120, so it stands. + assert _cfm_pad_frames(mel_frames=50, offset=50, noise_capacity=120, bucket_frames=16, disabled=False) == 14 + + +def test_padding_applies_with_cache_offset_when_it_fits(): + # offset 50 + mel 50 + pad 14 = 114 <= 200 -> pad stands. + assert _cfm_pad_frames(mel_frames=50, offset=50, noise_capacity=200, bucket_frames=16, disabled=False) == 14 + + +def test_steady_state_cache_width_saturates_on_bucket_grid(): + """Pin the real recurrence: grow, then trim to ``prompt_len + 100``. + + ``_decode_batch_once`` trims the estimator attention cache after every + decode, so the width saturates instead of growing without bound. What + matters for the graph cache is that the steady-state + ``(chunk_width, cache_width)`` pair settles on a single grid point. + """ + bucket, prompt_len = 16, 304 + cache_cap = prompt_len + 100 + width = prompt_len + shapes = set() + for _ in range(20): + pad = _cfm_pad_frames( + mel_frames=50, + offset=width, + noise_capacity=30000, + bucket_frames=bucket, + disabled=False, + ) + assert pad == 14 + shapes.add((50 + pad, width)) # (chunk width, cache width) as used + width = min(width + 50 + pad, cache_cap) + assert width == cache_cap + # (64, 304) -> (64, 368) -> (64, 404), then stable. + assert shapes == {(64, 304), (64, 368), (64, 404)} + + +def test_varied_chunk_lengths_collapse_onto_few_widths(): + """Bucketing exists for the varied-length calls, not the steady 50-frame one. + + The first/last chunk and the ``plan_token2wav_encode_slices`` splits land on + arbitrary lengths; those are the calls that would each become their own + capture shape. Pin that a realistic mix collapses onto a few padded widths. + """ + varied = (7, 12, 17, 25, 33, 50, 51, 64) + padded = { + mel + + _cfm_pad_frames( + mel_frames=mel, + offset=300, + noise_capacity=30000, + bucket_frames=16, + disabled=False, + ) + for mel in varied + } + assert len(padded) == 4, sorted(padded) # 16 / 32 / 48 / 64 + assert len(padded) < len(varied) + + +def test_zero_padded_frames_survives_an_integration_step(): + """The padded region must stay zero after ``x = x + dt * velocity``. + + Zeroing once before the CFM loop is not enough: the update touches every + column, so the padded region becomes non-zero again (0 -> 0.010 -> 0.020 + over the first steps). Pin the helper that re-zeroes it each step. + """ + mel, pad = 50, 14 + x = torch.zeros(2, 4, mel + pad) + x[..., :mel] = 0.05 + _zero_padded_frames(x, mel) + assert torch.all(x[..., mel:] == 0.0) + for _ in range(3): + x = x + 0.02 * torch.ones_like(x) # the integration step + _zero_padded_frames(x, mel) + assert torch.all(x[..., mel:] == 0.0) + assert torch.all(x[..., :mel] > 0.0) + + +def test_zero_padded_frames_is_a_noop_without_padding(): + x = torch.ones(1, 2, 32) + _zero_padded_frames(x, None) + _zero_padded_frames(x, 32) + assert torch.all(x == 1.0) diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index a0e7ed4db..edb813af3 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -523,7 +523,7 @@ def test_cfm_graph_receives_bfloat16_cache_in_compute_dtype(): seen_dtypes: list[torch.dtype] = [] class RecordingGraph: - def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out): + def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out, attn_mask=None): del time_embedding, cnn_cache seen_dtypes.append(att_cache.dtype) return estimator_input[:, :1], cnn_out, att_out @@ -1313,3 +1313,103 @@ def test_reference_voice_and_duplex_metadata_follow_request_lifecycle(): assert prompt_key not in model._runtime_prompts assert not Path(prompt_wav).exists() assert (prompt_cache_id, prompt_wav) not in model.backend._prompt_features + + +def _eager_cfm_wrapper(estimator: nn.Module) -> SimpleNamespace: + """A CFM graph wrapper that honours the contract but runs eagerly. + + ``_decode_cfm`` only takes the padded path when a wrapper is present, so + CPU tests that need that path (CUDA graphs are unavailable here) supply + this in place of the real wrapper. + """ + + def replay(estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out, attn_mask=None): + output = estimator.blocks_forward_chunk( + estimator_input, + time_emb, + attn_mask, + cnn_cache, + att_cache, + cnn_out, + att_out, + ) + return output, cnn_out, att_out + + return SimpleNamespace(enabled=True, replay=replay) + + +def _padded_decode_adapter(bucket_frames: int) -> tuple[BatchedToken2Wav, int, int]: + """Build an adapter over a real ``DiT`` that exercises the CFM path. + + Returns the adapter, the valid width, and the padding width. Every call + seeds the same way, so two adapters built here agree on weights and noise. + """ + from cosyvoice2.flow.decoder_dit import DiT + + torch.manual_seed(23) + estimator = DiT( + in_channels=4, + out_channels=1, + depth=2, + num_heads=2, + head_dim=2, + hidden_size=4, + ).eval() + token2wav = _FakeToken2Wav() + token2wav.flow.decoder.estimator = estimator + token2wav.flow.decoder.rand_noise = torch.randn(1, 1, 160) + adapter = BatchedToken2Wav(token2wav) + adapter._cfm_graph_wrapper = _eager_cfm_wrapper(estimator) + adapter._cfm_graph_bucket_frames = bucket_frames + mel_frames = 51 # the 16-frame grid pads this to 64 + pad_frames = (-(-mel_frames // 16) * 16 - mel_frames) if bucket_frames else 0 + return adapter, mel_frames, pad_frames + + +def test_padded_chunk_keeps_the_cross_chunk_caches_on_the_valid_boundary(): + """Caches the next chunk reads must not carry padding. + + The padded steps still run, so ``step_cnn``/``step_att`` come back with + padding-derived values in them. Unless those positions are cleared, the + following chunk consumes them as its keys and as its CNN left context. + This drives a padded chunk through the real ``DiT`` and checks both caches + at the boundary. + """ + adapter, mel_frames, pad_frames = _padded_decode_adapter(bucket_frames=16) + assert pad_frames == 13 + + mu = torch.randn(1, 1, mel_frames) + cond = torch.zeros(1, 1, mel_frames) + speakers = torch.randn(1, 1) + with torch.no_grad(): + x, cnn, att = adapter._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None) + + assert int(x.shape[2]) == mel_frames + assert int(att.shape[4]) == mel_frames + pad_frames + # These columns become the next chunk's keys. + assert torch.count_nonzero(att[..., mel_frames : mel_frames + pad_frames, :]) == 0 + # This tail becomes the next chunk's CNN left context. + assert torch.count_nonzero(cnn[..., -pad_frames:]) == 0 + + +def test_padding_does_not_change_the_valid_frames(): + """A padded chunk returns the same valid frames as an exact-width one. + + Masking takes the padded keys out of the attention and the padded CNN tail + never reaches the valid positions, so the only difference between the two + runs is the capture shape -- not the audio the chunk produces. + """ + padded, mel_frames, pad_frames = _padded_decode_adapter(bucket_frames=16) + exact, _, exact_pad = _padded_decode_adapter(bucket_frames=0) + assert pad_frames == 13 + assert exact_pad == 0 + + mu = torch.randn(1, 1, mel_frames, generator=torch.Generator().manual_seed(5)) + cond = torch.zeros(1, 1, mel_frames) + speakers = torch.randn(1, 1, generator=torch.Generator().manual_seed(7)) + with torch.no_grad(): + padded_x, _, _ = padded._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None) + exact_x, _, _ = exact._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None) + + assert int(padded_x.shape[2]) == int(exact_x.shape[2]) == mel_frames + assert torch.allclose(padded_x, exact_x, atol=1e-5) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 26daeff7b..959156b8f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -252,6 +252,37 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( wrapper._flush() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_capture_keeps_the_real_mask(monkeypatch: pytest.MonkeyPatch) -> None: + """Static buffers must be built from the real inputs, mask included. + + Capturing from zero-filled placeholders would bake the "nothing is masked" + branch into the graph, while replay copies the real mask into those very + buffers -- so the captured branch has to come from the real value. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4) + + inputs = _cfm_inputs(2, 12, 0) + frames = int(inputs[0].shape[2]) + mask = torch.ones(2, frames, frames, dtype=torch.bool, device="cuda") + mask[:, :, -2:] = False # two padded keys are masked out + + with torch.inference_mode(): + wrapper.replay(*inputs, mask) + assert wrapper._stats["captures"] == 1 + + static_inputs, _, _ = next(iter(wrapper._cache.values())) + static_mask = static_inputs[6] + assert static_mask is not None + assert torch.equal(static_mask, mask) + wrapper._flush() + + def test_importing_wrapper_does_not_resolve_platform() -> None: """Importing this module must not build the OmniPlatform singleton. @@ -354,7 +385,8 @@ def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.Monkey """A key that cannot round-trip is a property of one shape, not of the GPU. It must not disable the wrapper or retire the generation the way a capture - failure does. + failure does. `_tensors_from_key` is the fallback used when no real inputs + are available, so the probe drives `_capture` without them. """ pool = torch.cuda.graph_pool_handle() monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) @@ -374,7 +406,10 @@ def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.Monkey monkeypatch.setattr(wrapper_module, "_tensors_from_key", _reject_one_shape) with torch.inference_mode(): - wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0)) + key = ("estimator_step",) + tuple( + wrapper_module._tensor_signature(t) for t in _cfm_inputs(2, unbuildable_width, 0) + ) + assert wrapper._capture(key, None) is None assert wrapper.enabled is True assert wrapper._stats["captures"] == 0 @@ -382,11 +417,6 @@ def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.Monkey wrapper.replay(*_cfm_inputs(2, 12, 0)) assert wrapper._stats["captures"] == 1 - # and the rejected shape stays eager without retrying the capture - wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0)) - assert wrapper._stats["captures"] == 1 - assert wrapper._stats["eager"] == 2 - wrapper._flush() diff --git a/tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py b/tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py new file mode 100644 index 000000000..9b9bbf570 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the reference-audio normalization grid. + +Every request's reference audio is folded onto one sample rate and one fixed +length so the CFM attention cache origin (L0) is a single constant; that is +what collapsed the CFM CUDA-graph key space back from 58 shapes to a handful +and cut the capture storm behind the #6628 regression. These tests pin the +pure math so future edits cannot silently widen L0 again. +""" + +import logging +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + _REF_MAX_SECONDS, + _REF_TARGET_SAMPLE_RATE, + _normalize_reference, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TARGET = _REF_TARGET_SAMPLE_RATE +MAX_SAMPLES = int(_REF_MAX_SECONDS * TARGET) # 6 s @ 24 kHz + + +def test_mono_downmix_averages_channels_without_interleaving(): + """(2, T) stereo must become channel-mean mono, not a reshape interleave. + + reshape(-1) on (2, T) would alternate 1/0 samples; mean(dim=0) yields a + constant 0.5 waveform. The whole-output equality fails on any interleave. + """ + stereo = torch.zeros(2, 16) + stereo[0] = 1.0 + waveform, sr = _normalize_reference(stereo, TARGET) + assert sr == TARGET + assert waveform.shape == (MAX_SAMPLES,) + assert torch.all(waveform[:16] == 0.5) # mono body + assert torch.all(waveform[16:] == 0.0) # zero-padded tail + + +def test_resample_48k_to_24k_halves_sample_count(): + """Off-grid sample rates are resampled onto 24 kHz before padding.""" + wav_48k = torch.arange(24000, dtype=torch.float32) + waveform, sr = _normalize_reference(wav_48k, 48000) + assert sr == TARGET + assert waveform.shape == (MAX_SAMPLES,) + assert torch.all(waveform[:12000] > 0) # resampled body, then silence + assert torch.all(waveform[12000:] == 0) + + +def test_long_reference_is_truncated_to_6s(): + samples_7s = torch.arange(7 * TARGET, dtype=torch.float32) + waveform, sr = _normalize_reference(samples_7s, TARGET) + assert sr == TARGET + assert waveform.shape == (MAX_SAMPLES,) + assert torch.equal(waveform, samples_7s[:MAX_SAMPLES]) + + +def test_short_reference_is_zero_padded_to_6s(): + """The zero-pad is the load-bearing half of S1b. + + Truncation alone only caps long references; short ones would keep + distinct L0 values (verified experimentally: rtf 3.375 with truncation + only, 1.2 with zero-padding). Pin the padding behavior. + """ + samples_1s = torch.arange(TARGET, dtype=torch.float32) + waveform, sr = _normalize_reference(samples_1s, TARGET) + assert sr == TARGET + assert waveform.shape == (MAX_SAMPLES,) + assert torch.equal(waveform[:TARGET], samples_1s) + assert torch.all(waveform[TARGET:] == 0) + + +def test_exact_length_reference_passes_through_unchanged(): + samples_6s = torch.arange(MAX_SAMPLES, dtype=torch.float32) + waveform, sr = _normalize_reference(samples_6s, TARGET) + assert sr == TARGET + assert torch.equal(waveform, samples_6s) + + +def test_every_output_shares_one_length(): + """Whatever comes in, one request's L0 equals any other's.""" + for length in (TARGET // 4, TARGET, 3 * TARGET, MAX_SAMPLES + 1): + waveform, _ = _normalize_reference(torch.ones(length), TARGET) + assert waveform.numel() == MAX_SAMPLES + + +def test_empty_reference_stays_empty(): + """An empty waveform must not be padded into a valid-length reference. + + ``_materialize_runtime_prompt`` relies on numel() == 0 to raise + "empty_ref_audio" and reject the request; zero-padding would turn a + broken input into a silently accepted all-silence reference. + """ + waveform, sr = _normalize_reference(torch.zeros(0), TARGET) + assert waveform.numel() == 0 + assert sr == TARGET + + +def test_window_length_is_configurable(): + """Deployments can widen the window via ``ref_audio_max_seconds``.""" + samples_5s = torch.arange(5 * TARGET, dtype=torch.float32) + waveform, sr = _normalize_reference(samples_5s, TARGET, max_seconds=3.0) + assert sr == TARGET + assert waveform.numel() == 3 * TARGET + # A longer window keeps the whole 5s clip instead of truncating it. + waveform, _ = _normalize_reference(samples_5s, TARGET, max_seconds=8.0) + assert waveform.numel() == 8 * TARGET + assert torch.equal(waveform[: 5 * TARGET], samples_5s) + + +def test_truncation_emits_warning(caplog): + """Silent truncation would hide a user-visible behavior change. + + A bare ``caplog.at_level`` does not capture vLLM loggers (they do not + propagate to the root handlers the fixture taps), so attach the fixture + handler to the module logger directly -- same idiom as + ``tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py``. + """ + samples_7s = torch.arange(7 * TARGET, dtype=torch.float32) + target_logger = logging.getLogger("vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: + _normalize_reference(samples_7s, TARGET) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) + assert any("truncating" in record.getMessage() for record in caplog.records) + + +def test_soundfile_2d_output_is_transposed_before_the_normalizer(tmp_path): + """``sf.read`` gives (samples, channels); the normalizer wants the transpose. + + Without it a stereo default prompt looks like a huge channel count, the + downmix guard rejects it, and the caller falls back to the raw file -- so + that prompt keeps a second L0 value next to the normalized request + references. + """ + import soundfile as sf + + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + _read_reference_wav, + ) + + frames = TARGET + # Opposite-signed channels: the downmix is a true mean only when the two + # channels were aligned as rows, so a wrong layout cannot look correct. + stereo = torch.stack((torch.full((frames,), 0.5), torch.full((frames,), -0.5)), dim=1) + path = tmp_path / "stereo.wav" + sf.write(str(path), stereo.numpy(), TARGET, format="WAV") + + waveform, sample_rate_hz = _read_reference_wav(str(path)) + + assert sample_rate_hz == TARGET + assert waveform.shape == (2, frames) + normalized, target_sr = _normalize_reference(waveform, sample_rate_hz, max_seconds=1.0) + assert target_sr == TARGET + assert normalized.shape[-1] == TARGET + assert torch.allclose(normalized, torch.zeros_like(normalized), atol=1e-6) + + +def test_default_prompt_uses_the_configured_window(tmp_path): + """The default prompt has to honour ``ref_audio_max_seconds``. + + With a non-default window, a default prompt normalized at the fixed 6 s + constant would keep a different L0 from the request references, which is + the mismatch this normalization exists to remove. + """ + import tempfile + + import soundfile as sf + + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + MiniCPMO45Code2Wav, + _read_reference_wav, + ) + + frames = 2 * TARGET # shorter than the window, so it gets zero-padded + samples = torch.full((frames, 2), 0.25) + source = tmp_path / "default.wav" + sf.write(str(source), samples.numpy(), TARGET, format="WAV") + + window_seconds = 8.0 + with tempfile.TemporaryDirectory() as runtime_dir: + stub = SimpleNamespace( + _default_prompt_wav=str(source), + _default_prompt_id="HT_ref_audio", + _default_prompt_normalized=None, + _ref_max_seconds=window_seconds, + _runtime_prompt_dir=SimpleNamespace(name=runtime_dir), + ) + prompt_wav, _ = MiniCPMO45Code2Wav._normalized_default_prompt(stub) + # Read inside the block: the temporary directory is removed on exit. + normalized, sample_rate_hz = _read_reference_wav(prompt_wav) + + assert sample_rate_hz == TARGET + # The configured window, not the fixed 6 s default. + assert normalized.shape[-1] == int(window_seconds * TARGET) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index ca772afbd..5c26fc2c6 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -6,7 +6,6 @@ async_chunk: true session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. -active_stream_window: 4 duplex_session: # Matches Stage 0/1 max_num_seqs below. max_sessions: 4 @@ -30,6 +29,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 9600e5959..136dace7e 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -9,7 +9,6 @@ async_chunk: true session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. -active_stream_window: 4 duplex_session: # Matches Stage 0/1 max_num_seqs below. max_sessions: 4 @@ -33,6 +32,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index aae531218..9669a2768 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -9,7 +9,6 @@ async_chunk: true session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. -active_stream_window: 4 duplex_session: # Matches Stage 0/1 max_num_seqs below. max_sessions: 4 @@ -33,6 +32,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 00b1090f8..92065dd4f 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -19,7 +19,6 @@ async_chunk: true # (/v1/realtime?duplex=1). The default duplex_session caps sessions at 1, # matching Stage 0/1 max_num_seqs below. session_mode: duplex -active_stream_window: 1 trust_remote_code: true enable_prefix_caching: false @@ -40,6 +39,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 2bffeb49e..bec442282 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -46,6 +46,80 @@ def relpos_encode_token_budget( return max(lookahead + 1, min(int(cap), room)) +def _cfm_pad_frames( + *, + mel_frames: int, + offset: int, + noise_capacity: int, + bucket_frames: int, + disabled: bool, +) -> int: + """Frame padding that aligns one CFM call onto the capture-shape grid. + + Chunk lengths vary per request: the steady-state chunk is 50 mel frames, + but the first/last chunk and the ``plan_token2wav_encode_slices`` splits + land on arbitrary lengths, and every distinct length is its own CUDA-graph + capture shape. Padding the frame axis up to ``bucket_frames`` collapses + them onto one grid (the caller trims the output back). Measured on the + shipped config, 128 requests / concurrency 8: 61 captures and 1 cache + flush without bucketing versus 15 captures and 0 flushes with it, audio + RTF 1.317 -> 1.191 under otherwise identical conditions. + + The attention and CNN caches also advance by the padded width so their + shapes stay on the same grid; in practice that is bounded by the trim in + ``_decode_batch_once`` (the cache width saturates at ``prompt_len + + 100``) rather than by the decoder's ``rand_noise`` capacity. + + Returns 0 when bucketing does not apply: the ragged valid-lengths path, + graphs unavailable or already disabled by ``CFMGraphWrapper._disable``, + or padding that would overflow the decoder's noise buffer. + """ + if disabled or bucket_frames <= 1: + return 0 + pad = (bucket_frames - mel_frames % bucket_frames) % bucket_frames + if offset + mel_frames + pad > noise_capacity: + return 0 + return pad + + +def _zero_padded_frames(tensor: torch.Tensor, valid_frames: int | None) -> None: + """Keep the padded columns of ``tensor`` at zero, in place. + + Bucketing pads the frame axis so the capture shape stays fixed. The padded + columns carry no valid content: they are excluded from the attention by the + ``attn_mask`` built in ``_decode_batch_once``, and zeroing keeps them from + leaking back in. The integration step ``x = x + dt * velocity`` would + otherwise make the padded region non-zero again after the initial zeroing, + and the cross-chunk caches would read those values back. + """ + if valid_frames is None or valid_frames >= int(tensor.shape[-1]): + return + tensor[..., valid_frames:] = 0.0 + + +def _zero_padded_cnn_cache( + cnn_cache: torch.Tensor, + estimator: nn.Module, + pad_frames: int, +) -> None: + """Clear the cache positions that come from padded frames, in place. + + Each block's CNN cache holds the tail of its convolution output + (``new_cnn_cache = x[..., -causal_padding[0]:]``, inside ``stepaudio2``), + so when the padding sits at the chunk tail those positions are + padding-derived and would otherwise become the next chunk's left context. + Padding is at most ``pad_frames`` wide, so only the trailing positions that + can come from it are cleared; the valid part of the window is kept. + """ + for index, block in enumerate(estimator.blocks): + width = int(block.conv.block[1].causal_padding[0]) + if width <= 0: + continue + zero_from = max(0, width - pad_frames) + if zero_from < width: + cnn_cache[index][..., zero_from:] = 0.0 + + def plan_token2wav_encode_slices( num_frames: int, *, @@ -249,6 +323,17 @@ class BatchedToken2Wav(nn.Module): "CFM CUDA Graph is disabled on device type %s", flow_parameter.device.type if flow_parameter is not None else "unknown", ) + # mel-frame bucket size for the CFM CUDA Graph path. Pad each decode + # chunk up to a multiple of this many frames so the graph cache key + # space stays small (0 disables bucketing, e.g. when graphs are off). + self._cfm_graph_bucket_frames = ( + int(cfm_graph_cfg.get("bucket_frames", 0)) if self._cfm_graph_wrapper is not None else 0 + ) + if self._cfm_graph_bucket_frames > 1: + logger.info( + "CFM CUDA Graph bucketing enabled (bucket_frames=%d)", + self._cfm_graph_bucket_frames, + ) self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} def _hift_inference( @@ -398,6 +483,7 @@ class BatchedToken2Wav(nn.Module): att_cache: torch.Tensor | None, attn_mask: torch.Tensor | None = None, valid_lengths: list[int] | None = None, + valid_frames: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if self._trt_stepper is not None and valid_lengths is None: out, new_cnn, new_att = self._trt_stepper.step( @@ -413,6 +499,13 @@ class BatchedToken2Wav(nn.Module): time_embedding = estimator.t_embedder(time).unsqueeze(1) width = int(x.shape[-1]) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) + if valid_frames is not None and valid_frames < width: + # Every frame gets the speaker vector, so the padded columns would + # stay non-zero even with x/mu/cond zeroed -- and the estimator + # would then emit non-zero features there, leaving the + # chunk-boundary attention/CNN caches padding-derived. Zero them. + speaker_features = speaker_features.clone() + _zero_padded_frames(speaker_features, valid_frames) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) @@ -433,6 +526,7 @@ class BatchedToken2Wav(nn.Module): graph_att, cnn_out, att_out, + attn_mask, ) if valid_lengths is not None: if not hasattr(estimator, "in_proj"): @@ -566,6 +660,26 @@ class BatchedToken2Wav(nn.Module): estimator = decoder.estimator batch_size = int(mu.shape[0]) offset = int(att_cache.shape[4]) if att_cache is not None else 0 + mel_frames = int(mu.shape[2]) + pad_frames = _cfm_pad_frames( + mel_frames=mel_frames, + offset=offset, + noise_capacity=int(decoder.rand_noise.shape[2]), + bucket_frames=self._cfm_graph_bucket_frames, + disabled=( + valid_lengths is not None + or self._cfm_graph_wrapper is None + # `_disable` keeps the wrapper object alive, so check the flag + # too: padding under a disabled wrapper is pure overhead. + or not self._cfm_graph_wrapper.enabled + ), + ) + if pad_frames: + # Replicate rather than zero: a repeated last frame is a closer + # continuation of the chunk than silence, so the padded positions + # carry less of a step change into the attention and CNN caches. + mu = torch.nn.functional.pad(mu, (0, pad_frames), mode="replicate") + cond = torch.nn.functional.pad(cond, (0, pad_frames), mode="replicate") end = offset + int(mu.shape[2]) if end > int(decoder.rand_noise.shape[2]): raise RuntimeError( @@ -574,6 +688,13 @@ class BatchedToken2Wav(nn.Module): f'"available":{int(decoder.rand_noise.shape[2])}}}' ) x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + if pad_frames: + # The padded columns would otherwise carry real noise values. They + # are excluded from this chunk's attention by ``attn_mask`` below + # and held at zero here, so the padded region stays inert. Masking + # is what removes them; zeroing alone would not, since a zero row + # still occupies part of the softmax denominator. + x[:, :, mel_frames:] = 0.0 timeline = torch.linspace( 0, 1, @@ -603,6 +724,20 @@ class BatchedToken2Wav(nn.Module): device=mu.device, ) attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2) + elif pad_frames: + # Mask the padded keys instead of only zeroing their content: a + # zero-valued key/value pair still takes probability mass out of the + # softmax denominator, so the real frames keep attending to the + # padding unless it is explicitly excluded. + kv_len = int(mu.shape[2]) + offset + attn_mask = torch.ones( + 2 * batch_size, + int(mu.shape[2]), + kv_len, + dtype=torch.bool, + device=mu.device, + ) + attn_mask[:, :, mel_frames : mel_frames + pad_frames] = False next_cnn: list[torch.Tensor] = [] next_att_cache: torch.Tensor | None = None ragged_att_cache: list[torch.Tensor] | None = None @@ -622,10 +757,14 @@ class BatchedToken2Wav(nn.Module): att_cache=old_att, attn_mask=attn_mask, valid_lengths=valid_lengths, + valid_frames=mel_frames if pad_frames else None, ) + if pad_frames: + _zero_padded_cnn_cache(step_cnn, estimator, pad_frames) conditional, unconditional = estimate.split(batch_size, dim=0) velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional x = x + dt * velocity + _zero_padded_frames(x, mel_frames if pad_frames else None) time = time + dt if step + 1 < self.n_timesteps: dt = timeline[step + 2] - time[0] @@ -666,7 +805,15 @@ class BatchedToken2Wav(nn.Module): dtype=self._estimator_att_cache_dtype, ) next_att_cache[step].copy_(step_att) + if pad_frames: + # The padded steps ran but hold no valid content, and each + # cache entry becomes the next chunk's keys. Clear the + # padded columns so the boundary the next chunk reads is + # the valid one, in step with the mask built above. + next_att_cache[step][..., mel_frames : mel_frames + pad_frames, :] = 0.0 del step_att + if pad_frames: + x = x[:, :, :mel_frames] if ragged_att_cache is not None: return x, torch.stack(next_cnn), ragged_att_cache assert next_att_cache is not None diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index d737905fd..4db0c39bf 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py @@ -135,11 +135,13 @@ class HiFTGraphWrapper: return speech, cache_source -def _tensor_signature(value: torch.Tensor) -> tuple: +def _tensor_signature(value: torch.Tensor | None) -> tuple: + if value is None: + return (None,) return tuple(value.shape), str(value.dtype), str(value.device) -_DTYPE_MAP = {str(dtype): dtype for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.float64)} +_DTYPE_MAP = {str(dtype): dtype for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.float64, torch.bool)} def _memory_snapshot(device: torch.device) -> tuple[int, int] | None: @@ -231,7 +233,7 @@ class CFMGraphWrapper: return {**self._stats, "cache_size": len(self._cache)} def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor: - return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5]) + return self.graph_fn(args[0], args[1], args[6], args[2], args[3], args[4], args[5]) def _eager(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: self._stats["eager"] += 1 @@ -261,10 +263,21 @@ class CFMGraphWrapper: self.enabled = False self._flush() - def _capture(self, key: tuple) -> tuple | None: - """Capture a CUDA graph for the given key. Returns None on failure.""" + def _capture(self, key: tuple, inputs: tuple | None = None) -> tuple | None: + """Capture a CUDA graph for the given key. Returns None on failure. + + ``inputs`` should be the real tensors for this key whenever they are + available: the static buffers must be built from the values the graph + will actually run with. Zero-filled placeholders are only a fallback -- + capturing with them bakes the placeholder branch into the graph (e.g. an + attention mask that looks like "nothing is masked"). + """ try: - static_inputs = _tensors_from_key(key) + static_inputs = ( + tuple(None if tensor is None else tensor.detach().clone() for tensor in inputs) + if inputs is not None + else _tensors_from_key(key) + ) except KeyError: # An unsupported dtype is a property of this shape alone, so run it # eager and keep the other shapes on graphs. @@ -315,8 +328,9 @@ class CFMGraphWrapper: att_cache: torch.Tensor, cnn_out: torch.Tensor, att_out: torch.Tensor, + attn_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out) + inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out, attn_mask) self._stats["calls"] += 1 if not self.enabled or torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda": @@ -330,7 +344,7 @@ class CFMGraphWrapper: if entry is None: if len(self._cache) >= self.max_graphs: self._flush() - entry = self._capture(key) + entry = self._capture(key, inputs) if entry is None: return self._eager(inputs) self._cache[key] = entry @@ -339,7 +353,8 @@ class CFMGraphWrapper: static_inputs, static_output, graph = entry for static, current in zip(static_inputs, inputs, strict=True): - static.copy_(current) + if static is not None: + static.copy_(current) graph.replay() return ( static_output.detach().clone(), diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index f701e8b26..da2c0eb72 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -9,6 +9,7 @@ import os import tempfile from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass +from functools import lru_cache from hashlib import sha256 from pathlib import Path from typing import Any @@ -57,6 +58,92 @@ def _scalar(value: Any, default: Any = None) -> Any: return default if value is None else value +_REF_TARGET_SAMPLE_RATE = 24000 +# Overridable per deployment via ``ref_audio_max_seconds``. +_REF_MAX_SECONDS = 6.0 +# Channel ceiling for the optional (channels, samples) downmix guard. +_REF_MAX_CHANNELS = 8 + + +@lru_cache(maxsize=8) +def _get_resampler(orig_freq: int, new_freq: int): + """Cached anti-aliased resampler matching token2wav's own loader. + + ``torchaudio`` is imported lazily (it is not a hard runtime requirement of + this model on every platform) and the transform is built once per rate pair + instead of once per request, keeping it off the first-packet path. + """ + import torchaudio + + return torchaudio.transforms.Resample(orig_freq=orig_freq, new_freq=new_freq) + + +def _read_reference_wav(path: str) -> tuple[Any, int]: + """Read a WAV file for :func:`_normalize_reference`. + + ``soundfile`` returns ``(samples, channels)`` for multi-channel files + while the normalizer expects ``(channels, samples)``. Without the + transpose a stereo prompt looks like a huge channel count, the downmix + guard rejects it, and the caller silently falls back to the raw file -- + a second L0 value for the same content. + """ + waveform, sample_rate_hz = sf.read(path, dtype="float32", always_2d=False) + if getattr(waveform, "ndim", 1) > 1: + waveform = waveform.T + return waveform, int(sample_rate_hz) + + +def _normalize_reference( + ref_audio: Any, + sample_rate_hz: int, + *, + max_seconds: float = _REF_MAX_SECONDS, +) -> tuple[torch.Tensor, int]: + """Normalize a per-request reference waveform onto a shared grid. + + MiniCPM-o streaming uses the reference-audio length as the CFM attention + cache origin (L0). Unnormalized references give every request a distinct + L0, which explodes the CFM CUDA-graph cache key space (see #6628). Fold + every reference onto the same sample rate and a fixed length (truncate + long ones, zero-pad short ones) so L0 becomes one constant value. + + References longer than ``max_seconds`` are truncated with a warning; + the window is configurable via ``ref_audio_max_seconds``. + """ + tensor = torch.as_tensor(ref_audio, dtype=torch.float32) + if tensor.dim() > 1: + # (channels, samples) -> mono; plain reshape(-1) would interleave. + # Upstream flattens request references to 1-D, so anything else is a + # non-standard caller: reject layouts we cannot downmix unambiguously + # instead of silently averaging the waveform away. + if tensor.shape[0] > _REF_MAX_CHANNELS: + raise ValueError(f"reference audio must be 1-D or (channels, samples); got shape {tuple(tensor.shape)}") + tensor = tensor.mean(dim=0) + waveform = tensor.reshape(-1).cpu().contiguous() + if waveform.numel() == 0: + # Keep the caller's "empty_ref_audio" error path intact: an empty + # waveform must not be zero-padded into a valid-length reference. + return waveform, _REF_TARGET_SAMPLE_RATE + if sample_rate_hz != _REF_TARGET_SAMPLE_RATE: + # Match token2wav's own loader (torchaudio, anti-aliased): it only + # resamples when the stored rate differs, so this output is what the + # model consumes. + waveform = _get_resampler(sample_rate_hz, _REF_TARGET_SAMPLE_RATE)(waveform.view(1, -1)).view(-1).contiguous() + max_samples = max(1, int(max_seconds * _REF_TARGET_SAMPLE_RATE)) + if waveform.numel() > max_samples: + logger.warning( + "Reference audio is %.2fs, truncating to the %.2fs window (set ``ref_audio_max_seconds`` to keep more).", + waveform.numel() / _REF_TARGET_SAMPLE_RATE, + max_seconds, + ) + waveform = waveform[:max_samples].contiguous() + elif waveform.numel() < max_samples: + # Zero-pad short references so every request shares one L0; trailing + # silence has minimal style impact (verified via E3 WER/SIM). + waveform = torch.nn.functional.pad(waveform, (0, max_samples - waveform.numel())) + return waveform, _REF_TARGET_SAMPLE_RATE + + def _codec_tensor(value: Any, fallback: torch.Tensor) -> torch.Tensor: if isinstance(value, torch.Tensor): return value.reshape(-1).to(device=fallback.device, dtype=torch.long) @@ -181,7 +268,11 @@ class MiniCPMO45Code2Wav(nn.Module): self._cfm_graph_config = { "enabled": bool(extra.get("enable_cfm_graph", False)), "max_graphs": int(extra.get("cfm_max_graphs", 32)), + "bucket_frames": int(extra.get("cfm_graph_bucket_frames", 0)), } + self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS)) + if self._ref_max_seconds <= 0: + raise ValueError("MiniCPM-o Code2Wav ref_audio_max_seconds must be > 0") self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") @@ -192,6 +283,7 @@ class MiniCPMO45Code2Wav(nn.Module): raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") + self._default_prompt_normalized: tuple[str, str] | None = None @property def _default_prompt_wav(self) -> str: @@ -199,6 +291,46 @@ class MiniCPMO45Code2Wav(nn.Module): return str(self._prompt_wav_override) return str(Path(self.model_path) / "assets" / "HT_ref_audio.wav") + def _normalized_default_prompt(self) -> tuple[str, str]: + """Fold the shipped default prompt onto the request-reference grid. + + The shipped asset is 6.016 s and is loaded directly by token2wav, so a + request without a reference would otherwise keep a second L0 value a + couple of frames away from the normalized request references -- two + graph-key families instead of one (both use ``ref_audio_max_seconds``). + Returns ``(prompt_wav, + prompt_cache_id)``; if the asset cannot be read, the shipped path is + returned unchanged. + """ + if self._default_prompt_normalized is not None: + return self._default_prompt_normalized + source = self._default_prompt_wav + fallback = (source, self._default_prompt_id) + try: + waveform, sample_rate_hz = _read_reference_wav(source) + normalized, target_sr = _normalize_reference( + waveform, + sample_rate_hz, + max_seconds=self._ref_max_seconds, + ) + except Exception: + logger.warning("Could not normalize the default prompt %s; using it as-is", source) + self._default_prompt_normalized = fallback + return self._default_prompt_normalized + if normalized.numel() == 0: + self._default_prompt_normalized = fallback + return self._default_prompt_normalized + digest = sha256() + digest.update(normalized.numpy().tobytes()) + digest.update(str(target_sr).encode()) + cache_id = f"default-ref-{digest.hexdigest()[:24]}-{target_sr}" + path = Path(self._runtime_prompt_dir.name) / f"{cache_id}.wav" + if not path.is_file(): + sf.write(path, normalized.numpy(), target_sr, format="WAV") + logger.info("Default prompt normalized onto the %.2fs window: %s", self._ref_max_seconds, path) + self._default_prompt_normalized = (str(path), cache_id) + return self._default_prompt_normalized + def _extra_config(self) -> dict[str, Any]: model_config = getattr(self.vllm_config, "model_config", None) connector = getattr(model_config, "stage_connector_config", None) @@ -220,9 +352,13 @@ class MiniCPMO45Code2Wav(nn.Module): sample_rate: Any, ) -> tuple[str, _RuntimePrompt]: sample_rate_hz = int(_scalar(sample_rate, 0)) - waveform = torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1).cpu().contiguous() if sample_rate_hz <= 0: raise _batch_error("invalid_ref_audio_sample_rate", sample_rate=sample_rate_hz) + waveform, sample_rate_hz = _normalize_reference( + ref_audio, + sample_rate_hz, + max_seconds=self._ref_max_seconds, + ) if waveform.numel() == 0: raise _batch_error("empty_ref_audio") if not bool(torch.isfinite(waveform).all().item()): @@ -288,9 +424,10 @@ class MiniCPMO45Code2Wav(nn.Module): if entry is not None: return entry.cache_id, entry.path, cache_key + default_wav, default_cache_id = self._normalized_default_prompt() return ( - str(_scalar(meta.get("prompt_cache_id"), self._default_prompt_id)), - str(_scalar(meta.get("prompt_wav"), self._default_prompt_wav)), + str(_scalar(meta.get("prompt_cache_id"), default_cache_id)), + str(_scalar(meta.get("prompt_wav"), default_wav)), None, ) diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py index 0d34c93a5..4346e83b6 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -75,10 +75,15 @@ def _graphable_estimator_step( cond, cnn_cache, att_cache, + valid_frames=None, ): """Run the CFM estimator body after host-backed timestep embedding.""" width = int(x.shape[-1]) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) + if valid_frames is not None and valid_frames < width: + # Match the CUDA path: padded columns must not carry the speaker vector. + speaker_features = speaker_features.clone() + speaker_features[:, :, valid_frames:] = 0.0 estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = backend._estimator_buffers(estimator, estimator_input, att_cache) old_cnn = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) @@ -108,6 +113,7 @@ def _patched_estimator_step( att_cache, attn_mask=None, valid_lengths=None, + valid_frames=None, ): assert _original_estimator_step is not None graph_runner = _backend_graph_runners.get(self) @@ -130,10 +136,16 @@ def _patched_estimator_step( att_cache=att_cache, attn_mask=attn_mask, valid_lengths=valid_lengths, + valid_frames=valid_frames, ) if (cnn_cache is None) != (att_cache is None): raise ValueError("estimator CNN and attention caches must both be present or absent") + # valid_frames only matters for padded inputs; skip forwarding it for + # unpadded calls so platform-owned graphable bodies that predate the + # padding feature keep working with the old signature. + graphable_kwargs = {"valid_frames": valid_frames} if valid_frames is not None else {} + # The upstream embedder creates a frequency tensor on the host. Keep it # outside capture while retaining the tensor-only estimator body in graph. time_embedding = estimator.t_embedder(time).unsqueeze(1) @@ -152,6 +164,7 @@ def _patched_estimator_step( cond=step_cond, cnn_cache=None, att_cache=None, + **graphable_kwargs, ), ) @@ -169,6 +182,7 @@ def _patched_estimator_step( cond=step_cond, cnn_cache=step_cnn, att_cache=step_att, + **graphable_kwargs, ), ) From e78a5d01f785795b59bdcbc9f68214a3ce8cd06b Mon Sep 17 00:00:00 2001 From: Linze Shi Date: Wed, 16 Sep 2026 01:30:16 -0700 Subject: [PATCH 108/209] [Frontend] Add MiniMax-H3 references and Ref2VA workflow (#7483) Signed-off-by: Linze-Shi --- apps/ComfyUI-vLLM-Omni/README.md | 74 ++ .../comfyui_vllm_omni/nodes.py | 36 +- .../comfyui_vllm_omni/utils/api_client.py | 78 +- .../comfyui_vllm_omni/utils/types.py | 6 + ...LM-Omni MiniMax-H3 Reference to Video.json | 786 ++++++++++++++++++ recipes/MiniMaxAI/MiniMax-H3.md | 11 +- .../comfyui/test_comfyui_integration.py | 57 +- .../comfyui/test_h3_reference_workflow.py | 115 +++ .../features/comfyui/test_video_references.py | 138 +++ 9 files changed, 1235 insertions(+), 66 deletions(-) create mode 100644 apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax-H3 Reference to Video.json create mode 100644 tests/e2e/features/comfyui/test_h3_reference_workflow.py create mode 100644 tests/e2e/features/comfyui/test_video_references.py diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md index b848a0716..8f445635d 100644 --- a/apps/ComfyUI-vLLM-Omni/README.md +++ b/apps/ComfyUI-vLLM-Omni/README.md @@ -139,6 +139,80 @@ You can configure per-stage sampling parameters for multi-stage models. > > Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect. +For MiniMax-H3 Ref2VA, **Video References** accepts up to 9 images (`image_1`–`image_9`), +3 videos (`video_1`–`video_3`), and 3 audio clips (`audio_1`–`audio_3`), with at most +12 connected inputs in total. Any mixture containing at least one image or video +is supported; empty and audio-only references are rejected. For example, you can +combine 6 images, 3 videos, and 3 audio clips in one request. + +Within each media type, references follow slot-number order, skipping unconnected +slots. For example, connecting `image_2` and `image_9` sends `image_2` as the first +image and `image_9` as the second. Each image slot uses the first image in its batch. +Existing connections to `image_1`, `image_2`, `audio_1`, `audio_2`, `video_1`, and +`video_2` remain valid in saved workflows. + +#### MiniMax-H3 Reference to Video + +Load [MiniMax-H3 Reference to Video](example_workflows/vLLM-Omni%20MiniMax-H3%20Reference%20to%20Video.json) +from **Templates → ComfyUI-vLLM-Omni**, or drag the JSON onto the canvas. +It adapts the [official ComfyUI R2V workflow](https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_r2v.json) +to remote vLLM-Omni execution. Model loading, sampling, and VAE decoding run on +the server; ComfyUI loads references and saves the returned video. + +Use an extension version containing the audio-preserving output fix from +[#7456](https://github.com/vllm-project/vllm-omni/pull/7456). Older versions can +drop the server's audio while decoding the response. + +1. Start a Ref2VA-capable service using the [MiniMax-H3 recipe](../../recipes/MiniMaxAI/MiniMax-H3.md). + Set **Generate Video** to its `/v1` URL and served model name. +2. Select your own image in **Load Image**. The template's filenames are placeholders + relative to ComfyUI's input directory; no sample assets or model weights are bundled. +3. Connect **Load Video** and **Load Audio** when needed. Duplicate the loaders to + fill more reference slots, or disconnect the image when using video-only input. + Keep the references connected to **Generate Video** and leave `frame` disconnected. +4. Match prompt tags such as ``, `